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#[proc_macro_attribute]
45pub fn batch_impl(attr: TokenStream, item: TokenStream) -> TokenStream {
46 expand_attr_macro(attr, item, true)
47}
48
49#[proc_macro_attribute]
63pub fn batch_impl_only(attr: TokenStream, item: TokenStream) -> TokenStream {
64 expand_attr_macro(attr, item, false)
65}
66
67fn 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#[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 while cursor.is_punct(';') {
117 cursor.bump();
118 }
119 if cursor.at_end() {
120 break;
121 }
122
123 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 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
184fn 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
221fn 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
231fn 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
240fn generate_compile_error(msg: &str) -> TokenStream2 {
242 quote! { compile_error!(#msg); }
243}