Fix SlotDerivation tests generation

OpenZeppelin · May 13, 2024 · ca59b7e · ca59b7e
1 parent da8990c
commit ca59b7e
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 20 deletions.
diff --git a/scripts/generate/templates/SlotDerivation.t.js b/scripts/generate/templates/SlotDerivation.t.js
@@ -5,17 +5,19 @@ const { TYPES } = require('./Slot.opts');
 const header = `\
 pragma solidity ^0.8.20;
 
-import {Test} from "forge-std/Test.sol";
+// solhint-disable func-name-mixedcase
 
+import {Test} from "forge-std/Test.sol";
+import {SymTest} from "halmos-cheatcodes/SymTest.sol";
 import {SlotDerivation} from "@openzeppelin/contracts/utils/SlotDerivation.sol";
 `;
 
 const array = `\
 bytes[] private _array;
 
-function testDeriveArray(uint256 length, uint256 offset) public {
-  length = bound(length, 1, type(uint256).max);
-  offset = bound(offset, 0, length - 1);
+function testSymbolicDeriveArray(uint256 length, uint256 offset) public {
+  vm.assume(length > 0);
+  vm.assume(offset < length);
 
   bytes32 baseSlot;
   assembly {
@@ -33,10 +35,10 @@ function testDeriveArray(uint256 length, uint256 offset) public {
 }
 `;
 
-const mapping = ({ type, name, isValueType }) => `\
+const mapping = ({ type, name }) => `\
 mapping(${type} => bytes) private _${type}Mapping;
 
-function testDeriveMapping${name}(${type} ${isValueType ? '' : 'memory'} key) public {
+function testSymbolicDeriveMapping${name}(${type} key) public {
   bytes32 baseSlot;
   assembly {
     baseSlot := _${type}Mapping.slot
@@ -52,10 +54,37 @@ function testDeriveMapping${name}(${type} ${isValueType ? '' : 'memory'} key) pu
 }
 `;
 
+const boundedMapping = ({ type, name }) => `\
+mapping(${type} => bytes) private _${type}Mapping;
+
+function testDeriveMapping${name}(${type} memory key) public {
+  _assertDeriveMapping${name}(key);
+}
+
+function check_DeriveMapping${name}() public {
+  _assertDeriveMapping${name}(svm.create${name}(256, "DeriveMapping${name}Input"));
+}
+
+function _assertDeriveMapping${name}(${type} memory key) internal {
+  bytes32 baseSlot;
+  assembly {
+      baseSlot := _${type}Mapping.slot
+  }
+
+  bytes storage derived = _${type}Mapping[key];
+  bytes32 derivedSlot;
+  assembly {
+      derivedSlot := derived.slot
+  }
+
+  assertEq(baseSlot.deriveMapping(key), derivedSlot);
+}
+`;
+
 // GENERATE
 module.exports = format(
   header.trimEnd(),
-  'contract SlotDerivationTest is Test {',
+  'contract SlotDerivationTest is Test, SymTest {',
   'using SlotDerivation for bytes32;',
   '',
   array,
@@ -68,6 +97,6 @@ module.exports = format(
         isValueType: type.isValueType,
       })),
     ),
-  ).map(type => mapping(type)),
+  ).map(type => (!type.isValueType ? boundedMapping(type) : mapping(type))),
   '}',
 );
diff --git a/test/utils/SlotDerivation.t.sol b/test/utils/SlotDerivation.t.sol
@@ -179,18 +179,6 @@ contract SlotDerivationTest is Test, SymTest {
         _assertDeriveMappingString(svm.createString(256, "DeriveMappingStringInput"));
     }
 
-    mapping(bytes => bytes) private _bytesMapping;
-
-    function testDeriveMappingBytes(bytes memory key) public {
-        _assertDeriveMappingBytes(key);
-    }
-
-    function check_DeriveMappingBytes() public {
-        _assertDeriveMappingBytes(svm.createBytes(256, "DeriveMappingBytesInput"));
-    }
-
-    /// Asserts
-
     function _assertDeriveMappingString(string memory key) internal {
         bytes32 baseSlot;
         assembly {
@@ -206,6 +194,16 @@ contract SlotDerivationTest is Test, SymTest {
         assertEq(baseSlot.deriveMapping(key), derivedSlot);
     }
 
+    mapping(bytes => bytes) private _bytesMapping;
+
+    function testDeriveMappingBytes(bytes memory key) public {
+        _assertDeriveMappingBytes(key);
+    }
+
+    function check_DeriveMappingBytes() public {
+        _assertDeriveMappingBytes(svm.createBytes(256, "DeriveMappingBytesInput"));
+    }
+
     function _assertDeriveMappingBytes(bytes memory key) internal {
         bytes32 baseSlot;
         assembly {