feat(sol-macro): expand getter functions' return types

crates/sol-macro/src/expand/mod.rs

@@ -23,13 +23,13 @@ mod event;
 mod function;
 mod r#struct;
 mod udt;
+mod var_def;
 
 /// The limit for the number of times to resolve a type.
 const RESOLVE_LIMIT: usize = 8;
 
 /// The [`sol!`][crate::sol!] expansion implementation.
-pub fn expand(mut ast: File) -> Result<TokenStream> {
-    ast::VisitMut::visit_file(&mut MutateAst, &mut ast);
+pub fn expand(ast: File) -> Result<TokenStream> {
     ExpCtxt::new(&ast).expand()
 }
 
@@ -94,10 +94,8 @@ 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),
-            // public variables have their own getter function
-            Item::Variable(_) | Item::Import(_) | Item::Pragma(_) | Item::Using(_) => {
-                Ok(TokenStream::new())
-            }
+            Item::Variable(var_def) => var_def::expand(self, var_def),
+            Item::Import(_) | Item::Pragma(_) | Item::Using(_) => Ok(TokenStream::new()),
         }
     }
 }
@@ -248,47 +246,6 @@ 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(_))
-                    ) {
-                        // TODO: Move this to `expand_item` as we need the context to expand struct
-                        // return types
-                        functions
-                            .push((i + 1, ItemFunction::from_variable_definition(var.clone())));
-                    }
-                }
-                _ => {}
-            }
-        }
-        for (i, function) in functions.into_iter().rev() {
-            items.insert(i, Item::Function(function));
-        }
-    }
-}
-
 // utils
 impl ExpCtxt<'_> {
     #[allow(dead_code)]

crates/sol-macro/src/expand/ty.rs

@@ -112,13 +112,13 @@ fn rec_expand_type(ty: &Type, tokens: &mut TokenStream) {
                 }
             }
         }
-        Type::Function(ref function) => {
-            let span = function.span();
-            quote_spanned! {span=>
-                ::alloy_sol_types::sol_data::Function
-            }
-        }
-        Type::Mapping(ref mapping) => panic!("mapping types are unsupported: {mapping:?}"),
+        Type::Function(ref function) => quote_spanned! {function.span()=>
+            ::alloy_sol_types::sol_data::Function
+        },
+        Type::Mapping(ref mapping) => quote_spanned! {mapping.span()=>
+            ::core::compile_error!("Mapping types are not supported here")
+        },
+
         Type::Custom(ref custom) => return custom.to_tokens(tokens),
     };
     tokens.extend(tts);

crates/sol-macro/src/expand/var_def.rs

