Edited at

Rustのcustom attributeチュートリアル


追記(2019/7/22)

コメント欄にてさらに詳細な解説を頂きましたので、そちらもご覧ください。


はじめに

custom attributeとはRustのprocedural macroの一種です。例えば

#[hello_world]

fn test() {
println!("Hello, world!");
}

こんな感じで#[hello_world]というマクロを定義して関数などに付けることができます。せっかくRust 2018で安定化されたのに、あまり具体的なチュートリアルがなさそうだったので簡単な使い方を書いてみます。

マクロ周りはnightly前提だったりすることも多いですが、ここではstable(1.36.0)の範囲内で書けるものを扱います。

また、今回作成したサンプルソースは以下に置いてあります。

https://github.com/dalance/custom_attribute_sample


題材

関数呼び出しの前後にメッセージを出力するcustom attributeを実装します。

すなわち

#[trace]

fn test() {
println!("Hello, world!");
}

// println出力
// Enter: test
// Hello, world!
// Exit: test

こんな感じですね。任意の関数に#[trace]を付けると、入ったときと出る直前にメッセージを表示します。


準備

custom attributeに限らずprocedural macroは今のところ独立したクレートが必要です。

新しいプロジェクトを作って

$ cargo new custom_attribute_sample

その中に(実際には別に中でなくてもいいですが)さらにプロジェクトを作ります。

$ cd custom_attribute_sample

$ cargo new trace --lib

traceクレートは、procedural macroであることを示すため、Cargo.tomlに以下を追加します。


trace/Cargo.toml

[lib]

proc-macro = true

また、本体側のCargo.tomlにはdependencesにtraceを追加しておきます。


Cargo.toml

[dependencies]

trace = { path = "./trace" }


最初のcustom attribute

trace/src/lib.rsに以下のように書きます。


trace/src/lib.rs

extern crate proc_macro;

use crate::proc_macro::TokenStream;

#[proc_macro_attribute]
pub fn trace(_attr: TokenStream, item: TokenStream) -> TokenStream {
let item = dbg!(item);
item
}


#[proc_macro_attribute]trace関数がcustom attributeであることを示すアトリビュートです。custom attributeの名前はこの関数名になります。

