diff --git a/src/utils/EnumerableKeyValueSetLib.sol b/src/utils/EnumerableKeyValueSetLib.sol new file mode 100644 index 0000000000..6f80db9144 --- /dev/null +++ b/src/utils/EnumerableKeyValueSetLib.sol @@ -0,0 +1,314 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.0; + +import {EnumerableSetLib} from "./EnumerableSetLib.sol"; + +library EnumerableKeyValueSetLib { + /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/ + /* CUSTOM ERRORS */ + /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/ + + /// @dev The index must be less than the length. + error IndexOutOfBounds(); + + /// @dev The value cannot be the zero sentinel. + error ValueIsZeroSentinel(); + + /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/ + /* CONSTANTS */ + /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/ + + /// @dev A sentinel value to denote the zero value in storage. + /// No elements can be equal to this value. + /// `uint72(bytes9(keccak256(bytes("_ZERO_SENTINEL"))))`. + uint256 private constant _ZERO_SENTINEL = 0xfbb67fda52d4bfb8bf; + + /// @dev The storage layout is given by: + /// ``` + /// mstore(0x04, _ENUMERABLE_WORD_SET_SLOT_SEED) + /// mstore(0x00, set.slot) + /// let rootSlot := keccak256(0x00, 0x24) + /// mstore(0x20, rootSlot) + /// mstore(0x00, value) + /// let positionSlot := keccak256(0x00, 0x40) + /// let valueSlot := add(rootSlot, sload(positionSlot)) + /// let valueInStorage := sload(valueSlot) + /// let lazyLength := sload(not(rootSlot)) + /// ``` + uint256 private constant _ENUMERABLE_WORD_SET_SLOT_SEED = 0x18fb5864; + + /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/ + /* STRUCTS */ + /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/ + + /// @dev An enumerable KeyValue (address to uint96) set in storage. + struct KeyValueSet { + uint256 _spacer; + } + + /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/ + /* GETTERS / SETTERS */ + /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/ + + /// @notice Gets the address key from the keyValuePair. It extracts the first 160 bits. + /// @dev The first 160 (20 bytes) are for the address key, the last 96 bits (12 bytes) are for the value. + function getAddressKey(bytes32 keyValuePair) internal pure returns (address key) { + /// @solidity memory-safe-assembly + assembly { + key := shr(96, keyValuePair) + } + } + + /// @notice Gets the value from the keyValuePair. It extracts the last 96 bits. + /// @dev The first 160 (20 bytes) are for the address key, the last 96 bits (12 bytes) are for the value. + function getValue(bytes32 keyValuePair) internal pure returns (uint256 value) { + /// @solidity memory-safe-assembly + assembly { + value := shr(160, shl(160, keyValuePair)) + } + } + + function joinKeyValue(address key, uint96 value) internal pure returns (bytes32 keyValuePair) { + /// @solidity memory-safe-assembly + assembly { + keyValuePair := or(shl(96, key), shr(160, shl(160, value))) + } + } + + /// @dev Returns the number of elements in the set. + function length(KeyValueSet storage set) internal view returns (uint256) { + return EnumerableSetLib.length(_toBytes32Set(set)); + } + + /// @dev Returns whether address `key` is in the set. + function contains(KeyValueSet storage set, address key) internal view returns (bool result) { + bytes32 rootSlot = _rootSlot(set); + /// @solidity memory-safe-assembly + assembly { + key := shr(96, shl(96, key)) + if eq(key, _ZERO_SENTINEL) { + mstore(0x00, 0xf5a267f1) // `ValueIsZeroSentinel()`. + revert(0x1c, 0x04) + } + if iszero(key) { key := _ZERO_SENTINEL } + for {} 1 {} { + if iszero(sload(not(rootSlot))) { + result := 1 + /// @dev We only check if the first 20 bytes are equal to the key. (left-padded address) + if eq(shr(96, sload(rootSlot)), key) { break } + if eq(shr(96, sload(add(rootSlot, 1))), key) { break } + if eq(shr(96, sload(add(rootSlot, 2))), key) { break } + result := 0 + break + } + mstore(0x20, rootSlot) + mstore(0x00, key) + result := iszero(iszero(sload(keccak256(0x00, 0x40)))) + break + } + } + } + + /// @dev Adds `keyValuePair` to the set. Returns whether address `key` was not in the set. + /// @dev The first 160 (20 bytes) are for the address key, the last 96 bits (12 bytes) are for the value. + function add(KeyValueSet storage set, bytes32 keyValuePair) internal returns (bool result) { + bytes32 rootSlot = _rootSlot(set); + /// @solidity memory-safe-assembly + assembly { + let key := shr(96, keyValuePair) + if eq(key, _ZERO_SENTINEL) { + mstore(0x00, 0xf5a267f1) // `ValueIsZeroSentinel()`. + revert(0x1c, 0x04) + } + if iszero(key) { + key := _ZERO_SENTINEL + keyValuePair := or(shl(96, key), keyValuePair) + } + for { let n := sload(not(rootSlot)) } 1 {} { + mstore(0x20, rootSlot) + if iszero(n) { + let v0 := shr(96, sload(rootSlot)) + if iszero(v0) { + sstore(rootSlot, keyValuePair) + result := 1 + break + } + if eq(v0, key) { break } + let v1 := shr(96, sload(add(rootSlot, 1))) + if iszero(v1) { + sstore(add(rootSlot, 1), keyValuePair) + result := 1 + break + } + if eq(v1, key) { break } + let v2 := shr(96, sload(add(rootSlot, 2))) + if iszero(v2) { + sstore(add(rootSlot, 2), keyValuePair) + result := 1 + break + } + if eq(v2, key) { break } + mstore(0x00, v0) + sstore(keccak256(0x00, 0x40), 1) + mstore(0x00, v1) + sstore(keccak256(0x00, 0x40), 2) + mstore(0x00, v2) + sstore(keccak256(0x00, 0x40), 3) + n := 7 + } + mstore(0x00, key) + let p := keccak256(0x00, 0x40) + if iszero(sload(p)) { + n := shr(1, n) + sstore(add(rootSlot, n), keyValuePair) + sstore(p, add(1, n)) + sstore(not(rootSlot), or(1, shl(1, add(1, n)))) + result := 1 + break + } + break + } + } + } + + /// @dev Removes `keyValuePair` from the set if the `key` is present. Returns whether address `key` was in the set. + function remove(KeyValueSet storage set, address key) internal returns (bool result) { + bytes32 rootSlot = _rootSlot(set); + /// @solidity memory-safe-assembly + assembly { + key := shr(96, shl(96, key)) + if eq(key, _ZERO_SENTINEL) { + mstore(0x00, 0xf5a267f1) // `ValueIsZeroSentinel()`. + revert(0x1c, 0x04) + } + if iszero(key) { key := _ZERO_SENTINEL } + for { let n := sload(not(rootSlot)) } 1 {} { + if iszero(n) { + result := 1 + if eq(shr(96, sload(rootSlot)), key) { + sstore(rootSlot, sload(add(rootSlot, 1))) + sstore(add(rootSlot, 1), sload(add(rootSlot, 2))) + sstore(add(rootSlot, 2), 0) + break + } + if eq(shr(96, sload(add(rootSlot, 1))), key) { + sstore(add(rootSlot, 1), sload(add(rootSlot, 2))) + sstore(add(rootSlot, 2), 0) + break + } + if eq(shr(96, sload(add(rootSlot, 2))), key) { + sstore(add(rootSlot, 2), 0) + break + } + result := 0 + break + } + mstore(0x20, rootSlot) + mstore(0x00, key) + let p := keccak256(0x00, 0x40) + let position := sload(p) + if iszero(position) { break } + n := sub(shr(1, n), 1) + if iszero(eq(sub(position, 1), n)) { + let lastValue := sload(add(rootSlot, n)) + sstore(add(rootSlot, sub(position, 1)), lastValue) + sstore(add(rootSlot, n), 0) + mstore(0x00, shr(96, lastValue)) + sstore(keccak256(0x00, 0x40), position) + } + sstore(not(rootSlot), or(shl(1, n), 1)) + sstore(p, 0) + result := 1 + break + } + } + } + + /// @dev Returns all of the values in the set. + /// Note: This can consume more gas than the block gas limit for large sets. + function values(KeyValueSet storage set) internal view returns (bytes32[] memory result) { + bytes32 rootSlot = _rootSlot(set); + /// @solidity memory-safe-assembly + assembly { + let zs := _ZERO_SENTINEL + let n := sload(not(rootSlot)) + result := mload(0x40) + let o := add(0x20, result) + for {} 1 {} { + if iszero(n) { + let v := sload(rootSlot) + if v { + n := 1 + mstore(o, or(mul(v, iszero(eq(shr(96, v), zs))), shr(160, shl(160, v)))) + v := sload(add(rootSlot, n)) + if v { + n := 2 + mstore( + add(o, 0x20), + or(mul(v, iszero(eq(shr(96, v), zs))), shr(160, shl(160, v))) + ) + v := sload(add(rootSlot, n)) + if v { + n := 3 + mstore( + add(o, 0x40), + or(mul(v, iszero(eq(shr(96, v), zs))), shr(160, shl(160, v))) + ) + } + } + } + break + } + n := shr(1, n) + for { let i := 0 } lt(i, n) { i := add(i, 1) } { + let v := sload(add(rootSlot, i)) + mstore( + add(o, shl(5, i)), + or(mul(v, iszero(eq(shr(96, v), zs))), shr(160, shl(160, v))) + ) + } + break + } + mstore(result, n) + mstore(0x40, add(o, shl(5, n))) + } + } + + /// @dev Returns the element at index `i` in the set. + function at(KeyValueSet storage set, uint256 i) internal view returns (bytes32 result) { + result = _rootSlot(set); + /// @solidity memory-safe-assembly + assembly { + result := sload(add(result, i)) + result := + or(mul(result, iszero(eq(shr(96, result), _ZERO_SENTINEL))), shr(160, shl(160, result))) + } + if (i >= length(set)) revert IndexOutOfBounds(); + } + + /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/ + /* PRIVATE HELPERS */ + /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/ + + /// @dev Casts to a Bytes32Set. + function _toBytes32Set(KeyValueSet storage s) + private + pure + returns (EnumerableSetLib.Bytes32Set storage c) + { + /// @solidity memory-safe-assembly + assembly { + c.slot := s.slot + } + } + + /// @dev Returns the root slot. + function _rootSlot(KeyValueSet storage s) private pure returns (bytes32 r) { + /// @solidity memory-safe-assembly + assembly { + mstore(0x04, _ENUMERABLE_WORD_SET_SLOT_SEED) + mstore(0x00, s.slot) + r := keccak256(0x00, 0x24) + } + } +} diff --git a/test/EnumerableKeyValueSetLib.t.sol b/test/EnumerableKeyValueSetLib.t.sol new file mode 100644 index 0000000000..773d0ed128 --- /dev/null +++ b/test/EnumerableKeyValueSetLib.t.sol @@ -0,0 +1,421 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.4; + +import "./utils/SoladyTest.sol"; +import {LibSort} from "../src/utils/LibSort.sol"; +import {LibPRNG} from "../src/utils/LibPRNG.sol"; + +import { + EnumerableSetLib, EnumerableKeyValueSetLib +} from "../src/utils/EnumerableKeyValueSetLib.sol"; + +contract EnumerableKeyValueSetLibTest is SoladyTest { + using EnumerableKeyValueSetLib for EnumerableKeyValueSetLib.KeyValueSet; + using LibPRNG for *; + + address private constant _ZERO_SENTINEL = 0x0000000000000000000000fbb67FDa52D4Bfb8Bf; + + EnumerableKeyValueSetLib.KeyValueSet keyValueSet; + EnumerableKeyValueSetLib.KeyValueSet keyValueSet2; + + function _createKeyValuePair(address key, uint96 value) + internal + pure + returns (bytes32 keyValuePair) + { + /// @solidity memory-safe-assembly + assembly { + keyValuePair := or(shl(96, key), shr(160, shl(160, value))) + } + } + + function testEnumerableKeyValueSetNoStorageCollision() public { + keyValueSet.add(_createKeyValuePair(address(1), 1)); + assertEq(keyValueSet2.contains(address(1)), false); + keyValueSet2.add(_createKeyValuePair(address(2), 2)); + assertEq(keyValueSet.contains(address(1)), true); + assertEq(keyValueSet2.contains(address(1)), false); + assertEq(keyValueSet.contains(address(2)), false); + keyValueSet.add(_createKeyValuePair(address(2), 2)); + assertEq(keyValueSet.contains(address(2)), true); + assertEq(keyValueSet2.contains(address(1)), false); + keyValueSet2.add(_createKeyValuePair(address(1), 1)); + assertEq(keyValueSet.contains(address(2)), true); + assertEq(keyValueSet2.contains(address(1)), true); + } + + function testEnumerableKeyValueSetBasic() public { + bytes32[] memory data = new bytes32[](5); + data[0] = _createKeyValuePair(address(0), 1); + data[1] = _createKeyValuePair(address(2), 2); + data[2] = _createKeyValuePair(address(3), 3); + data[3] = _createKeyValuePair(address(4), 4); + data[4] = _createKeyValuePair(address(5), 5); + + assertEq(keyValueSet.length(), 0); + assertEq(keyValueSet.contains(address(0)), false); + assertEq(keyValueSet.contains(address(2)), false); + assertEq(keyValueSet.contains(address(3)), false); + assertEq(keyValueSet.contains(address(4)), false); + assertEq(keyValueSet.contains(address(5)), false); + _assertKeyValueSetValues(data, 0); + + assertTrue(keyValueSet.add(_createKeyValuePair(address(0), 1))); + assertFalse(keyValueSet.add(_createKeyValuePair(address(0), 1))); + + assertEq(keyValueSet.length(), 1); + assertEq(keyValueSet.contains(address(0)), true); + assertEq(keyValueSet.contains(address(2)), false); + assertEq(keyValueSet.contains(address(3)), false); + assertEq(keyValueSet.contains(address(4)), false); + assertEq(keyValueSet.contains(address(5)), false); + _assertKeyValueSetValues(data, 1); + + assertTrue(keyValueSet.add(_createKeyValuePair(address(2), 2))); + assertFalse(keyValueSet.add(_createKeyValuePair(address(2), 2))); + + assertEq(keyValueSet.length(), 2); + assertEq(keyValueSet.contains(address(0)), true); + assertEq(keyValueSet.contains(address(2)), true); + assertEq(keyValueSet.contains(address(3)), false); + assertEq(keyValueSet.contains(address(4)), false); + assertEq(keyValueSet.contains(address(5)), false); + _assertKeyValueSetValues(data, 2); + + assertTrue(keyValueSet.add(_createKeyValuePair(address(3), 3))); + assertFalse(keyValueSet.add(_createKeyValuePair(address(3), 3))); + + assertEq(keyValueSet.length(), 3); + assertEq(keyValueSet.contains(address(0)), true); + assertEq(keyValueSet.contains(address(2)), true); + assertEq(keyValueSet.contains(address(3)), true); + assertEq(keyValueSet.contains(address(4)), false); + assertEq(keyValueSet.contains(address(5)), false); + _assertKeyValueSetValues(data, 3); + + assertTrue(keyValueSet.add(_createKeyValuePair(address(4), 4))); + assertFalse(keyValueSet.add(_createKeyValuePair(address(4), 4))); + + assertEq(keyValueSet.length(), 4); + assertEq(keyValueSet.contains(address(0)), true); + assertEq(keyValueSet.contains(address(2)), true); + assertEq(keyValueSet.contains(address(3)), true); + assertEq(keyValueSet.contains(address(4)), true); + assertEq(keyValueSet.contains(address(5)), false); + _assertKeyValueSetValues(data, 4); + + assertTrue(keyValueSet.add(_createKeyValuePair(address(5), 5))); + assertFalse(keyValueSet.add(_createKeyValuePair(address(5), 5))); + + assertEq(keyValueSet.length(), 5); + assertEq(keyValueSet.contains(address(0)), true); + assertEq(keyValueSet.contains(address(2)), true); + assertEq(keyValueSet.contains(address(3)), true); + assertEq(keyValueSet.contains(address(4)), true); + assertEq(keyValueSet.contains(address(5)), true); + _assertKeyValueSetValues(data, 5); + } + + function testEnumerableKeyValueSetBasic2() public { + bytes32[] memory data = new bytes32[](3); + data[0] = _createKeyValuePair(address(0), 1); + data[1] = _createKeyValuePair(address(2), 2); + data[2] = _createKeyValuePair(address(3), 3); + + keyValueSet.add(_createKeyValuePair(address(0), 1)); + keyValueSet.add(_createKeyValuePair(address(2), 2)); + _assertKeyValueSetValues(data, 2); + data[0] = _createKeyValuePair(address(2), 2); + + keyValueSet.remove(address(0)); + assertEq(keyValueSet.length(), 1); + _assertKeyValueSetValues(data, 1); + keyValueSet.remove(address(2)); + assertEq(keyValueSet.length(), 0); + _assertKeyValueSetValues(data, 0); + + keyValueSet.add(_createKeyValuePair(address(0), 1)); + keyValueSet.add(_createKeyValuePair(address(2), 2)); + data[0] = _createKeyValuePair(address(0), 1); + _assertKeyValueSetValues(data, 2); + + keyValueSet.remove(address(2)); + assertEq(keyValueSet.length(), 1); + _assertKeyValueSetValues(data, 1); + keyValueSet.remove(address(0)); + assertEq(keyValueSet.length(), 0); + _assertKeyValueSetValues(data, 0); + + keyValueSet.add(_createKeyValuePair(address(0), 1)); + keyValueSet.add(_createKeyValuePair(address(2), 2)); + keyValueSet.add(_createKeyValuePair(address(3), 3)); + _assertKeyValueSetValues(data, 3); + + keyValueSet.remove(address(3)); + assertEq(keyValueSet.length(), 2); + _assertKeyValueSetValues(data, 2); + keyValueSet.remove(address(2)); + assertEq(keyValueSet.length(), 1); + _assertKeyValueSetValues(data, 1); + keyValueSet.remove(address(0)); + assertEq(keyValueSet.length(), 0); + _assertKeyValueSetValues(data, 0); + + keyValueSet.add(_createKeyValuePair(address(0), 1)); + keyValueSet.add(_createKeyValuePair(address(2), 2)); + keyValueSet.add(_createKeyValuePair(address(3), 3)); + _assertKeyValueSetValues(data, 3); + + data[0] = _createKeyValuePair(address(2), 2); + data[1] = _createKeyValuePair(address(3), 3); + + keyValueSet.remove(address(0)); + assertEq(keyValueSet.length(), 2); + _assertKeyValueSetValues(data, 2); + + data[0] = _createKeyValuePair(address(3), 3); + + keyValueSet.remove(address(2)); + assertEq(keyValueSet.length(), 1); + _assertKeyValueSetValues(data, 1); + keyValueSet.remove(address(3)); + assertEq(keyValueSet.length(), 0); + _assertKeyValueSetValues(data, 0); + } + + function testEnumerableSetFuzz(uint256 n) public { + if (_randomChance(2)) { + _testEnumerableKeyValueSetFuzz(n); + } else { + if (_randomChance(2)) _testEnumerableKeyValueSetFuzz(); + } + } + + function _testEnumerableKeyValueSetFuzz(uint256 n) internal { + unchecked { + LibPRNG.PRNG memory prng; + prng.state = n; + uint256[] memory additions = new uint256[](prng.next() % 16); + uint256 mask = _randomChance(2) ? 7 : 15; + + for (uint256 i; i != additions.length; ++i) { + uint256 x = prng.next() & mask; + x = uint256(_createKeyValuePair(address(uint160(x)), uint96(x))); + additions[i] = x; + keyValueSet.add(bytes32(x)); + assertTrue(keyValueSet.contains(EnumerableKeyValueSetLib.getAddressKey(bytes32(x)))); + } + LibSort.sort(additions); + LibSort.uniquifySorted(additions); + assertEq(keyValueSet.length(), additions.length); + { + bytes32[] memory values = keyValueSet.values(); + _checkMemory(); + uint256[] memory valuesCasted = _toUints(values); + LibSort.sort(valuesCasted); + assertEq(valuesCasted, additions); + } + + uint256[] memory removals = new uint256[](prng.next() % 16); + for (uint256 i; i != removals.length; ++i) { + uint256 x = prng.next() & mask; + x = uint256(_createKeyValuePair(address(uint160(x)), uint96(x))); + removals[i] = x; + address key = EnumerableKeyValueSetLib.getAddressKey(bytes32(x)); + keyValueSet.remove(key); + assertFalse(keyValueSet.contains(key)); + } + LibSort.sort(removals); + LibSort.uniquifySorted(removals); + + { + uint256[] memory difference = LibSort.difference(additions, removals); + bytes32[] memory values = keyValueSet.values(); + _checkMemory(); + if (_randomChance(8)) _checkKeyValueSetValues(values); + uint256[] memory valuesCasted = _toUints(values); + LibSort.sort(valuesCasted); + assertEq(valuesCasted, difference); + } + } + } + + function _testEnumerableKeyValueSetFuzz() internal { + uint256[] memory s = _makeArray(0); + do { + bytes32 r = bytes32(_random()); + if (_randomChance(16)) _brutalizeMemory(); + if (_randomChance(2)) { + keyValueSet.add(r); + _addToArray(s, uint256(r)); + assertTrue(keyValueSet.contains(EnumerableKeyValueSetLib.getAddressKey(r))); + } else { + address key = EnumerableKeyValueSetLib.getAddressKey(r); + keyValueSet.remove(key); + _removeFromArray(s, uint256(r)); + assertFalse(keyValueSet.contains(key)); + } + if (_randomChance(16)) _brutalizeMemory(); + if (_randomChance(16)) { + _checkArraysSortedEq(_toUints(keyValueSet.values()), s); + assertEq(keyValueSet.length(), s.length); + } + if (s.length == 512) break; + } while (!_randomChance(8)); + _checkArraysSortedEq(_toUints(keyValueSet.values()), s); + } + + function _checkArraysSortedEq(uint256[] memory a, uint256[] memory b) internal { + LibSort.sort(a); + LibSort.sort(b); + assertEq(a, b); + } + + function _checkKeyValueSetValues(bytes32[] memory values) internal { + unchecked { + for (uint256 i; i != values.length; ++i) { + assertEq(keyValueSet.at(i), values[i]); + } + vm.expectRevert(EnumerableSetLib.IndexOutOfBounds.selector); + keyValueSetAt(_bound(_random(), values.length, type(uint256).max)); + } + } + + function testEnumerableKeyValueSetRevertsOnSentinel(uint256) public { + do { + address key = _randomAddress(); + if (_randomChance(32)) { + key = _ZERO_SENTINEL; + } + bytes32 a = _createKeyValuePair(key, uint96(_random())); + uint256 r = _random() % 3; + if (r == 0) { + if (key == _ZERO_SENTINEL) { + vm.expectRevert(EnumerableSetLib.ValueIsZeroSentinel.selector); + } + this.addToKeyValueSet(a); + } + if (r == 1) { + if (key == _ZERO_SENTINEL) { + vm.expectRevert(EnumerableSetLib.ValueIsZeroSentinel.selector); + } + this.keyValueSetContains(a); + } + if (r == 2) { + if (key == _ZERO_SENTINEL) { + vm.expectRevert(EnumerableSetLib.ValueIsZeroSentinel.selector); + } + this.removeFromKeyValueSet(a); + } + } while (!_randomChance(2)); + } + + function addToKeyValueSet(bytes32 a) public returns (bool) { + return keyValueSet.add(a); + } + + function keyValueSetContains(bytes32 a) public view returns (bool) { + return keyValueSet.contains(EnumerableKeyValueSetLib.getAddressKey(a)); + } + + function removeFromKeyValueSet(bytes32 a) public returns (bool) { + return keyValueSet.remove(EnumerableKeyValueSetLib.getAddressKey(a)); + } + + function keyValueSetAt(uint256 i) public view returns (bytes32) { + return keyValueSet.at(i); + } + + function _toUints(address[] memory a) private pure returns (uint256[] memory result) { + /// @solidity memory-safe-assembly + assembly { + result := a + } + } + + function _toUints(int256[] memory a) private pure returns (uint256[] memory result) { + /// @solidity memory-safe-assembly + assembly { + result := a + } + } + + function _toUints(bytes32[] memory a) private pure returns (uint256[] memory result) { + /// @solidity memory-safe-assembly + assembly { + result := a + } + } + + function _makeArray(uint256 size, uint256 maxCap) + internal + pure + returns (uint256[] memory result) + { + /// @solidity memory-safe-assembly + assembly { + result := mload(0x40) + mstore(result, size) + mstore(0x40, add(result, shl(5, add(maxCap, 1)))) + } + } + + function _makeArray(uint256 size) internal pure returns (uint256[] memory result) { + require(size <= 512, "Size too big."); + result = _makeArray(size, 512); + } + + function _addToArray(uint256[] memory a, uint256 x) internal pure { + /// @solidity memory-safe-assembly + assembly { + let exists := 0 + let n := mload(a) + for { let i := 0 } lt(i, n) { i := add(i, 1) } { + let o := add(add(a, 0x20), shl(5, i)) + if eq(shr(96, mload(o)), shr(96, x)) { + exists := 1 + break + } + } + if iszero(exists) { + n := add(n, 1) + mstore(add(a, shl(5, n)), x) + mstore(a, n) + } + } + } + + function _removeFromArray(uint256[] memory a, uint256 x) internal pure { + /// @solidity memory-safe-assembly + assembly { + let n := mload(a) + for { let i := 0 } lt(i, n) { i := add(i, 1) } { + let o := add(add(a, 0x20), shl(5, i)) + if eq(shr(96, mload(o)), shr(96, x)) { + mstore(o, mload(add(a, shl(5, n)))) + mstore(a, sub(n, 1)) + break + } + } + } + } + + function _assertKeyValueSetValues(bytes32[] memory expected, uint256 length) internal { + uint256 originalLength; + /// @solidity memory-safe-assembly + assembly { + originalLength := mload(expected) + mstore(expected, length) + } + _assertKeyValueSetValues(expected); + /// @solidity memory-safe-assembly + assembly { + mstore(expected, originalLength) + } + } + + function _assertKeyValueSetValues(bytes32[] memory expected) internal { + assertEq(keyValueSet.values(), expected); + _checkKeyValueSetValues(expected); + } +}