Rustのproc macroを書いて、構造体定義からElasticsearchのmapping定義を導出する

弊社はGISデータ（地理空間情報）を専門的に取り扱うのですが、**Elasticsearchは実はGISデータも取り扱うことができる**ため、検索が必要な業務システムを作成する際には頻繁に利用します。

ただ、システム構築の際に、RDBMSなどを利用する場合はORMを利用することで、プログラム内部で利用する型とデータベースのテーブル定義をマッピングさせることができますが、Elasticsearchを利用する際にはあまりそういうツールはなさそうで、型定義からElasticsearch定義に変換するためのプログラムを個別に作る必要がありました。

最近、弊社ではRustを利用していく機運が高まっているため、将来的に「**構造体定義からElasticsearchのmapping定義を作成するproc macro**」があると便利だろうなーと思い、ちょっと書いてみました。

今回は特定のユースケースで利用したかっただけなので、かなり雑に書いており、Rustの全ての基本型に対応しているわけではなかったり、クレートとして公開できるようなものではなかったりするのですが、備忘のために残しておきます。
（将来的にはちゃんとしたものにしていきたい。）

https://github.com/nokonoko1203/es-mapping-macro-sample

## フォルダ構成

Rustのプロジェクトの作成方法などは他の記事に譲ります。
以下のようなフォルダ構成になりました。

```bash
❯ tree . -L 3
.
├── Cargo.lock
├── Cargo.toml
├── examples
│   └── basic_usage.rs
├── macros
│   ├── Cargo.toml
│   └── src
│       ├── lib.rs
│       └── types.rs
├── src
│   ├── lib.rs
│   └── traits
│       ├── document.rs
│       ├── mapping.rs
│       └── mod.rs
└── tests
    └── mapping.rs

14 directories, 12 files
```

今回はproc macroを作成していきますが、マクロの部分などはプロジェクト内部に別クレートを作成する必要があるようです。
クレート名は`macros`（名称は自由）としました。

