Rust 過程宏 proc-macro 是個啥

2023-06-25 21:01:19

定義一個 procedural macro

新建一個 lib 型別的 crate:

cargo new hello-macro --lib

procedural macros 只能在 proc-macro 型別的 crate 內定義,所以需要修改 Cargo.toml:

[lib]
proc-macro = true

刪除 src/lib.rs 裡的全部內容,然後定義第一個過程宏(procedural macro):

use proc_macro::TokenStream;

#[proc_macro]
pub fn hello_proc(input: TokenStream) -> TokenStream {
    input
}

目前它的作用跟下面這個宣告宏(declarative macro) 是等價的:

#[macro_export]
macro_rules! hello_macro {
    (
        $($tt: tt)*
    ) => {
        $($tt)*
    };
}

就是把所有傳入的 token 全部都原樣返回. TokenStream 相當於宣告宏裡的 $($tt: tt)*,
一連串的 token(TokenTree)
全部放到了一個 stream(其實內部就是個 Vec<TokenTree>) 裡

pub enum TokenTree {
    Group(Group), // [...], {...}, (...)
    Ident(Ident), // 函數名, struct 名等
    Punct(Punct), // 各種符號: + - * / ; &
    Literal(Literal), // 各種字面值: 123 'a' "hello" 
}

其中 Ident, PunctLiteral 都屬於單個的 token,
Group 是被三種括號(() [] {})包裹起來的 tokens

測試一下, 修改程式碼

#[proc_macro]
pub fn hello_proc(input: TokenStream) -> TokenStream {
    for tt in input.into_iter() {
        println!("tt: {:#?}", tt);
    }

    TokenStream::new()
}

然後

cargo new hello # 新建 bin 型別的 crate
cd hello
cargo add --path ../hello-macro # 新增我們的過程宏依賴

然後在 src/main.rs 裡呼叫 hello_proc

use hello_macro::hello_proc;

fn main() {
    hello_proc! {
        let a=8;[1,2,] {1+2 "hello world"}
    }
}

build 一下

cargo build
tt: Ident {
    ident: "let",
    span: #0 bytes(514..517),
}
tt: Ident {
    ident: "a",
    span: #0 bytes(518..519),
}
tt: Punct {
    ch: '=',
    spacing: Alone,
    span: #0 bytes(519..520),
}
tt: Literal {
    kind: Integer,
    symbol: "8",
    suffix: None,
    span: #0 bytes(520..521),
}
tt: Punct {
    ch: ';',
    spacing: Alone,
    span: #0 bytes(521..522),
}
tt: Group {
    delimiter: Bracket,
    stream: TokenStream [
        Literal {
            kind: Integer,
            symbol: "1",
            suffix: None,
            span: #0 bytes(523..524),
        },
        Punct {
            ch: ',',
            spacing: Alone,
            span: #0 bytes(524..525),
        },
        Literal {
            kind: Integer,
            symbol: "2",
            suffix: None,
            span: #0 bytes(525..526),
        },
        Punct {
            ch: ',',
            spacing: Alone,
            span: #0 bytes(526..527),
        },
    ],
    span: #0 bytes(522..528),
}
tt: Group {
    delimiter: Brace,
    stream: TokenStream [
        Literal {
            kind: Integer,
            symbol: "1",
            suffix: None,
            span: #0 bytes(530..531),
        },
        Punct {
            ch: '+',
            spacing: Alone,
            span: #0 bytes(531..532),
        },
        Literal {
            kind: Integer,
            symbol: "2",
            suffix: None,
            span: #0 bytes(532..533),
        },
        Literal {
            kind: Str,
            symbol: "hello world",
            suffix: None,
            span: #0 bytes(534..547),
        },
    ],
    span: #0 bytes(529..548),
}

能幹啥

過程宏的入參是一連串的 tokens, 這些都是編譯器在進行語法分析之前的 tokens, 而且我們可以在過程宏的函數裡執行復雜的邏輯, 且是在編譯期執行, 因此我們可以對這些 tokens 做任何事情, 比如定義一套新的語法,解析其它語言等等

甚至我可以在過程宏函數內執行一些毫不相干的程式碼,比如挖礦。這是一些惡意的過程宏可能會做的事情

Builder Pattern

先看需求:

derive_struct! {
    struct Foo {}
}

// derive_struct 展開後變成下面的程式碼
struct Foo {}
struct FooBuilder{}

