diff --git a/src/data-structures/Hashmap.huff b/src/data-structures/Hashmap.huff index a4517a7d..c6da7294 100644 --- a/src/data-structures/Hashmap.huff +++ b/src/data-structures/Hashmap.huff @@ -62,6 +62,49 @@ 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 + // [, slot, key1, key2, key3] + mstore // [key1, key2, key3] + + // next byte + 0x20 add // [ + 32, key1, key2, key3] + swap1 dup2 // [ + 32, key1, + 32, key2, key3] + mstore // [ + 32, key2, key3] + + 0x40 // [0x40, + 32, key2, key3] + // [, 0x40, + 32, key2, key3] + sha3 // [slot1, + 32, key2, key3] + + // concat the first two keys + dup2 // [ + 32, slot1, + 32, key2, key3] put slot1 in memory + mstore // [ + 32, key2, key3] + + // put key2 in memory, before slot1 + swap1 // [key2, + 32, key3] + // [, key2, + 32, key3] + mstore // [key3] + + 0x40 // [0x40, + 32, key3] + // [, 0x40, + 32, key3] + sha3 // [slot2, + 32, key3] + + // concat with the third key + swap1 // [ + 32, slot2, key3] put slot2 in memory + mstore // [key3] + + // put key3 in memory, before slot2 + // [, key3] + mstore // [] + + // Hash the data, generating the final slot3 + 0x40 // [0x40] + // [, 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 +126,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() // [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 +153,10 @@ GET_SLOT_FROM_KEYS_2D() // [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() // [slot, value] + sstore // [] +} \ No newline at end of file diff --git a/test/data-structures/Hashmap.t.sol b/test/data-structures/Hashmap.t.sol index 927594ac..e0b53377 100644 --- a/test/data-structures/Hashmap.t.sol +++ b/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); + } } diff --git a/test/data-structures/mocks/HashmapWrappers.huff b/test/data-structures/mocks/HashmapWrappers.huff index 448f3c04..c4b79630 100644 --- a/test/data-structures/mocks/HashmapWrappers.huff +++ b/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,6 +98,27 @@ 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 @@ -57,8 +126,12 @@ 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() }