Skip to main content

cellbook_macros/
lib.rs

1use proc_macro::TokenStream;
2use quote::{format_ident, quote};
3use syn::visit_mut::VisitMut;
4use syn::{DeriveInput, Expr, ExprLit, FnArg, ItemFn, Lit, Meta, MetaNameValue, parse_macro_input};
5
6/// Adds `ctx` prefix to context macro calls.
7struct CtxInjector;
8
9impl VisitMut for CtxInjector {
10    fn visit_macro_mut(&mut self, mac: &mut syn::Macro) {
11        let path = &mac.path;
12        let is_context_macro = path.is_ident("store")
13            || path.is_ident("storev")
14            || path.is_ident("load")
15            || path.is_ident("loadv")
16            || path.is_ident("remove")
17            || path.is_ident("consume")
18            || path.is_ident("consumev");
19
20        if is_context_macro {
21            let tokens = &mac.tokens;
22            mac.tokens = quote! { ctx, #tokens };
23        }
24    }
25}
26
27/// Marks an async function as a cellbook cell.
28///
29/// The macro:
30/// - Adds a `ctx: CellContext` parameter
31/// - Generates a `#[no_mangle]` wrapper for FFI
32/// - Registers the cell with inventory
33///
34/// ```ignore
35/// #[cell]
36/// async fn my_cell() -> Result<()> {
37///     store!(data)?;
38///     Ok(())
39/// }
40/// ```
41#[proc_macro_attribute]
42pub fn cell(_attr: TokenStream, item: TokenStream) -> TokenStream {
43    let mut input = parse_macro_input!(item as ItemFn);
44
45    let fn_name = input.sig.ident.clone();
46    let fn_name_str = fn_name.to_string();
47    let wrapper_name = format_ident!("__cellbook_cell_{}", fn_name_str);
48    let line = fn_name.span().start().line as u32;
49
50    CtxInjector.visit_item_fn_mut(&mut input);
51
52    let ctx_param: FnArg = syn::parse_quote!(ctx: &::cellbook::CellContext);
53    input.sig.inputs.insert(0, ctx_param);
54
55    let fn_vis = &input.vis;
56    let fn_sig = &input.sig;
57    let fn_block = &input.block;
58    let fn_attrs = &input.attrs;
59
60    let expanded = quote! {
61        #(#fn_attrs)*
62        #fn_vis #fn_sig #fn_block
63
64        #[doc(hidden)]
65        #[unsafe(no_mangle)]
66        pub fn #wrapper_name(
67            store_fn: fn(&str, Vec<u8>, &str),
68            load_fn: fn(&str) -> Option<(Vec<u8>, String)>,
69            remove_fn: fn(&str) -> Option<(Vec<u8>, String)>,
70            list_fn: fn() -> Vec<(String, String)>,
71        ) -> ::cellbook::futures::future::BoxFuture<'static, ::std::result::Result<(), Box<dyn ::std::error::Error + Send + Sync>>> {
72            let ctx = ::cellbook::CellContext::new(store_fn, load_fn, remove_fn, list_fn);
73            Box::pin(async move {
74                #fn_name(&ctx)
75                    .await
76                    .map_err(|e| -> Box<dyn ::std::error::Error + Send + Sync> { e.into() })
77            })
78        }
79
80        ::cellbook::inventory::submit!(::cellbook::CellInfo {
81            name: #fn_name_str,
82            func: #wrapper_name,
83            line: #line,
84        });
85    };
86
87    TokenStream::from(expanded)
88}
89
90/// Marks an async function as the required cellbook init entrypoint.
91///
92/// The macro:
93/// - Keeps the function as-is (arbitrary function name)
94/// - Exports `__cellbook_get_cells`
95/// - Exports `__cellbook_get_init`
96///
97/// ```ignore
98/// #[init]
99/// async fn setup() -> Result<()> {
100///     Ok(())
101/// }
102/// ```
103#[proc_macro_attribute]
104pub fn init(_attr: TokenStream, item: TokenStream) -> TokenStream {
105    let input = parse_macro_input!(item as ItemFn);
106    let fn_name = input.sig.ident.clone();
107    let fn_name_str = fn_name.to_string();
108    let wrapper_name = format_ident!("__cellbook_init_{}", fn_name_str);
109    let line = fn_name.span().start().line as u32;
110
111    let fn_vis = &input.vis;
112    let fn_sig = &input.sig;
113    let fn_block = &input.block;
114    let fn_attrs = &input.attrs;
115
116    let expanded = quote! {
117        #(#fn_attrs)*
118        #fn_vis #fn_sig #fn_block
119
120        #[doc(hidden)]
121        #[unsafe(no_mangle)]
122        pub fn #wrapper_name() -> ::cellbook::futures::future::BoxFuture<'static, ::std::result::Result<(), Box<dyn ::std::error::Error + Send + Sync>>> {
123            Box::pin(async move {
124                #fn_name()
125                    .await
126                    .map_err(|e| -> Box<dyn ::std::error::Error + Send + Sync> { e.into() })
127            })
128        }
129
130        #[unsafe(no_mangle)]
131        pub extern "Rust" fn __cellbook_get_cells() -> Vec<(
132            String,
133            u32,
134            fn(
135                fn(&str, Vec<u8>, &str),
136                fn(&str) -> Option<(Vec<u8>, String)>,
137                fn(&str) -> Option<(Vec<u8>, String)>,
138                fn() -> Vec<(String, String)>,
139            ) -> ::cellbook::futures::future::BoxFuture<'static, ::std::result::Result<(), Box<dyn ::std::error::Error + Send + Sync>>>
140        )> {
141            ::cellbook::registry::cells()
142                .into_iter()
143                .map(|c| (c.name.to_string(), c.line, c.func))
144                .collect()
145        }
146
147        #[unsafe(no_mangle)]
148        pub extern "Rust" fn __cellbook_get_init() -> (
149            String,
150            u32,
151            fn() -> ::cellbook::futures::future::BoxFuture<'static, ::std::result::Result<(), Box<dyn ::std::error::Error + Send + Sync>>>
152        ) {
153            (#fn_name_str.to_string(), #line, #wrapper_name)
154        }
155    };
156
157    TokenStream::from(expanded)
158}
159
160/// Derive `cellbook::StoreSchema` with a version set by `#[store_schema(version = N)]`.
161#[proc_macro_derive(StoreSchema, attributes(store_schema))]
162pub fn derive_store_schema(item: TokenStream) -> TokenStream {
163    let input = parse_macro_input!(item as DeriveInput);
164    let ident = input.ident;
165    let generics = input.generics;
166    let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
167    let mut version: Option<u32> = None;
168
169    for attr in &input.attrs {
170        if !attr.path().is_ident("store_schema") {
171            continue;
172        }
173
174        let parsed = match attr
175            .parse_args_with(syn::punctuated::Punctuated::<Meta, syn::Token![,]>::parse_terminated)
176        {
177            Ok(v) => v,
178            Err(e) => return e.to_compile_error().into(),
179        };
180
181        for meta in parsed {
182            let Meta::NameValue(MetaNameValue { path, value, .. }) = meta else {
183                return syn::Error::new_spanned(attr, "expected #[store_schema(version = <u32>)]")
184                    .to_compile_error()
185                    .into();
186            };
187
188            if !path.is_ident("version") {
189                return syn::Error::new_spanned(path, "unknown store_schema key")
190                    .to_compile_error()
191                    .into();
192            }
193
194            let Expr::Lit(ExprLit {
195                lit: Lit::Int(lit_int),
196                ..
197            }) = value
198            else {
199                return syn::Error::new_spanned(value, "version must be an integer literal")
200                    .to_compile_error()
201                    .into();
202            };
203
204            match lit_int.base10_parse::<u32>() {
205                Ok(v) => version = Some(v),
206                Err(e) => {
207                    return syn::Error::new_spanned(lit_int, e).to_compile_error().into();
208                }
209            }
210        }
211    }
212
213    let Some(version) = version else {
214        return syn::Error::new_spanned(
215            &ident,
216            "missing #[store_schema(version = <u32>)] for #[derive(StoreSchema)]",
217        )
218        .to_compile_error()
219        .into();
220    };
221
222    let expanded = quote! {
223        impl #impl_generics ::cellbook::StoreSchema for #ident #ty_generics #where_clause {
224            const VERSION: u32 = #version;
225        }
226    };
227    TokenStream::from(expanded)
228}