Add Hashmap slot+3keys (3D) macros

src/data-structures/Hashmap.huff

@@ -62,6 +62,47 @@
     sha3                // [key]
 }
 
+/// @notice Calculate the slot from three keys
+#define macro GET_SLOT_FROM_KEYS_3D(mem_ptr) = takes(4) returns (1) {
+    // Input stack: [slot, key1, key2, key3]
+    // Load the data into memory
+    <mem_ptr>           // [<mem_ptr>, slot, key1, key2, key3]
+    mstore              // [key1, key2, key3]
+
+    // next byte
+    <mem_ptr> 0x20 add  // [<mem_ptr> + 32, key1, key2, key3]
+    mstore              // [key2, key3]
+
+    0x40                // [0x40, key2, key3]
+    <mem_ptr>           // [<mem_ptr>, 0x40, key2, key3]
+    sha3                // [slot1, key2, key3]
+
+    // concat the first two keys
+    <mem_ptr> 0x20 add  // [<mem_ptr> + 32, slot1, key2, key3] put slot1 in memory
+    mstore              // [key2, key3]
+
+    // put key2 in memory, before slot1
+    <mem_ptr>           // [<mem_ptr>, key2, key3]
+    mstore              // [key3]
+
+    0x40                // [0x40, key3]
+    <mem_ptr>           // [<mem_ptr>, 0x40, key3]
+    sha3                // [slot2, key3]
+
+    // concat with the third key
+    <mem_ptr> 0x20 add  // [<mem_ptr> + 32, slot2, key3] put slot2 in memory
+    mstore              // [key3]
+
+    // put key3 in memory, before slot2
+    <mem_ptr>           // [<mem_ptr>, key3]
+    mstore              // []
+
+    // Hash the data, generating the final slot3
+    0x40                // [0x40]
+    <mem_ptr>           // [<mem_ptr>, 0x40]
+    sha3                // [slot3]
+}
+
 /// @notice Load an element onto the stack from a key
 #define macro LOAD_ELEMENT(mem_ptr) = takes(1) returns(1) {
     // Input stack: [key]
@@ -83,6 +124,13 @@
     sload                            // [value]
 }
 
+/// @notice Load an element onto the stack from a slot and three keys
+#define macro LOAD_ELEMENT_FROM_KEYS_3D(mem_ptr) = takes(4) returns(1) {
+    // Input stack: [slot, key1, key2, key3]
+    GET_SLOT_FROM_KEYS_3D(<mem_ptr>) // [slot]
+    sload                            // [value]
+}
+
 /// @notice Store an element from a key
 #define macro STORE_ELEMENT(mem_ptr) = takes(2) returns(0) {
     // Input stack: [key, value]
@@ -103,3 +151,10 @@
     GET_SLOT_FROM_KEYS_2D(<mem_ptr>)    // [slot, value]
     sstore                              // []
 }
+
+/// @notice Store an element from a slot and three keys
+#define macro STORE_ELEMENT_FROM_KEYS_3D(mem_ptr) = takes(5) returns (0) {
+    // Input stack: [slot, key1, key2, key3, value]
+    GET_SLOT_FROM_KEYS_3D(<mem_ptr>)    // [slot, value]
+    sstore                              // []
+}

test/data-structures/Hashmap.t.sol

@@ -8,8 +8,23 @@ import {HuffDeployer} from "foundry-huff/HuffDeployer.sol";
 interface Hashmap {
   function loadElement(bytes32) external view returns (bytes32);
   function loadElementFromKeys(bytes32, bytes32) external returns (bytes32);
+  function loadElementFromKeys2D(bytes32, bytes32, bytes32) external returns (bytes32);
+  function loadElementFromKeys3D(bytes32, bytes32, bytes32, bytes32) external returns (bytes32);
   function storeElement(bytes32 key, bytes32 value) external;
   function storeElementFromKeys(bytes32 key1, bytes32 key2, bytes32 value) external;
+  function storeElementFromKeys2D(
+    bytes32 slot,
+    bytes32 key1,
+    bytes32 key2,
+    bytes32 value
+  ) external;
+  function storeElementFromKeys3D(
+    bytes32 slot,
+    bytes32 key1,
+    bytes32 key2,
+    bytes32 key3,
+    bytes32 value
+  ) external;
 }
 
 contract HashmapTest is Test {
@@ -24,7 +39,7 @@ contract HashmapTest is Test {
     hmap = Hashmap(config.deploy("data-structures/Hashmap"));
   }
 
-  /// @notice Test getting a vlue for a key
+  /// @notice Test getting a value for a key
   function testGetKey(bytes32 key) public {
     bytes32 element = hmap.loadElement(key);
     assertEq(element, bytes32(0));
@@ -49,4 +64,29 @@ contract HashmapTest is Test {
     hmap.storeElementFromKeys(key_one, key_two, value);
     assertEq(hmap.loadElementFromKeys(key_one, key_two), value);
   }
+
+  /// @notice Test set with slot and 2 keys
+  function testSetKeys2D(
+    bytes32 slot, 
+    bytes32 key_one, 
+    bytes32 key_two, 
+    bytes32 value
+  ) public {
+    assertEq(hmap.loadElementFromKeys2D(slot, key_one, key_two), bytes32(0));
+    hmap.storeElementFromKeys2D(slot, key_one, key_two, value);
+    assertEq(hmap.loadElementFromKeys2D(slot, key_one, key_two), value);
+  }
+
+  /// @notice Test set with slot and 3 keys
+  function testSetKeys3D(
+    bytes32 slot, 
+    bytes32 key_one, 
+    bytes32 key_two,
+    bytes32 key_three, 
+    bytes32 value
+  ) public {
+    assertEq(hmap.loadElementFromKeys3D(slot, key_one, key_two, key_three), bytes32(0));
+    hmap.storeElementFromKeys3D(slot, key_one, key_two, key_three, value);
+    assertEq(hmap.loadElementFromKeys3D(slot, key_one, key_two, key_three), value);
+  }
 }

test/data-structures/mocks/HashmapWrappers.huff

@@ -1,18 +1,45 @@
-// Returns the slot for the given key.
+// Returns the slot value for the given key.
 // sig: 0x437b8ad6
 #define function loadElement(bytes32) view returns (bytes32)
 
-// Returns the slot by hashing the given keys.
+// Returns the slot value by hashing the given keys.
 // sig: 0xf268de10
 #define function loadElementFromKeys(bytes32, bytes32) view returns (bytes32)
 
+// Returns the slot value by hashing the given slot and two keys.
+// sig: 0xef5b4768
+#define function loadElementFromKeys2D(bytes32, bytes32, bytes32) view returns (bytes32)
+
+// Returns the slot value by hashing the given slot and three keys.
+// sig: 0x0cc703f5
+#define function loadElementFromKeys3D(bytes32, bytes32, bytes32, bytes32) view returns (bytes32)
+
 // Stores the value for the given key.
 // sig: 0x376caf9f
-#define function storeElement(bytes32 key, bytes32 value) view returns ()
+#define function storeElement(bytes32 key, bytes32 value) nonpayable returns ()
 
 // Stores the value for the given keys
 // sig: 0x2fdb44d8
-#define function storeElementFromKeys(bytes32 key1, bytes32 key2, bytes32 value) view returns ()
+#define function storeElementFromKeys(bytes32 key1, bytes32 key2, bytes32 value) nonpayable returns ()
+
+// Stores the value given a slot and two keys
+// sig: 0x64fab984
+#define function storeElementFromKeys2D(
+    bytes32 slot,
+    bytes32 key1,
+    bytes32 key2,
+    bytes32 value
+) nonpayable returns ()
+
+// Stores the value given a slot and three keys
+// sig: 0xd3b12314
+#define function storeElementFromKeys3D(
+    bytes32 slot,
+    bytes32 key1,
+    bytes32 key2,
+    bytes32 key3,
+    bytes32 value
+) nonpayable returns ()
 
 #define constant LOCATION = FREE_STORAGE_POINTER()
 
