Skip to main content

batch_impl/
lib.rs

1use proc_macro::TokenStream;
2use proc_macro2::{Ident, TokenStream as TokenStream2, TokenTree};
3use quote::quote;
4use syn::{parse_macro_input, ItemTrait};
5
6mod types;
7mod apply;
8mod parse;
9mod codegen;
10mod preprocess;
11
12use codegen::generate_impl;
13use parse::{parse_item, Cursor};
14use types::{reset_fresh_counter, Op};
15
16/// 为 trait 批量生成 `impl` 块的属性宏。
17///
18/// 在 trait 定义上标注 `#[batch_impl(...)]`,宏参数中的每个 impl-spec 都会
19/// 为该 trait 生成一个对应的 `impl` 块。
20///
21/// ## 语法
22///
23/// ```text
24/// #[batch_impl( impl-spec [, impl-spec]* [{ body }]? )]
25/// ```
26///
27/// impl-spec 由三部分组成(均可省略后半部分):
28/// - `<impl-泛型>` — `impl` 块的泛型参数
29/// - `Trait名<trait-泛型>` — trait 的泛型参数与关联类型绑定
30/// - 目标类型 — 用 `[]` 包裹表示并列,用 `^`/`-` 表示泛型应用
31///
32/// ## 示例
33///
34/// ```ignore
35/// #[batch_impl(usize, isize)]
36/// trait Numeric {}
37///
38/// #[batch_impl(<T> Vec<T>)]
39/// trait Collection {}
40///
41/// #[batch_impl(<T> FromValue<T> i32 { fn wrap(_: T) -> Self { 0 } })]
42/// trait FromValue<T> { fn wrap(val: T) -> Self; }
43/// ```
44#[proc_macro_attribute]
45pub fn batch_impl(attr: TokenStream, item: TokenStream) -> TokenStream {
46    expand_attr_macro(attr, item, true)
47}
48
49/// 与 `#[batch_impl]` 相同,但丢弃 trait 定义本身,只输出 `impl` 块。
50///
51/// 用于 trait 已在别处定义、只需批量生成 impl 的场景。
52/// 语法与 `#[batch_impl]` 完全一致。
53///
54/// ## 示例
55///
56/// ```ignore
57/// trait Greet { fn hello(&self) -> &str; }
58///
59/// #[batch_impl_only(usize #hello{"hi"})]
60/// trait Greet {}  // 此 trait 定义被丢弃,不影响已有的定义
61/// ```
62#[proc_macro_attribute]
63pub fn batch_impl_only(attr: TokenStream, item: TokenStream) -> TokenStream {
64    expand_attr_macro(attr, item, false)
65}
66
67/// 两个属性宏的共享实现
68fn expand_attr_macro(attr: TokenStream, item: TokenStream, include_trait: bool) -> TokenStream {
69    reset_fresh_counter();
70    let trait_item = parse_macro_input!(item as ItemTrait);
71    let trait_name = trait_item.ident.clone();
72    let attr_vec = TokenStream2::from(attr).into_iter().collect::<Vec<_>>();
73    let mut cursor = Cursor::new(&attr_vec);
74    let expanded = match preprocess::expand_tokens(&mut cursor, &trait_item) {
75        Ok(tokens) => tokens,
76        Err(err) => return err.into(),
77    };
78    cursor = Cursor::new(&expanded);
79    let trait_name_ts = quote![#trait_name];
80    let is_unsafe = trait_item.unsafety.is_some();
81    let start_trait = if include_trait { Some(trait_item) } else { None };
82    let impls = parse_batch_trait_entry(
83        &mut cursor, Op::Comma, &trait_name_ts, &trait_name,
84        is_unsafe, start_trait,
85    );
86    impls.into()
87}
88
89/// 对已声明的 trait 批量生成 `impl` 块的函数式宏。
90///
91/// 语法:`unsafe? Trait路径: impl-specs;`,以 `;` 分隔多个 trait 段。
92/// 每段的 `:` 之后是 DSL 表达式,与 `#[batch_impl]` 接受相同的语法。
93///
94/// ## 示例
95///
96/// ```ignore
97/// trait A {}
98/// trait B<T> {}
99/// mod foo { pub trait C {} }
100///
101/// batch_trait!(
102///     A: usize, isize;
103///     B: <T> B<T> Vec<T>;
104///     foo::C: u32;
105///     unsafe UnsafeTrait: usize
106/// );
107/// ```
108#[proc_macro]
109pub fn batch_trait(input: TokenStream) -> TokenStream {
110    reset_fresh_counter();
111    let tokens = TokenStream2::from(input).into_iter().collect::<Vec<_>>();
112    let mut cursor = Cursor::new(&tokens);
113    let mut result = quote![];
114    loop {
115        // 跳过前导 `;`(允许连续多个分号、尾随分号)
116        while cursor.is_punct(';') {
117            cursor.bump();
118        }
119        if cursor.at_end() {
120            break;
121        }
122
123        // `unsafe` 前缀:标记该段所有 impl 为 unsafe impl
124        let is_unsafe = if matches!(cursor.peek(), Some(TokenTree::Ident(id)) if *id == "unsafe") {
125            cursor.bump();
126            true
127        } else {
128            false
129        };
130
131        // 收集 trait 路径(遇到 `<>` 深度为 0 的 `:` 停止;`::` 路径分隔符一并收集)
132        let path_start = cursor.pos();
133        let mut depth = 0i32;
134        while let Some(token) = cursor.peek() {
135            match token {
136                TokenTree::Punct(p) if p.as_char() == '<' => {
137                    depth += 1;
138                    cursor.bump();
139                }
140                TokenTree::Punct(p) if p.as_char() == '>' => {
141                    depth -= 1;
142                    cursor.bump();
143                }
144                TokenTree::Punct(p) if p.as_char() == ':' && depth == 0 => {
145                    if matches!(cursor.peek_at(1), Some(TokenTree::Punct(p2)) if p2.as_char() == ':') {
146                        cursor.bump();
147                        cursor.bump();
148                    } else {
149                        break;
150                    }
151                }
152                _ => cursor.bump(),
153            }
154        }
155        let trait_path = cursor.slice_since(path_start);
156        if trait_path.is_empty() {
157            result.extend(generate_compile_error("batch_trait! 中期望 trait 名称"));
158            break;
159        }
160        let trait_full_path = match extract_trait_path(trait_path) {
161            Ok(path) => path,
162            Err(e) => { result.extend(e); break; }
163        };
164        let trait_last_ident = match extract_last_ident(trait_path) {
165            Ok(ident) => ident,
166            Err(e) => { result.extend(e); break; }
167        };
168        if !cursor.is_punct(':') {
169            result.extend(generate_compile_error(
170                "batch_trait! 中期望 ':' 分隔 trait 名称和 impl-specs",
171            ));
172            break;
173        }
174        cursor.bump();
175        let impl_code = parse_batch_trait_entry(
176            &mut cursor, Op::Semi, &trait_full_path,
177            trait_last_ident, is_unsafe, None,
178        );
179        result.extend(impl_code);
180    }
181    result.into()
182}
183
184/// 共享驱动:从游标解析 impl-specs,展开并列列表,生成 impl 块。
185///
186/// `top_level` 控制顶层优先级:
187/// - `Op::Comma` 用于 `#[batch_impl]`(整个参数按 `,` 分隔)
188/// - `Op::Semi` 用于 `batch_trait!` 的单段 specs(按 `,` 分隔,遇到 `;` 段落边界停止)
189///
190/// 展开阶段通过 BFS 工作清单把 `Ty::Array`(并列列表)逐层摊平为叶子 `Ty`,
191/// 再对每个叶子调用 `generate_impl` 生成对应的 impl 块。
192fn parse_batch_trait_entry(
193    cursor: &mut Cursor,
194    top_level: Op,
195    trait_full_path: &TokenStream2,
196    trait_last_ident: &Ident,
197    is_unsafe_trait: bool,
198    start_trait: Option<ItemTrait>,
199) -> TokenStream2 {
200    let mut tys = vec![];
201    while let Some(ty) = parse_item(cursor, top_level, Some(trait_last_ident)) {
202        let mut queue = vec![ty];
203        while let Some(item) = queue.pop() {
204            match item.expand() {
205                Ok(expanded) => {
206                    for e in expanded.into_iter().rev() {
207                        queue.push(e);
208                    }
209                }
210                Err(leaf) => tys.push(leaf),
211            }
212        }
213    }
214    let mut impls = start_trait.map_or(quote![], |t| quote![#t]);
215    for t in tys {
216        impls.extend(generate_impl(t, trait_full_path, is_unsafe_trait));
217    }
218    impls
219}
220
221/// 从 trait 路径 token 序列中提取完整路径(用于输出 `impl ... for`)
222fn extract_trait_path(trait_path: &[TokenTree]) -> Result<TokenStream2, TokenStream2> {
223    let path: TokenStream2 = trait_path.iter().cloned().collect();
224    if path.is_empty() {
225        Err(generate_compile_error("batch_trait! 中期望 trait 名称"))
226    } else {
227        Ok(path)
228    }
229}
230
231/// 从 trait 路径 token 序列中提取最后一个标识符(用作 trait_name 匹配)
232fn extract_last_ident(trait_path: &[TokenTree]) -> Result<&Ident, TokenStream2> {
233    trait_path
234        .iter()
235        .filter_map(|tt| if let TokenTree::Ident(id) = tt { Some(id) } else { None })
236        .next_back()
237        .ok_or_else(|| generate_compile_error("batch_trait! 中期望标识符作为 trait 名称"))
238}
239
240/// 构造 `compile_error!(msg)` 用于编译期报错
241fn generate_compile_error(msg: &str) -> TokenStream2 {
242    quote! { compile_error!(#msg); }
243}