RustでTypeScriptのOmitなどの型ユーティリティを実装してみた
はじめに
Rustを書いていると、たまにTypeScriptのOmitなどの型ユーティリティがほしいなと思う場面があります。
Rust本体にはそのような型ユーティリティはないので、勉強がてらにクレートを作ってみました
本記事は作ったクレートの簡単な紹介と実装について書いていきます。
使い方
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![;]>,
}
quote
はsyn
の構造体などからコンパイラが扱えるソースコードのトークンであるTokenStream
に変換してくれます。
let tokens = quote! {
struct #struct_name {
#(#fields),*
}
};
tokens.into::<TokenStream>()
実装の概要について
使用するクレートの概要についてわかったところで、omit
を使って実装の概要について説明していきます。
まず、最初にクレートの種類をproc-macro
にする必要があります。
[lib]
proc-macro = true
次にlib.rs
にomit
のマクロ関数を定義します。
#[proc_macro_attribute]
pub fn omit(attr: TokenStream, item: TokenStream) -> TokenStream {
// omitted
}
マクロ関数はTokenStream
を受け取って任意の構文木のTokenStream
を作成して返すことで、新しいソースコードを生成するイメージです。
TokenStream
はさきほど説明したとおり、コンパイラが扱えるソースコードのトークン情報です。
今回はomit
の新しい構造体のTokenStream
を作成して返す処理を実装します。
構造体の作成の流れは次のようになります。
-
syn::parse_macro_input!()
を使ってTokenStream
からsyn::ItemStruct
などに変換し、属性や構造体の情報を扱えるようにする - 新しい構造体に対応した
syn::ItemStruct
を作成する -
quote
を使って2で作成したsyn::ItemStruct
からTokenStream
に変換して返す
attribute
マクロの場合、関数は2つの引数を受け取ります。
1つ目は属性情報のTokenStream
、2つ目は構造体(や関数など)のTokenStream
です。
属性情報というのは次の例でいうとomit(...)
で囲っているNewS, [a, b]
の部分になります。
構造体のTokenStream
はstrcut S { ... }
の部分になります。
#[omit(NewS, [a, b])]
struct S {
// omitted
}
実装の詳細について
概要についてわかったところで、次にomit
の詳細について説明してきます。
#[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
トレイトを実装します。
#[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
した構造体に代入するくらいです。
omit
とpick
はフィルター時の判定が反転するだけなので、その判定を切り替えられるようにAttributeType::Omit|Pick
を用意しています。(正直いけていない感があるが、あんまりよい実装方思いつかずこうしています)
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-workshopのderive
マクロをやってみたんですが、attribute
マクロはやったことがなかったのでちょうどよい機会にやってみました。
実際に実装してみるとマクロに対する解像度がさらに上がったのでとてもよかったです。
また、普段からあんまりASTを触らないのもあって、RustのASTはちょっと分かりづらいなと思っていましたが今回でまた少し理解が深まりました。
実装時は、どのようなASTになるのかをAST Explorerを使って確認していました。
AST Explorer
もsyn
クレートを使ってパースしているので、構造体がそのままわかって大変助かりました。
次はfunction like
マクロを使ってなにかを作ってみようと思っています。
たとえばsql! { SELECT id, name from users where id = 1 };
みたいなのを書いたら、クエリを実行する・結果を構造体マッピングするメソッドをもつ構造体を生成できたら便利そうだなと思ったりしています。
Discussion