@@ -33,6 +60,27 @@
     0x20 0x00 return
 }
 
+// Get the value for the given slot and two keys
+#define macro GET_FROM_KEYS_2D() = takes(0) returns(0) {
+    0x44 calldataload
+    0x24 calldataload
+    0x04 calldataload
+    LOAD_ELEMENT_FROM_KEYS_2D(0x00)
+    0x00 mstore
+    0x20 0x00 return
+}
+
+// Get the value for the given slot and three keys
+#define macro GET_FROM_KEYS_3D() = takes(0) returns(0) {
+    0x64 calldataload
+    0x44 calldataload
+    0x24 calldataload
+    0x04 calldataload
+    LOAD_ELEMENT_FROM_KEYS_3D(0x00)
+    0x00 mstore
+    0x20 0x00 return
+}
+
 // Store the value for the given key.
 #define macro STORE() = takes(0) returns(0) {
     0x24 calldataload
@@ -50,15 +98,40 @@
     stop
 }
 
+// Store the value for the given slot and two keys.
+#define macro STORE_FROM_KEYS_2D() = takes(0) returns(0) {
+    0x64 calldataload
+    0x44 calldataload
+    0x24 calldataload
+    0x04 calldataload
+    STORE_ELEMENT_FROM_KEYS_2D(0x00)
+    stop
+}
+
+// Store the value for the given slot and three keys.
+#define macro STORE_FROM_KEYS_3D() = takes(0) returns(0) {
+    0x84 calldataload
+    0x64 calldataload
+    0x44 calldataload
+    0x24 calldataload
+    0x04 calldataload
+    STORE_ELEMENT_FROM_KEYS_3D(0x00)
+    stop
+}
+
 // Main Macro - The contract entrypoint
 #define macro MAIN() = takes(0) returns (0) {
     // Identify which function is being called using the 4 byte function signature
     pc calldataload 0xE0 shr
 
     dup1 __FUNC_SIG(loadElement)            eq load_element jumpi
     dup1 __FUNC_SIG(loadElementFromKeys)    eq load_element_from_keys jumpi
+    dup1 __FUNC_SIG(loadElementFromKeys2D)  eq load_element_from_keys_2d jumpi
+    dup1 __FUNC_SIG(loadElementFromKeys3D)  eq load_element_from_keys_3d jumpi
     dup1 __FUNC_SIG(storeElement)           eq store_element jumpi
     dup1 __FUNC_SIG(storeElementFromKeys)   eq store_element_from_keys jumpi
+    dup1 __FUNC_SIG(storeElementFromKeys2D) eq store_element_from_keys_2d jumpi
+    dup1 __FUNC_SIG(storeElementFromKeys3D) eq store_element_from_keys_3d jumpi
 
     // Revert if otherwise
     0x00 dup1 revert
@@ -67,8 +140,16 @@
         GET()
     load_element_from_keys:
         GET_FROM_KEYS()
+    load_element_from_keys_2d:
+        GET_FROM_KEYS_2D()
+    load_element_from_keys_3d:
+        GET_FROM_KEYS_3D()
     store_element:
         STORE()
     store_element_from_keys:
         STORE_FROM_KEYS()
+    store_element_from_keys_2d:
+        STORE_FROM_KEYS_2D()
+    store_element_from_keys_3d:
+        STORE_FROM_KEYS_3D()
 }