diff --git a/README.md b/README.md index 79ccd33..ab99baf 100644 --- a/README.md +++ b/README.md @@ -14,6 +14,7 @@ The Red-Black Tree implementation is based on [Solady's RedBlackTreeLib](https:/ - **Complex Value Support**: Stores `ValueLib.Value` structs containing any size of data. In this example, it shows 6 slots. - **Red-Black Tree Properties**: Maintains balanced tree structure ensuring O(log n) operations - **Zero Value Protection**: Prevents insertion of zero keys to maintain tree integrity +- **Dynamic Resizing**: Pre-allocate storage capacity to optimize gas costs for future insertions (available in `RedBlackTreeWithResizeLib`) ### Gas Performance @@ -38,15 +39,19 @@ The Red-Black Tree implementation shows **90% gas reduction** for hot insertions ``` src/ -├── RedBlackTreeKV.sol # Main KV store implementation -├── MappingKV.sol # Traditional mapping implementation for comparison +├── RedBlackTreeKV.sol # Main KV store implementation +├── RedBlackTreeWithResizeKV.sol # Enhanced KV store with resize capabilities +├── MappingKV.sol # Traditional mapping implementation for comparison └── lib/ - ├── RedBlackTreeLib.sol # Red-Black Tree data structure library - └── Value.sol # Value struct definition + ├── RedBlackTreeLib.sol # Red-Black Tree data structure library + ├── RedBlackTreeWithResizeLib.sol # Enhanced Red-Black Tree with resize capabilities + └── Value.sol # Value struct definition test/ -├── RedBlackTreeKV.t.sol # Comprehensive unit tests -└── RedBlackTreeKVGas.t.sol # Gas benchmark tests +├── RedBlackTreeKV.t.sol # Comprehensive unit tests +├── RedBlackTreeWithResizeKV.t.sol # Tests for resize-enabled KV store +├── RedBlackTreeWithResize.sol # Red-Black Tree library tests with resize functionality +└── RedBlackTreeKVGas.t.sol # Gas benchmark tests ``` ## Value Structure @@ -95,6 +100,54 @@ kv.deleteValue(1); uint256[] memory keys = kv.values(); ``` +### Enhanced Operations with Resize + +```solidity +RedBlackTreeWithResizeKV kvWithResize = new RedBlackTreeWithResizeKV(); + +// Pre-allocate storage for 1000 elements to optimize gas costs +kvWithResize.resize(1000); + +// Check current storage metrics +uint256 capacity = kvWithResize.getStorageSize(); // Returns 1000 +uint256 size = kvWithResize.getSize(); // Returns current number of elements + +// Set values with optimized gas costs +kvWithResize.setValue(1, value); +kvWithResize.setValue(2, anotherValue); + +// All other operations work the same +ValueLib.Value memory retrieved = kvWithResize.getValue(1); +kvWithResize.deleteValue(1); +uint256[] memory keys = kvWithResize.values(); +``` + +### WithResize Feature + +The `RedBlackTreeWithResizeLib` provides enhanced functionality for pre-allocating storage capacity to optimize gas costs: + +```solidity +import "./lib/RedBlackTreeWithResizeLib.sol"; + +RedBlackTreeWithResizeLib.Tree tree; + +// Pre-allocate storage for 1000 elements to save gas on future insertions +RedBlackTreeWithResizeLib.resize(tree, 1000); + +// Check current storage capacity +uint256 capacity = RedBlackTreeWithResizeLib.storageSize(tree); + +// Insert values with optimized gas costs +RedBlackTreeWithResizeLib.insert(tree, 42); +``` + +#### Resize Benefits + +- **Pre-allocation**: Reserve storage slots in advance to avoid costly allocations during insertion +- **Gas Optimization**: Significantly reduces gas costs for insertions when capacity is pre-allocated +- **Flexible Sizing**: Dynamically adjust storage capacity based on expected usage patterns +- **Auto-increment**: Storage capacity automatically increases when inserting beyond current capacity + ### Constraints - **No Zero Keys**: Keys cannot be 0 (will revert with `ValueIsEmpty()`) @@ -117,6 +170,11 @@ forge test # Run unit tests only forge test --match-contract RedBlackTreeKVTest +forge test --match-contract RedBlackTreeWithResizeKVTest + +# Run resize library tests +forge test --match-contract RedBlackTreeLibTest +forge test --match-contract RedBlackTreeWithResizeLibTest # Run gas benchmarks forge test --match-contract MappingGasTest -vv @@ -126,6 +184,8 @@ forge test --match-contract MappingGasTest -vv The test suite includes: +**RedBlackTreeKV Tests:** + - ✅ Basic CRUD operations - ✅ Edge cases (zero keys, max values) - ✅ Multiple value operations @@ -134,14 +194,40 @@ The test suite includes: - ✅ Fuzz testing with random inputs - ✅ Gas consumption benchmarks +**RedBlackTreeWithResizeKV Tests:** + +- ✅ All basic KV operations with resize functionality +- ✅ Storage capacity management and pre-allocation +- ✅ Resize operations (expand, shrink, same capacity) +- ✅ Data preservation during resize operations +- ✅ Error handling for invalid resize operations +- ✅ Resize with mixed insert/delete operations +- ✅ Fuzz testing for resize capacity variations + +**RedBlackTreeWithResizeLib Tests:** + +- ✅ Core Red-Black Tree operations with resize support +- ✅ Storage size tracking and capacity management +- ✅ Tree balancing with dynamic storage allocation +- ✅ Error cases for tree size limits and invalid operations + ## Use Cases This implementation is ideal for: -- **Order Books**: Frequent order insertions and cancellations -- **Gaming**: Player inventory systems with item trading -- **DeFi Protocols**: Position management with frequent updates -- **NFT Marketplaces**: Listing and delisting operations +- **Order Books**: Frequent order insertions and cancellations with predictable volume patterns +- **Gaming**: Player inventory systems with item trading and known capacity requirements +- **DeFi Protocols**: Position management with frequent updates and batch operations +- **NFT Marketplaces**: Listing and delisting operations with seasonal volume spikes +- **Data Analytics**: Time-series data storage with known dataset sizes +- **Batch Processing**: Systems that process large batches of data with predictable patterns + +**When to use RedBlackTreeWithResizeKV:** + +- You know the approximate maximum number of elements in advance +- You need to optimize for gas costs during high-frequency operations +- Your application has predictable usage patterns or batch processing requirements +- You want to avoid gas spikes during storage allocation ## License diff --git a/src/RedBlackTreeKV.sol b/src/RedBlackTreeKV.sol index 3c58d02..8296dc0 100644 --- a/src/RedBlackTreeKV.sol +++ b/src/RedBlackTreeKV.sol @@ -6,7 +6,8 @@ import {ValueLib} from "./lib/Value.sol"; // A simple RBT kv example. It is gas-efficient if frequently add & remove. contract RedBlackTreeKV { - uint256 private constant _DATA_SLOT_SEED = 0xdeadbeef; // Arbitrary unique seed + // This seed is derived from the bytes("RedBlackTreeKV") + uint256 private constant _DATA_SLOT_SEED = 0xc8a834a090c3fa550a8f4fd4767f0a9c3eb1fa1ce9154cab77c4b35ad48e385c; // Arbitrary unique seed uint256 private constant _SLOTS_PER_POSITION = ValueLib.SLOTS_PER_POSITION; // Dense slots per value RedBlackTreeLib.Tree public tree; diff --git a/src/RedBlackTreeWithResizeKV.sol b/src/RedBlackTreeWithResizeKV.sol new file mode 100644 index 0000000..c8f773c --- /dev/null +++ b/src/RedBlackTreeWithResizeKV.sol @@ -0,0 +1,143 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.4; + +import {console} from "forge-std/Test.sol"; +import {RedBlackTreeWithResizeLib} from "./lib/RedBlackTreeWithResizeLib.sol"; +import {ValueLib} from "./lib/Value.sol"; + +// A simple RBT kv example. It is gas-efficient if frequently add & remove. +contract RedBlackTreeWithResizeKV { + // This seed is derived from the bytes("RedBlackTreeKV") + uint256 private constant _DATA_SLOT_SEED = 0xc8a834a090c3fa550a8f4fd4767f0a9c3eb1fa1ce9154cab77c4b35ad48e385c; // Arbitrary unique seed + uint256 private constant _SLOTS_PER_POSITION = ValueLib.SLOTS_PER_POSITION; // Dense slots per value + + RedBlackTreeWithResizeLib.Tree public tree; + + using RedBlackTreeWithResizeLib for RedBlackTreeWithResizeLib.Tree; + + function setValue(uint256 key, ValueLib.Value memory value) public { + tree.insert(key); + bytes32 ptr = tree.find(key); + (, uint256 index) = _unpack(ptr); // Extract dense node index + uint256 dataBase = _dataBase(); + uint256 valueSlot = dataBase + index * _SLOTS_PER_POSITION; + /// @solidity memory-safe-assembly + assembly { + sstore(valueSlot, mload(value)) + sstore(add(valueSlot, 1), mload(add(value, 0x20))) + sstore(add(valueSlot, 2), mload(add(value, 0x40))) + sstore(add(valueSlot, 3), mload(add(value, 0x60))) + sstore(add(valueSlot, 4), mload(add(value, 0x80))) + sstore(add(valueSlot, 5), mload(add(value, 0xa0))) + } + } + + function getValue(uint256 key) public view returns (ValueLib.Value memory) { + bytes32 ptr = tree.find(key); + (, uint256 index) = _unpack(ptr); // Extract dense node index + uint256 dataBase = _dataBase(); + uint256 valueSlot = dataBase + index * _SLOTS_PER_POSITION; + ValueLib.Value memory value; + /// @solidity memory-safe-assembly + assembly { + mstore(value, sload(valueSlot)) + mstore(add(value, 0x20), sload(add(valueSlot, 1))) + mstore(add(value, 0x40), sload(add(valueSlot, 2))) + mstore(add(value, 0x60), sload(add(valueSlot, 3))) + mstore(add(value, 0x80), sload(add(valueSlot, 4))) + mstore(add(value, 0xa0), sload(add(valueSlot, 5))) + } + return value; + } + + function deleteValue(uint256 key) public { + bytes32 ptr = tree.find(key); + (, uint256 deletedIndex) = _unpack(ptr); // Extract dense node index + uint256 dataBase = _dataBase(); + uint256 deletedSlot = dataBase + deletedIndex * _SLOTS_PER_POSITION; + + uint256 treeSize = tree.size(); + uint256 lastIndex = treeSize; + bool needsCopy = (deletedIndex != lastIndex); + uint256 lastSlot = dataBase + lastIndex * _SLOTS_PER_POSITION; + + tree.remove(key); + + if (needsCopy) { + /// @solidity memory-safe-assembly + assembly { + /// @dev 6 = _SLOTS_PER_POSITION + for { let i := 0 } lt(i, 6) { i := add(i, 1) } { sstore(add(deletedSlot, i), sload(add(lastSlot, i))) } + } + } + } + + function resize(uint256 newCapacity) public { + uint256 oldCapacity = tree.storageSize(); + uint256 size = tree.size(); + if (newCapacity < size) { + revert RedBlackTreeWithResizeLib.InvalidStorageSize(); + } + tree.resize(newCapacity); + if (newCapacity > oldCapacity) { + uint256 dataBase = _dataBase(); + for (uint256 index = oldCapacity + 1; index <= newCapacity; index++) { + uint256 valueSlot = dataBase + index * _SLOTS_PER_POSITION; + bytes32 slot; + assembly { + slot := sload(valueSlot) + sstore(valueSlot, 1) + sstore(add(valueSlot, 1), 1) + sstore(add(valueSlot, 2), 1) + sstore(add(valueSlot, 3), 1) + sstore(add(valueSlot, 4), 1) + sstore(add(valueSlot, 5), 1) + } + } + } else if (newCapacity < oldCapacity) { + uint256 dataBase = _dataBase(); + for (uint256 index = newCapacity + 1; index <= oldCapacity; index++) { + uint256 valueSlot = dataBase + index * _SLOTS_PER_POSITION; + assembly { + sstore(valueSlot, 0) + sstore(add(valueSlot, 1), 0) + sstore(add(valueSlot, 2), 0) + sstore(add(valueSlot, 3), 0) + sstore(add(valueSlot, 4), 0) + sstore(add(valueSlot, 5), 0) + } + } + } + } + + // Helper: Compute data base slot (keccak-based, like tree nodes but unique seed) + function _dataBase() internal pure returns (uint256 slotBase) { + /// @solidity memory-safe-assembly + assembly { + mstore(0x20, tree.slot) + mstore(0x00, _DATA_SLOT_SEED) + slotBase := keccak256(0x00, 0x40) + } + } + + // Helper: Unpack pointer (copied from lib private _unpack for convenience) + function _unpack(bytes32 ptr) internal pure returns (uint256 nodes, uint256 key) { + /// @solidity memory-safe-assembly + assembly { + nodes := shl(32, shr(32, ptr)) // _NODES_SLOT_SHIFT == 32 + key := and(0x7fffffff, ptr) // _BITMASK_KEY == (1 << 31) - 1 + } + } + + function values() public view returns (uint256[] memory) { + return tree.values(); + } + + function getStorageSize() public view returns (uint256) { + return tree.storageSize(); + } + + function getSize() public view returns (uint256) { + return tree.size(); + } +} diff --git a/src/lib/RedBlackTreeWithResizeLib.sol b/src/lib/RedBlackTreeWithResizeLib.sol new file mode 100644 index 0000000..154f526 --- /dev/null +++ b/src/lib/RedBlackTreeWithResizeLib.sol @@ -0,0 +1,793 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.4; + +/// @notice Library for managing a red-black-tree in storage. +/// @author Solady (https://github.com/vectorized/solady/blob/main/src/utils/RedBlackTreeLib.sol) +/// @author Modified from BokkyPooBahsRedBlackTreeLibrary (https://github.com/bokkypoobah/BokkyPooBahsRedBlackTreeLibrary) +/// @dev This implementation does not support the zero (i.e. empty) value. +/// This implementation supports up to 2147483647 values. +/// @notice Added function "resize" and "storageSize" to the library. +library RedBlackTreeWithResizeLib { + /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/ + /* CUSTOM ERRORS */ + /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/ + + /// @dev The value cannot be zero. + error ValueIsEmpty(); + + /// @dev Cannot insert a value that already exists. + error ValueAlreadyExists(); + + /// @dev Cannot remove a value that does not exist. + error ValueDoesNotExist(); + + /// @dev The pointer is out of bounds. + error PointerOutOfBounds(); + + /// @dev The tree is full. + error TreeIsFull(); + + /// @dev The storage size is invalid. + error InvalidStorageSize(); + + /// @dev `bytes4(keccak256(bytes("ValueAlreadyExists()")))`. + uint256 internal constant ERROR_VALUE_ALREADY_EXISTS = 0xbb33e6ac; + + /// @dev `bytes4(keccak256(bytes("ValueDoesNotExist()")))`. + uint256 internal constant ERROR_VALUE_DOES_NOT_EXISTS = 0xb113638a; + + /// @dev `bytes4(keccak256(bytes("PointerOutOfBounds()")))`. + uint256 internal constant ERROR_POINTER_OUT_OF_BOUNDS = 0xccd52fbc; + + /// @dev `bytes4(keccak256(bytes("TreeIsFull()")))`. + uint256 internal constant ERROR_TREE_IS_FULL = 0xed732d0c; + + /// @dev `bytes4(keccak256(bytes("InvalidStorageSize()")))`. + uint256 internal constant ERROR_INVALID_STORAGE_SIZE = 0x3cd9b236; + + /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/ + /* STRUCTS */ + /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/ + + /// @dev A red-black-tree in storage. + struct Tree { + uint256 _spacer; + } + + // Custom storage: + // ``` + // mstore(0x20, tree.slot) + // mstore(0x00, _NODES_SLOT_SEED) + // let nodes := shl(_NODES_SLOT_SHIFT, keccak256(0x00, 0x40)) + // + // let root := shr(128, sload(nodes)) + // let totalNodes := and(sload(nodes), _BITMASK_KEY) + // + // let nodePacked := sload(or(nodes, nodeIndex)) + // let nodeLeft := and(nodePacked, _BITMASK_KEY) + // let nodeRight := and(shr(_BITPOS_RIGHT, nodePacked), _BITMASK_KEY) + // let nodeParent := and(shr(_BITPOS_PARENT, nodePacked), _BITMASK_KEY) + // let nodeRed := and(shr(_BITPOS_RED, nodePacked), 1) + // + // let nodeValue := shr(_BITPOS_PACKED_VALUE, nodePacked) + // if iszero(nodeValue) { + // nodeValue := sload(or(_BIT_FULL_VALUE_SLOT, or(nodes, nodeIndex))) + // } + // ``` + // + // Bits Layout of the Root Index Slot: + // - [0..30] `totalNodes` (actual size) + // - [32..61] `storageCapacity` (allocated storage size) + // - [128..159] `rootNodeIndex` + // + // Bits Layout of a Node: + // - [0..30] `leftChildIndex` + // - [31..61] `rightChildIndex` + // - [62..92] `parentIndex` + // - [93] `isRed` + // - [96..255] `nodePackedValue` + + uint256 private constant _NODES_SLOT_SEED = 0x1dc27bb5462fdadcb; + uint256 private constant _NODES_SLOT_SHIFT = 32; + uint256 private constant _BITMASK_KEY = (1 << 31) - 1; + uint256 private constant _BITPOS_LEFT = 0; + uint256 private constant _BITPOS_RIGHT = 31; + uint256 private constant _BITPOS_PARENT = 31 * 2; + uint256 private constant _BITPOS_RED = 31 * 3; + uint256 private constant _BITMASK_RED = 1 << (31 * 3); + uint256 private constant _BITPOS_PACKED_VALUE = 96; + uint256 private constant _BITMASK_PACKED_VALUE = (1 << 160) - 1; + uint256 private constant _BIT_FULL_VALUE_SLOT = 1 << 31; + uint256 private constant _BITPOS_STORAGE_CAPACITY = 32; + + /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/ + /* OPERATIONS */ + /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/ + + /// @dev Returns the number of unique values in the tree. + function size(Tree storage tree) internal view returns (uint256 result) { + uint256 nodes = _nodes(tree); + /// @solidity memory-safe-assembly + assembly { + result := and(sload(nodes), _BITMASK_KEY) + } + } + + /// @dev Returns the allocated storage capacity of the tree. + /// This represents the pre-allocated storage size, which may be larger than the actual size. + function storageSize(Tree storage tree) internal view returns (uint256 result) { + uint256 nodes = _nodes(tree); + /// @solidity memory-safe-assembly + assembly { + let rootSlot := sload(nodes) + result := and(shr(_BITPOS_STORAGE_CAPACITY, rootSlot), _BITMASK_KEY) + } + } + + /// @dev Resizes the storage capacity of the tree. + /// This pre-allocates storage slots to save gas on future insertions. + /// The new capacity must be at least as large as the current size. + function resize(Tree storage tree, uint256 newCapacity) internal { + uint256 err = tryResize(tree, newCapacity); + if (err != 0) _revert(err); + } + + /// @dev Resizes the storage capacity of the tree. + /// Returns a non-zero error code upon failure instead of reverting. + function tryResize(Tree storage tree, uint256 newCapacity) internal returns (uint256 err) { + uint256 currentSize = size(tree); + uint256 oldCapacity = storageSize(tree); + if (newCapacity > _BITMASK_KEY) return ERROR_TREE_IS_FULL; + if (newCapacity < currentSize) return ERROR_INVALID_STORAGE_SIZE; + + uint256 nodes = _nodes(tree); + /// @solidity memory-safe-assembly + assembly { + let rootSlot := sload(nodes) + // Update storage capacity in root slot + let newRootSlot := + or( + and(rootSlot, not(shl(_BITPOS_STORAGE_CAPACITY, _BITMASK_KEY))), + shl(_BITPOS_STORAGE_CAPACITY, newCapacity) + ) + sstore(nodes, newRootSlot) + + // Pre-allocate storage slots for larger capacity + let start := add(currentSize, 1) + let end := newCapacity + for { let index := start } iszero(gt(index, end)) { index := add(index, 1) } { + sstore(or(nodes, index), 1) + sstore(or(_BIT_FULL_VALUE_SLOT, or(nodes, index)), 1) + } + + // Release storage slots for smaller capacity + start := add(newCapacity, 1) + end := oldCapacity + for { let index := start } iszero(gt(index, end)) { index := add(index, 1) } { + sstore(or(nodes, index), 0) + sstore(or(_BIT_FULL_VALUE_SLOT, or(nodes, index)), 0) + } + } + return 0; + } + + /// @dev Returns an array of all the values in the tree in ascending sorted order. + /// WARNING! This function can exhaust the block gas limit if the tree is big. + /// It is intended for usage in off-chain view functions. + function values(Tree storage tree) internal view returns (uint256[] memory result) { + uint256 nodes = _nodes(tree); + /// @solidity memory-safe-assembly + assembly { + function visit(current_) { + if iszero(current_) { leave } // If the current node is null, leave. + current_ := or(mload(0x00), current_) // Current node's storage slot. + let packed_ := sload(current_) + visit(and(packed_, _BITMASK_KEY)) // Visit left child. + let value_ := shr(_BITPOS_PACKED_VALUE, packed_) // Current value. + if iszero(value_) { value_ := sload(or(current_, _BIT_FULL_VALUE_SLOT)) } + mstore(mload(0x20), value_) // Append the value to `results`. + mstore(0x20, add(0x20, mload(0x20))) // Advance the offset into `results`. + visit(and(shr(_BITPOS_RIGHT, packed_), _BITMASK_KEY)) // Visit right child. + } + result := mload(0x40) + let rootPacked := sload(nodes) + mstore(result, and(rootPacked, _BITMASK_KEY)) // Length of `result`. + mstore(0x00, nodes) // Cache the nodes pointer in scratch space. + mstore(0x20, add(result, 0x20)) // Cache the offset into `results` in scratch space. + mstore(0x40, add(mload(0x20), shl(5, mload(result)))) // Allocate memory. + visit(shr(128, rootPacked)) // Start the tree traversal from the root node. + } + } + + /// @dev Returns a pointer to the value `x`. + /// If the value `x` is not in the tree, the returned pointer will be empty. + function find(Tree storage tree, uint256 x) internal view returns (bytes32 result) { + (uint256 nodes,, uint256 key) = _find(tree, x); + result = _pack(nodes, key); + } + + /// @dev Returns a pointer to the nearest value to `x`. + /// In a tie-breaker, the returned pointer will point to the smaller value. + /// If the tree is empty, the returned pointer will be empty. + function nearest(Tree storage tree, uint256 x) internal view returns (bytes32 result) { + (uint256 nodes, uint256 cursor, uint256 key) = _find(tree, x); + unchecked { + if (cursor == uint256(0)) return result; // Nothing found -- empty tree. + if (key != uint256(0)) return _pack(nodes, key); // Exact match. + bytes32 a = _pack(nodes, cursor); + uint256 aValue = value(a); + bytes32 b = x < aValue ? prev(a) : next(a); + if (b == bytes32(0)) return a; // Only node found. + uint256 bValue = value(b); + uint256 aDist = x < aValue ? aValue - x : x - aValue; + uint256 bDist = x < bValue ? bValue - x : x - bValue; + return (aDist == bDist ? aValue < bValue : aDist < bDist) ? a : b; + } + } + + /// @dev Returns a pointer to the nearest value lesser or equal to `x`. + /// If there is no value lesser or equal to `x`, the returned pointer will be empty. + function nearestBefore(Tree storage tree, uint256 x) internal view returns (bytes32 result) { + (uint256 nodes, uint256 cursor, uint256 key) = _find(tree, x); + if (cursor == uint256(0)) return result; // Nothing found -- empty tree. + if (key != uint256(0)) return _pack(nodes, key); // Exact match. + bytes32 a = _pack(nodes, cursor); + return value(a) < x ? a : prev(a); + } + + /// @dev Returns a pointer to the nearest value greater or equal to `x`. + /// If there is no value greater or equal to `x`, the returned pointer will be empty. + function nearestAfter(Tree storage tree, uint256 x) internal view returns (bytes32 result) { + (uint256 nodes, uint256 cursor, uint256 key) = _find(tree, x); + if (cursor == uint256(0)) return result; // Nothing found -- empty tree. + if (key != uint256(0)) return _pack(nodes, key); // Exact match. + bytes32 a = _pack(nodes, cursor); + return value(a) > x ? a : next(a); + } + + /// @dev Returns whether the value `x` exists. + function exists(Tree storage tree, uint256 x) internal view returns (bool result) { + (,, uint256 key) = _find(tree, x); + result = key != 0; + } + + /// @dev Inserts the value `x` into the tree. + /// Reverts if the value `x` already exists. + function insert(Tree storage tree, uint256 x) internal { + uint256 err = tryInsert(tree, x); + if (err != 0) _revert(err); + } + + /// @dev Inserts the value `x` into the tree. + /// Returns a non-zero error code upon failure instead of reverting + /// (except for reverting if `x` is an empty value). + function tryInsert(Tree storage tree, uint256 x) internal returns (uint256 err) { + (uint256 nodes, uint256 cursor, uint256 key) = _find(tree, x); + err = _update(nodes, cursor, key, x, 0); + } + + /// @dev Removes the value `x` from the tree. + /// Reverts if the value does not exist. + function remove(Tree storage tree, uint256 x) internal { + uint256 err = tryRemove(tree, x); + if (err != 0) _revert(err); + } + + /// @dev Removes the value `x` from the tree. + /// Returns a non-zero error code upon failure instead of reverting + /// (except for reverting if `x` is an empty value). + function tryRemove(Tree storage tree, uint256 x) internal returns (uint256 err) { + (uint256 nodes,, uint256 key) = _find(tree, x); + err = _update(nodes, 0, key, 0, 1); + } + + /// @dev Removes the value at pointer `ptr` from the tree. + /// Reverts if `ptr` is empty (i.e. value does not exist), + /// or if `ptr` is out of bounds. + /// After removal, `ptr` may point to another existing value. + /// For safety, do not reuse `ptr` after calling remove on it. + function remove(bytes32 ptr) internal { + uint256 err = tryRemove(ptr); + if (err != 0) _revert(err); + } + + /// @dev Removes the value at pointer `ptr` from the tree. + /// Returns a non-zero error code upon failure instead of reverting. + function tryRemove(bytes32 ptr) internal returns (uint256 err) { + (uint256 nodes, uint256 key) = _unpack(ptr); + err = _update(nodes, 0, key, 0, 1); + } + + /// @dev Returns the value at pointer `ptr`. + /// If `ptr` is empty, the result will be zero. + function value(bytes32 ptr) internal view returns (uint256 result) { + if (ptr == bytes32(0)) return result; + /// @solidity memory-safe-assembly + assembly { + let packed := sload(ptr) + result := shr(_BITPOS_PACKED_VALUE, packed) + if iszero(result) { result := sload(or(ptr, _BIT_FULL_VALUE_SLOT)) } + } + } + + /// @dev Returns a pointer to the smallest value in the tree. + /// If the tree is empty, the returned pointer will be empty. + function first(Tree storage tree) internal view returns (bytes32 result) { + result = _end(tree, _BITPOS_LEFT); + } + + /// @dev Returns a pointer to the largest value in the tree. + /// If the tree is empty, the returned pointer will be empty. + function last(Tree storage tree) internal view returns (bytes32 result) { + result = _end(tree, _BITPOS_RIGHT); + } + + /// @dev Returns the pointer to the next largest value. + /// If there is no next value, or if `ptr` is empty, + /// the returned pointer will be empty. + function next(bytes32 ptr) internal view returns (bytes32 result) { + result = _step(ptr, _BITPOS_LEFT, _BITPOS_RIGHT); + } + + /// @dev Returns the pointer to the next smallest value. + /// If there is no previous value, or if `ptr` is empty, + /// the returned pointer will be empty. + function prev(bytes32 ptr) internal view returns (bytes32 result) { + result = _step(ptr, _BITPOS_RIGHT, _BITPOS_LEFT); + } + + /// @dev Returns whether the pointer is empty. + function isEmpty(bytes32 ptr) internal pure returns (bool result) { + result = ptr == bytes32(0); + } + + /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/ + /* PRIVATE HELPERS */ + /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/ + + /// @dev Unpacks the pointer `ptr` to its components. + function _unpack(bytes32 ptr) private pure returns (uint256 nodes, uint256 key) { + /// @solidity memory-safe-assembly + assembly { + nodes := shl(_NODES_SLOT_SHIFT, shr(_NODES_SLOT_SHIFT, ptr)) + key := and(_BITMASK_KEY, ptr) + } + } + + /// @dev Packs `nodes` and `key` into a single pointer. + function _pack(uint256 nodes, uint256 key) private pure returns (bytes32 result) { + /// @solidity memory-safe-assembly + assembly { + result := mul(or(nodes, key), iszero(iszero(key))) + } + } + + /// @dev Returns the pointer to either end of the tree. + function _end(Tree storage tree, uint256 L) private view returns (bytes32 result) { + uint256 nodes = _nodes(tree); + /// @solidity memory-safe-assembly + assembly { + result := shr(128, sload(nodes)) + if result { + for {} 1 {} { + let packed := sload(or(nodes, result)) + let left := and(shr(L, packed), _BITMASK_KEY) + if iszero(left) { break } + result := left + } + } + } + result = _pack(nodes, uint256(result)); + } + + /// @dev Step the pointer `ptr` forwards or backwards. + function _step(bytes32 ptr, uint256 L, uint256 R) private view returns (bytes32 result) { + if (ptr == bytes32(0)) return ptr; + (uint256 nodes, uint256 target) = _unpack(ptr); + /// @solidity memory-safe-assembly + assembly { + let packed := sload(ptr) + for { result := and(shr(R, packed), _BITMASK_KEY) } 1 {} { + if iszero(result) { + result := and(shr(_BITPOS_PARENT, packed), _BITMASK_KEY) + for {} 1 {} { + if iszero(result) { break } + packed := sload(or(nodes, result)) + if iszero(eq(target, and(shr(R, packed), _BITMASK_KEY))) { break } + target := result + result := and(shr(_BITPOS_PARENT, packed), _BITMASK_KEY) + } + break + } + for {} 1 {} { + packed := sload(or(nodes, result)) + let left := and(shr(L, packed), _BITMASK_KEY) + if iszero(left) { break } + result := left + } + break + } + } + result = _pack(nodes, uint256(result)); + } + + /// @dev Inserts or delete the value `x` from the tree. + function _update(uint256 nodes, uint256 cursor, uint256 key, uint256 x, uint256 mode) + private + returns (uint256 err) + { + /// @solidity memory-safe-assembly + assembly { + function getKey(packed_, bitpos_) -> index_ { + index_ := and(_BITMASK_KEY, shr(bitpos_, packed_)) + } + + function setKey(packed_, bitpos_, key_) -> result_ { + result_ := or(and(not(shl(bitpos_, _BITMASK_KEY)), packed_), shl(bitpos_, key_)) + } + + function rotate(nodes_, key_, L, R) { + let packed_ := sload(or(nodes_, key_)) + let cursor_ := getKey(packed_, R) + let parent_ := getKey(packed_, _BITPOS_PARENT) + let cursorPacked_ := sload(or(nodes_, cursor_)) + let cursorLeft_ := getKey(cursorPacked_, L) + + if cursorLeft_ { + let s_ := or(nodes_, cursorLeft_) + sstore(s_, setKey(sload(s_), _BITPOS_PARENT, key_)) + } + + for {} 1 {} { + if iszero(parent_) { + mstore(0x00, cursor_) + break + } + let s_ := or(nodes_, parent_) + let parentPacked_ := sload(s_) + if eq(key_, getKey(parentPacked_, L)) { + sstore(s_, setKey(parentPacked_, L, cursor_)) + break + } + sstore(s_, setKey(parentPacked_, R, cursor_)) + break + } + packed_ := setKey(packed_, R, cursorLeft_) + sstore(or(nodes_, key_), setKey(packed_, _BITPOS_PARENT, cursor_)) + cursorPacked_ := setKey(cursorPacked_, _BITPOS_PARENT, parent_) + sstore(or(nodes_, cursor_), setKey(cursorPacked_, L, key_)) + } + + function insert(nodes_, cursor_, key_, x_) -> err_ { + if key_ { + err_ := ERROR_VALUE_ALREADY_EXISTS + leave + } + + let rootData := mload(0x20) + let totalNodes_ := and(shr(128, rootData), _BITMASK_KEY) + let storageCapacity_ := and(shr(_BITPOS_STORAGE_CAPACITY, shr(128, rootData)), _BITMASK_KEY) + totalNodes_ := add(totalNodes_, 1) + if gt(totalNodes_, _BITMASK_KEY) { + err_ := ERROR_TREE_IS_FULL + leave + } + + // Check if we need to auto-increment storage capacity + // If storageSize < size, increment storageSize to maintain storageSize >= size constraint + if lt(storageCapacity_, totalNodes_) { + storageCapacity_ := totalNodes_ + if gt(storageCapacity_, _BITMASK_KEY) { + err_ := ERROR_TREE_IS_FULL + leave + } + } + + let newData := or(shl(128, totalNodes_), shl(_BITPOS_STORAGE_CAPACITY, shl(128, storageCapacity_))) + mstore(0x20, newData) + + { + let packed_ := or(_BITMASK_RED, shl(_BITPOS_PARENT, cursor_)) + let nodePointer_ := or(nodes_, totalNodes_) + + for {} 1 {} { + if iszero(gt(x_, _BITMASK_PACKED_VALUE)) { + packed_ := or(shl(_BITPOS_PACKED_VALUE, x_), packed_) + break + } + sstore(or(nodePointer_, _BIT_FULL_VALUE_SLOT), x_) + break + } + sstore(nodePointer_, packed_) + + for {} 1 {} { + if iszero(cursor_) { + mstore(0x00, totalNodes_) + break + } + let s_ := or(nodes_, cursor_) + let cPacked_ := sload(s_) + let cValue_ := shr(_BITPOS_PACKED_VALUE, cPacked_) + if iszero(cValue_) { cValue_ := sload(or(s_, _BIT_FULL_VALUE_SLOT)) } + if iszero(lt(x_, cValue_)) { + sstore(s_, setKey(cPacked_, _BITPOS_RIGHT, totalNodes_)) + break + } + sstore(s_, setKey(cPacked_, _BITPOS_LEFT, totalNodes_)) + break + } + } + + // Insert fixup workflow: + + key_ := totalNodes_ + let BR := _BITMASK_RED + for {} iszero(eq(key_, mload(0x00))) {} { + let packed_ := sload(or(nodes_, key_)) + let parent_ := getKey(packed_, _BITPOS_PARENT) + let parentPacked_ := sload(or(nodes_, parent_)) + if iszero(and(BR, parentPacked_)) { break } + + let grandParent_ := getKey(parentPacked_, _BITPOS_PARENT) + let grandParentPacked_ := sload(or(nodes_, grandParent_)) + + let R := mul(eq(parent_, getKey(grandParentPacked_, 0)), _BITPOS_RIGHT) + let L := xor(R, _BITPOS_RIGHT) + + let c_ := getKey(grandParentPacked_, R) + let cPacked_ := sload(or(nodes_, c_)) + if iszero(and(BR, cPacked_)) { + if eq(key_, getKey(parentPacked_, R)) { + key_ := parent_ + rotate(nodes_, key_, L, R) + parent_ := getKey(sload(or(nodes_, key_)), _BITPOS_PARENT) + parentPacked_ := sload(or(nodes_, parent_)) + } + sstore(or(nodes_, parent_), and(parentPacked_, not(BR))) + let s_ := or(nodes_, grandParent_) + sstore(s_, or(sload(s_), BR)) + rotate(nodes_, grandParent_, R, L) + break + } + sstore(or(nodes_, parent_), and(parentPacked_, not(BR))) + sstore(or(nodes_, c_), and(cPacked_, not(BR))) + sstore(or(nodes_, grandParent_), or(grandParentPacked_, BR)) + key_ := grandParent_ + } + let root_ := or(nodes_, mload(0x00)) + sstore(root_, and(sload(root_), not(BR))) + } + + function removeFixup(nodes_, key_) { + let BR := _BITMASK_RED + for {} iszero(eq(key_, mload(0x00))) {} { + let packed_ := sload(or(nodes_, key_)) + if and(BR, packed_) { break } + + let parent_ := getKey(packed_, _BITPOS_PARENT) + let parentPacked_ := sload(or(nodes_, parent_)) + + let R := mul(eq(key_, getKey(parentPacked_, 0)), _BITPOS_RIGHT) + let L := xor(R, _BITPOS_RIGHT) + + let cursor_ := getKey(parentPacked_, R) + let cursorPacked_ := sload(or(nodes_, cursor_)) + + if and(BR, cursorPacked_) { + sstore(or(nodes_, cursor_), and(cursorPacked_, not(BR))) + sstore(or(nodes_, parent_), or(parentPacked_, BR)) + rotate(nodes_, parent_, L, R) + cursor_ := getKey(sload(or(nodes_, parent_)), R) + cursorPacked_ := sload(or(nodes_, cursor_)) + } + + let cursorLeft_ := getKey(cursorPacked_, L) + let cursorLeftPacked_ := sload(or(nodes_, cursorLeft_)) + let cursorRight_ := getKey(cursorPacked_, R) + let cursorRightPacked_ := sload(or(nodes_, cursorRight_)) + + if iszero(and(BR, or(cursorLeftPacked_, cursorRightPacked_))) { + sstore(or(nodes_, cursor_), or(cursorPacked_, BR)) + key_ := parent_ + continue + } + + if iszero(and(BR, cursorRightPacked_)) { + sstore(or(nodes_, cursorLeft_), and(cursorLeftPacked_, not(BR))) + sstore(or(nodes_, cursor_), or(cursorPacked_, BR)) + rotate(nodes_, cursor_, R, L) + cursor_ := getKey(sload(or(nodes_, parent_)), R) + cursorPacked_ := sload(or(nodes_, cursor_)) + cursorRight_ := getKey(cursorPacked_, R) + cursorRightPacked_ := sload(or(nodes_, cursorRight_)) + } + + parentPacked_ := sload(or(nodes_, parent_)) + // forgefmt: disable-next-item + sstore(or(nodes_, cursor_), xor(cursorPacked_, and(BR, xor(cursorPacked_, parentPacked_)))) + sstore(or(nodes_, parent_), and(parentPacked_, not(BR))) + sstore(or(nodes_, cursorRight_), and(cursorRightPacked_, not(BR))) + rotate(nodes_, parent_, L, R) + break + } + sstore(or(nodes_, key_), and(sload(or(nodes_, key_)), not(BR))) + } + + function replaceParent(nodes_, parent_, a_, b_) { + if iszero(parent_) { + mstore(0x00, a_) + leave + } + let s_ := or(nodes_, parent_) + let p_ := sload(s_) + let t_ := iszero(eq(b_, getKey(p_, _BITPOS_LEFT))) + sstore(s_, setKey(p_, mul(t_, _BITPOS_RIGHT), a_)) + } + + // In `remove`, the parent of the null value (index 0) may be temporarily set + // to a non-zero value. This is an optimization that unifies the removal cases. + function remove(nodes_, key_) -> err_ { + if gt(key_, and(shr(128, mload(0x20)), _BITMASK_KEY)) { + err_ := ERROR_POINTER_OUT_OF_BOUNDS + leave + } + if iszero(key_) { + err_ := ERROR_VALUE_DOES_NOT_EXISTS + leave + } + + let cursor_ := key_ + { + let packed_ := sload(or(nodes_, key_)) + let left_ := getKey(packed_, _BITPOS_LEFT) + let right_ := getKey(packed_, _BITPOS_RIGHT) + if mul(left_, right_) { + for { cursor_ := right_ } 1 {} { + let cursorLeft_ := getKey(sload(or(nodes_, cursor_)), _BITPOS_LEFT) + if iszero(cursorLeft_) { break } + cursor_ := cursorLeft_ + } + } + } + + let cursorPacked_ := sload(or(nodes_, cursor_)) + let probe_ := getKey(cursorPacked_, _BITPOS_LEFT) + probe_ := getKey(cursorPacked_, mul(iszero(probe_), _BITPOS_RIGHT)) + + let yParent_ := getKey(cursorPacked_, _BITPOS_PARENT) + let probeSlot_ := or(nodes_, probe_) + sstore(probeSlot_, setKey(sload(probeSlot_), _BITPOS_PARENT, yParent_)) + replaceParent(nodes_, yParent_, probe_, cursor_) + + if iszero(eq(cursor_, key_)) { + let packed_ := sload(or(nodes_, key_)) + replaceParent(nodes_, getKey(packed_, _BITPOS_PARENT), cursor_, key_) + + let leftSlot_ := or(nodes_, getKey(packed_, _BITPOS_LEFT)) + sstore(leftSlot_, setKey(sload(leftSlot_), _BITPOS_PARENT, cursor_)) + + let rightSlot_ := or(nodes_, getKey(packed_, _BITPOS_RIGHT)) + sstore(rightSlot_, setKey(sload(rightSlot_), _BITPOS_PARENT, cursor_)) + + // Copy `left`, `right`, `red` from `key_` to `cursor_`. + // forgefmt: disable-next-item + sstore(or(nodes_, cursor_), xor(cursorPacked_, + and(xor(packed_, cursorPacked_), sub(shl(_BITPOS_PACKED_VALUE, 1), 1)))) + + let t_ := cursor_ + cursor_ := key_ + key_ := t_ + } + + if iszero(and(_BITMASK_RED, cursorPacked_)) { removeFixup(nodes_, probe_) } + + // Remove last workflow: + + let last_ := and(shr(128, mload(0x20)), _BITMASK_KEY) + let capacity_ := and(shr(_BITPOS_STORAGE_CAPACITY, shr(128, mload(0x20))), _BITMASK_KEY) + let lastPacked_ := sload(or(nodes_, last_)) + let lastValue_ := shr(_BITPOS_PACKED_VALUE, lastPacked_) + let lastFullValue_ := 0 + if iszero(lastValue_) { + lastValue_ := sload(or(_BIT_FULL_VALUE_SLOT, or(nodes_, last_))) + lastFullValue_ := lastValue_ + } + + let cursorValue_ := shr(_BITPOS_PACKED_VALUE, sload(or(nodes_, cursor_))) + let cursorFullValue_ := 0 + if iszero(cursorValue_) { + cursorValue_ := sload(or(_BIT_FULL_VALUE_SLOT, or(nodes_, cursor_))) + cursorFullValue_ := cursorValue_ + } + + if iszero(eq(lastValue_, cursorValue_)) { + sstore(or(nodes_, cursor_), lastPacked_) + if iszero(eq(lastFullValue_, cursorFullValue_)) { + sstore(or(_BIT_FULL_VALUE_SLOT, or(nodes_, cursor_)), lastFullValue_) + } + for { let lastParent_ := getKey(lastPacked_, _BITPOS_PARENT) } 1 {} { + if iszero(lastParent_) { + mstore(0x00, cursor_) + break + } + let s_ := or(nodes_, lastParent_) + let p_ := sload(s_) + let t_ := iszero(eq(last_, getKey(p_, _BITPOS_LEFT))) + sstore(s_, setKey(p_, mul(t_, _BITPOS_RIGHT), cursor_)) + break + } + let lastRight_ := getKey(lastPacked_, _BITPOS_RIGHT) + if lastRight_ { + let s_ := or(nodes_, lastRight_) + sstore(s_, setKey(sload(s_), _BITPOS_PARENT, cursor_)) + } + let lastLeft_ := getKey(lastPacked_, _BITPOS_LEFT) + if lastLeft_ { + let s_ := or(nodes_, lastLeft_) + sstore(s_, setKey(sload(s_), _BITPOS_PARENT, cursor_)) + } + } + // save both capacity and total nodes + let newData := or(shl(128, sub(last_, 1)), shl(_BITPOS_STORAGE_CAPACITY, shl(128, capacity_))) + mstore(0x20, newData) + } + + mstore(0x00, codesize()) // Zeroize the first 0x10 bytes. + mstore(0x10, sload(nodes)) + + for {} 1 {} { + if iszero(mode) { + err := insert(nodes, cursor, key, x) + break + } + err := remove(nodes, key) + break + } + + sstore(nodes, mload(0x10)) + } + } + + /// @dev Returns the pointer to the `nodes` for the tree. + function _nodes(Tree storage tree) private pure returns (uint256 nodes) { + /// @solidity memory-safe-assembly + assembly { + mstore(0x20, tree.slot) + mstore(0x00, _NODES_SLOT_SEED) + nodes := shl(_NODES_SLOT_SHIFT, keccak256(0x00, 0x40)) + } + } + + /// @dev Finds `x` in `tree`. The `key` will be zero if `x` is not found. + function _find(Tree storage tree, uint256 x) private view returns (uint256 nodes, uint256 cursor, uint256 key) { + if (x == uint256(0)) _revert(0xc94f1877); // `ValueIsEmpty()`. + /// @solidity memory-safe-assembly + assembly { + mstore(0x20, tree.slot) + mstore(0x00, _NODES_SLOT_SEED) + nodes := shl(_NODES_SLOT_SHIFT, keccak256(0x00, 0x40)) + // Layout scratch space so that `mload(0x00) == 0`, `mload(0x01) == _BITPOS_RIGHT`. + mstore(0x01, _BITPOS_RIGHT) // `_BITPOS_RIGHT` is 31. + for { let probe := shr(128, sload(nodes)) } probe {} { + cursor := probe + let nodePacked := sload(or(nodes, probe)) + let nodeValue := shr(_BITPOS_PACKED_VALUE, nodePacked) + if iszero(nodeValue) { nodeValue := sload(or(or(nodes, probe), _BIT_FULL_VALUE_SLOT)) } + if eq(nodeValue, x) { + key := cursor + break + } + probe := and(shr(mload(gt(x, nodeValue)), nodePacked), _BITMASK_KEY) + } + } + } + + /// @dev Helper to revert `err` efficiently. + function _revert(uint256 err) private pure { + /// @solidity memory-safe-assembly + assembly { + mstore(0x00, err) + revert(0x1c, 0x04) + } + } +} diff --git a/test/RedBlackTreeKVGas.t.sol b/test/RedBlackTreeKVGas.t.sol index dd89c34..b93be15 100644 --- a/test/RedBlackTreeKVGas.t.sol +++ b/test/RedBlackTreeKVGas.t.sol @@ -2,16 +2,19 @@ pragma solidity ^0.8.30; import {Test, console} from "forge-std/Test.sol"; +import {RedBlackTreeWithResizeKV} from "../src/RedBlackTreeWithResizeKV.sol"; import {RedBlackTreeKV} from "../src/RedBlackTreeKV.sol"; import {MappingKV} from "../src/MappingKV.sol"; import {ValueLib} from "../src/lib/Value.sol"; contract MappingGasTest is Test { uint256 private constant _INSERT_COUNT = 1000; + RedBlackTreeWithResizeKV public redBlockTreeWithResizeKV; RedBlackTreeKV public redBlockTreeKV; MappingKV public mappingKV; function setUp() public { + redBlockTreeWithResizeKV = new RedBlackTreeWithResizeKV(); redBlockTreeKV = new RedBlackTreeKV(); mappingKV = new MappingKV(); } @@ -208,4 +211,93 @@ contract MappingGasTest is Test { assertEq(value.slot6, address(uint160(i + 1))); } } + + function test_rbt_with_resize_fully_insert() public { + redBlockTreeWithResizeKV.setValue(uint256(keccak256(abi.encodePacked("valueMAX"))), generateValue(UINT256_MAX)); + redBlockTreeWithResizeKV.setValue( + uint256(keccak256(abi.encodePacked("valueMAX-1"))), generateValue(UINT256_MAX - 1) + ); + redBlockTreeWithResizeKV.deleteValue(uint256(keccak256(abi.encodePacked("valueMAX")))); + uint256 gasBefore; + uint256 gasAfter; + uint256 totalColdInsertGas = 0; + uint256 totalHotInsertGas = 0; + uint256 totalDeleteGas = 0; + // each loop: insert 2 values and delete 1, first insert is hot, second is cold + for (uint256 i = 1; i <= _INSERT_COUNT * 2; i += 2) { + uint256 orderId1 = uint256(keccak256(abi.encodePacked("value", i))); + uint256 orderId2 = uint256(keccak256(abi.encodePacked("value", i + 1))); + gasBefore = gasleft(); + redBlockTreeWithResizeKV.setValue(orderId1, generateValue(i)); + gasAfter = gasleft(); + totalHotInsertGas += gasBefore - gasAfter; + gasBefore = gasleft(); + redBlockTreeWithResizeKV.setValue(orderId2, generateValue(i + 1)); + gasAfter = gasleft(); + totalColdInsertGas += gasBefore - gasAfter; + gasBefore = gasleft(); + redBlockTreeWithResizeKV.deleteValue(orderId1); + gasAfter = gasleft(); + totalDeleteGas += gasBefore - gasAfter; + } + console.log("_INSERT_COUNT", _INSERT_COUNT); + console.log("avgColdInsertGas", totalColdInsertGas / _INSERT_COUNT); + console.log("avgHotInsertGas", totalHotInsertGas / _INSERT_COUNT); + console.log("avgDeleteGas", totalDeleteGas / _INSERT_COUNT); + for (uint256 i = 1; i <= _INSERT_COUNT * 2; i += 2) { + uint256 orderId2 = uint256(keccak256(abi.encodePacked("value", i + 1))); + ValueLib.Value memory value = redBlockTreeWithResizeKV.getValue(orderId2); + assertEq(value.slot1, bytes32(uint256(i + 1))); + assertEq(value.slot2, uint256(i + 1)); + assertEq(value.slot3, int256(i + 1)); + assertEq(value.slot4, bytes32(uint256(i + 1))); + assertEq(value.slot5, uint256(i + 1)); + assertEq(value.slot6, address(uint160(i + 1))); + } + } + + function test_rbt_with_resize_fully_insert_hot() public { + redBlockTreeWithResizeKV.setValue(uint256(keccak256(abi.encodePacked("valueMAX"))), generateValue(UINT256_MAX)); + redBlockTreeWithResizeKV.setValue( + uint256(keccak256(abi.encodePacked("valueMAX-1"))), generateValue(UINT256_MAX - 1) + ); + redBlockTreeWithResizeKV.deleteValue(uint256(keccak256(abi.encodePacked("valueMAX")))); + redBlockTreeWithResizeKV.resize(1000); + uint256 gasBefore; + uint256 gasAfter; + uint256 totalColdInsertGas = 0; + uint256 totalHotInsertGas = 0; + uint256 totalDeleteGas = 0; + // each loop: insert 2 values and delete 1, first insert is hot, second is cold + for (uint256 i = 1; i <= _INSERT_COUNT * 2; i += 2) { + uint256 orderId1 = uint256(keccak256(abi.encodePacked("value", i))); + uint256 orderId2 = uint256(keccak256(abi.encodePacked("value", i + 1))); + gasBefore = gasleft(); + redBlockTreeWithResizeKV.setValue(orderId1, generateValue(i)); + gasAfter = gasleft(); + totalHotInsertGas += gasBefore - gasAfter; + gasBefore = gasleft(); + redBlockTreeWithResizeKV.setValue(orderId2, generateValue(i + 1)); + gasAfter = gasleft(); + totalColdInsertGas += gasBefore - gasAfter; + gasBefore = gasleft(); + redBlockTreeWithResizeKV.deleteValue(orderId1); + gasAfter = gasleft(); + totalDeleteGas += gasBefore - gasAfter; + } + console.log("_INSERT_COUNT", _INSERT_COUNT); + console.log("avgColdInsertGas", totalColdInsertGas / _INSERT_COUNT); + console.log("avgHotInsertGas", totalHotInsertGas / _INSERT_COUNT); + console.log("avgDeleteGas", totalDeleteGas / _INSERT_COUNT); + for (uint256 i = 1; i <= _INSERT_COUNT * 2; i += 2) { + uint256 orderId2 = uint256(keccak256(abi.encodePacked("value", i + 1))); + ValueLib.Value memory value = redBlockTreeWithResizeKV.getValue(orderId2); + assertEq(value.slot1, bytes32(uint256(i + 1))); + assertEq(value.slot2, uint256(i + 1)); + assertEq(value.slot3, int256(i + 1)); + assertEq(value.slot4, bytes32(uint256(i + 1))); + assertEq(value.slot5, uint256(i + 1)); + assertEq(value.slot6, address(uint160(i + 1))); + } + } } diff --git a/test/RedBlackTreeWithResize.sol b/test/RedBlackTreeWithResize.sol new file mode 100644 index 0000000..5258497 --- /dev/null +++ b/test/RedBlackTreeWithResize.sol @@ -0,0 +1,743 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.4; + +import {console} from "forge-std/Test.sol"; +import "./utils/SoladyTest.sol"; +import {LibSort} from "../src/lib/LibSort.sol"; +import {LibPRNG} from "../src/lib/LibPRNG.sol"; +import {RedBlackTreeWithResizeLib} from "../src/lib/RedBlackTreeWithResizeLib.sol"; + +contract RedBlackTreeWithResizeLibTest is SoladyTest { + using RedBlackTreeWithResizeLib for *; + using LibPRNG for *; + + RedBlackTreeWithResizeLib.Tree tree; + RedBlackTreeWithResizeLib.Tree tree2; + + function testRedBlackTreeInsertBenchStep() public { + unchecked { + LibPRNG.PRNG memory prng = LibPRNG.PRNG(123); + uint256 n = 128; + uint256 m = (1 << 160) - 1; + for (uint256 i; i != n; ++i) { + uint256 r = 1 | (prng.next() & m); + tree.insert(r); + } + _testIterateTree(); + } + } + + function testRedBlackTreeInsertBenchUint160() public { + unchecked { + LibPRNG.PRNG memory prng = LibPRNG.PRNG(123); + uint256 n = 128; + uint256[] memory a = _makeArray(n); + uint256 m = (1 << 160) - 1; + for (uint256 i; i != n; ++i) { + uint256 r = 1 | (prng.next() & m); + a[i] = r; + tree.insert(r); + } + } + } + + function testRedBlackTreeBenchUint160() public { + unchecked { + LibPRNG.PRNG memory prng = LibPRNG.PRNG(123); + uint256 n = 128; + uint256[] memory a = _makeArray(n); + uint256 m = (1 << 160) - 1; + for (uint256 i; i != n; ++i) { + uint256 r = 1 | (prng.next() & m); + a[i] = r; + tree.insert(r); + } + prng.shuffle(a); + for (uint256 i; i != n; ++i) { + tree.remove(a[i]); + } + assertEq(tree.size(), 0); + } + } + + function testRedBlackTreeInsertBenchUint256() public { + unchecked { + LibPRNG.PRNG memory prng = LibPRNG.PRNG(123); + uint256 n = 128; + uint256[] memory a = _makeArray(n); + for (uint256 i; i != n; ++i) { + uint256 r = 1 | prng.next(); + a[i] = r; + tree.insert(r); + } + } + } + + function testRedBlackTreeBenchUint256() public { + unchecked { + LibPRNG.PRNG memory prng = LibPRNG.PRNG(123); + uint256 n = 128; + uint256[] memory a = _makeArray(n); + for (uint256 i; i != n; ++i) { + uint256 r = 1 | prng.next(); + a[i] = r; + tree.insert(r); + } + prng.shuffle(a); + for (uint256 i; i != n; ++i) { + tree.remove(a[i]); + } + assertEq(tree.size(), 0); + } + } + + function testRedBlackTreeInsertAndRemove(uint256) public { + unchecked { + for (uint256 t; t < 2; ++t) { + _testRedBlackTreeInsertAndRemove(); + } + } + } + + function _testRemoveAndInsertBack(uint256[] memory a, uint256 n, uint256 t) internal { + unchecked { + uint256 choice = a[_random() % n]; + bytes32 ptr = tree.find(choice); + bool exists = !ptr.isEmpty(); + if (exists) { + assertEq(ptr.value(), choice); + _brutalizeScratchSpace(); + ptr.remove(); + if (_randomChance(4)) { + _brutalizeScratchSpace(); + tree.tryRemove(choice); + } + assertTrue(tree.find(choice).isEmpty()); + assertFalse(tree.exists(choice)); + } + if (t != 0) { + _testRemoveAndInsertBack(a, n, t - 1); + } + if (exists) { + _brutalizeScratchSpace(); + tree.insert(choice); + if (_randomChance(4)) { + _brutalizeScratchSpace(); + tree.tryInsert(choice); + } + assertFalse(tree.find(choice).isEmpty()); + assertTrue(tree.exists(choice)); + } + } + } + + function _testIterateTree() internal view { + bytes32 ptr = tree.first(); + uint256 prevValue; + while (!ptr.isEmpty()) { + uint256 v = ptr.value(); + assertTrue(prevValue < v); + prevValue = v; + ptr = ptr.next(); + } + assertEq(ptr.next().value(), 0); + + ptr = tree.last(); + prevValue = 0; + while (!ptr.isEmpty()) { + uint256 v = ptr.value(); + assertTrue(prevValue == 0 || prevValue > v); + prevValue = v; + ptr = ptr.prev(); + } + assertEq(ptr.prev().value(), 0); + } + + function _testRedBlackTreeInsertAndRemove() internal { + uint256 n = _random() % (_randomChance(128) ? 32 : 8); + uint256[] memory a = _fillTree(n); + + LibSort.sort(a); + LibSort.uniquifySorted(a); + assertEq(a.length, n); + assertEq(tree.size(), n); + + assertEq(tree2.size(), 0); + + unchecked { + uint256 i; + bytes32 ptr = tree.first(); + while (!ptr.isEmpty()) { + assertEq(a[i++], ptr.value()); + ptr = ptr.next(); + } + assertEq(ptr.next().value(), 0); + } + + unchecked { + uint256 i = n; + bytes32 ptr = tree.last(); + while (!ptr.isEmpty()) { + assertEq(a[--i], ptr.value()); + ptr = ptr.prev(); + } + assertEq(ptr.prev().value(), 0); + } + + _testIterateTree(); + + LibPRNG.PRNG memory prng = LibPRNG.PRNG(_random()); + prng.shuffle(a); + + unchecked { + uint256 m = n < 8 ? 4 : n; + for (uint256 i; i != n; ++i) { + _brutalizeScratchSpace(); + tree.remove(a[i]); + assertEq(tree.size(), n - i - 1); + if (_random() % m == 0) { + _testIterateTree(); + } + } + } + assertEq(tree.size(), 0); + + unchecked { + if (_randomChance(2)) { + for (uint256 i; i != n; ++i) { + assertTrue(tree.find(a[i]).isEmpty()); + } + } + assertTrue(tree.first().isEmpty()); + assertEq(tree.first().value(), 0); + assertTrue(tree.last().isEmpty()); + assertEq(tree.last().value(), 0); + } + + assertEq(tree2.size(), 0); + } + + function testRedBlackTreeInsertAndRemove2(uint256) public { + unchecked { + uint256 n = _randomChance(2) ? 16 : 32; + uint256[] memory candidates = _makeArray(n); + for (uint256 i; i != n; ++i) { + candidates[i] = _bound(_random(), 1, type(uint256).max); + } + uint256[] memory records = _makeArray(0); + uint256 mode = 0; + for (uint256 t = _random() % 32 + 1; t != 0; --t) { + uint256 r = candidates[_random() % n]; + bytes32 ptr = tree.find(r); + if (mode == 0) { + if (ptr.isEmpty()) { + _brutalizeScratchSpace(); + tree.insert(r); + _addToArray(records, r); + } + } else { + if (!ptr.isEmpty()) { + _brutalizeScratchSpace(); + tree.remove(r); + _removeFromArray(records, r); + } + } + if (_randomChance(3)) mode = _random() % 2; + } + LibSort.sort(records); + assertEq(tree.size(), records.length); + + assertEq(tree2.size(), 0); + + { + uint256 i = 0; + bytes32 ptr = tree.first(); + while (!ptr.isEmpty()) { + assertEq(records[i++], ptr.value()); + ptr = ptr.next(); + } + assertEq(ptr.next().value(), 0); + } + } + } + + 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(mload(o), 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(mload(o), x) { + mstore(o, mload(add(a, shl(5, n)))) + mstore(a, sub(n, 1)) + break + } + } + } + } + + function testRedBlackTreeInsertAndRemove3() public { + unchecked { + uint256 m = type(uint256).max; + for (uint256 i; i < 256; ++i) { + _brutalizeScratchSpace(); + tree.insert(m - i); + assertEq(tree.size(), i + 1); + } + for (uint256 i; i < 256; ++i) { + tree2.insert(i + 1); + assertEq(tree2.size(), i + 1); + } + for (uint256 i; i < 256; ++i) { + assertTrue(tree.exists(m - i)); + assertFalse(tree.exists(i + 1)); + assertTrue(tree2.exists(i + 1)); + assertFalse(tree2.exists(m - i)); + } + bytes32[] memory ptrs = new bytes32[](256); + for (uint256 i; i < 256; ++i) { + bytes32 ptr = tree.find(m - i); + _brutalizeScratchSpace(); + ptr.remove(); + // assertTrue(ptr.value() != m - i); + ptrs[i] = ptr; + assertEq(tree.size(), 256 - (i + 1)); + } + for (uint256 i; i < 256; ++i) { + // assertEq(ptrs[i].value(), 0); + vm.expectRevert(RedBlackTreeWithResizeLib.PointerOutOfBounds.selector); + _brutalizeScratchSpace(); + this.remove(ptrs[i]); + } + for (uint256 i; i < 256; ++i) { + _brutalizeScratchSpace(); + tree2.remove(i + 1); + assertEq(tree2.size(), 256 - (i + 1)); + } + } + } + + function find(uint256 x) public view { + tree.find(x); + } + + function insert(uint256 x) public { + tree.insert(x); + } + + function remove(bytes32 ptr) public { + ptr.remove(); + } + + function testRedBlackTreeInsertOneGas() public { + unchecked { + for (uint256 i; i != 1; ++i) { + tree.insert(i + 1); + } + } + } + + function testRedBlackTreeInsertTwoGas() public { + unchecked { + for (uint256 i; i != 2; ++i) { + tree.insert(i + 1); + } + } + } + + function testRedBlackTreeInsertThreeGas() public { + unchecked { + for (uint256 i; i != 3; ++i) { + tree.insert(i + 1); + } + } + } + + function testRedBlackTreeInsertTenGas() public { + unchecked { + for (uint256 i; i != 10; ++i) { + tree.insert(i + 1); + } + } + } + + function testRedBlackTreeValues() public { + testRedBlackTreeValues(3); + } + + function testRedBlackTreeValues(uint256 n) public { + unchecked { + n = n & 7; + while (true) { + uint256[] memory values = new uint256[](n); + for (uint256 i; i != n; ++i) { + values[i] = 1 | _random(); + _brutalizeScratchSpace(); + tree.tryInsert(values[i]); + } + LibSort.sort(values); + LibSort.uniquifySorted(values); + uint256[] memory retrieved = tree.values(); + _checkMemory(); + assertEq(retrieved, values); + n = values.length; + if (_random() & 1 == 0) { + LibPRNG.PRNG memory prng = LibPRNG.PRNG(_random()); + prng.shuffle(values); + for (uint256 i; i != n; ++i) { + _brutalizeScratchSpace(); + tree.tryRemove(values[i]); + } + assertEq(tree.values(), new uint256[](0)); + if (_random() & 1 == 0) { + n += _random() & 15; + continue; + } + } + break; + } + } + } + + function testRedBlackTreeRejectsEmptyValue() public { + vm.expectRevert(RedBlackTreeWithResizeLib.ValueIsEmpty.selector); + this.insert(0); + vm.expectRevert(RedBlackTreeWithResizeLib.ValueDoesNotExist.selector); + this.remove(0); + vm.expectRevert(RedBlackTreeWithResizeLib.ValueIsEmpty.selector); + this.find(0); + } + + function testRedBlackTreeRemoveViaPointer() public { + tree.insert(1); + tree.insert(2); + + bytes32 ptr = tree.find(1); + ptr.remove(); + ptr.remove(); + + vm.expectRevert(RedBlackTreeWithResizeLib.PointerOutOfBounds.selector); + this.remove(ptr); + + ptr = bytes32(0); + vm.expectRevert(RedBlackTreeWithResizeLib.ValueDoesNotExist.selector); + this.remove(ptr); + } + + function testRedBlackTreeTryInsertAndRemove() public { + tree.tryInsert(1); + tree.tryInsert(2); + assertEq(tree.size(), 2); + tree.tryInsert(1); + assertEq(tree.size(), 2); + tree.tryRemove(2); + assertEq(tree.size(), 1); + tree.tryRemove(2); + assertEq(tree.size(), 1); + } + + function testRedBlackTreeTreeFullReverts() public { + tree.insert(1); + bytes32 ptr = tree.find(1); + /// @solidity memory-safe-assembly + assembly { + ptr := shl(32, shr(32, ptr)) + sstore(ptr, or(sload(ptr), sub(shl(31, 1), 1))) + } + vm.expectRevert(RedBlackTreeWithResizeLib.TreeIsFull.selector); + this.insert(2); + assertEq(tree.size(), 2 ** 31 - 1); + } + + function testRedBlackTreePointers() public { + assertTrue(tree.find(1).isEmpty()); + assertTrue(tree.find(2).isEmpty()); + + tree.insert(1); + tree.insert(2); + + assertFalse(tree.find(1).isEmpty()); + assertFalse(tree.find(2).isEmpty()); + + assertTrue(tree.find(1).prev().isEmpty()); + assertFalse(tree.find(1).next().isEmpty()); + + assertFalse(tree.find(2).prev().isEmpty()); + assertTrue(tree.find(2).next().isEmpty()); + + assertEq(tree.find(1).next(), tree.find(2)); + assertEq(tree.find(1), tree.find(2).prev()); + + assertTrue(tree.find(1).prev().isEmpty()); + assertTrue(tree.find(1).prev().prev().isEmpty()); + assertTrue(tree.find(1).prev().next().isEmpty()); + + assertTrue(tree.find(2).next().isEmpty()); + assertTrue(tree.find(2).next().next().isEmpty()); + assertTrue(tree.find(2).next().prev().isEmpty()); + + assertEq(tree.first(), tree.find(1)); + assertEq(tree.last(), tree.find(2)); + + assertTrue(tree.find(3).isEmpty()); + } + + function testRedBlackTreeNearest(uint256) public { + assertEq(tree.nearest(1), bytes32(0)); + uint256[] memory a = _fillTree(_random() % 8); + uint256 x = _bound(_random(), 1, type(uint256).max); + (uint256 nearestIndex, bool found) = _nearestIndex(a, x); + if (found) { + assertEq(tree.nearest(x).value(), a[nearestIndex]); + } else { + assertEq(tree.nearest(x), bytes32(0)); + } + } + + function _nearestIndex(uint256[] memory a, uint256 x) internal pure returns (uint256 nearestIndex, bool found) { + unchecked { + uint256 nearestValue = type(uint256).max; + uint256 nearestDist = type(uint256).max; + uint256 n = a.length; + for (uint256 i; i != n; ++i) { + uint256 y = a[i]; + uint256 dist = x < y ? y - x : x - y; + if (dist < nearestDist || (dist == nearestDist && y < nearestValue)) { + nearestIndex = i; + nearestValue = y; + nearestDist = dist; + found = true; + } + } + } + } + + function testRedBlackTreeNearestBefore(uint256) public { + assertEq(tree.nearestBefore(1), bytes32(0)); + uint256[] memory a = _fillTree(_random() % 8); + uint256 x = _bound(_random(), 1, type(uint256).max); + (uint256 nearestIndexBefore, bool found) = _nearestIndexBefore(a, x); + if (found) { + assertEq(tree.nearestBefore(x).value(), a[nearestIndexBefore]); + } else { + assertEq(tree.nearestBefore(x), bytes32(0)); + } + } + + function _nearestIndexBefore(uint256[] memory a, uint256 x) + internal + pure + returns (uint256 nearestIndex, bool found) + { + unchecked { + uint256 nearestDist = type(uint256).max; + uint256 n = a.length; + for (uint256 i; i != n; ++i) { + uint256 y = a[i]; + if (y > x) continue; + uint256 dist = x - y; + if (dist < nearestDist) { + nearestIndex = i; + nearestDist = dist; + found = true; + } + } + } + } + + function testRedBlackTreeNearestAfter(uint256) public { + assertEq(tree.nearestAfter(1), bytes32(0)); + uint256[] memory a = _fillTree(_random() % 8); + uint256 x = _bound(_random(), 1, type(uint256).max); + (uint256 nearestIndexAfter, bool found) = _nearestIndexAfter(a, x); + if (found) { + assertEq(tree.nearestAfter(x).value(), a[nearestIndexAfter]); + } else { + assertEq(tree.nearestAfter(x), bytes32(0)); + } + } + + function _nearestIndexAfter(uint256[] memory a, uint256 x) + internal + pure + returns (uint256 nearestIndex, bool found) + { + unchecked { + uint256 nearestDist = type(uint256).max; + uint256 n = a.length; + for (uint256 i; i != n; ++i) { + uint256 y = a[i]; + if (y < x) continue; + uint256 dist = y - x; + if (dist < nearestDist) { + nearestIndex = i; + nearestDist = dist; + found = true; + } + } + } + } + + function _fillTree(uint256 n) internal returns (uint256[] memory a) { + a = _makeArray(n); + unchecked { + for (uint256 i; i != n;) { + uint256 r = _bound(_random(), 1, type(uint256).max); + if (tree.find(r).isEmpty()) { + a[i++] = r; + _brutalizeScratchSpace(); + tree.insert(r); + } + if (_randomChance(4)) { + _testRemoveAndInsertBack(a, i, (3 + i >> 2)); + } + } + } + } + + function testRedBlackTreeStorageSize() public { + // Initial storage size should be 0 + assertEq(tree.storageSize(), 0); + assertEq(tree.size(), 0); + + // Insert one element + tree.insert(1); + assertEq(tree.size(), 1); + assertEq(tree.storageSize(), 1); + assertTrue(tree.storageSize() >= tree.size()); + + // Insert more elements + tree.insert(2); + tree.insert(3); + assertEq(tree.size(), 3); + assertEq(tree.storageSize(), 3); + assertTrue(tree.storageSize() >= tree.size()); + + // Test that storage size is never less than actual size + uint256 currentStorageSize = tree.storageSize(); + assertTrue(currentStorageSize >= 3); + } + + function testRedBlackTreeResize() public { + // Test resize on empty tree + tree.resize(10); + assertEq(tree.storageSize(), 10); + assertEq(tree.size(), 0); + + // Insert elements and check storage is not exceeded + tree.insert(1); + tree.insert(2); + tree.insert(3); + assertEq(tree.size(), 3); + assertEq(tree.storageSize(), 10); + + // Resize to larger capacity + tree.resize(20); + assertEq(tree.storageSize(), 20); + assertEq(tree.size(), 3); + + // Insert more elements + for (uint256 i = 4; i <= 15; i++) { + tree.insert(i); + } + assertEq(tree.size(), 15); + assertEq(tree.storageSize(), 20); + } + + function testRedBlackTreeResizeErrorCases() public { + // Insert some elements + tree.insert(1); + tree.insert(2); + tree.insert(3); + assertEq(tree.size(), 3); + + // Try to resize to capacity smaller than current size + vm.expectRevert(RedBlackTreeWithResizeLib.InvalidStorageSize.selector); + this.resize(2); + + // Try to resize to maximum capacity + vm.expectRevert(RedBlackTreeWithResizeLib.TreeIsFull.selector); + this.resize(2 ** 31); + + // Valid resize should still work + tree.resize(5); + assertEq(tree.storageSize(), 5); + } + + function testRedBlackTreeTryResize() public { + // Test tryResize success cases + assertEq(tree.tryResize(10), 0); + assertEq(tree.storageSize(), 10); + + tree.insert(1); + tree.insert(2); + assertEq(tree.tryResize(15), 0); + assertEq(tree.storageSize(), 15); + + // Test tryResize error cases + uint256 err = tree.tryResize(1); + assertTrue(err != 0); // Should return InvalidStorageSize error + + err = tree.tryResize(2 ** 31); + assertTrue(err != 0); // Should return TreeIsFull error + } + + function testRedBlackTreeResizeWithRemoval() public { + // Pre-allocate and fill + tree.resize(10); + for (uint256 i = 1; i <= 8; i++) { + tree.insert(i); + } + assertEq(tree.size(), 8); + assertEq(tree.storageSize(), 10); + + // Remove some elements + tree.remove(1); + tree.remove(2); + assertEq(tree.size(), 6); + assertEq(tree.storageSize(), 10); // Storage size shouldn't change + + // Should still be able to resize + tree.resize(15); + assertEq(tree.storageSize(), 15); + assertEq(tree.size(), 6); + } + + function resize(uint256 newCapacity) public { + tree.resize(newCapacity); + } +} diff --git a/test/RedBlackTreeWithResizeKV.t.sol b/test/RedBlackTreeWithResizeKV.t.sol new file mode 100644 index 0000000..5845e48 --- /dev/null +++ b/test/RedBlackTreeWithResizeKV.t.sol @@ -0,0 +1,402 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.30; + +import {Test, console} from "forge-std/Test.sol"; +import {RedBlackTreeWithResizeKV} from "../src/RedBlackTreeWithResizeKV.sol"; +import {ValueLib} from "../src/lib/Value.sol"; +import {RedBlackTreeWithResizeLib} from "../src/lib/RedBlackTreeWithResizeLib.sol"; + +contract RedBlackTreeWithResizeKVTest is Test { + RedBlackTreeWithResizeKV public redBlockTreeKV; + + function setUp() public { + redBlockTreeKV = new RedBlackTreeWithResizeKV(); + } + + function generateValue(uint256 value) internal pure returns (ValueLib.Value memory) { + return ValueLib.Value({ + slot1: bytes32(uint256(value)), + slot2: uint256(value), + slot3: int256(value), + slot4: bytes32(uint256(value)), + slot5: uint256(value), + slot6: address(uint160(value)) + }); + } + + function test_setValue_and_getValue() public { + uint256 key = 1; + ValueLib.Value memory expectedValue = generateValue(100); + + redBlockTreeKV.setValue(key, expectedValue); + ValueLib.Value memory actualValue = redBlockTreeKV.getValue(key); + + assertEq(actualValue.slot1, expectedValue.slot1); + assertEq(actualValue.slot2, expectedValue.slot2); + assertEq(actualValue.slot3, expectedValue.slot3); + assertEq(actualValue.slot4, expectedValue.slot4); + assertEq(actualValue.slot5, expectedValue.slot5); + assertEq(actualValue.slot6, expectedValue.slot6); + } + + function test_setValue_overwrite() public { + uint256 key = 1; + ValueLib.Value memory value1 = generateValue(100); + ValueLib.Value memory value2 = generateValue(200); + + redBlockTreeKV.setValue(key, value1); + vm.expectRevert(RedBlackTreeWithResizeLib.ValueAlreadyExists.selector); + redBlockTreeKV.setValue(key, value2); + } + + function test_deleteValue() public { + uint256 key = 1; + ValueLib.Value memory value = generateValue(100); + + redBlockTreeKV.setValue(key, value); + redBlockTreeKV.deleteValue(key); + + uint256[] memory values = redBlockTreeKV.values(); + assertEq(values.length, 0); + } + + function test_multiple_values() public { + uint256 key1 = 1; + uint256 key2 = 2; + uint256 key3 = 3; + + ValueLib.Value memory value1 = generateValue(100); + ValueLib.Value memory value2 = generateValue(200); + ValueLib.Value memory value3 = generateValue(300); + + redBlockTreeKV.setValue(key1, value1); + redBlockTreeKV.setValue(key2, value2); + redBlockTreeKV.setValue(key3, value3); + + ValueLib.Value memory actualValue1 = redBlockTreeKV.getValue(key1); + ValueLib.Value memory actualValue2 = redBlockTreeKV.getValue(key2); + ValueLib.Value memory actualValue3 = redBlockTreeKV.getValue(key3); + + assertEq(actualValue1.slot2, 100); + assertEq(actualValue2.slot2, 200); + assertEq(actualValue3.slot2, 300); + + uint256[] memory values = redBlockTreeKV.values(); + assertEq(values.length, 3); + } + + function test_delete_middle_value() public { + uint256 key1 = 1; + uint256 key2 = 2; + uint256 key3 = 3; + + ValueLib.Value memory value1 = generateValue(100); + ValueLib.Value memory value2 = generateValue(200); + ValueLib.Value memory value3 = generateValue(300); + + redBlockTreeKV.setValue(key1, value1); + redBlockTreeKV.setValue(key2, value2); + redBlockTreeKV.setValue(key3, value3); + + redBlockTreeKV.deleteValue(key2); + + ValueLib.Value memory actualValue1 = redBlockTreeKV.getValue(key1); + ValueLib.Value memory actualValue3 = redBlockTreeKV.getValue(key3); + + assertEq(actualValue1.slot2, 100); + assertEq(actualValue3.slot2, 300); + + uint256[] memory values = redBlockTreeKV.values(); + assertEq(values.length, 2); + } + + function test_values_empty() public view { + uint256[] memory values = redBlockTreeKV.values(); + assertEq(values.length, 0); + } + + function test_values_ordering() public { + uint256[] memory keys = new uint256[](5); + keys[0] = 50; + keys[1] = 30; + keys[2] = 70; + keys[3] = 20; + keys[4] = 60; + + for (uint256 i = 0; i < keys.length; i++) { + redBlockTreeKV.setValue(keys[i], generateValue(keys[i])); + } + + uint256[] memory values = redBlockTreeKV.values(); + assertEq(values.length, 5); + + for (uint256 i = 0; i < values.length - 1; i++) { + assertLt(values[i], values[i + 1]); + } + } + + function test_edge_case_zero_key() public { + uint256 key = 0; + ValueLib.Value memory value = generateValue(100); + + vm.expectRevert(RedBlackTreeWithResizeLib.ValueIsEmpty.selector); + redBlockTreeKV.setValue(key, value); + } + + function test_edge_case_max_key() public { + uint256 key = type(uint256).max; + ValueLib.Value memory value = generateValue(100); + + redBlockTreeKV.setValue(key, value); + ValueLib.Value memory actualValue = redBlockTreeKV.getValue(key); + + assertEq(actualValue.slot2, 100); + } + + function test_large_dataset() public { + uint256 count = 100; + + for (uint256 i = 1; i <= count; i++) { + redBlockTreeKV.setValue(i, generateValue(i * 10)); + } + + for (uint256 i = 1; i <= count; i++) { + ValueLib.Value memory value = redBlockTreeKV.getValue(i); + assertEq(value.slot2, i * 10); + } + + uint256[] memory values = redBlockTreeKV.values(); + assertEq(values.length, count); + } + + function test_mixed_operations() public { + redBlockTreeKV.setValue(1, generateValue(100)); + redBlockTreeKV.setValue(2, generateValue(200)); + redBlockTreeKV.setValue(3, generateValue(300)); + + redBlockTreeKV.deleteValue(2); + + redBlockTreeKV.setValue(4, generateValue(400)); + redBlockTreeKV.setValue(5, generateValue(500)); + + redBlockTreeKV.deleteValue(1); + redBlockTreeKV.deleteValue(3); + + ValueLib.Value memory value4 = redBlockTreeKV.getValue(4); + ValueLib.Value memory value5 = redBlockTreeKV.getValue(5); + + assertEq(value4.slot2, 400); + assertEq(value5.slot2, 500); + + uint256[] memory values = redBlockTreeKV.values(); + assertEq(values.length, 2); + } + + function test_fuzz_setValue_getValue(uint256 key, uint256 value) public { + vm.assume(key != 0); + vm.assume(value != 0); + + ValueLib.Value memory testValue = generateValue(value); + redBlockTreeKV.setValue(key, testValue); + + ValueLib.Value memory retrievedValue = redBlockTreeKV.getValue(key); + + assertEq(retrievedValue.slot1, testValue.slot1); + assertEq(retrievedValue.slot2, testValue.slot2); + assertEq(retrievedValue.slot3, testValue.slot3); + assertEq(retrievedValue.slot4, testValue.slot4); + assertEq(retrievedValue.slot5, testValue.slot5); + assertEq(retrievedValue.slot6, testValue.slot6); + } + + function test_fuzz_multiple_operations(uint256[10] memory keys, uint256[10] memory values) public { + for (uint256 i = 0; i < keys.length; i++) { + vm.assume(keys[i] != 0); + vm.assume(values[i] != 0); + for (uint256 j = i + 1; j < keys.length; j++) { + vm.assume(keys[i] != keys[j]); + } + } + + for (uint256 i = 0; i < keys.length; i++) { + redBlockTreeKV.setValue(keys[i], generateValue(values[i])); + } + + for (uint256 i = 0; i < keys.length; i++) { + ValueLib.Value memory retrievedValue = redBlockTreeKV.getValue(keys[i]); + assertEq(retrievedValue.slot2, values[i]); + } + } + + function test_resize_expand_capacity() public { + uint256 initialCapacity = redBlockTreeKV.getStorageSize(); + uint256 newCapacity = initialCapacity + 50; + + redBlockTreeKV.resize(newCapacity); + + assertEq(redBlockTreeKV.getStorageSize(), newCapacity); + assertEq(redBlockTreeKV.getSize(), 0); + } + + function test_resize_shrink_capacity() public { + redBlockTreeKV.setValue(1, generateValue(100)); + redBlockTreeKV.setValue(2, generateValue(200)); + + assertEq(redBlockTreeKV.getSize(), 2); + assertEq(redBlockTreeKV.getStorageSize(), 2); + + uint256 currentSize = redBlockTreeKV.getSize(); + uint256 newCapacity = currentSize + 10; + + redBlockTreeKV.resize(newCapacity); + + assertEq(redBlockTreeKV.getStorageSize(), newCapacity); + assertEq(redBlockTreeKV.getSize(), currentSize); + + ValueLib.Value memory value1 = redBlockTreeKV.getValue(1); + ValueLib.Value memory value2 = redBlockTreeKV.getValue(2); + assertEq(value1.slot2, 100); + assertEq(value2.slot2, 200); + } + + function test_resize_same_capacity() public { + redBlockTreeKV.setValue(1, generateValue(100)); + + uint256 currentCapacity = redBlockTreeKV.getStorageSize(); + uint256 currentSize = redBlockTreeKV.getSize(); + + redBlockTreeKV.resize(currentCapacity); + + assertEq(redBlockTreeKV.getStorageSize(), currentCapacity); + assertEq(redBlockTreeKV.getSize(), currentSize); + + ValueLib.Value memory value = redBlockTreeKV.getValue(1); + assertEq(value.slot2, 100); + } + + function test_resize_invalid_capacity_too_small() public { + redBlockTreeKV.setValue(1, generateValue(100)); + redBlockTreeKV.setValue(2, generateValue(200)); + redBlockTreeKV.setValue(3, generateValue(300)); + + uint256 currentSize = redBlockTreeKV.getSize(); + uint256 invalidCapacity = currentSize - 1; + + vm.expectRevert(RedBlackTreeWithResizeLib.InvalidStorageSize.selector); + redBlockTreeKV.resize(invalidCapacity); + } + + function test_resize_preserve_data_after_expansion() public { + for (uint256 i = 1; i <= 10; i++) { + redBlockTreeKV.setValue(i, generateValue(i * 100)); + } + + uint256 currentCapacity = redBlockTreeKV.getStorageSize(); + uint256 newCapacity = currentCapacity + 100; + + redBlockTreeKV.resize(newCapacity); + + assertEq(redBlockTreeKV.getStorageSize(), newCapacity); + assertEq(redBlockTreeKV.getSize(), 10); + + for (uint256 i = 1; i <= 10; i++) { + ValueLib.Value memory value = redBlockTreeKV.getValue(i); + assertEq(value.slot2, i * 100); + } + } + + function test_resize_preserve_data_after_shrinking() public { + for (uint256 i = 1; i <= 5; i++) { + redBlockTreeKV.setValue(i, generateValue(i * 100)); + } + + uint256 currentSize = redBlockTreeKV.getSize(); + uint256 newCapacity = currentSize + 5; + + redBlockTreeKV.resize(newCapacity); + + assertEq(redBlockTreeKV.getStorageSize(), newCapacity); + assertEq(redBlockTreeKV.getSize(), currentSize); + + for (uint256 i = 1; i <= 5; i++) { + ValueLib.Value memory value = redBlockTreeKV.getValue(i); + assertEq(value.slot2, i * 100); + } + } + + function test_resize_after_deletions() public { + for (uint256 i = 1; i <= 10; i++) { + redBlockTreeKV.setValue(i, generateValue(i * 100)); + } + + redBlockTreeKV.deleteValue(5); + redBlockTreeKV.deleteValue(8); + redBlockTreeKV.deleteValue(2); + + uint256 currentSize = redBlockTreeKV.getSize(); + uint256 newCapacity = currentSize + 20; + + redBlockTreeKV.resize(newCapacity); + + assertEq(redBlockTreeKV.getStorageSize(), newCapacity); + assertEq(redBlockTreeKV.getSize(), currentSize); + + uint256[] memory values = redBlockTreeKV.values(); + assertEq(values.length, 7); + } + + function test_resize_zero_capacity() public { + uint256 currentSize = redBlockTreeKV.getSize(); + + if (currentSize == 0) { + redBlockTreeKV.resize(0); + assertEq(redBlockTreeKV.getStorageSize(), 0); + } else { + vm.expectRevert(RedBlackTreeWithResizeLib.InvalidStorageSize.selector); + redBlockTreeKV.resize(0); + } + } + + function test_resize_large_capacity() public { + uint256 largeCapacity = 1000; + + redBlockTreeKV.resize(largeCapacity); + + assertEq(redBlockTreeKV.getStorageSize(), largeCapacity); + assertEq(redBlockTreeKV.getSize(), 0); + + redBlockTreeKV.setValue(1, generateValue(100)); + redBlockTreeKV.setValue(50, generateValue(5000)); + + assertEq(redBlockTreeKV.getSize(), 2); + + ValueLib.Value memory value1 = redBlockTreeKV.getValue(1); + ValueLib.Value memory value2 = redBlockTreeKV.getValue(50); + assertEq(value1.slot2, 100); + assertEq(value2.slot2, 5000); + } + + function test_fuzz_resize_capacity(uint256 newCapacity) public { + vm.assume(newCapacity > 0 && newCapacity < 1000); + + redBlockTreeKV.setValue(1, generateValue(100)); + redBlockTreeKV.setValue(2, generateValue(200)); + + uint256 currentSize = redBlockTreeKV.getSize(); + + if (newCapacity >= currentSize) { + redBlockTreeKV.resize(newCapacity); + + assertEq(redBlockTreeKV.getStorageSize(), newCapacity); + assertEq(redBlockTreeKV.getSize(), currentSize); + + ValueLib.Value memory value1 = redBlockTreeKV.getValue(1); + ValueLib.Value memory value2 = redBlockTreeKV.getValue(2); + assertEq(value1.slot2, 100); + assertEq(value2.slot2, 200); + } else { + vm.expectRevert(RedBlackTreeWithResizeLib.InvalidStorageSize.selector); + redBlockTreeKV.resize(newCapacity); + } + } +}