分析一下, 我們需要給傳入的 struct 加一個 Builder. 如果用「宣告式宏」來做, 怎樣才能把一個 ident(Foo) 變成
另一個 ident(FooBuilder) 呢? 好像沒有辦法(如果你知道的話, 請一定告訴我). 那麼我們用過程宏呢, 我們可以取得 ident(Foo),
也可以定義新的 ident(FooBuilder), 理論上完全 OK.

來,讓我們在不借助第三方庫的情況下試一下

#[proc_macro]
pub fn derive_struct(mut input: TokenStream) -> TokenStream {
   let mut iter = input.clone().into_iter();

   assert_eq!(iter.next().unwrap().to_string().as_str(), "struct");

   let Some(proc_macro::TokenTree::Ident(ident)) = iter.next() else {
       panic!("parse struct identifier error");
   };

   let builder: TokenStream = format!(
           "struct {}{} {}", 
           ident, "Builder", "{}"
       )
       .parse().unwrap();

   input.extend(builder.into_iter());

   input
}

測試程式碼 main.rs

use hello_macro::derive_struct;

derive_struct! {
    struct Foo {
        a: u8,
    }
}

fn main() {}

檢視展開後的程式碼

# 安裝 cargo-expand
# cargo install cargo-expand
cargo expand

展開後的程式碼:

struct Foo {
    a: u8,
}
struct FooBuilder {}

我們目前只解析了最簡單形式的 struct, 如果要再複雜一些, 比如帶泛型和 meta data, 那麼解析起來就會麻煩很多。
幸運的是我們可以藉助 syn 來代替我們手動 parse,
這篇文章 中所有 Metavariables 都能用 syn 來解析,
我們現在需要解析出 ItemStruct 就夠了

在 hello-macro 目錄下新增依賴:

cargo add syn --features full # syn::Item 需要 full feature

然後修改 derive_struct:

#[proc_macro]
pub fn derive_struct(mut input: TokenStream) -> TokenStream {
    let item_struct: syn::ItemStruct = syn::parse(input.clone()).unwrap();

    let ident = item_struct.ident;

    let builder: TokenStream = format!(
            "struct {}{} {}", 
            ident, "Builder", "{}"
        )
        .parse().unwrap();

    input.extend(builder.into_iter());

    input
}

TokenStremsyn::Item 簡單了,那反方向解析有沒有方便使用的 crate 呢?
有, quote

新增依賴

cargo add quote

修改我們的 derive_struct:

#[proc_macro]
pub fn derive_struct(input: TokenStream) -> TokenStream {
    let item_struct: syn::ItemStruct = syn::parse(input.clone()).unwrap();

    let vis = &item_struct.vis;
    let ident = quote::format_ident!("{}Builder", item_struct.ident);
    let generics = &item_struct.generics;

    quote! {
        #item_struct

        #vis struct #ident #generics {}
    }
    .into()
}

quote::quote 是一個「宣告式宏」, 它的內部其實是將 (# $var:ident) 替換為 var.to_tokens()(需要 var 的型別實現 ToTokens trait),
#(#var)* 的用法也跟宣告式宏類似

繼續改進:

#[proc_macro]
pub fn derive_struct(input: TokenStream) -> TokenStream {
    let mut item_struct: syn::ItemStruct = syn::parse(input.clone()).unwrap();

    let attr: syn::Attribute = syn::parse_quote! {
        #[derive(Default)]
    };

    if item_struct.attrs.iter().all(|x| {
        x.to_token_stream().to_string() != attr.to_token_stream().to_string()
    }) {
        item_struct.attrs.push(attr);
    }

    item_struct.generics.make_where_clause();

    let vis = &item_struct.vis;
    let generics = &item_struct.generics; // <T: Default>
    let generic_where_clause = &generics.where_clause;

    let mut generic_params = generics.params.clone();
    generic_params = generic_params.into_iter().filter_map(|mut v| {
        match &mut v {
            syn::GenericParam::Lifetime(_) => None,
            syn::GenericParam::Type(ty) => {
                ty.bounds.clear();
                ty.attrs.clear();
                Some(v)
            },
            syn::GenericParam::Const(c) => {
                let ident = c.ident.clone();
                Some(syn::parse_quote! {
                    #ident
                })
            },
        }
    }).collect();

    // println!("generics: {}", generics.to_token_stream());
    // println!("generic_params: {}", generic_params.to_token_stream());
    // println!("generic_where_clause: {}", generic_where_clause.to_token_stream());

    let ident = &item_struct.ident;
    let builder_ident = quote::format_ident!("{}Builder", item_struct.ident);
    let fields = &item_struct.fields;

    let syn::Fields::Named(_) = fields else {
        panic!("struct with unnamed fields like `struct Foo(String);` is not supported.");
    };

    let field_ident: Vec<syn::Ident> = fields.iter().map(|f|f.ident.clone().unwrap()).collect();
    let field_ty: Vec<syn::Type> = fields.iter().map(|f|f.ty.clone()).collect();

    quote! {
        #item_struct

        impl #generics #ident <#generic_params> {
            pub fn builder() -> #builder_ident <#generic_params>{
                Default::default()
            }
        }

        #[derive(Default)]
        #vis struct #builder_ident #generics {
            inner: #ident <#generic_params>,
        }

        impl #generics #builder_ident <#generic_params> {
            pub fn build(self) -> #ident <#generic_params> {
                self.inner
            }
            #(
                pub fn #field_ident(mut self, #field_ident: #field_ty) -> Self {
                    self.inner.#field_ident = #field_ident;
                    self
                }
            )*
        }
    }
    .into()
}

