sg_basic_whitelist_derive/
lib.rs

1use proc_macro::TokenStream;
2use quote::quote;
3use syn::{parse_macro_input, AttributeArgs, DataEnum, DeriveInput};
4
5/// Merges the variants of two enums.
6///
7/// Adapted from cw-plus-plus which is adapted from DAO DAO:
8/// https://github.com/larry0x/cw-plus-plus/blob/main/packages/ownable/derive/src/lib.rs
9fn merge_variants(metadata: TokenStream, left: TokenStream, right: TokenStream) -> TokenStream {
10    use syn::Data::Enum;
11
12    // parse metadata
13    let args = parse_macro_input!(metadata as AttributeArgs);
14    if let Some(first_arg) = args.first() {
15        return syn::Error::new_spanned(first_arg, "macro takes no arguments")
16            .to_compile_error()
17            .into();
18    }
19
20    // parse the left enum
21    let mut left: DeriveInput = parse_macro_input!(left);
22    let Enum(DataEnum {
23        variants,
24        ..
25    }) = &mut left.data else {
26        return syn::Error::new(left.ident.span(), "only enums can accept variants")
27            .to_compile_error()
28            .into();
29    };
30
31    // parse the right enum
32    let right: DeriveInput = parse_macro_input!(right);
33    let Enum(DataEnum {
34        variants: to_add,
35        ..
36    }) = right.data else {
37        return syn::Error::new(left.ident.span(), "only enums can provide variants")
38            .to_compile_error()
39            .into();
40    };
41
42    // insert variants from the right to the left
43    variants.extend(to_add.into_iter());
44
45    quote! { #left }.into()
46}
47
48/// Append basic-whitelist-related query message variant(s) to an enum.
49///
50/// For example, apply the `sg_basic_whitelist_query` macro to the following enum:
51///
52/// ```rust
53/// #[sg_basic_whitelist_query]
54/// #[cw_serde]
55/// #[derive(QueryResponses)]
56/// enum QueryMsg {
57///     #[returns(FooResponse)]
58///     Foo {},
59///     #[returns(BarResponse)]
60///     Bar {},
61/// }
62/// ```
63///
64/// Is equivalent to:
65///
66/// ```rust
67/// #[cw_serde]
68/// #[derive(QueryResponses)]
69/// enum QueryMsg {
70///     #[returns(bool)]
71///     IncludesAddress { address: String },
72///     #[returns(FooResponse)]
73///     Foo {},
74///     #[returns(BarResponse)]
75///     Bar {},
76/// }
77/// ```
78///
79/// Note: `#[sg_basic_whitelist_query]` must be applied _before_ `#[cw_serde]`.
80#[proc_macro_attribute]
81pub fn sg_basic_whitelist_query(metadata: TokenStream, input: TokenStream) -> TokenStream {
82    merge_variants(
83        metadata,
84        input,
85        quote! {
86            enum Right {
87                /// Query the contract's basic whitelist info
88                #[returns(bool)]
89                IncludesAddress { address: String },
90            }
91        }
92        .into(),
93    )
94}