Rustでコンパイラプラグインを書いてみる

rustでコンパイラプラグインを書きくなったため色々試行錯誤した記録

1.22.0-nightly で確認
nightly API使ってる(?)ので要注意

準備

$ cargo new sample-plugin

Cargo.toml

[package]
name = "sample_plugin"
version = "0.1.0"
authors = ["UEKI Masaru <[e-mail]>"]

[lib]
name = "sample_plugin"
path = "src/lib.rs"
plugin = true

[dependencies]

plugin = trueが無いとエラーが出る。

文字列を返す

まずは単純に文字列を返す例

src/lib.rs

#![crate_type="dylib"]
#![feature(plugin_registrar, rustc_private)]

extern crate syntax;
extern crate syntax_pos;
extern crate rustc;
extern crate rustc_plugin;

use syntax::ast;
use syntax::tokenstream::TokenTree;
use syntax::ext::base::{ExtCtxt, MacResult, MacEager};
use syntax::ext::build::AstBuilder;
use syntax_pos::Span;
use syntax_pos::symbol::Symbol;
use rustc_plugin::Registry;

fn expand_sample(cx: &mut ExtCtxt, sp: Span, _args: &[TokenTree])
                 -> Box<MacResult + 'static> {
    MacEager::expr(cx.expr_lit(sp, ast::LitKind::Str(Symbol::intern("hoge"), ast::StrStyle::Cooked)))
}

#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_macro("sample", expand_sample);
}

expand_sample()がプラグイン本体。plugin_registrar()ではexpand_sampleをsample!()として登録している。

プラグインではAST(抽象構文木)を構築して返す必要がある。
ここでは"hoge"という文字リテラルをexpr_lit()を使って構築している。

テストを書いて試してみる

tests/sample_plugin_tests.rs

#![feature(plugin)]
#![plugin(sample_plugin)]

#[test] fn sample_plugin_test() {
    let hoge = sample!();
    assert_eq!(hoge, "hoge");
}

let hoge = sample!();の部分がコンパイル時にlet hoge = "hoge;"と置き換わる想定。

$ cargo testで確かに動作していることが確認できる。

ブロックを定義する

次はブロックを定義してみる。今度はsample!()が

{ "hoge" }

と置き換えられる想定。

src/lib.rs

fn expand_sample(cx: &mut ExtCtxt, sp: Span, _args: &[TokenTree])
                 -> Box<MacResult + 'static> {
    let expr = cx.expr_lit(sp, ast::LitKind::Str(Symbol::intern("hoge"), ast::StrStyle::Cooked));
    let stmt = cx.stmt_expr(expr);
    MacEager::expr(cx.expr_block(cx.block(sp, vec!(stmt))))
}

$ cargo test で確認できる。

letとメソッド呼出

今度はsample!()が

{
  let h = "hoge".to_string();
  h
}

と置き換えられるようにしてみる。

src/lib.rs

fn expand_sample(cx: &mut ExtCtxt, sp: Span, _args: &[TokenTree])
                 -> Box<MacResult + 'static> {
    let ident_h = Ident::from_str("h");
    let hoge_lit = cx.expr_lit(sp, ast::LitKind::Str(Symbol::intern("hoge"), ast::StrStyle::Cooked));

    let stmt_let_h = cx.stmt_let(sp, false, ident_h, cx.expr_method_call(sp, hoge_lit, Ident::from_str("to_string"), vec!()));

    let stmt_ident = cx.stmt_expr(cx.expr_ident(sp, ident_h));
    MacEager::expr(cx.expr_block(cx.block(sp, vec!(stmt_let_h, stmt_ident))))
}

expr_method_call()によって"hoge"に対してto_string()を呼出し、stmt_let()により結果をhに束縛している。

テストはそのままでOK

別プラグインの呼出

試しに別のプラグインを呼び出してみたが難なく動いた。
今回は

{
  let h = "hoge".to_string();
  h.push_str(include_str!("fuga.txt"));
  h
}

と置き換わる想定

src/fuga.txt

fuga

src/lib.rs

fn expand_sample(cx: &mut ExtCtxt, sp: Span, _args: &[TokenTree])
                 -> Box<MacResult + 'static> {
    let fuga_str = include_str!("fuga.txt");

    let ident_h = Ident::from_str("h");
    let hoge_lit = cx.expr_lit(sp, ast::LitKind::Str(Symbol::intern("hoge"), ast::StrStyle::Cooked));

    let stmt_let_h = cx.stmt_let(sp, true, ident_h, cx.expr_method_call(sp, hoge_lit, Ident::from_str("to_string"), vec!()));

    let fuga_lit = cx.expr_lit(sp, ast::LitKind::Str(Symbol::intern(&fuga_str[0..4]), ast::StrStyle::Cooked));

    let stmt_push_str = cx.stmt_expr(cx.expr_method_call(sp, cx.expr_ident(sp,ident_h), Ident::from_str("push_str"),
                                                         vec!(fuga_lit)));

    let stmt_ident = cx.stmt_expr(cx.expr_ident(sp, ident_h));
    MacEager::expr(cx.expr_block(cx.block(sp, vec!(stmt_let_h, stmt_push_str, stmt_ident))))
}

本当はinclude_str!()ではなくexpand_include_str()を呼びたかったのだが、argsの構築方法がわからなかったため上のような形になった。

追記(2017/10/18): include_str!()直呼びで問題ないと思ったが、パスがlib.rsからの相対パスになってしまうため、本来の想定とはやはり異なる。

tests/sample_plugin_tests.rs

#![feature(plugin)]
#![plugin(sample_plugin)]

#[test] fn sample_plugin_test() {
    let hoge = sample!();
    assert_eq!(hoge, "hogefuga");
}

quoteマクロを使う

いままではASTを関数に地道に構築してきたがquoteマクロを使うともっと簡単にASTを作ることが出来る。
例えば「letとメソッド呼出」の例であれば

src/lib.rs

#![feature(plugin_registrar, quote, rustc-private)] // <- quote追加

fn expand_sample(cx: &mut ExtCtxt, sp: Span, _args: &[TokenTree])
                 -> Box<MacResult + 'static> {
    MacEager::expr(quote_expr!(cx, {
        let h = "hoge".to_string();
        h
    }))
}

このようにそのまま書き下すことが出来る。quoteマクロの使い方についてはrustソースコードの「src/test/run-pass-fulldeps/quote-tokens.rs」を見てもらうと感じがつかめると思う。

おわりに

ドキュメントではASTの作り方についてはあまり解説されていないので、rustソースのsrc/libsyntax/ext/build.rsにある関数でgrepしたりして試行錯誤した。
この記事ではコンパイラプラグインでなくても出来るような例しかあげられなかったが、コンパイラプラグイン作成の入り口として役に立てれば幸いである。