Skip to main content

basm_sdk_community_macros/
lib.rs

1use darling::{ast::NestedMeta, FromMeta};
2use proc_macro::TokenStream;
3use quote::{quote};
4use syn::{parse_macro_input, ItemFn, Signature, Type};
5
6#[derive(Debug, FromMeta)]
7struct EntrypointArgs {
8    #[darling(default="default_json_deserializer_fn")]
9    secret_deserializer: syn::TypePath,
10    #[darling(default="default_json_deserializer_fn")]
11    input_deserializer: syn::TypePath,
12    #[darling(default="default_json_serializer_fn")]
13    output_serializer: syn::TypePath,
14}
15
16fn default_json_serializer_fn() -> syn::TypePath {
17    syn::parse_str("serde_json::to_string").unwrap()
18}
19fn default_json_deserializer_fn() -> syn::TypePath {
20    syn::parse_str("serde_json::from_str").unwrap()
21}
22
23// Extracts the context types from the function signature.
24fn extract_context_types(sig: &Signature) -> Result<(Type, Type), syn::Error> {
25    // 1. Get the first function argument.
26    let first_arg = sig.inputs.first().ok_or_else(|| {
27        syn::Error::new_spanned(sig, "Function must have at least one argument")
28    })?;
29
30    // 2. Check if the first argument is a typed argument.
31    let pat_type = if let syn::FnArg::Typed(pt) = first_arg {
32        pt
33    } else {
34        return Err(syn::Error::new_spanned(
35            first_arg,
36            "Expected a typed argument",
37        ));
38    };
39
40    // 3. Ensure the argument's type is a path type, e.g., `my_crate::Context<A, B>`.
41    let type_path = if let syn::Type::Path(tp) = &*pat_type.ty {
42        tp
43    } else {
44        return Err(syn::Error::new_spanned(
45            &pat_type.ty,
46            "Argument type must be a path like `Context<A, B>`",
47        ));
48    };
49
50    // 4. Get the last part of the path (the type name and its generics).
51    let last_segment = type_path.path.segments.last().ok_or_else(|| {
52        syn::Error::new_spanned(&type_path.path, "Type path must have at least one segment")
53    })?;
54
55    // 5. Check that the type is `Context`.
56    if last_segment.ident != "Context" {
57        return Err(syn::Error::new_spanned(
58            &last_segment.ident,
59            "Expected the argument type to be `Context`",
60        ));
61    }
62
63    // 6. Extract the angle-bracketed generic arguments, e.g., `<A, B>`.
64    let angle_bracketed_args =
65        if let syn::PathArguments::AngleBracketed(args) = &last_segment.arguments {
66            args
67        } else {
68            return Err(syn::Error::new_spanned(
69                last_segment,
70                "Expected angle-bracketed generic arguments like `<A, B>`",
71            ));
72        };
73
74    // 7. Collect all the *type* arguments from the generics.
75    let generic_types: Vec<&Type> = angle_bracketed_args
76        .args
77        .iter()
78        .filter_map(|arg| if let syn::GenericArgument::Type(ty) = arg { Some(ty) } else { None })
79        .collect();
80
81    // 8. There should be either 1 or 2 type arguments.
82    if generic_types.len() < 1 || generic_types.len() > 2 {
83        return Err(syn::Error::new_spanned(
84            &angle_bracketed_args.args,
85            "Expected 1 or 2 type arguments in `Context<A, B>`",
86        ));
87    }
88
89    if generic_types.len() == 1 {
90        return Ok((
91            generic_types[0].clone(),
92            syn::parse_quote!(()),
93        ));
94    } else {
95        return Ok((
96            generic_types[0].clone(),
97            generic_types[1].clone(),
98        ));
99    }
100}
101
102/// Serves as entrypoint for a function that can be called from the host environment.
103#[proc_macro_attribute]
104pub fn bky_entrypoint(args: TokenStream, input: TokenStream) -> TokenStream {
105    let attr_args = match NestedMeta::parse_meta_list(args.into()) {
106        Ok(v) => v,
107        Err(e) => return e.into_compile_error().into(),
108    };
109    let EntrypointArgs {
110        input_deserializer,
111        output_serializer,
112        secret_deserializer
113    } = match EntrypointArgs::from_list(&attr_args) {
114        Ok(v) => v,
115        Err(e) => return e.write_errors().into(),
116    };
117
118    let input = parse_macro_input!(input as ItemFn);
119    let sig = &input.sig;
120
121    let (input_type, secret_type) = match extract_context_types(sig) {
122        Ok((input_type, secret_type)) => (input_type, secret_type),
123        Err(e) => return e.into_compile_error().into(),
124    };
125
126    let fn_name = &sig.ident.clone();
127
128    quote! {
129        #[unsafe(no_mangle)]
130        pub extern "C" fn #fn_name(
131            input_ptr: basm_sdk_community::memory::FatPointer,
132            secret_ptr: basm_sdk_community::memory::FatPointer
133        ) -> basm_sdk_community::memory::FatPointer {
134            #input
135
136            std::panic::set_hook(Box::new(|p| {
137                basm_sdk_community::host_log!("{}", p);
138                std::process::abort();
139            }));
140
141            let input_data = unsafe { input_ptr.copy_data() };
142            let secret_data = unsafe { secret_ptr.copy_data() };
143
144            let input_str = str::from_utf8(&input_data).unwrap();
145            let input: #input_type = #input_deserializer(input_str)
146                .expect("Failed to deserialize input JSON");
147
148            let secrets_str = str::from_utf8(&secret_data).unwrap();
149            let secrets: #secret_type = #secret_deserializer(secrets_str)
150                .expect("Failed to deserialize secrets JSON");
151
152            let ctx = Context {
153                input: input,
154                secrets: secrets
155            };
156
157            let output = #fn_name(ctx);
158
159            let x = #output_serializer(&output)
160                .expect("Failed to serialize output into JSON");
161            unsafe { basm_sdk_community::memory::leak_to_shared_memory(&x.as_bytes()) }
162        }
163    }.into()
164}