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}