こちらの記事などを参考にしました！
[[Rust] Procedural Macroの仕組みと実装方法](https://zenn.dev/tak_iwamoto/articles/890771ea5b8ad3)


## 利用方法

以下のように、「構造体にderiveして利用する」ことを目指します。

- `examples/basic_usage.rs`

```rust
use es_mapping::traits::{document::EsDocument, mapping::EsMap};
use macros::EsMapping;
use serde::{Deserialize, Serialize};
use serde_json::json;

#[derive(Serialize, Deserialize, EsMapping)]
struct User {
    #[serde(rename = "userId")]
    #[es(type = "keyword")]
    user_id: i32,
    #[es(type = "text")]
    name: String,
    #[es(type = "text")]
    age: i32,
    year: i32,
}

impl EsDocument for User {}

fn main() {
    let mapping = User::generate_mapping();
    println!("Mapping: {}", json!(mapping));
}
```

## 必要なプログラムを実装

`lib.rs`や`mod.rs`などは省略しますので、リポジトリの方を見てください。
また`Cargo.toml`も同様です。

重要そうなファイルのみ記載していきます。

まずはトレイトを定義していきます。
`EsMap`トレイトは`#[derive(EsMapping)]`のように構造体にマクロが利用された場合に、その構造体に自動的に実装されるメソッドを定義します。
つまり、`EsMapping`をderiveすると`EsMap`の実装が強制され、`generate_mapping()`メソッドが生えます。

- `src/traits/mapping.rs`
```rust
pub trait EsMap {
    fn generate_mapping() -> serde_json::Value;
}
```

次に、マクロ本体を定義します。

- `macros/src/lib.rs`
```rust
mod types;

use proc_macro::TokenStream;
use proc_macro2::TokenStream as TokenStream2;
use quote::quote;
use syn::{parse_macro_input, DeriveInput, LitStr};
use types::EsType;

#[proc_macro_derive(EsMapping, attributes(es))]
pub fn es_mapping_derive(token: TokenStream) -> TokenStream {
    let derive_input = &parse_macro_input!(token as DeriveInput);

    let fields = match &derive_input.data {
        syn::Data::Struct(ref data) => &data.fields,
        _ => {
            todo!("Only struct is supported");
        }
    };

    let mut es_fields: Vec<TokenStream2> = Vec::new();

    for field in fields {
        // field name
        let ident = field.ident.as_ref().unwrap();
        let ident_str = ident.to_string();

        // field type
        // todo: What to do when renaming with serde
        let ty = &field.ty;
        let es_type_text = rust_type_to_es_type(ty).as_str();
        let mut es_type = quote! { #es_type_text };

        // if the field has es attribute, override the type
        for attr in &field.attrs {
            let _ = attr.parse_nested_meta(|meta| {
                if meta.path.is_ident("type") {
                    let name: LitStr = meta.value()?.parse()?;
                    es_type = es_type_text_to_quote(name.value().as_str());
                }

                Ok(())
            });
        }
        es_fields.push(quote! {
            #ident_str: {
                "type": #es_type
            }
        });
    }

    let name = &derive_input.ident;
    let expanded: TokenStream2 = quote! {
        impl EsMap for #name {
            fn generate_mapping() -> serde_json::Value {
                serde_json::json!({
                    "mappings":{
                        "properties": {
                            #( #es_fields ),*
                        }
                    }
                })
            }
        }
    };

    TokenStream::from(expanded)
}

fn es_type_text_to_quote(es_type: &str) -> TokenStream2 {
    match es_type {
        "text" => quote! { "text" },
        "keyword" => quote! { "keyword" },
        "long" => quote! { "long" },
        "boolean" => quote! { "boolean" },
        "unsigned_long" => quote! { "unsigned_long" },
        "double" => quote! { "double" },
        "object" => quote! { "object" },
        "nested" => quote! { "nested" },
        "geo_point" => quote! { "geo_point" },
        "geo_shape" => quote! { "geo_shape" },
        _ => quote! { "text" },
    }
}

fn rust_type_to_es_type(ty: &syn::Type) -> EsType {
    match ty {
        syn::Type::Path(type_path) if type_path.path.is_ident("String") => EsType::Text,
        syn::Type::Path(type_path) if type_path.path.is_ident("bool") => EsType::Boolean,
        syn::Type::Path(type_path)
            if type_path.path.is_ident("i8")
                || type_path.path.is_ident("i16")
                || type_path.path.is_ident("i32")
                || type_path.path.is_ident("i64")
                || type_path.path.is_ident("i128")
                || type_path.path.is_ident("isize") =>
        {
            EsType::Number
        }
        syn::Type::Path(type_path)
            if type_path.path.is_ident("u8")
                || type_path.path.is_ident("u16")
                || type_path.path.is_ident("u32")
                || type_path.path.is_ident("u64")
                || type_path.path.is_ident("u128")
                || type_path.path.is_ident("usize") =>
        {
            EsType::UnsignedNumber
        }
        syn::Type::Path(type_path)
            if type_path.path.is_ident("f32") || type_path.path.is_ident("f64") =>
        {
            EsType::Double
        }
        syn::Type::Path(type_path) if type_path.path.is_ident("char") => EsType::Keyword,
        // todo: handle object
        syn::Type::Path(type_path)
            if type_path.path.is_ident("IndexMap") || type_path.path.is_ident("HashMap") =>
        {
            EsType::Object
        }
        // todo: handle nested
        syn::Type::Path(type_path) if type_path.path.is_ident("Vec") => EsType::Nested,
        syn::Type::Array(_) => EsType::Nested,
        // todo: handle Option<T>
        syn::Type::Path(type_path) if type_path.path.is_ident("Option") => EsType::Text,
        // todo: handle geo_point
        syn::Type::Path(type_path) if type_path.path.is_ident("GeoPoint") => EsType::GeoPoint,
        // todo: handle geo_shape
        syn::Type::Path(type_path) if type_path.path.is_ident("GeoShape") => EsType::GeoShape,
        _ => EsType::Text,
    }
}
```

基本的にマクロは「ソースコードを解析してASTに変換し、コードを付与して返す」ようなことをするため、読みづらいコードになりがちですが、今回な単純なことしかしていないので、比較的読みやすいかとは思います。
順に解説していきます。

まず`es_mapping_derive`関数では以下のように定義しています。

```rust
#[proc_macro_derive(EsMapping, attributes(es))]
pub fn es_mapping_derive(token: TokenStream) -> TokenStream {
    let derive_input = &parse_macro_input!(token as DeriveInput);

    let fields = match &derive_input.data {
        syn::Data::Struct(ref data) => &data.fields,
        _ => {
            todo!("Only struct is supported");
        }
    };
```

proc_macro_deriveマクロを利用して、マクロがどのように呼び出されるかを定義します。
今回は`EsMapping`という名前が定義されています。
また、`attributes(es)`の部分は、構造体のフィールドに対して`#[es(type = "keyword")]`のようにesという属性が利用できるようになることを示しています。

`es_mapping_derive`関数ではトークン（コンパイルされるソースコードを示すものの最小単位）のストリームを受け取り、加工されたトークンを返します。

`TokenStream`がソースコード本体のようなものなので、ソースコード内に記載された構造体の情報を受け取ることができます。

ここでは、トークンから構造体かどうかチェックしています。
enumなどでも利用することができますが、今回は構造体に絞って対応しました。

```rust
    let mut es_fields: Vec<TokenStream2> = Vec::new();

    for field in fields {
        // field name
        let ident = field.ident.as_ref().unwrap();
        let ident_str = ident.to_string();

        // field type
        // todo: What to do when renaming with serde
        let ty = &field.ty;
        let es_type_text = rust_type_to_es_type(ty).as_str();
        let mut es_type = quote! { #es_type_text };

        //後述
        ...
    }
```

上のコードでは構造体から属性名や型などを取り出しています。
途中で`rust_type_to_es_type`や`es_type_text_to_quote`という関数を利用しています。

- `rust_type_to_es_type`
```rust
fn rust_type_to_es_type(ty: &syn::Type) -> EsType {
    match ty {
        syn::Type::Path(type_path) if type_path.path.is_ident("String") => EsType::Text,
        syn::Type::Path(type_path) if type_path.path.is_ident("bool") => EsType::Boolean,
        syn::Type::Path(type_path)
            if type_path.path.is_ident("i8")
                || type_path.path.is_ident("i16")
                || type_path.path.is_ident("i32")
                || type_path.path.is_ident("i64")
                || type_path.path.is_ident("i128")
                || type_path.path.is_ident("isize") =>
        {
            EsType::Number
        }
        syn::Type::Path(type_path)
            if type_path.path.is_ident("u8")
                || type_path.path.is_ident("u16")
                || type_path.path.is_ident("u32")
                || type_path.path.is_ident("u64")
                || type_path.path.is_ident("u128")
                || type_path.path.is_ident("usize") =>
        {
            EsType::UnsignedNumber
        }
        syn::Type::Path(type_path)
            if type_path.path.is_ident("f32") || type_path.path.is_ident("f64") =>
        {
            EsType::Double
        }
        syn::Type::Path(type_path) if type_path.path.is_ident("char") => EsType::Keyword,
        // todo: handle object
        syn::Type::Path(type_path)
            if type_path.path.is_ident("IndexMap") || type_path.path.is_ident("HashMap") =>
        {
            EsType::Object
        }
        // todo: handle nested
        syn::Type::Path(type_path) if type_path.path.is_ident("Vec") => EsType::Nested,
        syn::Type::Array(_) => EsType::Nested,
        // todo: handle Option<T>
        syn::Type::Path(type_path) if type_path.path.is_ident("Option") => EsType::Text,
        // todo: handle geo_point
        syn::Type::Path(type_path) if type_path.path.is_ident("GeoPoint") => EsType::GeoPoint,
        // todo: handle geo_shape
        syn::Type::Path(type_path) if type_path.path.is_ident("GeoShape") => EsType::GeoShape,
        _ => EsType::Text,
    }
}
```

まずは`rust_type_to_es_type`でRustの型を独自に定義した`EsType`という型に変換していきます。
型定義は`macros/src/types.rs`で行っています。

- `macros/src/types.rs`
```rust
#[derive(Debug)]
pub enum EsType {
    Text,
    Keyword,
    Number,
    Boolean,
    UnsignedNumber,
    Double,
    Object,
    Nested,
    GeoPoint,
    GeoShape,
}

impl EsType {
    pub fn as_str(&self) -> &'static str {
        match self {
            EsType::Text => "text",
            EsType::Keyword => "keyword",
            EsType::Number => "long",
            EsType::Boolean => "boolean",
            EsType::UnsignedNumber => "unsigned_long",
            EsType::Double => "double",
            EsType::Object => "object",
            EsType::Nested => "nested",
            EsType::GeoPoint => "geo_point",
            EsType::GeoShape => "geo_shape",
        }
    }
}

pub struct GeoPoint {
    pub lat: f64,
    pub lon: f64,
}

pub struct GeoShape {
    pub points: Vec<GeoPoint>,
}
```

`EsType`は`as_str()`というメソッドを持っており、Elasticsearchのマッピング定義で利用する文字列を出力します。

これを以下のように`quote!`マクロでTokenStreamに変換します。
`#`をつけて利用することで、変数も埋め込むことができます。

```rust
        let ty = &field.ty;
        let es_type_text = rust_type_to_es_type(ty).as_str();
        let mut es_type = quote! { #es_type_text };
```

この処理で、Rustの型に応じたデフォルトのElasticsearchの型を定義していきました。

次のコートでは、属性ごとにElasticsearchの型を上書きするコードを書いていきます。

```rust
    let mut es_fields: Vec<TokenStream2> = Vec::new();

    for field in fields {
        // ...前述

        // if the field has es attribute, override the type
        for attr in &field.attrs {
            let _ = attr.parse_nested_meta(|meta| {
                if meta.path.is_ident("type") {
                    let name: LitStr = meta.value()?.parse()?;
                    es_type = es_type_text_to_quote(name.value().as_str());
                }

                Ok(())
            });
        }
        es_fields.push(quote! {
            #ident_str: {
                "type": #es_type
            }
        });
    }
```
 
esアトリビュートがついている場合は、`es_type_text_to_quote`関数を利用し、Elasticsearchの型情報を上書きしていきます。

- `es_type_text_to_quote`
```rust
fn es_type_text_to_quote(es_type: &str) -> TokenStream2 {
    match es_type {
        "text" => quote! { "text" },
        "keyword" => quote! { "keyword" },
        "long" => quote! { "long" },
        "boolean" => quote! { "boolean" },
        "unsigned_long" => quote! { "unsigned_long" },
        "double" => quote! { "double" },
        "object" => quote! { "object" },
        "nested" => quote! { "nested" },
        "geo_point" => quote! { "geo_point" },
        "geo_shape" => quote! { "geo_shape" },
        _ => quote! { "text" },
    }
}
```

最後に、`EsMap`の`generate_mappinng()`メソッドを強制的に実装させるコードを書きます。

```rust
    let name = &derive_input.ident;
    let expanded: TokenStream2 = quote! {
        impl EsMap for #name {
            fn generate_mapping() -> serde_json::Value {
                serde_json::json!({
                    "mappings":{
                        "properties": {
                            #( #es_fields ),*
                        }
                    }
                })
            }
        }
    };

    TokenStream::from(expanded)
```

できたので、利用してみましょう。

- `examples/basic_usage.rs`
```rust
use es_mapping::traits::{document::EsDocument, mapping::EsMap};
use macros::EsMapping;
use serde::{Deserialize, Serialize};
use serde_json::json;

#[derive(Serialize, Deserialize, EsMapping)]
struct User {
    #[serde(rename = "userId")]
    #[es(type = "keyword")]
    user_id: i32,
    #[es(type = "text")]
    name: String,
    #[es(type = "text")]
    age: i32,
    year: i32,
}

impl EsDocument for User {}

fn main() {
    let mapping = User::generate_mapping();
    println!("Mapping: {}", json!(mapping));
}
```

以下を実行すると、以下のようなマッピング定義が出てくると思います！

```bash
❯ cargo run --example basic_usage
Mapping: {"mappings":{"properties":{"age":{"type":"text"},"name":{"type":"text"},"user_id":{"type":"keyword"},"year":{"type":"long"}}}}
```

## 終わりに

ということで、Rustでマクロを書いていきました。
思ったよりも簡単にかけましたが、冒頭述べた通り、雑に書いているので、nestedやobjectなどの型にちゃんと対応しようと思うとコード量がもっと増えると思いますし、黒魔術感が出てくるのでコードリーディングも大変そうだなとは思いました。

適度に使う分には便利だと思うので、用法容量守って書いていきたいですね。