🦍

RustでTypeScriptのOmitなどの型ユーティリティを実装してみた

はじめに

Rustを書いていると、たまにTypeScriptのOmitなどの型ユーティリティがほしいなと思う場面があります。
Rust本体にはそのような型ユーティリティはないので、勉強がてらにクレートを作ってみました

本記事は作ったクレートの簡単な紹介と実装について書いていきます。

https://github.com/skanehira/type-utilities-rs

使い方

attributeマクロでomitなどを提供しているので、次のように使います。

use type_utilities_rs::omit;

// Create a new struct `NewS` with omitted field `b`
#[omit(NewS, [b])]
struct S {
    a: i32,
    b: &str,
}
// `NewS` will only have field `a`
let _ = NewS { a: 1 };

シンタックスはomit(新しい構造体名, [省略するフィールド, ...])という感じになります。
上記の例では、フィールドbを省略したNewSという構造体が生成されます。

struct NewS {
    a: i32,
}

他にもpartialなどがあります。

use type_utilities_rs::partial;

// Create a new struct `NewS` with all fields optional
#[partial(NewS)]
struct S<'a> {
  a: i32,
  b: &'a str,
  c: f64,
}

// `NewS` will have all fields optional
let _ = NewS { a: Some(1), b: Some("hello"), c: Some(1.5) };

現時点では実装予定・実装済みのマクロは次になります。

  • Omit
  • Pick
  • Partial
  • Required
  • Exclude
  • Extract

インターフェイスの検討

TypeScriptの型ユーティリティをRustで実現方法について、次のようなインターフェイスを検討していました。

1つはderiveマクロを使う場合です。

#[derive(Omitter)]
struct S<'a> {
    #[omit]
    a: i3,
    b: &'a str,
    c: f64
}

deriveマクロはattributeマクロと比べて、省略したいフィールドに#[omit]を記述する必要があり冗長なのと、あたらしい構造体名をどこで指定するかといった問題があります。
#[derive(Omit(NewS))]みたいな書き方ができればよいですが(もしかしてできる…?)

次にfunction likeマクロを使う場合です。

omit!(S, NewS, [a, b]);

こちらもattributeマクロと比べて、生成もとの構造体名を指定しないといけないので、すこし不便かなと思いました。

上記2つを鑑みて、最終的に冒頭で紹介したattributeマクロを使ったインターフェイスにしました。

使用するクレートについて

Rustで構文木を扱う際に外せない次のクレート使いました。

[dependencies]
quote = "1.0.35"
syn = { version = "2.0.48", features = ["full"] }

synはRustの構文木のデータをより簡単に扱えるようにするためのクレートです。
たとえばstructを表すsyn::ItemStructがあり、そこから構造体名やフィールド名と型などの情報を取得できます。

pub struct ItemStruct {
    pub attrs: Vec<Attribute>,
    pub vis: Visibility, 
    pub struct_token: Token![struct],
    pub ident: Ident,
    pub generics: Generics,
    pub fields: Fields,
    pub semi_token: Option<Token![;]>,
}

quotesynの構造体などからコンパイラが扱えるソースコードのトークンであるTokenStreamに変換してくれます。

let tokens = quote! {
    struct #struct_name {
        #(#fields),*
    }
};

tokens.into::<TokenStream>()

実装の概要について

使用するクレートの概要についてわかったところで、omitを使って実装の概要について説明していきます。

まず、最初にクレートの種類をproc-macroにする必要があります。

Cargo.toml
[lib]
proc-macro = true

次にlib.rsomitのマクロ関数を定義します。

lib.rs
#[proc_macro_attribute]
pub fn omit(attr: TokenStream, item: TokenStream) -> TokenStream {
    // omitted
}

マクロ関数はTokenStreamを受け取って任意の構文木のTokenStreamを作成して返すことで、新しいソースコードを生成するイメージです。
TokenStreamはさきほど説明したとおり、コンパイラが扱えるソースコードのトークン情報です。
今回はomitの新しい構造体のTokenStreamを作成して返す処理を実装します。

構造体の作成の流れは次のようになります。

  1. syn::parse_macro_input!()を使ってTokenStreamからsyn::ItemStructなどに変換し、属性や構造体の情報を扱えるようにする
  2. 新しい構造体に対応したsyn::ItemStructを作成する
  3. quoteを使って2で作成したsyn::ItemStructからTokenStreamに変換して返す

attributeマクロの場合、関数は2つの引数を受け取ります。
1つ目は属性情報のTokenStream、2つ目は構造体(や関数など)のTokenStreamです。

属性情報というのは次の例でいうとomit(...)で囲っているNewS, [a, b]の部分になります。
構造体のTokenStreamstrcut S { ... }の部分になります。

#[omit(NewS, [a, b])]
struct S {
    // omitted
}

実装の詳細について

概要についてわかったところで、次にomitの詳細について説明してきます。

lib.rs
#[proc_macro_attribute]
pub fn omit(attr: TokenStream, item: TokenStream) -> TokenStream {
    let attr = parse_macro_input!(attr as StructAttribute);
    let item = parse_macro_input!(item as ItemStruct);

    let new_item = omit_or_pick(attr, item.clone(), attribute::AttributeType::Omit);

    quote! {
        #item

        #new_item
    }
    .into()
}

まず最初に、属性の情報を扱うためにTokenStreamをパースします。

parse_macro_input!(attr as StructAttribute);

今回のような属性情報を表現した構造体がsynにないため、StructAttributeという構造体を用意してsyn::parse::Parseトレイトを実装します。

