axiom_sdk_derive/
lib.rs

1use input::{impl_flatten_and_raw_input, impl_new_struct};
2use proc_macro::TokenStream;
3use syn::{parse_macro_input, DeriveInput, FnArg, Ident, ItemFn, ItemStruct, Type, Visibility};
4
5extern crate proc_macro;
6extern crate syn;
7#[macro_use]
8extern crate quote;
9
10mod input;
11
12#[proc_macro_attribute]
13#[allow(non_snake_case)]
14/// Derive the `AxiomComputeInput` trait for a struct.
15/// The struct must be named `_Input`. Ex: `ExampleInput`.
16/// All the fields of the struct must implement `RawInput` from `axiom_circuit::input::raw_input`,
17/// which has already been implemented for most primitive Rust types and Ethers types.
18pub fn AxiomComputeInput(_args: TokenStream, input: TokenStream) -> TokenStream {
19    let input_clone = input.clone();
20    let ast = parse_macro_input!(input_clone as ItemStruct);
21    // let name = ast.ident.clone();
22    let new_struct = impl_new_struct(&ast);
23    if let Err(err) = new_struct {
24        return err.into();
25    }
26    let new_struct = new_struct.unwrap();
27    let new_derive_input: TokenStream = new_struct.clone().into();
28    let new_ast = parse_macro_input!(new_derive_input as DeriveInput);
29    let flatten = impl_flatten_and_raw_input(&new_ast);
30    quote! {
31        #[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
32        #ast
33        #new_struct
34        #flatten
35    }
36    .into()
37}
38
39#[proc_macro_attribute]
40#[allow(non_snake_case)]
41pub fn AxiomCompute(_attr: TokenStream, item: TokenStream) -> TokenStream {
42    let input_fn = parse_macro_input!(item as ItemFn);
43
44    let _fn_name = &input_fn.sig.ident;
45
46    let mut trait_fn = input_fn.clone();
47    trait_fn.sig.ident = Ident::new("compute", input_fn.sig.ident.span());
48    trait_fn.vis = Visibility::Inherited;
49
50    let fourth_input_type = if let Some(FnArg::Typed(pat_type)) = input_fn.sig.inputs.iter().nth(3)
51    {
52        match &*pat_type.ty {
53            Type::Path(type_path) => type_path
54                .path
55                .segments
56                .first()
57                .expect("Type path should have at least one segment")
58                .ident
59                .to_string(),
60            _ => panic!("Unsupported type format for the fourth input"),
61        }
62    } else {
63        panic!("The function does not have a fourth input");
64    };
65
66    if !fourth_input_type.ends_with("CircuitInput") {
67        panic!("The assigned input name must end with `CircuitInput`");
68    }
69
70    let input_name = fourth_input_type
71        .trim_end_matches("CircuitInput")
72        .to_string()
73        + "Input";
74    let input_name_ident = Ident::new(&input_name, input_fn.sig.ident.span());
75
76    quote! {
77        impl crate::compute::AxiomComputeFn for #input_name_ident {
78            #trait_fn
79        }
80    }
81    .into()
82}
axiom_sdk_derive/lib.rs

axiom_sdk_derive/
lib.rs