引数はTokenStreamです。itemの方にアトリビュートが付いた関数全体のトークン列が渡されます。_attrは今回は使いませんが、アトリビュートの引数(#[trace(Arg)]と書いた場合のArg部分)を受け取ることができます。

特に処理はせず、単にdbg!で出力してみます。

ちなみにRust 2018では基本的にextern crate不要になっていますがproc_macroにはまだ必要です。

src/main.rsを以下のようにすれば、実際に使ってみることができます。


src/main.rs

use trace::trace;

#[trace]
fn test() {
println!("Hello, world!");
}

fn main() {
test()
}


これで実行すると

[Running 'cargo run']

Compiling trace v0.1.0
Compiling custom_attribute_sample v0.1.0
[trace/src/lib.rs:9] item = TokenStream [
Ident {
ident: "fn",
span: #0 bytes(28..30),
},
Ident {
ident: "test",
span: #0 bytes(31..35),
},
Group {
delimiter: Parenthesis,
stream: TokenStream [],
span: #0 bytes(35..37),
},
Group {
delimiter: Brace,
stream: TokenStream [
Ident {
ident: "println",
span: #0 bytes(44..51),
},
Punct {
ch: '!',
spacing: Alone,
span: #0 bytes(51..52),
},
Group {
delimiter: Parenthesis,
stream: TokenStream [
Literal { lit: Str_(Hello, world!), suffix: None, span: Span { lo: BytePos(53), hi: BytePos(68), ctxt: #0 } },
],
span: #0 bytes(52..69),
},
Punct {
ch: ';',
spacing: Alone,
span: #0 bytes(69..70),
},
],
span: #0 bytes(38..72),
},
]
Finished dev [unoptimized + debuginfo] target(s) in 0.39s
Running `target/debug/custom_attribute_sample`
Hello, world!
[Finished running. Exit status: 0]

となって、確かに関数のトークン列を表示できています。

この表示はコンパイル中に行われるので、出来たバイナリを実行しても表示されません。


TokenStreamをパースする

単に表示するだけだとあまり意味がないので、操作することを考えます。TokenStreamの出力を見れば分かるように、このままでは関数名とか引数といった意味のある単位で扱うことができません。このTokenStreamをパースして抽象構文木(AST)の形に変換するのがsynクレートです。trace/Cargo.tomlに追加します。


trace/Cargo.toml

[dependencies]

syn = { version = "0.15.18", features = ["full", "extra-traits"] }

featuresのうちfullは全ての構文要素を扱うためのフィーチャで、extra-traitsはASTのDebugトレイトなどを実装するフィーチャです。

これでTokenStreamをパースしてASTとして扱うことができるようになります。


trace/src/lib.rs

extern crate proc_macro;

use crate::proc_macro::TokenStream;
use syn::{parse_macro_input, ItemFn};

#[proc_macro_attribute]
pub fn trace(_attr: TokenStream, item: TokenStream) -> TokenStream {
let ret = item.clone();
let ast = parse_macro_input!(item as ItemFn);
dbg!(ast);
ret
}


parse_macro_input!がパース用のマクロです。引数のasはキャストしているわけではなく(わかりにくい…)マクロが提供している構文で、何にパースするかを示します。今は関数をパースしたいので関数を表すItemFnになります。

実行するとこうなります。

[Running 'cargo run']

Compiling trace v0.1.0
Compiling custom_attribute_sample v0.1.0
[trace/src/lib.rs:18] ast = ItemFn {
attrs: [],
vis: Inherited,
constness: None,
asyncness: None,
unsafety: None,
abi: None,
ident: Ident {
ident: "test",
span: #0 bytes(31..35),
},
decl: FnDecl {
fn_token: Fn,
generics: Generics {
lt_token: None,
params: [],
gt_token: None,
where_clause: None,
},
paren_token: Paren,
inputs: [],
variadic: None,
output: Default,
},

// 省略

}
Finished dev [unoptimized + debuginfo] target(s) in 9.76s
Running `target/debug/custom_attribute_sample`
Hello, world!
[Finished running. Exit status: 0]

例えばast.identで関数名が取れますし、引数や型引数の情報も取ることができます。


コード生成

ここまではdbg!での表示でしたが、いよいよ埋め込むコードを生成します。コード生成で使うクレートはquoteです。


trace/Cargo.toml

[dependencies]

syn = { version = "0.15.18", features = ["full", "extra-traits"] }
quote = "0.6.9"

これを使ってコード生成した例が以下になります。


trace/src/lib.rs

extern crate proc_macro;

use crate::proc_macro::TokenStream;
use quote::quote;
use syn::{parse_macro_input, ItemFn, Stmt};

#[proc_macro_attribute]
pub fn trace(_attr: TokenStream, item: TokenStream) -> TokenStream {
let mut ast = parse_macro_input!(item as ItemFn);

let new_stmt = quote! {
println!("trace");
};
let new_stmt: TokenStream = new_stmt.into();
let new_stmt = parse_macro_input!(new_stmt as Stmt);

ast.block.stmts.clear();
ast.block.stmts.push(new_stmt);

let gen = quote! {
#ast
};
gen.into()
}


quote!マクロで囲んだ中身が実際に展開されるコードになります。このコードは単なる文字列ではなくRustのコードとして扱われるので、例えば閉じていない括弧のように構文的に変な記述をするとコンパイルエラーとなります。逆に構文さえ満たしていれば内容はどんなものでも大丈夫です。

quote!で囲んだものはinto()TokenStreamに変換して使いますが、TokenStreamのままではASTに挿入したりできないので再度パースします。

    let new_stmt = quote! {

println!("trace");
};
let new_stmt: TokenStream = new_stmt.into();
let new_stmt = parse_macro_input!(new_stmt as Stmt);

ここでは;で終端されたStatementなのでそれを表すStmtにパースします。

これを元の関数のASTに挿入します。

    ast.block.stmts.clear();

ast.block.stmts.push(new_stmt);

ast.block.stmtsが関数内のStatementを保持するVec<Stmt>になっているので一旦消してから目的のものを入れています。

最後に全体を再度TokenStreamに変換して出力します。quote!内で#astのように#を付けることでquote!の外の変数を参照することができます。

    let gen = quote! {

#ast
};
gen.into()

これを実行すると

[Running 'cargo run']

Compiling trace v0.1.0
Compiling custom_attribute_sample v0.1.0
Finished dev [unoptimized + debuginfo] target(s) in 1.16s
Running `target/debug/custom_attribute_sample`
trace
[Finished running. Exit status: 0]

となって、確かに関数の中身がprintln!("trace");に書き換わったことがわかります。


仕上げ

ここまでできればあとは細かい実装だけです。


trace/src/lib.rs

extern crate proc_macro;

use crate::proc_macro::TokenStream;
use quote::quote;
use syn::{parse_macro_input, ItemFn, Stmt};

#[proc_macro_attribute]
pub fn trace(_attr: TokenStream, item: TokenStream) -> TokenStream {
let mut ast = parse_macro_input!(item as ItemFn);
let ident = &ast.ident;

let enter = quote! {
println!("Enter: {}", stringify!(#ident) );
};
let enter: TokenStream = enter.into();
let enter = parse_macro_input!(enter as Stmt);

let mut body = quote! {};
for s in &ast.block.stmts {
body = quote! {
#body
#s
};
}
let body = quote! {
let body = || { #body };
};
let body: TokenStream = body.into();
let body = parse_macro_input!(body as Stmt);

let exit = quote! {
{
let ret = body();
println!("Exit: {}", stringify!(#ident) );
ret
}
};
let exit: TokenStream = exit.into();
let exit = parse_macro_input!(exit as Stmt);

ast.block.stmts.clear();
ast.block.stmts.push(enter);
ast.block.stmts.push(body);
ast.block.stmts.push(exit);

let gen = quote! {
#ast
};

gen.into()
}


関数に入ったときに出力するメッセージはこのように生成しています。

    let enter = quote! {

println!("Enter: {}", stringify!(#ident) );
};
let enter: TokenStream = enter.into();
let enter = parse_macro_input!(enter as Stmt);

stringify!マクロは引数をそのまま文字列リテラル化するマクロです。

今回は元の関数の本体も必要なので、一旦bodyに集めます。

    let mut body = quote! {};

for s in &ast.block.stmts {
body = quote! {
#body
#s
};
}

さらに集めたbodyをクロージャに入れます。

    let body = quote! {

let body = || { #body };
};
let body: TokenStream = body.into();
let body = parse_macro_input!(body as Stmt);

次は関数を出るときのメッセージを生成する部分です。bodyをクロージャに入れた理由は、単に末尾にprintln!を挿入しただけでは関数が早期リターンした場合に対応できないためです。なのでlet ret = body();で戻り値を受けて、println!、戻り値を返す、の順になります。

    let exit = quote! {

{
let ret = body();
println!("Exit: {}", stringify!(#ident) );
ret
}
};
let exit: TokenStream = exit.into();
let exit = parse_macro_input!(exit as Stmt);

最後にここまでで作った3つのStmtを元の関数に挿入して完成です。

    ast.block.stmts.clear();

ast.block.stmts.push(enter);
ast.block.stmts.push(body);
ast.block.stmts.push(exit);

let gen = quote! {
#ast
};

gen.into()

これを実行すると

[Running 'cargo run']

Compiling trace v0.1.0
Compiling custom_attribute_sample v0.1.0
Finished dev [unoptimized + debuginfo] target(s) in 1.12s
Running `target/debug/custom_attribute_sample`
Enter: test
Hello, world!
Exit: test
[Finished running. Exit status: 0]

ということで目的のcustom attributeを実装出来ました。