@@ -0,0 +1,105 @@
+//! State variable ([`VariableDefinition`]) expansion.
+
+use super::ExpCtxt;
+use ast::{ItemFunction, ParameterList, Spanned, Type, VariableDeclaration, VariableDefinition};
+use proc_macro2::TokenStream;
+use syn::{Error, Result};
+
+/// Expands a [`VariableDefinition`].
+///
+/// See [`ItemFunction::from_variable_definition`].
+pub(super) fn expand(cx: &ExpCtxt<'_>, var_def: &VariableDefinition) -> Result<TokenStream> {
+    // only expand public or external state variables
+    if !var_def
+        .attributes
+        .visibility()
+        .map_or(false, |v| v.is_public() || v.is_external())
+    {
+        return Ok(TokenStream::new())
+    }
+
+    let mut function = ItemFunction::from_variable_definition(var_def.clone());
+    expand_returns(cx, &mut function)?;
+    super::function::expand(cx, &function)
+}
+
+/// Expands return-position custom types.
+fn expand_returns(cx: &ExpCtxt<'_>, f: &mut ItemFunction) -> Result<()> {
+    let returns = f
+        .returns
+        .as_mut()
+        .expect("generated getter function with no returns");
+    let ret = returns.returns.first_mut().unwrap();
+    if !ret.ty.has_custom_simple() {
+        return Ok(())
+    }
+
+    let mut ty = &ret.ty;
+
+    // resolve if custom
+    if let Type::Custom(name) = ty {
+        ty = cx.custom_type(name);
+    }
+
+    // skip if not tuple with complex types
+    let Type::Tuple(tup) = ty else { return Ok(()) };
+    if !tup.types.iter().any(type_is_complex) {
+        return Ok(())
+    }
+
+    // retain only non-complex types
+    // TODO: assign return types' names from the original struct
+    let mut new_returns = ParameterList::new();
+    for p in tup.types.pairs() {
+        let (ty, comma) = p.into_tuple();
+        if !type_is_complex(ty) {
+            new_returns.push_value(VariableDeclaration::new(ty.clone()));
+            if let Some(comma) = comma {
+                new_returns.push_punct(*comma);
+            }
+        }
+    }
+
+    // all types were complex, Solidity doesn't accept this
+    if new_returns.is_empty() {
+        return Err(Error::new(f.name().span(), "invalid state variable type"))
+    }
+
+    returns.returns = new_returns;
+    Ok(())
+}
+
+/// Returns `true` if a type is complex for the purposes of state variable
+/// getters.
+///
+/// Technically tuples are also complex if they contain complex types, but only
+/// at the first "depth" level.
+///
+/// Here, `mapA` is fine but `mapB` throws an error; you can test that pushing
+/// and reading to `mapA` works fine (last checked at Solc version `0.8.21`):
+///
+/// ```solidity
+/// contract Complex {
+///     struct A {
+///         B b;
+///     }
+///     struct B {
+///         uint[] arr;
+///     }
+///
+///     mapping(uint => A) public mapA;
+///
+///     function pushValueA(uint idx, uint val) public {
+///         mapA[idx].b.arr.push(val);
+///     }
+///
+///     mapping(uint => B) public mapB;
+///
+///     function pushValueB(uint idx, uint val) public {
+///         mapB[idx].arr.push(val);
+///     }
+/// }
+/// ```
+fn type_is_complex(ty: &Type) -> bool {
+    matches!(ty, Type::Mapping(_) | Type::Array(_))
+}

crates/sol-types/tests/sol.rs

@@ -182,6 +182,44 @@ fn function_names() {
     assert!(!call.is___f());
 }
 
+#[test]
+fn getters() {
+    // modified from https://docs.soliditylang.org/en/latest/contracts.html#getter-functions
+    sol! {
+        struct Data {
+            uint a;
+            bytes3 b;
+            uint[3] c;
+            uint[] d;
+            bytes e;
+        }
+        mapping(uint => mapping(bool => Data[])) public data1;
+        mapping(uint => mapping(bool => Data)) public data2;
+
+        mapping(bool => mapping(address => uint256[])[])[][] public nestedMapArray;
+    }
+
+    assert_eq!(data1Call::SIGNATURE, "data1(uint256,bool,uint256)");
+    let _ = data1Return {
+        _0: U256::ZERO,
+        _1: [0, 0, 0],
+        _2: vec![],
+    };
+
+    assert_eq!(data2Call::SIGNATURE, "data2(uint256,bool)");
+    let _ = data2Return {
+        _0: U256::ZERO,
+        _1: [0, 0, 0],
+        _2: vec![],
+    };
+
+    assert_eq!(
+        nestedMapArrayCall::SIGNATURE,
+        "nestedMapArray(uint256,uint256,bool,uint256,address,uint256)"
+    );
+    let _ = nestedMapArrayReturn { _0: U256::ZERO };
+}
+
 #[test]
 fn abigen_sol_multicall() {
     sol!("../syn-solidity/tests/contracts/Multicall.sol");

crates/sol-types/tests/ui/type.rs

@@ -740,6 +740,14 @@ sol! {
     }
 }
 
+sol! {
+    mapping(bool => mapping(address => uint256[])[])[][] public nestedMapArray;
+}
+
+sol! {
+    mapping(mapping(int => int) => int) public mapKeyOfMap;
+}
+
 sol! {
     function mappings(mapping(uint256 a => bool b), mapping(bool => bool) x);
 }

