feat(sol-macro): expand getter functions for public state variables (#…

…218) * feat(sol-macro): expand getter functions for public state variables * chore: clippy * fixes
alloy-rs · Aug 2, 2023 · 0ff43ed · 0ff43ed
1 parent 43a14af
commit 0ff43ed
Show file tree

Hide file tree

Showing 21 changed files with 468 additions and 67 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -38,6 +38,7 @@ syn = "2.0"
 cfg-if = "1.0.0"
 derive_more = "0.99"
 hex-literal = "0.4"
+paste = "1.0"
 strum = { version = "0.25", features = ["derive"] }
 num_enum = "0.6"
 thiserror = "1.0"

diff --git a/crates/json-abi/src/abi.rs b/crates/json-abi/src/abi.rs
@@ -57,9 +57,9 @@ impl JsonAbi {
             constructor: self.constructor.as_ref(),
             fallback: self.fallback.as_ref(),
             receive: self.receive.as_ref(),
-            functions: self.functions.values().flatten(),
-            events: self.events.values().flatten(),
-            errors: self.errors.values().flatten(),
+            functions: self.functions(),
+            events: self.events(),
+            errors: self.errors(),
         }
     }
 

diff --git a/crates/sol-macro/src/expand/contract.rs b/crates/sol-macro/src/expand/contract.rs
@@ -51,7 +51,7 @@ pub(super) fn expand(cx: &ExpCtxt<'_>, contract: &ItemContract) -> Result<TokenS
     let d_attrs: Vec<Attribute> = attr::derives(&attrs).cloned().collect();
     for item in body {
         match item {
-            Item::Function(function) => functions.push(function),
+            Item::Function(function) if function.name.is_some() => functions.push(function),
             Item::Error(error) => errors.push(error),
             Item::Event(event) => events.push(event),
             _ => {}

diff --git a/crates/sol-macro/src/expand/mod.rs b/crates/sol-macro/src/expand/mod.rs
@@ -28,7 +28,8 @@ mod udt;
 const RESOLVE_LIMIT: usize = 8;
 
 /// The [`sol!`][crate::sol!] expansion implementation.
-pub fn expand(ast: File) -> Result<TokenStream> {
+pub fn expand(mut ast: File) -> Result<TokenStream> {
+    ast::VisitMut::visit_file(&mut MutateAst, &mut ast);
     ExpCtxt::new(&ast).expand()
 }
 
@@ -93,11 +94,10 @@ impl<'ast> ExpCtxt<'ast> {
             Item::Function(function) => function::expand(self, function),
             Item::Struct(strukt) => r#struct::expand(self, strukt),
             Item::Udt(udt) => udt::expand(self, udt),
-            Item::Variable(_) => {
-                // TODO: Expand getter function for public variables
+            // public variables have their own getter function
+            Item::Variable(_) | Item::Import(_) | Item::Pragma(_) | Item::Using(_) => {
                 Ok(TokenStream::new())
             }
-            Item::Import(_) | Item::Pragma(_) | Item::Using(_) => Ok(TokenStream::new()),
         }
     }
 }
@@ -200,7 +200,9 @@ impl ExpCtxt<'_> {
             }
 
             for (i, &function) in functions.iter().enumerate() {
-                let old_name = function.name();
+                let Some(old_name) = function.name.as_ref() else {
+                    continue
+                };
                 let new_name = format!("{old_name}_{i}");
                 if let Some(other) = all_orig_names.iter().find(|x| x.0 == new_name) {
                     let msg = format!(
@@ -246,8 +248,47 @@ impl<'ast> Visit<'ast> for ExpCtxt<'ast> {
     }
 }
 
+struct MutateAst;
+
+impl<'ast> ast::VisitMut<'ast> for MutateAst {
+    fn visit_file(&mut self, file: &'ast mut File) {
+        Self::visit_items(&mut file.items);
+        ast::visit_mut::visit_file(self, file);
+    }
+
+    fn visit_item_contract(&mut self, contract: &'ast mut ast::ItemContract) {
+        Self::visit_items(&mut contract.body);
+        ast::visit_mut::visit_item_contract(self, contract);
+    }
+}
+
+impl MutateAst {
+    #[allow(clippy::single_match)]
+    fn visit_items(items: &mut Vec<Item>) {
+        // add a getter function for each public variable
+        let mut functions = Vec::new();
+        for (i, item) in items.iter().enumerate() {
+            match item {
+                Item::Variable(var) => {
+                    if matches!(
+                        var.attributes.visibility(),
+                        Some(ast::Visibility::Public(_) | ast::Visibility::External(_))
+                    ) {
+                        functions.push((i + 1, ItemFunction::from_variable_definition(var)))
+                    }
+                }
+                _ => {}
+            }
+        }
+        for (i, function) in functions.into_iter().rev() {
+            items.insert(i, Item::Function(function));
+        }
+    }
+}
+
 // utils
 impl ExpCtxt<'_> {
+    #[allow(dead_code)]
     fn get_item(&self, name: &SolPath) -> &Item {
         match self.try_get_item(name) {
             Some(item) => item,

diff --git a/crates/sol-macro/src/expand/struct.rs b/crates/sol-macro/src/expand/struct.rs
@@ -40,12 +40,12 @@ pub(super) fn expand(cx: &ExpCtxt<'_>, s: &ItemStruct) -> Result<TokenStream> {
         .map(|f| (expand_type(&f.ty), f.name.as_ref().unwrap()))
         .unzip();
 
-    let eip712_encode_type_fns: TokenStream = expand_encode_type_fns(fields, name);
+    let eip712_encode_type_fns = expand_encode_type_fns(fields, name);
 
     let tokenize_impl = expand_tokenize_func(fields.iter());
 
     let encode_data_impl = match fields.len() {
-        0 => unreachable!(),
+        0 => unreachable!("struct with zero fields"),
         1 => {
             let VariableDeclaration { ty, name, .. } = fields.first().unwrap();
             let ty = expand_type(ty);
@@ -135,14 +135,11 @@ pub(super) fn expand(cx: &ExpCtxt<'_>, s: &ItemStruct) -> Result<TokenStream> {
     Ok(tokens)
 }
 
-fn expand_encode_type_fns(
-    fields: &ast::Parameters<syn::token::Semi>,
-    name: &ast::SolIdent,
-) -> TokenStream {
+fn expand_encode_type_fns(fields: &ast::FieldList, name: &ast::SolIdent) -> TokenStream {
     let components_impl = expand_eip712_components(fields);
     let root_type_impl = fields.eip712_signature(name.as_string());
 
-    let encode_type_impl_opt: Option<TokenStream> = if fields.iter().any(|f| f.ty.is_custom()) {
+    let encode_type_impl_opt = if fields.iter().any(|f| f.ty.is_custom()) {
         None
     } else {
         Some(quote! {
@@ -165,7 +162,7 @@ fn expand_encode_type_fns(
     }
 }
 
-fn expand_eip712_components(fields: &ast::Parameters<syn::token::Semi>) -> TokenStream {
+fn expand_eip712_components(fields: &ast::FieldList) -> TokenStream {
     let bits: Vec<TokenStream> = fields
         .iter()
         .filter(|f| f.ty.is_custom())
@@ -179,7 +176,7 @@ fn expand_eip712_components(fields: &ast::Parameters<syn::token::Semi>) -> Token
         .collect();
 
     if bits.is_empty() {
-        quote! { ::alloy_sol_types::private::Vec::with_capacity(0) }
+        quote! { ::alloy_sol_types::private::Vec::new() }
     } else {
         quote! {
             let mut components = ::alloy_sol_types::private::Vec::new();

diff --git a/crates/sol-macro/src/expand/ty.rs b/crates/sol-macro/src/expand/ty.rs
@@ -112,8 +112,8 @@ fn rec_expand_type(ty: &Type, tokens: &mut TokenStream) {
                 }
             }
         }
-        Type::Function(ref _function) => todo!(),
-        Type::Mapping(ref _mapping) => todo!(),
+        Type::Function(ref function) => todo!("function types: {function:?}"),
+        Type::Mapping(ref mapping) => todo!("mapping types: {mapping:?}"),
         Type::Custom(ref custom) => return custom.to_tokens(tokens),
     };
     tokens.extend(tts);
@@ -161,15 +161,26 @@ pub(super) fn type_base_data_size(cx: &ExpCtxt<'_>, ty: &Type) -> usize {
             .map(|ty| type_base_data_size(cx, ty))
             .sum(),
 
-        Type::Custom(name) => match cx.get_item(name) {
-            Item::Enum(_) => 32,
-            Item::Struct(strukt) => strukt
+        Type::Custom(name) => match cx.try_get_item(name) {
+            Some(Item::Enum(_)) => 32,
+            Some(Item::Error(error)) => error
+                .parameters
+                .types()
+                .map(|ty| type_base_data_size(cx, ty))
+                .sum(),
+            Some(Item::Event(event)) => event
+                .parameters
+                .iter()
+                .map(|p| type_base_data_size(cx, &p.ty))
+                .sum(),
+            Some(Item::Struct(strukt)) => strukt
                 .fields
                 .types()
                 .map(|ty| type_base_data_size(cx, ty))
                 .sum(),
-            Item::Udt(udt) => type_base_data_size(cx, &udt.ty),
-            _ => unreachable!(),
+            Some(Item::Udt(udt)) => type_base_data_size(cx, &udt.ty),
+            Some(item) => panic!("Invalid item in param list: {item:?}"),
+            None => 0,
         },
 
         // not applicable
@@ -197,12 +208,20 @@ pub(super) fn can_derive_default(cx: &ExpCtxt<'_>, ty: &Type) -> bool {
 
         Type::Custom(name) => match cx.try_get_item(name) {
             Some(Item::Enum(_)) => false,
+            Some(Item::Error(error)) => error
+                .parameters
+                .types()
+                .all(|ty| can_derive_default(cx, ty)),
+            Some(Item::Event(event)) => event
+                .parameters
+                .iter()
+                .all(|p| can_derive_default(cx, &p.ty)),
             Some(Item::Struct(strukt)) => {
                 strukt.fields.types().all(|ty| can_derive_default(cx, ty))
             }
             Some(Item::Udt(udt)) => can_derive_default(cx, &udt.ty),
-            Some(_) => unreachable!(),
-            None => false,
+            Some(item) => panic!("Invalid item in param list: {item:?}"),
+            _ => false,
         },
 
         _ => true,
@@ -227,12 +246,21 @@ pub(super) fn can_derive_builtin_traits(cx: &ExpCtxt<'_>, ty: &Type) -> bool {
 
         Type::Custom(name) => match cx.try_get_item(name) {
             Some(Item::Enum(_)) => true,
-            Some(Item::Struct(strukt)) => {
-                strukt.fields.types().all(|ty| can_derive_default(cx, ty))
-            }
-            Some(Item::Udt(udt)) => can_derive_default(cx, &udt.ty),
-            Some(_) => unreachable!(),
-            None => false,
+            Some(Item::Error(error)) => error
+                .parameters
+                .types()
+                .all(|ty| can_derive_builtin_traits(cx, ty)),
+            Some(Item::Event(event)) => event
+                .parameters
+                .iter()
+                .all(|p| can_derive_builtin_traits(cx, &p.ty)),
+            Some(Item::Struct(strukt)) => strukt
+                .fields
+                .types()
+                .all(|ty| can_derive_builtin_traits(cx, ty)),
+            Some(Item::Udt(udt)) => can_derive_builtin_traits(cx, &udt.ty),
+            Some(item) => panic!("Invalid item in param list: {item:?}"),
+            _ => false,
         },
 
         _ => true,

diff --git a/crates/sol-macro/src/expand/udt.rs b/crates/sol-macro/src/expand/udt.rs
@@ -11,8 +11,10 @@ pub(super) fn expand(cx: &ExpCtxt<'_>, udt: &ItemUdt) -> Result<TokenStream> {
         name, ty, attrs, ..
     } = udt;
 
-    let (_sol_attrs, mut attrs) = crate::attr::SolAttrs::parse(attrs)?;
-    cx.type_derives(&mut attrs, Some(ty), true);
+    // TODO: Uncomment after migrating `define_udt!`
+    let _ = cx;
+    // let (_sol_attrs, mut attrs) = crate::attr::SolAttrs::parse(attrs)?;
+    // cx.type_derives(&mut attrs, Some(ty), true);
 
     let ty = expand_type(ty);
     let tokens = quote! {

diff --git a/crates/sol-macro/src/json.rs b/crates/sol-macro/src/json.rs
@@ -78,8 +78,8 @@ fn expand_abi(name: &Ident, abi: JsonAbi) -> Result<TokenStream> {
     let udvts = udvts.iter().map(expand_udvt);
 
     let structs = structs.iter().map(expand_struct);
-    let events = abi.events.values().flatten().map(expand_event);
-    let errors = abi.errors.values().flatten().map(expand_error);
+    let events = abi.events().map(expand_event);
+    let errors = abi.errors().map(expand_error);
 
     let constructor = abi
         .constructor
@@ -93,10 +93,9 @@ fn expand_abi(name: &Ident, abi: JsonAbi) -> Result<TokenStream> {
         .receive
         .as_ref()
         .map(|r| AbiFunction::Receive.expand(r.state_mutability));
-    let functions =
-        abi.functions.values().flatten().map(|f| {
-            AbiFunction::Function(&f.name, &f.inputs, &f.outputs).expand(f.state_mutability)
-        });
+    let functions = abi
+        .functions()
+        .map(|f| AbiFunction::Function(&f.name, &f.inputs, &f.outputs).expand(f.state_mutability));
 
     let tokens = quote! {
         interface #name {

diff --git a/crates/sol-types/tests/sol.rs b/crates/sol-types/tests/sol.rs
@@ -163,7 +163,7 @@ fn empty_call() {
 }
 
 #[test]
-fn abigen_sol() {
+fn abigen_sol_multicall() {
     sol!("../syn-solidity/tests/contracts/Multicall.sol");
 
     sol! {
@@ -275,14 +275,21 @@ fn abigen_sol() {
 
 #[test]
 #[cfg(feature = "json")]
-fn abigen_json() {
-    sol!(Contract, "../json-abi/tests/abi/LargeArray.json");
+fn abigen_json_large_array() {
+    sol!(LargeArray, "../json-abi/tests/abi/LargeArray.json");
     assert_eq!(
-        Contract::callWithLongArrayCall::SIGNATURE,
+        LargeArray::callWithLongArrayCall::SIGNATURE,
         "callWithLongArray(uint64[128])"
     );
 }
 
+// TODO
+// #[test]
+// #[cfg(feature = "json")]
+// fn abigen_json_seaport() {
+//     sol!(Seaport, "../json-abi/tests/abi/Seaport.json");
+// }
+
 #[test]
 fn eip712_encode_type_nesting() {
     sol! {

diff --git a/crates/syn-solidity/Cargo.toml b/crates/syn-solidity/Cargo.toml
@@ -18,6 +18,7 @@ all-features = true
 rustdoc-args = ["--cfg", "docsrs"]
 
 [dependencies]
+paste.workspace = true
 proc-macro2.workspace = true
 quote.workspace = true
 syn = { workspace = true, features = ["extra-traits"] }