attribute.rs
#[derive(Debug)]
pub(crate) struct StructAttribute {
    // omitに定義した新しい構造体名
    pub(crate) name: Ident,
    // 省略したいフィールド名
    pub(crate) fields: HashMap<Ident, ()>,
}

impl Parse for StructAttribute {
    // `input`は`omit(...)`の`(...)`内のトークン情報
    fn parse(input: ParseStream) -> syn::Result<Self> {
        // 構造体名を取得する
        // `omit(NewS, ..)`の場合`NewS`になる
        let name = input.parse()?;

        let mut fields = HashMap::new();
        // 次のトークンが`,`じゃない場合は処理を切り上げる
        // つまり`omit(NewS)`の場合、フィールドを省略しない
        if !input.peek(Token![,]) {
            return Ok(StructAttribute { name, fields });
        }

        // `,`は特に使わないのでトークンを消費する
        input.parse::<Token![,]>()?;

        // `omit(NewS, [...])`の`[...]`内のトークンを取得
        let content;
        bracketed!(content in input);

        // トークンがなくなるまで、指定しているフィールド名をパースする
        while !content.is_empty() {
            fields.insert(content.parse()?, ());
            if content.is_empty() {
                break;
            }
            // `[a, b, ..]`というふうに`,`で区切られているため`,`を消費する
            content.parse::<Token![,]>()?;
        }

        // `omit(NewS, [])`を許容しない方針にしているため、
        // フィールド名が指定されていない場合はエラーとする
        if fields.is_empty() {
            return Err(syn::Error::new(
                content.span(),
                "Attribute must have at least one field",
            ));
        }

        Ok(StructAttribute { name, fields })
    }
}

これでparse_macro_input!()StructAttributeを使えるようになります。

続けて、構造体もパースします。構造体はsyn::ItemStructがあるのでそれを使います。

let item = parse_macro_input!(item as ItemStruct);

これで属性の情報と構造体の情報が揃ったので、omit_or_pick()を使って新しい構造体を生成していきます。

let new_item = omit_or_pick(attr, item.clone(), attribute::AttributeType::Omit);

といっても基本的にフィールドをフィルターしてcloneした構造体に代入するくらいです。

omitpickはフィルター時の判定が反転するだけなので、その判定を切り替えられるようにAttributeType::Omit|Pickを用意しています。(正直いけていない感があるが、あんまりよい実装方思いつかずこうしています)

refine.rs
pub(super) fn omit_or_pick(
    attr: StructAttribute,
    mut item: ItemStruct,
    attr_type: AttributeType,
) -> ItemStruct {
    // 指定した新しい構造体名を更新
    item.ident = syn::Ident::new(&attr.name.to_string(), item.ident.span());

    // Rustの場合は次も構造体になるが、対応しない方針としているためその場合はスキップする
    // - `struct S;`
    // - `struct S(T);`
    let is_tuple_or_unit = matches!(item.fields, syn::Fields::Unnamed(_) | syn::Fields::Unit);

    if !is_tuple_or_unit {
        // 構造体のフィールドをフィルターする
        let fields = item
            .fields
            .into_iter()
            .filter(|field| {
                // `AttributeType::Omit|Pick`でフィールドを省略/選択を指定する
                let should_pick = matches!(attr_type, AttributeType::Pick);
                if let Some(ref ident) = field.ident {
                    if attr.fields.contains_key(ident) {
                        return should_pick;
                    }
                };
                !should_pick
            })
            .collect();

        // フィルターしたフィールドを更新
        item.fields = syn::Fields::Named(syn::FieldsNamed {
            brace_token: syn::token::Brace::default(),
            named: fields,
        });
    }

    item
}

これであたらしい構造体を用意できたので、最後にquote!{}を使ってTokenStreamに変換します。

quote! {
    #item

    #new_item
}
.into()

quote!{}の戻り値をさらにinto()している理由ですが、
quote!{}が返すのはproc_macro2::TokenStreamとなっていて、こちらはTokenStreamのラッパーとなっています。
互換性はあってInto<T>を実装しているので、それを使ってTokenStreamに変換しています。

これでomitの実装は終わりです。

試しに、次のように使ってみると

#[omit(NewS, [b])]
struct S {
    a: i32,
    b: &'static str,
}

cargo expandで確認すると最終的に次のコードが吐かれるのが確認できます。

struct S {
    a: i32,
    b: &'static str,
}
struct NewS {
    a: i32,
}

omitの制約について

実は今のomitの実装だとジェネリクスを使った場合にうまく動かないことがあります。
具体的には次のケースの場合、ジェネリクスBが残ってしまいます。
これは今後対応予定です。

#[ommit(NewS, [b])]
struct S<A, B> {
    a: A,
    b: B,
}

// 生成される構造体に`B`が残ってしまう
struct NewS<A, B> {
    a: A,
}

まとめ

以前に仲間内でproc-macro-workshopderiveマクロをやってみたんですが、attributeマクロはやったことがなかったのでちょうどよい機会にやってみました。
実際に実装してみるとマクロに対する解像度がさらに上がったのでとてもよかったです。

また、普段からあんまりASTを触らないのもあって、RustのASTはちょっと分かりづらいなと思っていましたが今回でまた少し理解が深まりました。

実装時は、どのようなASTになるのかをAST Explorerを使って確認していました。
AST Explorersynクレートを使ってパースしているので、構造体がそのままわかって大変助かりました。

次はfunction likeマクロを使ってなにかを作ってみようと思っています。
たとえばsql! { SELECT id, name from users where id = 1 };みたいなのを書いたら、クエリを実行する・結果を構造体マッピングするメソッドをもつ構造体を生成できたら便利そうだなと思ったりしています。

FRAIMテックブログ

Discussion