crates/sol-types/tests/ui/type.stderr

@@ -64,27 +64,29 @@ error: enum has too many variants
 476 |     enum TooBigEnum {
     |          ^^^^^^^^^^
 
-error: proc macro panicked
-   --> tests/ui/type.rs:737:1
+error: Mapping types are not supported here
+   --> tests/ui/type.rs:739:9
     |
-737 | / sol! {
-738 | |     struct Mappings {
-739 | |         mapping(mapping(a b => c d) e => mapping(f g => h i) j) map;
-740 | |     }
-741 | | }
-    | |_^
+739 |         mapping(mapping(a b => c d) e => mapping(f g => h i) j) map;
+    |         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+error: Mapping types are not supported here
+   --> tests/ui/type.rs:748:13
     |
-    = help: message: mapping types are unsupported: TypeMapping { key: Type::TypeMapping { key: Type::Custom([SolIdent("a")]), key_name: Some(SolIdent("b")), value: Type::Custom([SolIdent("c")]), value_name: Some(SolIdent("d")) }, key_name: Some(SolIdent("e")), value: Type::TypeMapping { key: Type::Custom([SolIdent("f")]), key_name: Some(SolIdent("g")), value: Type::Custom([SolIdent("h")]), value_name: Some(SolIdent("i")) }, value_name: Some(SolIdent("j")) }
+748 |     mapping(mapping(int => int) => int) public mapKeyOfMap;
+    |             ^^^^^^^^^^^^^^^^^^^
 
-error: proc macro panicked
-   --> tests/ui/type.rs:743:1
+error: Mapping types are not supported here
+   --> tests/ui/type.rs:752:23
     |
-743 | / sol! {
-744 | |     function mappings(mapping(uint256 a => bool b), mapping(bool => bool) x);
-745 | | }
-    | |_^
+752 |     function mappings(mapping(uint256 a => bool b), mapping(bool => bool) x);
+    |                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+error: Mapping types are not supported here
+   --> tests/ui/type.rs:752:53
     |
-    = help: message: mapping types are unsupported: TypeMapping { key: Type::Uint(Some(256)), key_name: Some(SolIdent("a")), value: Type::Bool, value_name: Some(SolIdent("b")) }
+752 |     function mappings(mapping(uint256 a => bool b), mapping(bool => bool) x);
+    |                                                     ^^^^^^^^^^^^^^^^^^^^^
 
 error[E0412]: cannot find type `bytes_` in this scope
    --> tests/ui/type.rs:205:9
@@ -122,6 +124,12 @@ error[E0412]: cannot find type `int_256` in this scope
 210 |         int_256 f;
     |         ^^^^^^^ not found in this scope
 
+error[E0412]: cannot find type `a` in this scope
+   --> tests/ui/type.rs:739:25
+    |
+739 |         mapping(mapping(a b => c d) e => mapping(f g => h i) j) map;
+    |                         ^ not found in this scope
+
 error[E0277]: the trait bound `(Address, Address, alloy_sol_types::sol_data::String, Bool, alloy_sol_types::sol_data::Bytes, alloy_sol_types::sol_data::FixedBytes<1>, alloy_sol_types::sol_data::FixedBytes<2>, alloy_sol_types::sol_data::FixedBytes<3>, alloy_sol_types::sol_data::FixedBytes<4>, alloy_sol_types::sol_data::FixedBytes<5>, alloy_sol_types::sol_data::FixedBytes<6>, alloy_sol_types::sol_data::FixedBytes<7>, alloy_sol_types::sol_data::FixedBytes<8>, alloy_sol_types::sol_data::FixedBytes<9>, alloy_sol_types::sol_data::FixedBytes<10>, alloy_sol_types::sol_data::FixedBytes<11>, alloy_sol_types::sol_data::FixedBytes<12>, alloy_sol_types::sol_data::FixedBytes<13>, alloy_sol_types::sol_data::FixedBytes<14>, alloy_sol_types::sol_data::FixedBytes<15>, alloy_sol_types::sol_data::FixedBytes<16>, alloy_sol_types::sol_data::FixedBytes<17>, alloy_sol_types::sol_data::FixedBytes<18>, alloy_sol_types::sol_data::FixedBytes<19>, alloy_sol_types::sol_data::FixedBytes<20>, alloy_sol_types::sol_data::FixedBytes<21>, alloy_sol_types::sol_data::FixedBytes<22>, alloy_sol_types::sol_data::FixedBytes<23>, alloy_sol_types::sol_data::FixedBytes<24>, alloy_sol_types::sol_data::FixedBytes<25>, alloy_sol_types::sol_data::FixedBytes<26>, alloy_sol_types::sol_data::FixedBytes<27>, alloy_sol_types::sol_data::FixedBytes<28>, alloy_sol_types::sol_data::FixedBytes<29>, alloy_sol_types::sol_data::FixedBytes<30>, alloy_sol_types::sol_data::FixedBytes<31>, alloy_sol_types::sol_data::FixedBytes<32>, Int<256>, Int<8>, Int<16>, Int<24>, Int<32>, Int<40>, Int<48>, Int<56>, Int<64>, Int<72>, Int<80>, Int<88>, Int<96>, Int<104>, Int<112>, Int<120>, Int<128>, Int<136>, Int<144>, Int<152>, Int<160>, Int<168>, Int<176>, Int<184>, Int<192>, Int<200>, Int<208>, Int<216>, Int<224>, Int<232>, Int<240>, Int<248>, Int<256>, Uint<256>, Uint<8>, Uint<16>, Uint<24>, Uint<32>, Uint<40>, Uint<48>, Uint<56>, Uint<64>, Uint<72>, Uint<80>, Uint<88>, Uint<96>, Uint<104>, Uint<112>, Uint<120>, Uint<128>, Uint<136>, Uint<144>, Uint<152>, Uint<160>, Uint<168>, Uint<176>, Uint<184>, Uint<192>, Uint<200>, Uint<208>, Uint<216>, Uint<224>, Uint<232>, Uint<240>, Uint<248>, Uint<256>): SolStruct` is not satisfied
    --> tests/ui/type.rs:3:1
     |
@@ -139,6 +147,7 @@ error[E0277]: the trait bound `(Address, Address, alloy_sol_types::sol_data::Str
               ArrayTypes
               TupleTypes
               CustomTypes
+              Mappings
               FunctionTypes
     = note: required for `(Address, ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ...)` to implement `SolType`
 note: required by a bound in `Encodable`
@@ -165,6 +174,7 @@ error[E0277]: the trait bound `(Address, Address, alloy_sol_types::sol_data::Str
               ArrayTypes
               TupleTypes
               CustomTypes
+              Mappings
               FunctionTypes
     = note: required for `(Address, ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ..., ...)` to implement `SolType`
     = note: this error originates in the macro `sol` (in Nightly builds, run with -Z macro-backtrace for more info)

crates/sol-types/tests/ui/var_def.rs

@@ -0,0 +1,33 @@
+use alloy_sol_types::sol;
+
+// OK
+sol! {
+    struct Simple {
+        uint a;
+    }
+
+    mapping(int => Simple) public simpleMap;
+}
+
+// Not OK
+sol! {
+    struct Complex1 {
+        uint[] a;
+    }
+
+    mapping(int => Complex1) public complexMap;
+}
+
+// OK
+sol! {
+    struct DoubleComplex {
+        Complex2 a;
+    }
+    struct Complex2 {
+        uint[] a;
+    }
+
+    mapping(int => DoubleComplex) public complexMap;
+}
+
+fn main() {}

crates/sol-types/tests/ui/var_def.stderr

@@ -0,0 +1,5 @@
+error: invalid state variable type
+  --> tests/ui/var_def.rs:18:37
+   |
+18 |     mapping(int => Complex1) public complexMap;
+   |                                     ^^^^^^^^^^