Zenn
🦀

[Rust] Procedural Macros を実装するときに準備していること

2024/02/19に公開
Rust
#
tech
#
macro
tech

はじめに

今回は僕が Rust の Procedural Macros を実装する際に準備/整備してよかったと思うことを紹介してみようと思います。
前回 Rust のマクロについて、実際に使われているコードの紹介をしましたが、こういったキャッチアップが必要だったのも僕自身が Rust のマクロを記述する必要があるためでした。

https://zenn.dev/linnefromice/articles/when-to-use-macros-in-rust

マクロの記述自体も少しキャッチアップが必要ですが、その品質を維持しつつ拡張していくことが最初は結構大変でした (数日見ないと元の設計もそのマクロによる生成コードも記憶から薄れてしまっていたので...)。
そこで今回はこういった問題にアプローチするためのいくつかのことを紹介しようと思います。

設計/実装

エントリポイントとマクロロジックの分離

Procedural Macros を実装する際に、まずエントリポイントとなる proc_macro 属性を付与した関数と、実際に入力 TokenStream から出力 TokenStream を生成するロジックのための関数を分離し、エントリポイントとなる関数はロジックのための関数を呼び出すのみとします。外部に公開するマクロ関数がなんであるかをわかりやすくするという利点もありますが、後述するテストなどにもこの分離は効果的です。

// src/lib.rs
#[proc_macro]
pub fn generate_something(input: TokenStream) -> TokenStream {
    functions::generate_something(input)
}

// src/functions.rs
pub fn generate_something(input: TokenStream) -> TokenStream {
    let args = syn::parse_macro_input!(input as GenerateSomethingArgs);
    generate_something_internal(args).into()
}

僕の場合ではモジュールごとわけることが多いです。たくさんのマクロ関数を公開するパッケージの場合、以下のようにしてどういったマクロのロジックを格納しているのか把握しやすくするためにモジュール分離を行います。

ex

root/src
|- func.rs // ex: 便利関数を提供するマクロ
|- storage.rs // ex: ストレージに関するマクロ
...
L lib.rs // エントリーポイント

syn,proc_macro2 ベースにする

前セクションでエントリポイントとロジックを分離することを挙げましたが、ロジック側で syn, proc_macro2 クレートをベースにロジック部分を実装していくのがオススメです。syn,quote,proc_macro2 などマクロに関する便利クレートがあり、基本的にこれを駆使してマクロコードを生成します。これらのクレートに関する細かい解説は行いませんが、簡単に触れておきます。

  • proc_macro2
    • dtolnay/proc-macro2
    • Procedural Macros 用の proc_macro クレートのラッパーであり、マクロではないコードでも宣言可能な上、後述する syn, quote クレートとの連携を可能にします。
    • 特に proc_macro2::TokenStream が syn, quote との連携を可能にするトークン表現の構造体であり proc_macro::TokenStream に変換可能です。
  • syn
    • dtolnay/syn: Parser for Rust source code
    • Rust のソースコードから構文木を解析する
    • Struct, Enum, Function, Trait などの要素を表現する構造体を保持し、それらを解析可能にしている
    • このクレートをベースに構文木を再構築して、目的のコードを生成するマクロを構築する
  • quote
    • dtolnay/quote: Rust quasi-quoting
    • quote::quote! マクロによって、引数に入れた Rust のコードの文字列を、構文木に変換することができる
      • つまり直接 Rust コードを書いて、構文木を扱うための proc_macro2::TokenStream の構造体に落とし込むことができます

proc_macro2::TokenStream ベースで扱うようにすることで syn,quote と組み合わせてすぐに通常のコーディング通りに扱うことができ、proc_macro 自体を意識する必要はなくなります。そしてロジック関数の最後に quote::quote! や proc_macro2::TokenStream をリターンすることで、エントリポイント側では .into() で容易に proc_macro::TokenStream で出力できます。

pub fn generate_something(input: TokenStream) -> TokenStream {
    let args = syn::parse_macro_input!(input as GenerateSomethingArgs);
    generate_something_internal(args).into()
}
fn generate_something_internal(args: TimerTaskArgs) -> proc_macro2::TokenStream {
  ...

  quote! {
    fn something() -> String { ... }
  }
}

parse_macro_input! による入力パラメータの構造体化

既に登場していますが解説していなかった syn::parse_macro_input! 部分を取り上げます。

let args = syn::parse_macro_input!(input as GenerateSomethingArgs);

このマクロは "入力パラメータを構造化する" ことができます。つまり、マクロの引数に指定したパラメータによる proc_macro::TokenStream を適切な構造体に変換することができます。この一行以降、入力パラメータは親しみのある構造体で扱うことができ、開発効率と可読性の向上に直結します。

