bolt_attribute_bolt_component/
lib.rs

1use proc_macro::TokenStream;
2
3use quote::quote;
4use syn::{
5    parse_macro_input, parse_quote, Attribute, DeriveInput, Lit, Meta, MetaList, MetaNameValue,
6    NestedMeta,
7};
8
9use bolt_utils::add_bolt_metadata;
10use heck::ToSnakeCase;
11
12/// This Component attribute is used to automatically generate the seed and size functions
13///
14/// The component_id can be used to define the seed used to generate the PDA which stores the component data.
15/// The macro also adds the InitSpace and Default derives to the struct.
16///
17/// #[component]
18/// #[derive(Default, Copy)]
19/// pub struct Position {
20///     pub x: i64,
21///     pub y: i64,
22///     pub z: i64,
23/// }
24/// ```
25#[proc_macro_attribute]
26pub fn component(attr: TokenStream, item: TokenStream) -> TokenStream {
27    let mut input = parse_macro_input!(item as DeriveInput);
28    let mut component_id_value = None;
29    let mut delegate_set = false;
30
31    if !attr.is_empty() {
32        let attr_meta = parse_macro_input!(attr as Meta);
33        delegate_set = is_delegate_set(&attr_meta);
34        component_id_value = match attr_meta {
35            Meta::Path(_) => None,
36            Meta::NameValue(meta_name_value) => extract_component_id(&meta_name_value),
37            Meta::List(meta_list) => {
38                if !delegate_set {
39                    delegate_set = is_delegate_set(&Meta::List(meta_list.clone()));
40                }
41                find_component_id_in_list(meta_list)
42            }
43        };
44    }
45
46    let component_id_value = component_id_value.unwrap_or_else(|| "".to_string());
47
48    let additional_macro: Attribute = parse_quote! { #[account] };
49    let additional_derives: Attribute = parse_quote! { #[derive(InitSpace)] };
50    input.attrs.push(additional_derives);
51
52    let new_fn = define_new_fn(&input);
53
54    add_bolt_metadata(&mut input);
55
56    let name = &input.ident;
57
58    let snake_case_name = name.to_string().to_snake_case();
59    let component_name = syn::Ident::new(&snake_case_name, input.ident.span());
60
61    let bolt_program = if delegate_set {
62        quote! {
63            #[delegate(#name)]
64            #[bolt_program(#name)]
65            pub mod #component_name {
66                use super::*;
67            }
68        }
69    } else {
70        quote! {
71            #[bolt_program(#name)]
72            pub mod #component_name {
73                use super::*;
74            }
75        }
76    };
77
78    let expanded = quote! {
79        #bolt_program
80
81        #additional_macro
82        #input
83
84        #new_fn
85
86        #[automatically_derived]
87        impl ComponentTraits for #name {
88            fn seed() -> &'static [u8] {
89                #component_id_value.as_bytes()
90            }
91
92            fn size() -> usize {
93                8 + <#name>::INIT_SPACE
94            }
95        }
96
97    };
98    expanded.into()
99}
100
101/// Create a fn `new` to initialize the struct without bolt_metadata field
102fn define_new_fn(input: &DeriveInput) -> proc_macro2::TokenStream {
103    let struct_name = &input.ident;
104    let init_struct_name = syn::Ident::new(&format!("{}Init", struct_name), struct_name.span());
105
106    if let syn::Data::Struct(ref data) = input.data {
107        if let syn::Fields::Named(ref fields) = data.fields {
108            // Generate fields for the init struct
109            let init_struct_fields = fields.named.iter().map(|f| {
110                let name = &f.ident;
111                let ty = &f.ty;
112                quote! { pub #name: #ty }
113            });
114
115            // Generate struct initialization code using the init struct
116            let struct_init_fields = fields.named.iter().map(|f| {
117                let name = &f.ident;
118                quote! { #name: init_struct.#name }
119            });
120
121            // Generate the new function and the init struct
122            let gen = quote! {
123                // Define a new struct to hold initialization parameters
124                pub struct #init_struct_name {
125                    #(#init_struct_fields),*
126                }
127
128                impl #struct_name {
129                    pub fn new(init_struct: #init_struct_name) -> Self {
130                        Self {
131                            #(#struct_init_fields,)*
132                            bolt_metadata: BoltMetadata::default(),
133                        }
134                    }
135                }
136            };
137            return gen;
138        }
139    }
140    quote! {}
141}
142
143fn is_delegate_set(meta: &Meta) -> bool {
144    match meta {
145        Meta::Path(path) => path.is_ident("delegate"),
146        Meta::List(meta_list) => meta_list.nested.iter().any(|nested_meta| {
147            if let NestedMeta::Meta(Meta::Path(path)) = nested_meta {
148                path.is_ident("delegate")
149            } else {
150                false
151            }
152        }),
153        _ => false,
154    }
155}
156
157fn extract_component_id(meta_name_value: &MetaNameValue) -> Option<String> {
158    if meta_name_value.path.is_ident("component_id") {
159        if let Lit::Str(lit) = &meta_name_value.lit {
160            return Some(lit.value());
161        }
162    }
163    None
164}
165
166fn find_component_id_in_list(meta_list: MetaList) -> Option<String> {
167    meta_list.nested.into_iter().find_map(|nested_meta| {
168        if let NestedMeta::Meta(Meta::NameValue(meta_name_value)) = nested_meta {
169            extract_component_id(&meta_name_value)
170        } else {
171            None
172        }
173    })
174}