fix: returns working with struct mappings (#103)

Closes issue #94
defi-wonderland · Dec 21, 2021 · 27fe7a5 · 27fe7a5
1 parent 9561ca5
commit 27fe7a5
Show file tree

Hide file tree

Showing 5 changed files with 76 additions and 6 deletions.
diff --git a/src/factories/smock-contract.ts b/src/factories/smock-contract.ts
@@ -10,7 +10,7 @@ import { ProgrammableFunctionLogic, SafeProgrammableContract } from '../logic/pr
 import { ObservableVM } from '../observable-vm';
 import { Sandbox } from '../sandbox';
 import { ContractCall, FakeContract, MockContractFactory, ProgrammableContractFunction, ProgrammedReturnValue } from '../types';
-import { fromFancyAddress, impersonate, toFancyAddress, toHexString } from '../utils';
+import { convertPojoToStruct, fromFancyAddress, impersonate, isPojo, toFancyAddress, toHexString } from '../utils';
 import { getStorageLayout } from '../utils/storage';
 
 export async function createFakeContract<Contract extends BaseContract>(
@@ -115,12 +115,20 @@ function getFunctionEncoder(contractInterface: ethers.utils.Interface, sighash:
     // if it is a fallback function, return simplest encoder
     return (values) => values;
   } else {
-    const fnFragment = contractInterface.getFunction(sighash);
     return (values) => {
+      const fnFragment = contractInterface.getFunction(sighash);
       try {
         return contractInterface.encodeFunctionResult(fnFragment, [values]);
       } catch {
-        return contractInterface.encodeFunctionResult(fnFragment, values);
+        try {
+          return contractInterface.encodeFunctionResult(fnFragment, values);
+        } catch (err) {
+          if (isPojo(values)) {
+            return contractInterface.encodeFunctionResult(fnFragment, convertPojoToStruct(values, fnFragment));
+          }
+
+          throw err;
+        }
       }
     };
   }

diff --git a/src/utils/serdes.ts b/src/utils/serdes.ts
@@ -1,3 +1,5 @@
+import { ethers } from 'ethers';
+
 const proto = Object.prototype;
 const gpo = Object.getPrototypeOf;
 
@@ -19,6 +21,32 @@ export function convertStructToPojo(struct: any): object {
   return obj;
 }
 
+export function convertPojoToStruct(value: Record<string, unknown>, fnFragment: ethers.utils.FunctionFragment): unknown[] {
+  const parsedValue = {
+    [fnFragment.name]: value,
+  };
+  const parsedFnFragment: Partial<ethers.utils.ParamType> = {
+    name: fnFragment.name,
+    components: fnFragment.outputs!,
+  };
+
+  return convertPojoToStructRecursive(parsedValue, [parsedFnFragment])[0];
+}
+
+export function convertPojoToStructRecursive(value: any, fnFragments: Partial<ethers.utils.ParamType>[]): unknown[][] {
+  let res: unknown[][] = [];
+
+  fnFragments.forEach((item) => {
+    if (item.components) {
+      res.push(convertPojoToStructRecursive(value[item.name!], item.components));
+    } else {
+      res.push(value[item.name!]);
+    }
+  });
+
+  return res;
+}
+
 export function getObjectAndStruct(obj1: unknown, obj2: unknown): [object, unknown[]] | undefined {
   if (isPojo(obj1) && isStruct(obj2)) {
     return [obj1 as object, obj2 as unknown[]];

diff --git a/test/contracts/programmable-function-logic/Returner.sol b/test/contracts/programmable-function-logic/Returner.sol
@@ -30,6 +30,9 @@ contract Returner {
 
   function getStructFixedSize() public returns (StructFixedSize memory _out1) {}
 
+  mapping(uint256 => StructFixedSize) public structMap;
+  mapping(uint256 => StructMixedNested) public nestedStructMap;
+
   struct StructDynamicSize {
     bytes valBytes;
     string valString;
@@ -47,6 +50,13 @@ contract Returner {
 
   function getStructMixedSize() public returns (StructMixedSize memory _out1) {}
 
+  struct StructMixedNested {
+    bool valRootBoolean;
+    string valRootString;
+    StructFixedSize valStructFixedSize;
+    StructDynamicSize valStructDynamicSize;
+  }
+
   struct StructNested {
     StructFixedSize valStructFixedSize;
     StructDynamicSize valStructDynamicSize;

diff --git a/test/unit/programmable-function-logic/type-handling.spec.ts b/test/unit/programmable-function-logic/type-handling.spec.ts
@@ -1,4 +1,5 @@
 import { FakeContract, smock } from '@src';
+import { convertStructToPojo } from '@src/utils';
 import { Returner } from '@typechained';
 import chai, { expect } from 'chai';
 import { BigNumber, utils } from 'ethers';
@@ -82,6 +83,29 @@ describe('ProgrammableFunctionLogic: Type Handling', () => {
         expect(result.valBytes32).to.equal(expected.valBytes32);
       });
 
+      it('should be able to return a struct inside a mapping', async () => {
+        const expected = STRUCT_FIXED_SIZE_EXAMPLE;
+
+        fake.structMap.returns(expected);
+
+        const result = convertStructToPojo(await fake.callStatic.structMap(BigNumber.from(1)));
+        expect(result).to.deep.equal(expected);
+      });
+
+      it('should be able to return a nested struct inside a mapping', async () => {
+        const expected = {
+          valRootString: 'Random',
+          valStructFixedSize: STRUCT_FIXED_SIZE_EXAMPLE,
+          valStructDynamicSize: STRUCT_DYNAMIC_SIZE_EXAMPLE,
+          valRootBoolean: true,
+        };
+
+        fake.nestedStructMap.returns(expected);
+
+        const result = convertStructToPojo(await fake.callStatic.nestedStructMap(BigNumber.from(1)));
+        expect(result).to.deep.equal(expected);
+      });
+
       it('should be able to return a struct with dynamic size values', async () => {
         const expected = STRUCT_DYNAMIC_SIZE_EXAMPLE;
         fake.getStructDynamicSize.returns(expected);

diff --git a/test/utils/constants.ts b/test/utils/constants.ts
@@ -8,12 +8,12 @@ export const BYTES_EXAMPLE =
   '0x56785678567856785678567856785678567856785678567856785678567856785678567856785678567856785678567856785678567856785678567856785678';
 
 export const STRUCT_FIXED_SIZE_EXAMPLE = {
-  valUint256: utils.parseUnits('1'),
-  valBoolean: true,
   valBytes32: BYTES32_EXAMPLE,
+  valBoolean: true,
+  valUint256: utils.parseUnits('1'),
 };
 
 export const STRUCT_DYNAMIC_SIZE_EXAMPLE = {
-  valBytes: BYTES_EXAMPLE,
   valString: 'hola',
+  valBytes: BYTES_EXAMPLE,
 };