struct GenerateSomethingArgs {
    func_name: syn::LitStr,
    request_args: syn::Type,
    response: Option<syn::Type>,
}
impl Parse for GenerateSomethingArgs {
    fn parse(input: ParseStream) -> Result<Self> {
        let func_name: syn::LitStr = input.parse()?;
        input.parse::<syn::Token![,]>()?;
        let request_args: syn::Type = input.parse()?;
        let response = if input.peek(syn::Token![,]) {
            input.parse::<syn::Token![,]>()?;
            Some(input.parse()?)
        } else {
            None
        };
        Ok(GenerateSomethingArgs {
            func_name,
            request_args,
            response,
        })
    }
}
pub fn generate_something(input: TokenStream) -> TokenStream {
    let args = syn::parse_macro_input!(input as GenerateSomethingArgs);
    generate_something_internal(args).into()
}

// ↓

generate_something!("hello", String);
generate_something!("register_account", Account, RegistrationResult);

構造体に変換するための Parser の記述方法について、少しキャッチアップが必要ですが、慣れるとすぐに記述できるようになります。カンマなどの区切り文字扱いと、変換する syn のタイプ (LitXxx etc) さえわかれば対応可能になります。parse メソッドの引数である ParseStream 型の値について、parse メソッドで指定した型にトークンを変換しストリームの位置を1つ進めたり、peek 関数でストリームの次のトークンを判定したりします。

syn::parse - Rust
ParseBuffer in syn::parse - Rust

テスト

Procedural Macros のテストコードは通常の Rust コードをテストするより難しいことが多いです。シンプルな関数や構造体を生成する場合、あるいは Derive Macros の場合はユニットテストで書ける可能性が高いですが、複数の関数を生成したりする場合は困難です。とはいえ、何も検証のためのコードが書けないのは保守も難しくなるため、いくつか代替となるテストを作ります。

Snapshots tests の活用

Snapshot tests を用いて、マクロコードの出力を検証します。この Snapshot tests を行うために insta クレートがオススメです。

https://docs.rs/insta/latest/insta/

insta を利用して下記のようにテストコードを書き、

#[test]
fn test_snapshot_generate_something() {
    let input = quote! {"something", String, String};
    let args: syn::Result<GenerateSomethingArgs> = syn::parse2(input);
    let generated = generate_something_internal(args.unwrap());
    let formatted = RustFmt::default()
        .format_str(generated.to_string())
        .expect("rustfmt failed");
    assert_snapshot!("snapshot__generate_something", formatted);
}

このテストを実行すると .snap というファイルが生成され、そのファイルにマクロコードによって展開されたコードが保存されます。

cargo test

---
source: src/functions.rs
assertion_line: 58
expression: formatted
---
fn something(request: String) -> String {
    unimplemented!()
}

もし .snap に変更前のマクロから展開されたコードがある場合には、以下のような実装修正を行い、

fn something(request: String) -> String {
-  todo!()
+  unimplemented!()
}

再度テストを実行することで、下記のように検証することができます。

test functions::test::test_snapshot_generate_something ... FAILED

failures:

---- functions::test::test_snapshot_generate_something stdout ----
━━━━ Snapshot Summary ━━━━
Snapshot file: src/snapshots/macros_3__functions__test__snapshot__generate_something.snap
Snapshot: snapshot__generate_something
Source: src/functions.rs:58
──────────────────────────
Expression: formatted
──────────────────────────
-old snapshot
+new results
────────────┬─────────────
    0     0 │ fn something(request: String) -> String {
    1       │-    todo!()
          1 │+    unimplemented!()
    2     2 │ }
────────────┴─────────────

このように Snapshot tests を導入し活用することで、想定通りに展開されるマクロが作成できているか、意図した通りの改修ができているかを確認できます。

Integration tests の活用

次は Integration tests を活用します。Integration test で生成したマクロ関数を利用して、実際にテストを書くことができます。
Integration tests はルートフォルダから tests フォルダを作成して追加します。

https://doc.rust-lang.org/book/ch11-03-test-organization.html#integration-tests

// tests/functions.rs
mod functions_test {
    use macros_3::generate_something;

    generate_something!("hello", String, String);

    #[test]
    #[should_panic(expected = "not implemented")]
    fn test() {
        assert_eq!(hello("World.".to_string()), "hello");
    }
}

テストモジュール内でマクロを使用しコード生成することで、実際のテスト関数でその挙動をテストすることができます。

終わりに

また Rust のマクロに関する記事を書いてみましたが、今回は僕自身がマクロ実装を行なって得た気づきを紹介してみました。僕はマクロでも Procedural Macros を実装することが多いのでそこにフォーカスしましたが、考え方自体はどのマクロの実装パターンでも参考にできるかと思います。皆さんもマクロを書く機会があればぜひ参考にしてみてください。

Discussion