natom_macros/
lib.rs

1extern crate proc_macro;
2extern crate proc_macro2;
3extern crate quote;
4extern crate syn;
5
6use proc_macro::TokenStream;
7use proc_macro2::Span;
8use quote::{quote, format_ident};
9use syn::{Attribute, parse_quote, ReturnType};
10
11macro_rules! parse_error {
12    ($($arg:tt)*) => {{
13        let error = syn::LitStr::new(&format!($($arg)*), Span::call_site());
14        quote!(compile_error!(#error);).into()
15    }}
16}
17
18#[derive(Hash, PartialEq, Eq, Clone, Debug)]
19enum AtomScope {
20    Initial,
21    Singleton,
22    Dependent,
23}
24
25impl AtomScope {
26    fn from(value: &str) -> Option<Self> {
27        match value {
28            "initial" => Some(Self::Initial),
29            "singleton" => Some(Self::Singleton),
30            "dependent" => Some(Self::Dependent),
31            _ => None,
32        }
33    }
34}
35
36impl ToString for AtomScope {
37    fn to_string(&self) -> String {
38        match self {
39            Self::Initial => "Initial".to_string(),
40            Self::Singleton => "Singleton".to_string(),
41            Self::Dependent => "Dependent".to_string(),
42        }
43    }
44}
45
46fn append_atom_doc(attrs: &mut Vec<Attribute>) {
47    let mut added_documentation: Vec<syn::Attribute> = vec![
48        parse_quote!(#[doc = "### Atom"]),
49        parse_quote!(#[doc = "This object is an atom. More documentation coming soon."]),
50    ];
51    attrs.append(&mut added_documentation);
52}
53
54fn parse_scope(item: TokenStream) -> Result<AtomScope, String> {
55    if item.is_empty() {
56        return Ok(AtomScope::Initial)
57    }
58
59    let Ok(item_ast) = syn::parse::<syn::Ident>(item) else {
60        return Err("malformed atom scope, expecting a single identifier.".into());
61    };
62    let Some(atom_scope) = AtomScope::from(&item_ast.to_string()) else {
63        return Err("invalid atom scope".into());
64    };
65
66    return Ok(atom_scope)
67}
68
69fn field_is_injected(field: &syn::Field) -> bool {
70    field.attrs.iter().find(|a| a.path().is_ident("inject")).is_some()
71}
72
73fn make_injections_arc(fields: &mut syn::Fields) {
74    fields.iter_mut()
75        .filter(|f| field_is_injected(f))
76        .for_each(|f| {
77            let ty = &f.ty;
78            f.ty = parse_quote!(std::sync::Arc<#ty>)
79        });
80}
81
82static mut AFABRIC_COMPTIME_REGISTERED_ATOMS: Vec<String> = Vec::new();
83
84fn atom_impl(atom_type: AtomScope, decl: syn::ItemStruct) -> TokenStream {
85    let mut decl = decl.clone();
86    let decl_name = &decl.ident;
87
88    append_atom_doc(&mut decl.attrs);
89
90    make_injections_arc(&mut decl.fields);
91
92    let injected_fields = decl.fields.iter()
93        .filter(|f| field_is_injected(f));
94    let injected_types: Vec<syn::Type> = injected_fields.clone().map(|f| f.ty.clone()).collect();
95
96    let gen_depends_on = quote!(
97        vec![
98            #(std::any::TypeId::of::<#injected_types>()),*
99        ]
100    );
101
102    let injections: Vec<syn::FieldValue> = injected_fields.clone().map(|f| {
103        let injected_type = &f.ty;
104        syn::FieldValue {
105            attrs: Vec::new(),
106            member: syn::Member::Named(f.ident.clone().unwrap()),
107            colon_token: Some(parse_quote!(:)),
108            expr: parse_quote!(
109                system.resolve(std::any::TypeId::of::<#injected_type>()).downcast_ref::<#injected_type>().unwrap()
110            )
111        }
112    }).collect();
113    let injections = quote!(
114        #(#injections),*
115    );
116
117    // remove all [inject] from the decl
118    decl.fields.iter_mut().for_each(|f| {
119        if field_is_injected(&f) {
120            f.attrs.retain(|a| !a.path().is_ident("inject"));
121        }
122    });
123
124    let atom_type_ident = format_ident!("{}", atom_type.to_string());
125
126    unsafe {
127        AFABRIC_COMPTIME_REGISTERED_ATOMS.push(decl_name.to_string());
128    }
129
130    let gen = quote!(
131        #decl
132
133        impl #decl_name {
134            #[doc(hidden)]
135            fn __atom_create<T: afabric::AtomRegistry>(system: afabric::Context<T>) -> #decl_name {
136                #decl_name {
137                    #injections
138                }
139            }
140
141            #[doc(hidden)]
142            fn __atom_meta() -> afabric::AtomMeta {
143                afabric::AtomMeta {
144                    type_id: std::any::TypeId::of::<#decl_name>(),
145                    atom_type: afabric::AtomType::#atom_type_ident,
146                    depends_on: #gen_depends_on,
147                }
148            }
149        }
150    );
151
152    gen.into()
153}
154
155/// Annotates a class to be an Atom.
156#[proc_macro_attribute]
157pub fn atom(attr: TokenStream, item: TokenStream) -> TokenStream {
158    let atom_scope_or_err = parse_scope(attr);
159    let Ok(atom_scope) = atom_scope_or_err else {
160        return parse_error!("{}", atom_scope_or_err.unwrap_err());
161    };
162    
163    let Ok(decl) = syn::parse::<syn::ItemStruct>(item) else {
164        return parse_error!("atom only applies to structs.");
165    };
166    
167    atom_impl(atom_scope, decl)
168}
169
170#[proc_macro_attribute]
171pub fn produces(_attr: TokenStream, item: TokenStream) -> TokenStream {
172    // we ignore attr for now, but will be used for named lol.
173    let Ok(func_impl) = syn::parse::<syn::ImplItemFn>(item) else {
174        return parse_error!("'produces' should only annotate functions on atoms.");
175    };
176
177    // ensure we have a specific return-type (not '()' ).
178    let ReturnType::Type(_, ret_type) = func_impl.clone().sig.output else {
179        return parse_error!("'produces' requires function to return discrete type, not '()'");
180    };
181
182    // ensure we have zero or one parameter, where the parameter must be {self, &self}.
183    let inputs_count = func_impl.clone().sig.inputs.iter().count();
184    if inputs_count > 1 {
185        return parse_error!("'produces' requires function to take zero or one argument, but got {}", inputs_count);
186    }
187
188    println!("produces {} ->", quote!(#ret_type));
189
190    quote!(#func_impl).into()
191}
192
193#[proc_macro]
194pub fn define_base_atoms(_item: TokenStream) -> TokenStream {
195    unsafe {
196        let get_atoms: Vec<proc_macro2::TokenStream> = AFABRIC_COMPTIME_REGISTERED_ATOMS.iter().map(|name| {
197            let name = format_ident!("{}", name);
198            quote!(
199                let meta = #name::__atom_meta();
200                if id == meta.type_id {
201                    return Some(meta)
202                }
203            )
204        }).collect();
205
206        let atom_meta: Vec<proc_macro2::TokenStream> = AFABRIC_COMPTIME_REGISTERED_ATOMS.iter().map(|name| {
207            let name = format_ident!("{}", name);
208            quote!(#name::__atom_meta())
209        }).collect();
210
211        quote!(
212            pub struct BaseAtomRegistry {}
213            impl afabric::AtomRegistry for BaseAtomRegistry {
214                fn get(id: std::any::TypeId) -> Option<afabric::AtomMeta> {
215                    #(#get_atoms)*
216    
217                    None
218                }
219
220                fn all() -> Vec<afabric::AtomMeta> {
221                    vec![
222                        #(#atom_meta),*
223                    ]
224                }
225            }
226        ).into()
227    }
228}
natom_macros/lib.rs

natom_macros/
lib.rs
