Rust 関数風マクロを軽くハンズオン


こちらの記事は [Rustマクロ冬期講習アドベントカレンダー](https://qiita.com/advent-calendar/2024/rust-macro) 19日目の記事です！

> アドカレまとめ記事はこちら！: [Rustマクロ作成チートシート！](https://qiita.com/namn1125/items/d601d2cbbf2cc7b5d3dc)

[前回](https://qiita.com/namn1125/items/b05ad80d324afcd2ed91)まででRustの手続きマクロで作れる3つのマクロ

- [関数風マクロ](https://qiita.com/namn1125/items/fc5aa71690fcb5109cf0)
- [属性風マクロ](https://qiita.com/namn1125/items/b7d9374e3d6d3f64ab65)
- [deriveマクロ](https://qiita.com/namn1125/items/b05ad80d324afcd2ed91)

について、最小構成を見てきました。

何もしないマクロならば三種の神器 [`syn`](https://docs.rs/syn/latest/syn/), [`quote`](https://docs.rs/quote/latest/quote/), [`proc-macro2`](https://docs.rs/proc-macro2/latest/proc_macro2/) を使うまでもないことを示したわけですが、今回からはちゃんと神器を使ってある程度実用的なマクロを書いてみます！

なお、題材については手続きマクロを書くことに特化するため、宣言マクロ等でも簡単に書けるかもしれない程度のものにしたいと思います。

:::note
**関数風マクロの定義・構造**

`Cargo.toml` に、 `[lib] proc-macro = true` という項目を追加した上で、 `src/lib.rs` に、

```rust:src/lib.rs
#[proc_macro]
pub fn マクロ名(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    // ...
}
```

というように `#[proc_macro]` を付けた関数を定義すると、 `マクロ名!()` みたいに呼び出せる関数マクロが定義されます。

三種の神器クレート [syn](https://docs.rs/syn/latest/syn/) [quote](https://docs.rs/quote/latest/quote/) [proc-macro2](https://docs.rs/proc-macro2/latest/proc_macro2/) を使って入力の [`input: proc_macro::TokenStream`](https://doc.rust-lang.org/beta/proc_macro/struct.TokenStream.html) を加工していくことで出力を作ります。
:::

# 関数風マクロ「 `each_expr!` 」

C言語高速化手法の一つに「 `for` 文のループアンローリング」というものがあります。これに着想を得て、というわけじゃないですが、「式のリスト」と「繰り返したい構文」を渡すことで、リストにある式の個数分構文を繰り返すマクロを作ってみます。

```rust:Rust
each_expr! {
    for V in [
        10;
        "hoge";
        true;
        if false { 0 } else { 1 };
    ] {
        println!("{:?}", V);
    }
};
```

上記のように書いた時、以下のように展開されることを目指します！

```rust:Rust
println!("{:?}", 10);
println!("{:?}", "hoge");
println!("{:?}", true);
println!("{:?}", if false { 0 } else { 1 });
```

# 手順 1. 骨組みを作る

[RTA](https://qiita.com/namn1125/items/fc5aa71690fcb5109cf0)と同様に準備していきます。

```bash
cargo new --lib each_expr
```

```diff_toml:Cargo.toml
[package]
name = "each_expr"
version = "0.1.0"
edition = "2021"

+[lib]
+proc-macro = true

[dependencies]
```

`proc_macro::TokenStream` と `proc_macro2::TokenStream` 等の名前衝突を防ぐため、それ以外にもデバッグ等をしやすくするため、 実装は `src/impls.rs` で行うこととし、 `src/lib.rs` はなるべくシンプルにします。

```plain:ディレクトリ構造
.
├── Cargo.lock
├── Cargo.toml
├── src
│   ├── impls.rs
│   ├── lib.rs
│   └── main.rs # マクロの動作確認に使用
└── target
```

```rust:src/lib.rs
mod impls;

use proc_macro::TokenStream;

#[proc_macro]
pub fn each_expr(input: TokenStream) -> TokenStream {
    todo!()
}
```

三種の神器を `cargo add` しておきます。

```bash
cargo add proc-macro2 quote
cargo add syn --features full --features extra-traits
```

`syn` の2つのfeaturesである `full` と `extra-traits` は便利なので入れています。

# 手順 2. 入力をパースする

今想定しているマクロは次のような構造をしています。

```rust:Rust
each_expr! {
    for 識別子 in [
        式のリスト;
    ] {
        繰り返す処理
    }
};
```

これをパースしてマクロ側で扱いやすいようにまとめます！

```rust:src/impls.rs
use proc_macro2::TokenStream;
use syn::{Expr, Ident};

pub struct EachExprInput {
    pub var: Ident, // 識別子
    pub exprs: Vec<Expr>, // 式のリスト
    pub repeat_target: TokenStream, // 繰り返す処理
}
```

`repeat_target` にはトークン木の集合である `TokenStream` をそのまま持つようにします。ただし、  特殊な`proc_macro` クレートではなく神器の方の `proc_macro2` の `TokenStream` です！

[`syn::parse_macro_input!`](https://docs.rs/syn/latest/syn/macro.parse_macro_input.html) マクロを利用し、

```rust:Rust
let input = parse_macro_input!(input as EachExprInput);
```

と書くことで入力がパースされるようにします。

`EachExprInput` に [`syn::parse::Parse`](https://docs.rs/syn/latest/syn/parse/trait.Parse.html) トレイトを実装することで上記のように書けるようになります。

```rust:src/impls.rs
use proc_macro2::TokenStream;
use syn::{
    braced, bracketed,
    parse::{Parse, ParseStream},
    Expr, Ident, Result, Token,
};

pub struct EachExprInput {
    pub var: Ident,
    pub exprs: Vec<Expr>,
    pub repeat_target: TokenStream,
}

impl Parse for EachExprInput {
    fn parse(input: ParseStream) -> Result<Self> {
        todo!()
    }
}
```

`todo!()` を埋めていきます。 `Parse` トレイトの実装方法ですが、入れ子構造にすることで実装していきます。

すなわち、 `EachExprInput` という構造体として受け取れる `for 識別子 in [...;] {}` という構文の内部について、一つ一つパースしていく感じです！

例えば `for` 部分をパースするとこんな感じです。パースできたらその値は不要なので `_for` としています。

```rust:Rust
let _for: syn::Token![for] = input.parse()?;
```

[`ParseStream::parse`](https://docs.rs/syn/latest/syn/parse/struct.ParseBuffer.html#method.parse) [^parsestream]メソッドは、`Parse` トレイトを実装している構造体へとマクロの入力をパースしていきます。

[^parsestream]: [`ParseStream<'a>` の正体は `&'a ParseBuffer<'a>` のエイリアス](https://docs.rs/syn/latest/syn/parse/type.ParseStream.html) ですが、わかりやすさと名前的に好きなので `ParseStream` にしています。

[`std::iter::Iterator::next`](https://doc.rust-lang.org/std/iter/trait.Iterator.html#tymethod.next) と似ており、パースに使われた分入力が消費されていくので、 `for 識別子 in [...;] {}` を構成するパーツを先頭から順にパースして消費していくことでいい感じにパースができます！

次は繰り返し変数に当たる変数名が来ます。 `syn::Ident` としてパースします。

```rust:Rust
let var: syn::Ident = input.parse()?;
```

こちらは出力に使用するので `var` という変数で持ちました。

この調子で、残りの入力もパースしていきます！全体像は以下のようになります。

```rust:Rust
use proc_macro2::TokenStream;
use syn::{
    braced, bracketed,
    parse::{Parse, ParseStream},
    Expr, Ident, Result, Token,
};

// 省略...

impl Parse for EachExprInput {
    fn parse(input: ParseStream) -> Result<Self> {
        // for
        let _for: Token![for] = input.parse()?;

        // V
        let var = input.parse()?;
        // in
        let _in: Token![in] = input.parse()?;

        // exprs in [..]
        let content;
        let _ = bracketed!(content in input);
        let exprs = content
            .parse_terminated(Expr::parse, Token![;])?
            .into_iter()
            .collect();

        // repeat_target in {..}
        let content;
        let _ = braced!(content in input);
        let repeat_target = TokenStream::parse(&content)?;

        Ok(Self {
            var,
            exprs,
            repeat_target,
        })
    }
}
```

- カッコ系のパースにはマクロを使うのが便利です！
  - `()`: [`syn::parenthesized`](https://docs.rs/syn/latest/syn/macro.parenthesized.html)
  - `{}`: [`syn::braced`](https://docs.rs/syn/latest/syn/macro.braced.html)
  - `[]`: [`syn::bracketed`](https://docs.rs/syn/latest/syn/macro.bracketed.html)

例えば以下のように書くことで `[]` 内の `ParseStream` を `content` にバインドできます。

```rust:Rust
// exprs in [..]
let content;
let _ = bracketed!(content in input);
```

そして [`ParseStream::parse_terminated`](https://docs.rs/syn/latest/syn/parse/struct.ParseBuffer.html#method.parse_terminated) を利用すると `1; 2; 3` みたいな区切り文字(英語でPunctuation)で区切られた表現をイテレータとして得ることができるので、これを利用して `Vec<Expr>` へとパースしています。

これで入力を受け取ることができました！実際に以下を書いてコンパイルが通るところまで確認しましょう。

<details><summary>src/impls.rs</summary>

```rust:src/impls.rs
use proc_macro2::TokenStream;
use syn::{
    braced, bracketed,
    parse::{Parse, ParseStream},
    Expr, Ident, Result, Token,
};

pub struct EachExprInput {
    pub var: Ident,
    pub exprs: Vec<Expr>,
    pub repeat_target: TokenStream,
}

impl Parse for EachExprInput {
    fn parse(input: ParseStream) -> Result<Self> {
        // for
        let _for: Token![for] = input.parse()?;

        // T
        let var = input.parse()?;
        // in
        let _in: Token![in] = input.parse()?;

        // exprs in [..]
        let content;
        let _ = bracketed!(content in input);
        let exprs = content
            .parse_terminated(Expr::parse, Token![;])?
            .into_iter()
            .collect();

        // repeat_target in {..}
        let content;
        let _ = braced!(content in input);
        let repeat_target = TokenStream::parse(&content)?;

        Ok(Self {
            var,
            exprs,
            repeat_target,
        })
    }
}
```
</details>

```rust:src/lib.rs
mod impls;
use impls::EachExprInput;

use proc_macro::TokenStream;
use syn::parse_macro_input;

#[proc_macro]
pub fn each_expr(input: TokenStream) -> TokenStream {
    let _input = parse_macro_input!(input as EachExprInput);

    TokenStream::new()
}
```

お試しは `src/main.rs` でできます！

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

fn main() {
    each_expr! {
        for x in [1; 2; 3] {
            println!("{}", x);
        }
    };
}
```

とりあえずコンパイルが通れば入力のパースは成功しています！

# 手順 3. 出力を組み立てる

出力に必要な情報は `EachExprInput` に集めることができているため、あとは出力するのみです！トークン木のうちグループは深くネストしている可能性があるので再帰を利用する必要があります。以下の方針で行きます！

- 式のリストに対して、式ごとに繰り返して出力の `TokenStream` を作ります。
  - [`quote::quote!`](https://docs.rs/quote/latest/quote/macro.quote.html) マクロや、 [`quote::ToTokens::to_token_stream`](https://docs.rs/quote/latest/quote/trait.ToTokens.html) メソッドを使うことで `TokenStream` を作れます！
- `repeat_target` は `TokenStream` になっています。
  - ★ `repeat_target` が持っているトークン木を一つ一つ見ていき
    - `Group` ( `()` などで囲まれたトークン木の枝) ならさらにその中身を再帰的に見る
    - `Ident` で、繰り返し変数 `var` と同じなら `Expr` に置き換える
    - それ以外ならそのまま出力する

```rust:src/impls.rs
use proc_macro2::{Group, TokenStream, TokenTree};
use quote::ToTokens;
use syn::{
    braced, bracketed,
    parse::{Parse, ParseStream},
    Expr, Ident, Result, Token,
};

// 省略

impl EachExprInput {
    pub fn render(&self) -> TokenStream {
        let Self {
            var,
            exprs,
            repeat_target,
        } = self;

        exprs
            // 各式について繰り返す
            .iter()
            // 平坦化する
            .flat_map(|expr| {
                // repeat_target は TokenStream
                // TokenStream は TokenTree のイテレータ
                repeat_target
                    .clone()
                    .into_iter()
                    .map(|tt| /* ★再帰的にTokenTreeを得る */)
            })
            // impl FromIterator<TokenStream> for TokenStream が実装されているので
            // Iterator<Item = TokenStream> -> TokenStream にできる
            .collect()
    }
}
```

★の再帰部分を実装すれば完成です！

```rust:src/impls.rs
fn render_rec(tt: &TokenTree, var: &Ident, expr: &Expr) -> TokenStream {
    match tt {
        // 再帰で中身も置換
        TokenTree::Group(group) => {
            let content = group
                .stream()
                .into_iter()
                .map(|tt| render_rec(&tt, var, expr))
                .collect();

            let mut new_group = Group::new(group.delimiter(), content);
            // Spanの設定。別な回で解説予定
            new_group.set_span(group.span());

            new_group.to_token_stream()
        }
        // 置換対象なら式に置換！
        TokenTree::Ident(ident) if ident == var => expr.to_token_stream(),
        // それ以外はそのまま出力
        t => t.to_token_stream(),
    }
}

impl EachExprInput {
    pub fn render(&self) -> TokenStream {
        let Self {
            var,
            exprs,
            repeat_target,
        } = self;

        exprs
            .iter()
            .flat_map(|expr| {
                repeat_target
                    .clone()
                    .into_iter()
                    .map(|tt| /* ★ */ render_rec(&tt, var, expr))
            })
            .collect()
    }
}
```

完成した `render` メソッドを `src/lib.rs` にて呼びます。

```rust:src/lib.rs
mod impls;
use impls::EachExprInput;

use proc_macro::TokenStream;
use syn::parse_macro_input;

#[proc_macro]
pub fn each_expr(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as EachExprInput);

    input.render()
        // ↓ proc_macro2::TokenStream から proc_macro::TokenStream に変換    
        .into()
}
```

最後に、目的としていた構文を入力に渡してエラーにならなければ完成です！

```rust:main.rs
use each_expr::each_expr;

fn main() {
    each_expr! {
        for V in [
            10;
            "hoge";
            true;
            if false { 0 } else { 1 };
        ] {
            println!("{:?}", V);
        }
    };
}
```

```bash:cargo expandしてみる
$ cargo expand --bin each_expr
    Checking each_expr v0.1.0 (/path/to/each_expr)
    Finished `dev` profile [unoptimized + debuginfo] target(s) in 0.07s

#![feature(prelude_import)]
#[prelude_import]
use std::prelude::rust_2021::*;
#[macro_use]
extern crate std;
use each_expr::each_expr;
fn main() {
    {
        ::std::io::_print(format_args!("{0:?}\n", 10));
    };
    {
        ::std::io::_print(format_args!("{0:?}\n", "hoge"));
    };
    {
        ::std::io::_print(format_args!("{0:?}\n", true));
    };
    {
        ::std::io::_print(format_args!("{0:?}\n", if false { 0 } else { 1 }));
    };
}
```

```plain:実行してみる
$ cargo run -q
10
"hoge"
true
1
```

無事完成です！

:::note
**出力のデバッグ**

実装を進めていると、謎のコンパイルエラーが発生してしまいました！

```plain
error: expected one of `(`, `[`, or `{`, found `"{:?}"`
  --> src/main.rs:11:22
   |
11 |             println!("{:?}", V);
   |                      ^^^^^^ expected one of `(`, `[`, or `{`
```

こういう時は `render_rec` が出力する `TokenStream` を [`dbg!`](https://doc.rust-lang.org/std/macro.dbg.html) 等でデバッグ出力すると原因がわかるかもしれません。

```rust:Rust
exprs
    .iter()
    .flat_map(|expr| {
        repeat_target.clone().into_iter().map(|tt| {
            let tmp = render_rec(&tt, var, expr);
            dbg!(&tmp);
            tmp
        })
    })
    .collect()
```

```plain
[src/impls.rs:85:21] &tmp = TokenStream [
    Ident {
        ident: "println",
        span: #0 bytes(192..199),
    },
]
[src/impls.rs:85:21] &tmp = TokenStream [
    Punct {
        ch: '!',
        spacing: Alone,
        span: #0 bytes(199..200),
    },
]
[src/impls.rs:85:21] &tmp = TokenStream [
    Literal {
        kind: Str,
        symbol: "{:?}",
        suffix: None,
        span: #0 bytes(201..207),
    },
    // 省略
]
```

グループの `()` を出力し忘れていたのが原因だったみたいです！実際に遭遇したエラーだったのですが無事に修正できました。

`syn::Expr` をデバッグ出力したい場合などもあるでしょう。 `syn` のfeaturesに `extra-traits` を含めたのは、このデバッグ出力を可能にするためだったりしました！(`TokenStream` のデバッグ出力では不要でしたが :sweat_smile:)

参考: https://docs.rs/syn/latest/syn/enum.Expr.html#impl-Debug-for-Expr
:::

# まとめ・所感

以下、全体像になります。

```toml:Cargo.toml
[package]
name = "each_expr"
version = "0.1.0"
edition = "2021"

[lib]
proc-macro = true

[dependencies]
proc-macro2 = "1.0.92"
quote = "1.0.37"
syn = { version = "2.0.91", features = ["full", "extra-traits"] }
```

```rust:src/impls.rs
use proc_macro2::{Group, TokenStream, TokenTree};
use quote::ToTokens;
use syn::{
    braced, bracketed,
    parse::{Parse, ParseStream},
    Expr, Ident, Result, Token,
};

pub struct EachExprInput {
    pub var: Ident,
    pub exprs: Vec<Expr>,
    pub repeat_target: TokenStream,
}

impl Parse for EachExprInput {
    fn parse(input: ParseStream) -> Result<Self> {
        // for
        let _for: Token![for] = input.parse()?;

        // T
        let var = input.parse()?;
        // in
        let _in: Token![in] = input.parse()?;

        // exprs in [..]
        let content;
        let _ = bracketed!(content in input);
        let exprs = content
            .parse_terminated(Expr::parse, Token![;])?
            .into_iter()
            .collect();

        // repeat_target in {..}
        let content;
        let _ = braced!(content in input);
        let repeat_target = TokenStream::parse(&content)?;

        Ok(Self {
            var,
            exprs,
            repeat_target,
        })
    }
}

fn render_rec(tt: &TokenTree, var: &Ident, expr: &Expr) -> TokenStream {
    match tt {
        TokenTree::Group(group) => {
            let content = group
                .stream()
                .into_iter()
                .map(|tt| render_rec(&tt, var, expr))
                .collect();

            let mut new_group = Group::new(group.delimiter(), content);
            new_group.set_span(group.span());

            new_group.to_token_stream()
        }
        TokenTree::Ident(ident) if ident == var => expr.to_token_stream(),
        t => t.to_token_stream(),
    }
}

impl EachExprInput {
    pub fn render(&self) -> TokenStream {
        let Self {
            var,
            exprs,
            repeat_target,
        } = self;

        exprs
            .iter()
            .flat_map(|expr| {
                repeat_target
                    .clone()
                    .into_iter()
                    .map(|tt| render_rec(&tt, var, expr))
            })
            .collect()
    }
}
```

```rust:src/lib.rs
mod impls;
use impls::EachExprInput;

use proc_macro::TokenStream;
use syn::parse_macro_input;

#[proc_macro]
pub fn each_expr(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as EachExprInput);

    input.render().into()
}
```

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

fn main() {
    each_expr! {
        for V in [
            10;
            "hoge";
            true;
            if false { 0 } else { 1 };
        ] {
            println!("{:?}", V);
        }
    };
}
```

- `syn` クレートを利用して入力をパース
- `to_token_stream` メソッドを使用して `proc-macro2::TokenStream` を組み立て

という手順2つを解説し、入力を元に出力をシンプルに決める関数風マクロを作りました！

実は今回のマクロは [proc-macro-workshop](https://github.com/dtolnay/proc-macro-workshop) の [`seq!`](https://github.com/dtolnay/proc-macro-workshop/tree/master/seq) マクロ課題の簡易版といったところなので、もっと応用的なことをやってみたい方はこの課題に取り組んでみてください。