Skip to main content

typhoon_account_macro/
lib.rs

1use {
2    keys::PrimaryKeys,
3    quote::{quote, ToTokens},
4    syn::{parse_macro_input, punctuated::Punctuated, spanned::Spanned, Error, Item, Path, Token},
5    typhoon_discriminator::DiscriminatorBuilder,
6};
7
8mod keys;
9
10fn has_derive(attrs: &[syn::Attribute], derive_name: &str) -> bool {
11    attrs
12        .iter()
13        .filter(|attr| attr.path().is_ident("derive"))
14        .filter_map(|attr| {
15            attr.parse_args_with(Punctuated::<Path, Token![,]>::parse_terminated)
16                .ok()
17        })
18        .flatten()
19        .any(|path| {
20            path.segments
21                .last()
22                .is_some_and(|segment| segment.ident == derive_name)
23        })
24}
25
26#[proc_macro_derive(AccountState, attributes(key, no_space))]
27pub fn derive_account(item: proc_macro::TokenStream) -> proc_macro::TokenStream {
28    let item = parse_macro_input!(item as Item);
29    let (attrs, name, generics, fields) = match item {
30        Item::Struct(ref item_struct) => (
31            &item_struct.attrs,
32            &item_struct.ident,
33            &item_struct.generics,
34            &item_struct.fields,
35        ),
36        _ => {
37            return Error::new(item.span(), "Invalid account type")
38                .into_compile_error()
39                .into()
40        }
41    };
42
43    let space_token = if attrs.iter().any(|a| a.path().is_ident("no_space")) {
44        None
45    } else {
46        Some(quote! {
47            impl #name {
48                pub const SPACE: usize = <#name as Discriminator>::DISCRIMINATOR.len() + core::mem::size_of::<#name>();
49            }
50        })
51    };
52    let (_, ty_generics, where_clause) = generics.split_for_impl();
53
54    let keys = match PrimaryKeys::try_from(fields) {
55        Ok(fields) => fields,
56        Err(err) => return err.to_compile_error().into(),
57    };
58    let seeded_trait = keys.split_for_impl(name);
59    let discriminator = DiscriminatorBuilder::new(&name.to_string()).build();
60    let account_strategy = if has_derive(attrs, "SchemaRead") {
61        quote!(
62            WincodeStrategy<
63                {
64                    matches!(
65                <Self as wincode::SchemaRead<'static, wincode::config::DefaultConfig>>::TYPE_META,
66                wincode::TypeMeta::Static { zero_copy: true, .. }
67            )
68                },
69            >
70        )
71    } else if has_derive(attrs, "BorshDeserialize") {
72        quote!(BorshStrategy)
73    } else {
74        quote!(BytemuckStrategy)
75    };
76
77    quote! {
78        impl CheckOwner for #name #ty_generics #where_clause {
79            #[inline(always)]
80            fn owned_by(owner: &Address) -> bool {
81                address_eq(owner, &crate::ID)
82            }
83        }
84
85        impl Discriminator for #name #ty_generics #where_clause {
86            const DISCRIMINATOR: &'static [u8] = &[#(#discriminator),*];
87        }
88
89        impl DataStrategy for #name #ty_generics #where_clause {
90            type Strategy = #account_strategy;
91        }
92
93        #space_token
94
95        #seeded_trait
96    }
97    .into_token_stream()
98    .into()
99}