目前的 derive_struct 已經可以支援下面這種 struct 了

derive_struct! {
    #[derive(Debug)]
    pub struct Bar<const N: usize, T: Default> {
        a: u8,
        b: String,
        c: T,
    }
}

派生宏

我們前面定義的過程宏 derive_struct 中文名叫「函數式宏」, 在這個場景下雖然能用, 但是每次都要把整個 struct 包裹起來,還是很麻煩的。這時 proc_macro_derive(中文叫「派生宏」) 就該出場了,
定義一個名為 Builder 的派生宏:

// attributes 可以加到 fields 上, 如果不需要可以不要這個 attributes
#[proc_macro_derive(Builder, attributes(attr1, attr2,))]
pub fn my_builder(input: TokenStream) -> TokenStream {
    let input: syn::DeriveInput = syn::parse(input).unwrap();
    let syn::Data::Struct(data) = input.data else {
        panic!("Sorry, we only support struct.");
    };

    let vis = input.vis;
    let generics = input.generics;
    let builder_ident = quote::format_ident!("{}Builder", input.ident);

    // input.attrs;
    // data.fields;

    quote! {
        #vis struct #builder_ident #generics {}
    }
    .into()
}

proc_macro_derive 是專門用來處理 derive 型別的過程宏的, 函數名可以隨意, input 引數是跟宏相關聯的某個 item, 在這裡它總是 enum, struct 或 union 其中的一種, 因為只有這
三種 item 可以標註 derive 屬性。函數返回值會被追加到 item 後面(「函數式宏」會完全替換掉原來的 TokenStream)

#[derive(Debug, Builder)]
struct Foo {
    a: u32,
    #[attr1]
    b: String,
    #[attr2(hello = world)]
    c: (u32, u32),
}
// struct FooBuilder {} // 會被追加到這裡

屬性宏

「屬性宏」的返回值也是會完全替換掉輸入的 item

#[proc_macro_attribute]
pub fn hello_attr(attr: TokenStream, item: TokenStream) -> TokenStream {
    // println!("hello_attr attr: {}, item: {}", attr, item);
    item
}
#[hello_attr(hello world)]
fn foo() {}

總結

  • 「過程宏」是比「宣告式宏」能力更強的一種宏,可以在編譯期執行復雜邏輯
  • 熟練寫「宣告式宏」對理解「過程宏」很有幫助,建議學習「過程宏」之前先學習好「宣告式宏」
  • 寫宏的時候多多參閱 The Rust Reference, 可以更深入地理解 Rust 語言
  • 在學習過程中,使用 proc-macro2, synquote 之前,建議先嚐試用 Rust 標準庫程式碼實現,這樣可以更好的理解這幾個庫
  • 寫宏的過程會強迫你對 Rust 語言的細節有更多的理解

關於 proc-macro2

https://crates.io/crates/proc-macro2

https://veykril.github.io/tlborm/proc-macros/third-party-crates.html

由於 proc_macro crate 是專門為 proc_macro 型別 crate 設計的,因此使它們可進行單元測試或從非 proc_macro 程式碼中存取它們幾乎是不可能的。鑑於此,proc-macro2 crate 模仿了原始 proc_macro crate 的 API,在 proc_macro crates 中充當包裝器,在非 proc_macro crates 中則可獨立使用。因此,建議針對 proc_macro 程式碼構建庫時,使用 proc-macro2 來進行構建,這將使這些庫可進行單元測試,這也是為什麼下面列出的 crate 取出和發射 proc-macro2::TokenStreams 的原因。當需要 proc_macro token stream 時,可以簡單地將 proc-macro2 token stream 轉換為 proc_macro 版本,反之亦然。