mtb_entity_slab_macros/
lib.rs

1use proc_macro::TokenStream;
2use proc_macro2::Span;
3use proc_macro_crate::{crate_name, FoundCrate};
4use quote::{quote, ToTokens};
5use syn::{
6    parse::Parse, punctuated::Punctuated, spanned::Spanned, Error, Expr, ExprLit, Generics, Ident,
7    Lit, Meta, MetaNameValue,
8};
9
10/// 自动为结构体/枚举实现 `IEntityAllocatable` 的属性宏。
11///
12/// 参数（均为可选，顺序不限）：
13/// - `policy = ...` 选择分配策略。可接受：
14///   - 短名称：`Policy128 | Policy256 | Policy512 | Policy1024 | Policy2048 | Policy4096`
15///   - 完整类型名（不需要带类型参数）：`EntityAllocPolicy256`
16///   - 整数字面量：`128 | 256 | 512 | 1024 | 2048 | 4096`
17///   默认：`Policy256`。
18/// - `ptrid = TypePath` 使用你自定义的 PtrID 类型（必须实现 `Copy + Eq`，并提供
19///   `From<PtrID<Self>>` 与 `From<Custom> for PtrID<Self>` 转换）。
20/// - `wrapper = Ident` 生成一个透明的 PtrID 包装新类型（字段为 `pub`，`Debug` 打印指针地址）。
21/// - `opaque_wrapper = Ident` 生成一个不透明的 PtrID 包装新类型（字段为私有，`Debug` 隐藏指针）。
22///
23/// 冲突与重复：
24/// - `ptrid` 不能与 `wrapper`/`opaque_wrapper` 同时出现；
25/// - `wrapper` 与 `opaque_wrapper` 不能同时出现；
26/// - 同一键名重复书写会报编译错误。
27///
28/// 示例：
29/// ```ignore
30/// use mtb_entity_slab::entity_allocatable;
31///
32/// #[entity_allocatable]
33/// struct A;
34///
35/// #[entity_allocatable(policy = Policy512)]
36/// struct B<T>(T);
37///
38/// #[entity_allocatable(wrapper = CPtr)]
39/// struct C;
40///
41/// #[entity_allocatable(opaque_wrapper = DPtr, policy = 1024)]
42/// enum D { X, Y }
43///
44/// #[entity_allocatable(ptrid = my_crate::MyPtrId)]
45/// struct E;
46/// ```
47///
48/// English brief: Attribute macro implementing `IEntityAllocatable`.
49/// Args (optional, order independent): `policy`, `ptrid`, `wrapper`, `opaque_wrapper`.
50/// See README for details and examples.
51#[proc_macro_attribute]
52pub fn entity_allocatable(attr: TokenStream, item: TokenStream) -> TokenStream {
53    let mut attr = syn::parse_macro_input!(attr as EntityAllocatableAttr);
54    let item = syn::parse_macro_input!(item as syn::DeriveInput);
55    // Only allow on struct or enum
56    match &item.data {
57        syn::Data::Struct(_) | syn::Data::Enum(_) => {}
58        _ => {
59            return syn::Error::new(
60                item.ident.span(),
61                "#[entity_allocatable] 只能用于 struct 或 enum",
62            )
63            .to_compile_error()
64            .into();
65        }
66    }
67
68    // Capture ident and generics for codegen
69    attr.ident = item.ident.clone();
70    attr.generics = item.generics.clone();
71
72    let output = quote! {
73        #item
74        #attr
75    };
76    TokenStream::from(output)
77}
78
79struct EntityAllocatableAttr {
80    crate_path: proc_macro2::TokenStream,
81    ident: Ident,
82    generics: syn::Generics,
83    policy: AllocPolicy,
84    ptrid_wrapper: PtrIDWrapper,
85}
86
87impl Parse for EntityAllocatableAttr {
88    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
89        // Defaults
90        let crate_path: proc_macro2::TokenStream = match crate_name("mtb-entity-slab") {
91            Ok(FoundCrate::Itself) => quote! { crate },
92            Ok(FoundCrate::Name(name)) => {
93                let ident = Ident::new(&name, proc_macro2::Span::call_site());
94                quote! { #ident }
95            }
96            Err(_) => quote! { mtb_entity_slab },
97        };
98
99        let mut ret = Self {
100            crate_path,
101            ident: syn::parse_quote! { __entity_placeholder }, // will be overwritten in entry
102            generics: syn::Generics::default(),
103            policy: AllocPolicy::Policy256,
104            ptrid_wrapper: PtrIDWrapper::Default,
105        };
106        // Parse comma-separated meta list: key = value, ...
107        let metas = Punctuated::<Meta, syn::token::Comma>::parse_terminated(input)?;
108        for meta in metas {
109            ret.parse_meta(meta)?;
110        }
111        Ok(ret)
112    }
113}
114
115impl EntityAllocatableAttr {
116    fn parse_meta(&mut self, meta: Meta) -> syn::Result<()> {
117        const UNSUPPORTED_ARG: &str =
118        "Unsupported argument; available arguments: policy=..., ptrid=..., wrapper=..., opaque_wrapper=...";
119
120        let Meta::NameValue(named) = meta else {
121            return Err(Error::new(meta.span(), UNSUPPORTED_ARG));
122        };
123        let Some(ident) = named.path.get_ident() else {
124            return Err(Error::new(named.path.span(), UNSUPPORTED_ARG));
125        };
126        match ident.to_string().as_str() {
127            "policy" => self.parse_policy(&named),
128            "ptrid" => self.parse_ptrid(&named),
129            "wrapper" => self.parse_wrapper(&named),
130            "opaque_wrapper" => self.parse_opaque_wrapper(&named),
131            _ => return Err(Error::new(named.span(), UNSUPPORTED_ARG)),
132        }
133    }
134
135    fn parse_policy(&mut self, named: &MetaNameValue) -> syn::Result<()> {
136        const POLICY_TRUE_KIND: &str =
137            "policy should be an identifier or integer, e.g., Policy256 or 256";
138        self.policy = match &named.value {
139            Expr::Path(ep) => {
140                let last = ep.path.segments.last();
141                let ident = match last {
142                    Some(last) => last.ident.to_string(),
143                    None => String::new(),
144                };
145                AllocPolicy::from_str(&ident, ep.span())
146            }
147            Expr::Lit(ExprLit {
148                lit: Lit::Int(i), ..
149            }) => {
150                let s = i.base10_digits();
151                AllocPolicy::from_str(s, i.span())
152            }
153            other => Err(Error::new(other.span(), POLICY_TRUE_KIND)),
154        }?;
155        Ok(())
156    }
157    fn parse_ptrid(&mut self, named: &MetaNameValue) -> syn::Result<()> {
158        const PTRID_CONFLICTS: &str = "`ptrid` conflicts with `wrapper`/`opaque_wrapper`";
159        const REPEATED_PTRID: &str = "`ptrid` specified multiple times";
160        const SHOULD_BE_PATH: &str = "external `ptrid` must be a type path, e.g., MyPtrID";
161        match self.ptrid_wrapper {
162            PtrIDWrapper::Transparent(_) | PtrIDWrapper::Opaque(_) => {
163                return Err(Error::new(named.span(), PTRID_CONFLICTS));
164            }
165            PtrIDWrapper::External(_) => {
166                return Err(Error::new(named.span(), REPEATED_PTRID));
167            }
168            _ => {}
169        }
170        let Expr::Path(ep) = &named.value else {
171            return Err(Error::new(named.value.span(), SHOULD_BE_PATH));
172        };
173        let ty: syn::Type = syn::parse_quote! { #ep };
174        self.ptrid_wrapper = PtrIDWrapper::External(ty);
175        Ok(())
176    }
177    fn parse_wrapper(&mut self, named: &MetaNameValue) -> Result<(), Error> {
178        const WRAPPER_CONFLICTS: &str = "`wrapper` conflicts with `ptrid`/`opaque_wrapper`";
179        const REPEATED_WRAPPER: &str = "`wrapper` specified multiple times";
180        const SHOULD_BE_INDENT: &str = "wrapper must be an identifier, e.g., wrapper = MyPtrID";
181        match self.ptrid_wrapper {
182            PtrIDWrapper::External(_) | PtrIDWrapper::Opaque(_) => {
183                return Err(Error::new(named.span(), WRAPPER_CONFLICTS));
184            }
185            PtrIDWrapper::Transparent(_) => {
186                return Err(Error::new(named.span(), REPEATED_WRAPPER));
187            }
188            _ => {}
189        }
190        let Expr::Path(ep) = &named.value else {
191            return Err(Error::new(named.value.span(), SHOULD_BE_INDENT));
192        };
193        if let Some(seg) = ep.path.segments.last() {
194            self.ptrid_wrapper = PtrIDWrapper::Transparent(seg.ident.clone());
195            Ok(())
196        } else {
197            Err(Error::new(ep.span(), SHOULD_BE_INDENT))
198        }
199    }
200    fn parse_opaque_wrapper(&mut self, named: &MetaNameValue) -> syn::Result<()> {
201        const WRAPPER_CONFLICTS: &str = "`opaque_wrapper` conflicts with `ptrid`/`wrapper`";
202        const REPEATED_WRAPPER: &str = "`opaque_wrapper` specified multiple times";
203        const SHOULD_BE_INDENT: &str =
204            "opaque_wrapper must be an identifier, e.g., opaque_wrapper = MyPtrID";
205        match self.ptrid_wrapper {
206            PtrIDWrapper::External(_) | PtrIDWrapper::Transparent(_) => {
207                return Err(Error::new(named.span(), WRAPPER_CONFLICTS));
208            }
209            PtrIDWrapper::Opaque(_) => {
210                return Err(Error::new(named.span(), REPEATED_WRAPPER));
211            }
212            _ => {}
213        }
214        let Expr::Path(ep) = &named.value else {
215            return Err(Error::new(named.value.span(), SHOULD_BE_INDENT));
216        };
217        if let Some(seg) = ep.path.segments.last() {
218            self.ptrid_wrapper = PtrIDWrapper::Opaque(seg.ident.clone());
219            Ok(())
220        } else {
221            Err(Error::new(ep.span(), SHOULD_BE_INDENT))
222        }
223    }
224}
225
226impl ToTokens for EntityAllocatableAttr {
227    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
228        let crate_path = &self.crate_path;
229        let ident = &self.ident;
230        let (impl_generics, ty_generics, where_clause) = self.generics.split_for_impl();
231        let entity_ty = quote! { #ident #ty_generics };
232        let ptrid_ty_tokens = self
233            .ptrid_wrapper
234            .ptrid_ty_tokens(crate_path, ident, &ty_generics);
235        let policy_ty_tokens = self.policy.to_tokens(crate_path);
236
237        tokens.extend(quote! {
238            impl #impl_generics #crate_path::IEntityAllocatable for #ident #ty_generics #where_clause {
239                type AllocatePolicyT = #policy_ty_tokens<#entity_ty>;
240                type PtrID = #ptrid_ty_tokens;
241            }
242        });
243
244        if let Some(wrapper) = self
245            .ptrid_wrapper
246            .define_wrapper(crate_path, ident, &self.generics)
247        {
248            tokens.extend(wrapper);
249        }
250    }
251}
252
253enum PtrIDWrapper {
254    External(syn::Type),
255    Transparent(Ident),
256    Opaque(Ident),
257    Default,
258}
259impl PtrIDWrapper {
260    fn define_wrapper(
261        &self,
262        crate_path: &proc_macro2::TokenStream,
263        ident: &Ident,
264        generics: &Generics,
265    ) -> Option<proc_macro2::TokenStream> {
266        use PtrIDWrapper::*;
267        let wrapper_ident = match self {
268            Transparent(id) | Opaque(id) => id,
269            _ => return None,
270        };
271        let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
272        let field_vis = if matches!(self, Transparent(_)) {
273            quote!(pub)
274        } else {
275            quote!()
276        };
277
278        // Debug body: opaque does not reveal address
279        let debug_body = if matches!(self, Opaque(_)) {
280            quote! { write!(f, concat!(stringify!(#wrapper_ident), "(<opaque>)")) }
281        } else {
282            quote! { write!(f, concat!(stringify!(#wrapper_ident), "({:p})"), self.0) }
283        };
284
285        Some(quote! {
286            #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
287            pub struct #wrapper_ident #impl_generics (#field_vis #crate_path::PtrID<#ident #ty_generics>) #where_clause;
288
289            impl #impl_generics std::fmt::Debug for #wrapper_ident #ty_generics #where_clause {
290                fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
291                    #debug_body
292                }
293            }
294            impl #impl_generics From<#crate_path::PtrID<#ident #ty_generics>> for #wrapper_ident #ty_generics #where_clause {
295                fn from(p: #crate_path::PtrID<#ident #ty_generics>) -> Self { Self(p) }
296            }
297            impl #impl_generics From<#wrapper_ident #ty_generics> for #crate_path::PtrID<#ident #ty_generics> #where_clause {
298                fn from(x: #wrapper_ident #ty_generics) -> Self { x.0 }
299            }
300        })
301    }
302    fn ptrid_ty_tokens(
303        &self,
304        crate_path: &proc_macro2::TokenStream,
305        ident: &Ident,
306        ty_generics: &syn::TypeGenerics,
307    ) -> proc_macro2::TokenStream {
308        use PtrIDWrapper::*;
309        match self {
310            External(ty) => quote! { #ty },
311            Transparent(wrapper_ident) | Opaque(wrapper_ident) => {
312                quote! { #wrapper_ident #ty_generics }
313            }
314            Default => quote! { #crate_path::PtrID<#ident #ty_generics> },
315        }
316    }
317}
318
319enum AllocPolicy {
320    Policy128,
321    Policy256,
322    Policy512,
323    Policy1024,
324    Policy2048,
325    Policy4096,
326}
327
328impl AllocPolicy {
329    const ERRMSG: &str =
330        "Unknown allocation policy. Expected one of: Policy{128|256|512|1024|2048|4096}";
331
332    fn to_tokens(&self, crate_path: &proc_macro2::TokenStream) -> proc_macro2::TokenStream {
333        match self {
334            AllocPolicy::Policy128 => quote! { #crate_path::EntityAllocPolicy128 },
335            AllocPolicy::Policy256 => quote! { #crate_path::EntityAllocPolicy256 },
336            AllocPolicy::Policy512 => quote! { #crate_path::EntityAllocPolicy512 },
337            AllocPolicy::Policy1024 => quote! { #crate_path::EntityAllocPolicy1024 },
338            AllocPolicy::Policy2048 => quote! { #crate_path::EntityAllocPolicy2048 },
339            AllocPolicy::Policy4096 => quote! { #crate_path::EntityAllocPolicy4096 },
340        }
341    }
342
343    fn from_str(s: &str, errpos: Span) -> syn::Result<Self> {
344        match s {
345            "128" | "Policy128" | "EntityAllocPolicy128" => Ok(AllocPolicy::Policy128),
346            "256" | "Policy256" | "EntityAllocPolicy256" => Ok(AllocPolicy::Policy256),
347            "512" | "Policy512" | "EntityAllocPolicy512" => Ok(AllocPolicy::Policy512),
348            "1024" | "Policy1024" | "EntityAllocPolicy1024" => Ok(AllocPolicy::Policy1024),
349            "2048" | "Policy2048" | "EntityAllocPolicy2048" => Ok(AllocPolicy::Policy2048),
350            "4096" | "Policy4096" | "EntityAllocPolicy4096" => Ok(AllocPolicy::Policy4096),
351            _ => Err(syn::Error::new(errpos, Self::ERRMSG)),
352        }
353    }
354}
mtb_entity_slab_macros/lib.rs

mtb_entity_slab_macros/
lib.rs
