feat: mock multidimensional array (#74)

defi-wonderland · Apr 29, 2024 · 2476e3e · 2476e3e
1 parent 9467f24
commit 2476e3e
Show file tree

Hide file tree

Showing 7 changed files with 324 additions and 29 deletions.
diff --git a/README.md b/README.md
@@ -113,6 +113,7 @@ greeter.set__greeting('Hola');
 - Please, note that if you want to mock `internal` functions, you **must** make them `virtual`. The tool will not generate mocks for internal functions that are not virtual.
 - Cannot set `private` variables and mock `private` functions.
 - Mocking of structs containing mappings is not supported.
+- Mocking of multi-dimensional arrays of structs is not supported.
 
 # Licensing
 

diff --git a/solidity/contracts/utils/ContractG.sol b/solidity/contracts/utils/ContractG.sol
@@ -53,7 +53,17 @@ contract ContractG {
 
   NestedStruct public nestedStruct;
 
-  // CommonStruct[][] public multidimensionalStruct;
+  CommonStruct[] public structArray;
+
+  CommonStruct[][] public twoDimensionalStruct;
+
+  CommonStruct[][][] public threeDimensionalStruct;
+
+  uint256[][] public twoDimensionalArray;
+
+  string[][] public twoDimensionalStringArray;
+
+  uint256[][][] public threeDimensionalArray;
 
   function setNestedStruct(NestedStruct memory _nestedStruct) public {
     nestedStruct = _nestedStruct;

diff --git a/solidity/test/utils/ContractG.t.sol b/solidity/test/utils/ContractG.t.sol
@@ -0,0 +1,230 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.0;
+
+import {Test} from 'forge-std/Test.sol';
+import {MockContractG} from 'test/smock/contracts/utils/MockContractG.sol';
+import {SmockHelper} from 'test/smock/SmockHelper.sol';
+import {ContractG} from 'contracts/utils/ContractG.sol';
+
+contract UnitMockContractG is Test, SmockHelper {
+  address internal _owner = makeAddr('owner');
+  MockContractG internal _contractTest;
+
+  function setUp() public {
+    vm.prank(_owner);
+
+    _contractTest = MockContractG(deployMock('TestContractG', type(MockContractG).creationCode, abi.encode()));
+  }
+
+  function test_Set_MyStructArray() public {
+    ContractG.CommonStruct[] memory _myStructArray = new ContractG.CommonStruct[](2);
+    _myStructArray[0] = ContractG.CommonStruct(10);
+    _myStructArray[1] = ContractG.CommonStruct(20);
+
+    _contractTest.set_structArray(_myStructArray);
+    (uint256 _value) = _contractTest.structArray(0);
+    assertEq(_value, 10);
+
+    (_value) = _contractTest.structArray(1);
+    assertEq(_value, 20);
+  }
+
+  function test_Set_TwoDimensionalArray() public {
+    uint256[][] memory _twoDimensionalStruct = new uint256[][](2);
+    _twoDimensionalStruct[0] = new uint256[](2);
+    _twoDimensionalStruct[1] = new uint256[](2);
+
+    _twoDimensionalStruct[0][0] = 10;
+    _twoDimensionalStruct[0][1] = 20;
+    _twoDimensionalStruct[1][0] = 30;
+    _twoDimensionalStruct[1][1] = 40;
+
+    _contractTest.set_twoDimensionalArray(_twoDimensionalStruct);
+    (uint256 _value) = _contractTest.twoDimensionalArray(0, 0);
+    assertEq(_value, 10);
+
+    (_value) = _contractTest.twoDimensionalArray(0, 1);
+    assertEq(_value, 20);
+
+    (_value) = _contractTest.twoDimensionalArray(1, 0);
+    assertEq(_value, 30);
+
+    (_value) = _contractTest.twoDimensionalArray(1, 1);
+    assertEq(_value, 40);
+  }
+
+  function test_Call_TwoDimensionalArray() public {
+    uint256[][] memory _twoDimensionalStruct = new uint256[][](2);
+    _twoDimensionalStruct[0] = new uint256[](2);
+    _twoDimensionalStruct[1] = new uint256[](2);
+
+    _twoDimensionalStruct[0][0] = 10;
+    _twoDimensionalStruct[0][1] = 20;
+    _twoDimensionalStruct[1][0] = 30;
+    _twoDimensionalStruct[1][1] = 40;
+
+    _contractTest.mock_call_twoDimensionalArray(0, 0, _twoDimensionalStruct[0][0]);
+    (uint256 _value) = _contractTest.twoDimensionalArray(0, 0);
+    assertEq(_value, 10);
+
+    _contractTest.mock_call_twoDimensionalArray(0, 1, _twoDimensionalStruct[0][1]);
+    (_value) = _contractTest.twoDimensionalArray(0, 1);
+    assertEq(_value, 20);
+
+    _contractTest.mock_call_twoDimensionalArray(1, 0, _twoDimensionalStruct[1][0]);
+    (_value) = _contractTest.twoDimensionalArray(1, 0);
+    assertEq(_value, 30);
+
+    _contractTest.mock_call_twoDimensionalArray(1, 1, _twoDimensionalStruct[1][1]);
+    (_value) = _contractTest.twoDimensionalArray(1, 1);
+    assertEq(_value, 40);
+  }
+
+  function test_Set_ThreeDimensionalArray() public {
+    uint256[][][] memory _threeDimensionalStruct = new uint256[][][](2);
+    _threeDimensionalStruct[0] = new uint256[][](2);
+    _threeDimensionalStruct[1] = new uint256[][](2);
+
+    _threeDimensionalStruct[0][0] = new uint256[](2);
+    _threeDimensionalStruct[0][1] = new uint256[](2);
+    _threeDimensionalStruct[1][0] = new uint256[](2);
+    _threeDimensionalStruct[1][1] = new uint256[](2);
+
+    _threeDimensionalStruct[0][0][0] = 10;
+    _threeDimensionalStruct[0][0][1] = 20;
+    _threeDimensionalStruct[0][1][0] = 30;
+    _threeDimensionalStruct[0][1][1] = 40;
+    _threeDimensionalStruct[1][0][0] = 50;
+    _threeDimensionalStruct[1][0][1] = 60;
+    _threeDimensionalStruct[1][1][0] = 70;
+    _threeDimensionalStruct[1][1][1] = 80;
+
+    _contractTest.set_threeDimensionalArray(_threeDimensionalStruct);
+    (uint256 _value) = _contractTest.threeDimensionalArray(0, 0, 0);
+    assertEq(_value, 10);
+
+    (_value) = _contractTest.threeDimensionalArray(0, 0, 1);
+    assertEq(_value, 20);
+
+    (_value) = _contractTest.threeDimensionalArray(0, 1, 0);
+    assertEq(_value, 30);
+
+    (_value) = _contractTest.threeDimensionalArray(0, 1, 1);
+    assertEq(_value, 40);
+
+    (_value) = _contractTest.threeDimensionalArray(1, 0, 0);
+    assertEq(_value, 50);
+
+    (_value) = _contractTest.threeDimensionalArray(1, 0, 1);
+    assertEq(_value, 60);
+
+    (_value) = _contractTest.threeDimensionalArray(1, 1, 0);
+    assertEq(_value, 70);
+
+    (_value) = _contractTest.threeDimensionalArray(1, 1, 1);
+    assertEq(_value, 80);
+  }
+
+  function test_Call_ThreeDimensionalArray() public {
+    uint256[][][] memory _threeDimensionalStruct = new uint256[][][](2);
+    _threeDimensionalStruct[0] = new uint256[][](2);
+    _threeDimensionalStruct[1] = new uint256[][](2);
+
+    _threeDimensionalStruct[0][0] = new uint256[](2);
+    _threeDimensionalStruct[0][1] = new uint256[](2);
+    _threeDimensionalStruct[1][0] = new uint256[](2);
+    _threeDimensionalStruct[1][1] = new uint256[](2);
+
+    _threeDimensionalStruct[0][0][0] = 10;
+    _threeDimensionalStruct[0][0][1] = 20;
+    _threeDimensionalStruct[0][1][0] = 30;
+    _threeDimensionalStruct[0][1][1] = 40;
+    _threeDimensionalStruct[1][0][0] = 50;
+    _threeDimensionalStruct[1][0][1] = 60;
+    _threeDimensionalStruct[1][1][0] = 70;
+    _threeDimensionalStruct[1][1][1] = 80;
+
+    _contractTest.mock_call_threeDimensionalArray(0, 0, 0, _threeDimensionalStruct[0][0][0]);
+    (uint256 _value) = _contractTest.threeDimensionalArray(0, 0, 0);
+    assertEq(_value, 10);
+
+    _contractTest.mock_call_threeDimensionalArray(0, 0, 1, _threeDimensionalStruct[0][0][1]);
+    (_value) = _contractTest.threeDimensionalArray(0, 0, 1);
+    assertEq(_value, 20);
+
+    _contractTest.mock_call_threeDimensionalArray(0, 1, 0, _threeDimensionalStruct[0][1][0]);
+    (_value) = _contractTest.threeDimensionalArray(0, 1, 0);
+    assertEq(_value, 30);
+
+    _contractTest.mock_call_threeDimensionalArray(0, 1, 1, _threeDimensionalStruct[0][1][1]);
+    (_value) = _contractTest.threeDimensionalArray(0, 1, 1);
+    assertEq(_value, 40);
+
+    _contractTest.mock_call_threeDimensionalArray(1, 0, 0, _threeDimensionalStruct[1][0][0]);
+    (_value) = _contractTest.threeDimensionalArray(1, 0, 0);
+    assertEq(_value, 50);
+
+    _contractTest.mock_call_threeDimensionalArray(1, 0, 1, _threeDimensionalStruct[1][0][1]);
+    (_value) = _contractTest.threeDimensionalArray(1, 0, 1);
+    assertEq(_value, 60);
+
+    _contractTest.mock_call_threeDimensionalArray(1, 1, 0, _threeDimensionalStruct[1][1][0]);
+    (_value) = _contractTest.threeDimensionalArray(1, 1, 0);
+    assertEq(_value, 70);
+
+    _contractTest.mock_call_threeDimensionalArray(1, 1, 1, _threeDimensionalStruct[1][1][1]);
+    (_value) = _contractTest.threeDimensionalArray(1, 1, 1);
+    assertEq(_value, 80);
+  }
+
+  function test_Set_TwoDimensionalStringArray() public {
+    string[][] memory _twoDimensionalStringArray = new string[][](2);
+    _twoDimensionalStringArray[0] = new string[](2);
+    _twoDimensionalStringArray[1] = new string[](2);
+
+    _twoDimensionalStringArray[0][0] = '10';
+    _twoDimensionalStringArray[0][1] = '20';
+    _twoDimensionalStringArray[1][0] = '30';
+    _twoDimensionalStringArray[1][1] = '40';
+
+    _contractTest.set_twoDimensionalStringArray(_twoDimensionalStringArray);
+    (string memory _value) = _contractTest.twoDimensionalStringArray(0, 0);
+    assertEq(_value, '10');
+
+    (_value) = _contractTest.twoDimensionalStringArray(0, 1);
+    assertEq(_value, '20');
+
+    (_value) = _contractTest.twoDimensionalStringArray(1, 0);
+    assertEq(_value, '30');
+
+    (_value) = _contractTest.twoDimensionalStringArray(1, 1);
+    assertEq(_value, '40');
+  }
+
+  function test_Call_TwoDimensionalStringArray() public {
+    string[][] memory _twoDimensionalStringArray = new string[][](2);
+    _twoDimensionalStringArray[0] = new string[](2);
+    _twoDimensionalStringArray[1] = new string[](2);
+
+    _twoDimensionalStringArray[0][0] = '10';
+    _twoDimensionalStringArray[0][1] = '20';
+    _twoDimensionalStringArray[1][0] = '30';
+    _twoDimensionalStringArray[1][1] = '40';
+
+    _contractTest.mock_call_twoDimensionalStringArray(0, 0, _twoDimensionalStringArray[0][0]);
+    (string memory _value) = _contractTest.twoDimensionalStringArray(0, 0);
+    assertEq(_value, '10');
+
+    _contractTest.mock_call_twoDimensionalStringArray(0, 1, _twoDimensionalStringArray[0][1]);
+    (_value) = _contractTest.twoDimensionalStringArray(0, 1);
+    assertEq(_value, '20');
+
+    _contractTest.mock_call_twoDimensionalStringArray(1, 0, _twoDimensionalStringArray[1][0]);
+    (_value) = _contractTest.twoDimensionalStringArray(1, 0);
+    assertEq(_value, '30');
+
+    _contractTest.mock_call_twoDimensionalStringArray(1, 1, _twoDimensionalStringArray[1][1]);
+    (_value) = _contractTest.twoDimensionalStringArray(1, 1);
+    assertEq(_value, '40');
+  }
+}
diff --git a/src/context.ts b/src/context.ts
@@ -205,12 +205,19 @@ export function arrayVariableContext(node: VariableDeclaration): ArrayVariableCo
   // Array type
   const arrayType: string = sanitizeParameterType(explicitTypeStorageLocation(node.typeString));
 
-  // Base type
-  const baseType: string = sanitizeParameterType(explicitTypeStorageLocation(node.vType['vBaseType'].typeString));
-
   // Struct flag
   const isStructArray: boolean = node.typeString.startsWith('struct ');
 
+  // Check if the array is multi-dimensional
+  const dimensionsQuantity = node.typeString.split('[]').length - 1;
+  const isMultiDimensional = dimensionsQuantity > 1;
+  const isMultiDimensionalStruct = isMultiDimensional && isStructArray;
+  const dimensions = [...Array(dimensionsQuantity).keys()];
+
+  // Base type
+  const baseTypeString = isMultiDimensional ? node.typeString.replace(/\[\]/g, '') : node.vType['vBaseType'].typeString;
+  const baseType = sanitizeParameterType(explicitTypeStorageLocation(baseTypeString));
+
   // Check if the variable is a struct and get its fields
   const structFields = extractStructFieldsNames(node.vType);
 
@@ -231,6 +238,8 @@ export function arrayVariableContext(node: VariableDeclaration): ArrayVariableCo
     },
     isInternal: isInternal,
     isStructArray: isStructArray,
+    isMultiDimensionalStruct,
+    dimensions,
   };
 }
 

