Rustのcustom attributeチュートリアル

# 追記（2019/7/22）
コメント欄にてさらに[詳細な解説](https://qiita.com/dalance/items/1911a775ee23f3e35d18#comment-bd2a67e53f8a61c07d6d)を頂きましたので、そちらもご覧ください。

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

```rust
#[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を実装します。
すなわち

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

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

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

# 準備
custom attributeに限らずprocedural macroは今のところ独立したクレートが必要です。
新しいプロジェクトを作って

```console
$ cargo new custom_attribute_sample
```

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

```console
$ cd custom_attribute_sample
$ cargo new trace --lib
```

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

```toml:trace/Cargo.toml
[lib]
proc-macro = true
```

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

```toml:Cargo.toml
[dependencies]
trace = { path = "./trace" }
```

# 最初のcustom attribute
trace/src/lib.rsに以下のように書きます。

```rust: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を以下のようにすれば、実際に使ってみることができます。

```rust:src/main.rs
use trace::trace;

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

fn main() {
    test()
}
```


これで実行すると

```console
[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に追加します。

```toml:trace/Cargo.toml
[dependencies]
syn = { version = "0.15.18", features = ["full", "extra-traits"] }
```

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

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

```rust: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`になります。

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

```console
[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です。

```toml:trace/Cargo.toml
[dependencies]
syn = { version = "0.15.18", features = ["full", "extra-traits"] }
quote = "0.6.9"
```

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

```rust: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に挿入したりできないので再度パースします。

```rust
    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に挿入します。

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

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

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

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

これを実行すると

```console
[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");`に書き換わったことがわかります。

# 仕上げ

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

```rust: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()
}
```

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

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

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

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

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

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

```rust
    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!`、戻り値を返す、の順になります。

```rust
    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`を元の関数に挿入して完成です。

```rust
    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()
```

これを実行すると

```console
[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を実装出来ました。