LoginSignup
28
18

More than 3 years have passed since last update.

@yagince

[Rust] proc_macro入門 CustomDeriveを書いてみる

Posted at

davidpdrsn/juniper-eager-loading: Library for avoiding N+1 query bugs with Juniper
↑の実装を読んでいて、CustomDeriveってこうやって書くのかぁ、と思ったので、proc_macroでCustomDriveを書いてみる為にちょっと試したので、メモ。

環境

  • Rust 1.41

構造体の名前を返すメソッドを定義するCustomDeriveを書いてみる

シンプルな例にする為に構造体名を文字列で返すメソッドを定義するCustomDriveを書いてみます。

Cargo.toml
[package]
name = "proc_macro_sample"
version = "0.1.0"
authors = ["yagince <xxxx@gmail.com>"]
edition = "2018"
publish = false

[lib]
proc-macro = true

[dependencies]
quote = "1.0.2"
syn = { version = "1.0.14", features = ["full", "extra-traits"] }

test

先にどんな感じになって欲しいかテストを書いてみます。

tests/test.rs
use proc_macro_sample::SelfName;

#[derive(SelfName)]
struct Hoge {}

#[test]
fn test_hoge() {
    let hoge = Hoge {};
    assert_eq!(hoge.self_name(), "Hoge");
}

derive(SelfName) すると self_name というメソッドが定義されるのを目指します。

実装

src/lib.rs
extern crate proc_macro;
use quote::quote;
use syn::{parse_macro_input, ItemStruct};

#[proc_macro_derive(SelfName)]
pub fn derive_self_name(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let item = parse_macro_input!(input as ItemStruct);
    let struct_name = item.ident;
    let gen = quote! {
        impl #struct_name {
            pub fn self_name(&self) -> &str {
                stringify!(#struct_name)
            }
        }
    };
    gen.into()
}
  • proc_macro_derive(SelfName) で derive(SelfName) できるようにする
  • synのparse_macro_inputマクロで syn::ItemStruct にパースする
    • TokenStreamそのままだと直感的に扱えないので、もうちょい扱いやすい状態に変換している
  • quoteのquoteマクロで定義したいメソッドを書く

ちなみに、TokenStreamとItemStructの状態はこんな状態。

[src/lib.rs:7] &input = TokenStream [
    Ident {
        ident: "struct",
        span: #0 bytes(79..85),
    },
    Ident {
        ident: "Hoge",
        span: #0 bytes(86..90),
    },
    Group {
        delimiter: Brace,
        stream: TokenStream [],
        span: #0 bytes(91..93),
    },
]
[src/lib.rs:9] &item = ItemStruct {
    attrs: [],
    vis: Inherited,
    struct_token: Struct,
    ident: Ident {
        ident: "Hoge",
        span: #0 bytes(86..90),
    },
    generics: Generics {
        lt_token: None,
        params: [],
        gt_token: None,
        where_clause: None,
    },
    fields: Named(
        FieldsNamed {
            brace_token: Brace,
            named: [],
        },
    ),
    semi_token: None,
}

構造体名はidentで取得できるし、fieldsからフィールド定義が取れるので、とても扱いやすい状態になっているように見えます。

オプションを設定できるようにしたい

何かしら設定を受け取れるようにしたいので、今回はシンプルにlowercaseにするオプションを設定出来る状態を目指すことにしました。

一旦設定を受け取ってみる

src/lib.rs
...

#[proc_macro_derive(SelfName, attributes(self_name))]
pub fn derive_self_name(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let item = parse_macro_input!(input as ItemStruct);
    let struct_name = dbg!(item).ident;

    let gen = quote! {
        impl #struct_name {
            pub fn self_name(&self) -> &str {
                stringify!(#struct_name)
            }
        }
    };

    gen.into()
}

attributes で self_name というオプションを受け取れるようにしてみた

test

tests/test.rs
#[derive(SelfName)]
#[self_name(lowercase)]
struct HogeLowercase {}

#[test]
fn test_hoge_lowercase() {
    let hoge = HogeLowercase {};
    assert_eq!(hoge.self_name(), "hogelowercase");
}

[src/lib.rs:17] item = ItemStruct {
    attrs: [
        Attribute {
            pound_token: Pound,
            style: Outer,
            bracket_token: Bracket,
            path: Path {
                leading_colon: None,
                segments: [
                    PathSegment {
                        ident: Ident {
                            ident: "self_name",
                            span: #0 bytes(186..195),
                        },
                        arguments: None,
                    },
                ],
            },
            tokens: TokenStream [
                Group {
                    delimiter: Parenthesis,
                    stream: TokenStream [
                        Ident {
                            ident: "lowercase",
                            span: #0 bytes(196..205),
                        },
                    ],
                    span: #0 bytes(195..206),
                },
            ],
        },
    ],
    vis: Inherited,
    struct_token: Struct,
    ident: Ident {
        ident: "HogeLowercase",
        span: #0 bytes(215..228),
    },
    generics: Generics {
        lt_token: None,
        params: [],
        gt_token: None,
        where_clause: None,
    },
    fields: Named(
        FieldsNamed {
            brace_token: Brace,
            named: [],
        },
    ),
    semi_token: None,
}

attrs というフィールドにAttribute型のデータが入っていて
path にattribute名
tokens にattributeの中身が入っています

このままだとちょっと扱いづらいですね。

Attributeを構造体にマッピングしたい