diff --git a/src/templates/partials/array-state-variable.hbs b/src/templates/partials/array-state-variable.hbs
@@ -1,27 +1,31 @@
-function set_{{setFunction.functionName}}({{setFunction.arrayType}} _{{setFunction.paramName}}) public {
-  {{#if isStructArray}}
-    for (uint256 _i; _i < _{{setFunction.paramName}}.length; ++_i) {
-      {{setFunction.functionName}}.push(_{{setFunction.paramName}}[_i]);
-    }
-  {{else}}
-    {{setFunction.functionName}} = _{{setFunction.paramName}};
-  {{/if}}
-}
+{{#unless isMultiDimensionalStruct}}
+  function set_{{setFunction.functionName}}({{setFunction.arrayType}} _{{setFunction.paramName}}) public {
+    {{#if isStructArray}}
+      for (uint256 _i; _i < _{{setFunction.paramName}}.length; ++_i) {
+        {{setFunction.functionName}}.push(_{{setFunction.paramName}}[_i]);
+      }
+    {{else}}
+      {{setFunction.functionName}} = _{{setFunction.paramName}};
+    {{/if}}
+  }
 
-{{#unless isInternal}}
-function mock_call_{{mockFunction.functionName}}(uint256 _index, {{mockFunction.baseType}} _value) public {
-  vm.mockCall(
-    address(this),
-    abi.encodeWithSignature('{{mockFunction.functionName}}(uint256)', _index),
-    abi.encode(
-      {{#if isStructArray}}
-        {{#each mockFunction.structFields}}
-          _value.{{this}}{{#unless @last}}, {{/unless}}
-        {{/each}}
-      {{else}}
-        _value
-      {{/if}}
-    )
-  );
-}
+  {{#unless isInternal}}
+  function mock_call_{{mockFunction.functionName}}({{#each dimensions}}uint256 _index{{@index}}, {{/each}} {{mockFunction.baseType}} _value) public {
+    vm.mockCall(
+      address(this),
+      abi.encodeWithSignature(
+        '{{mockFunction.functionName}}({{#each dimensions}}uint256{{#unless @last}},{{/unless}}{{/each}})', 
+        {{#each dimensions}}_index{{@index}}{{#unless @last}}, {{/unless}} {{/each}}),
+      abi.encode(
+        {{#if isStructArray}}
+          {{#each mockFunction.structFields}}
+            _value.{{this}}{{#unless @last}}, {{/unless}}
+          {{/each}}
+        {{else}}
+          _value
+        {{/if}}
+      )
+    );
+  }
+  {{/unless}}
 {{/unless}}
diff --git a/src/types.ts b/src/types.ts
@@ -70,6 +70,8 @@ export interface ArrayVariableContext {
   };
   isInternal: boolean;
   isStructArray: boolean;
+  isMultiDimensionalStruct: boolean;
+  dimensions: number[];
 }
 
 export interface StateVariableContext {