davidpdrsn/juniper-eager-loadingでは
https://github.com/davidpdrsn/juniper-eager-loading/blob/590b93ee4d/juniper-eager-loading-code-gen/src/derive_eager_loading/field_args.rs#L18-L27
bae::FromAttributesというのを使っているみたいですね。
bae - crates.io: Rust Package Registry　※同じ作者
juniper-eager-loading/derive_eager_loading.rs at 590b93ee4d0d7fe629b939bfc3f051458a98149d · davidpdrsn/juniper-eager-loading
EagarLoading::from_attributes(&attrs) で構造体にマッピングしているようです。

bae is a crate for proc macro authors, which simplifies parsing of attributes. It is heavily inspired by darling but has a significantly simpler API.

inspired by darling って書いてあるので darling も見てみます。

darling - crates.io: Rust Package Registry
https://github.com/TedDriggs/darling/blob/57c4ef486c/examples/fallible_read.rs
FromDeriveInput と Xx::from_derive_input で行けそうです。

今回は darling を使ってみることにしました。

Attributeを構造体にマッピングしてみる

darling::FromDeriveInput - Rust
from_derive_input は syn::DeriveInput 型を引数に取ります。
今は ItemStruct にパースしてるので、そのままだと渡せません。
なので、パースする型を DeriveInput に変えてみました。

#[proc_macro_derive(SelfName, attributes(self_name))]
pub fn derive_self_name(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let item = parse_macro_input!(input as DeriveInput);
    let struct_name = dbg!(item).ident;

    let gen = quote! {
        impl #struct_name {
            pub fn self_name(&self) -> &str {
                stringify!(#struct_name)
            }
        }
    };

    gen.into()
}

[src/lib.rs:17] item = DeriveInput {
    attrs: [
        Attribute {
            pound_token: Pound,
            style: Outer,
            bracket_token: Bracket,
            path: Path {
                leading_colon: None,
                segments: [
                    PathSegment {
                        ident: Ident {
                            ident: "self_name",
                            span: #0 bytes(186..195),
                        },
                        arguments: None,
                    },
                ],
            },
            tokens: TokenStream [
                Group {
                    delimiter: Parenthesis,
                    stream: TokenStream [
                        Ident {
                            ident: "lowercase",
                            span: #0 bytes(196..205),
                        },
                    ],
                    span: #0 bytes(195..206),
                },
            ],
        },
    ],
    vis: Inherited,
    ident: Ident {
        ident: "HogeLowercase",
        span: #0 bytes(215..228),
    },
    generics: Generics {
        lt_token: None,
        params: [],
        gt_token: None,
        where_clause: None,
    },
    data: Struct(
        DataStruct {
            struct_token: Struct,
            fields: Named(
                FieldsNamed {
                    brace_token: Brace,
                    named: [],
                },
            ),
            semi_token: None,
        },
    ),
}

あんまり変わってないのでDeriveInputで行けそうです。

構造体にマッピング

lowercase というオプションを受け取りたいので、以下のように構造体を定義。

#[derive(Debug, FromDeriveInput)]
#[darling(attributes(self_name))]
struct SelfNameOption {
    #[darling(default)]
    lowercase: bool,
}
  • self_name というattributeで受け取る
  • lowercaseオプションはdefaultでfalseで構わないので、darling(default) でdefaultが設定されるように指定

やってみます

    let item = parse_macro_input!(input as DeriveInput);
    let option = SelfNameOption::from_derive_input(&item).unwrap();
tests/test.rs
use proc_macro_sample::SelfName;

#[derive(SelfName)]
struct Hoge {}

#[test]
fn test_hoge() {
    let hoge = Hoge {};
    assert_eq!(hoge.self_name(), "Hoge");
}

#[derive(SelfName)]
#[self_name(lowercase)]
struct HogeLowercase {}

#[test]
fn test_hoge_lowercase() {
    let hoge = HogeLowercase {};
    assert_eq!(hoge.self_name(), "hogelowercase");
}

[src/lib.rs:21] option = SelfNameOption {
    lowercase: false,
}
[src/lib.rs:21] option = SelfNameOption {
    lowercase: true,
}

オプション無しの場合は false
オプションありの場合は true
になりました。

最終形がこちら

lib.rs
extern crate proc_macro;

use quote::quote;
use syn::{parse_macro_input, DeriveInput};
use darling::{FromDeriveInput};

#[derive(Debug, FromDeriveInput)]
#[darling(attributes(self_name))]
struct SelfNameOption {
    #[darling(default)]
    lowercase: bool,
}

#[proc_macro_derive(SelfName, attributes(self_name))]
pub fn derive_self_name(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let item = parse_macro_input!(input as DeriveInput);
    let option = SelfNameOption::from_derive_input(&item).unwrap();
    let struct_name = item.ident;
    let mut name_str = struct_name.to_string();

    if option.lowercase {
        name_str = name_str.to_lowercase();
    }

    let gen = quote! {
        impl #struct_name {
            pub fn self_name(&self) -> &str {
                #name_str
            }
        }
    };

    gen.into()
}

running 2 tests
test test_hoge ... ok
test test_hoge_lowercase ... ok

test result: ok. 2 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out

まとめ

  • proc_macro_derive でCustomDeriveは定義できる
  • TokenStreamのままだと扱いづらいので、synや周辺ライブラリで扱いやすい形にパースすると良い
  • quoteで直感的にTokenStreamを生成できる

TODO

28
18
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
28
18