diff --git a/src/HeapOrdering.sol b/src/HeapOrdering.sol index 75794c1b..e5f29e21 100644 --- a/src/HeapOrdering.sol +++ b/src/HeapOrdering.sol @@ -10,8 +10,9 @@ library HeapOrdering { } struct HeapArray { - Account[] accounts; // All the accounts. - uint256 size; // The size of the heap portion of the structure, should be less than accounts length, the rest is an unordered array. + mapping(uint256 => Account) accounts; // All the accounts. + uint128 arrayLength; // The length of the array represented by the `accounts` mapping. + uint128 size; // The size of the heap portion of the structure, should be less than `arrayLength`, the rest is an unordered array. mapping(address => uint256) indexOf; // A mapping from an address to an index in accounts. From index i, the parent index is (i-1)/2, the left child index is 2*i+1 and the right child index is 2*i+2. } @@ -45,7 +46,7 @@ library HeapOrdering { uint256 size = _heap.size; uint256 newSize = computeSize(size, _maxSortedUsers); - if (size != newSize) _heap.size = newSize; + if (size != newSize) _heap.size = uint128(newSize); if (formerValue != newValue) { if (newValue == 0) remove(_heap, newSize, _id, formerValue); @@ -181,14 +182,15 @@ library HeapOrdering { if (_id == address(0)) revert AddressIsZero(); // Put the account at the end of accounts. - uint256 accountsLength = _heap.accounts.length; - _heap.accounts.push(Account(_id, _value)); + uint256 accountsLength = _heap.arrayLength; + _heap.accounts[accountsLength] = Account(_id, _value); + _heap.arrayLength = uint128(accountsLength + 1); _heap.indexOf[_id] = accountsLength; // Move the account at the end of the heap and restore the invariant. swap(_heap, _size, accountsLength); shiftUp(_heap, _size); - _heap.size = computeSize(_size + 1, _maxSortedUsers); + _heap.size = uint128(computeSize(_size + 1, _maxSortedUsers)); } /// @notice Decreases the amount of an account in the `_heap`. @@ -231,7 +233,7 @@ library HeapOrdering { else { swap(_heap, _size, index); shiftUp(_heap, _size); - _heap.size = computeSize(_size + 1, _maxSortedUsers); + _heap.size = uint128(computeSize(_size + 1, _maxSortedUsers)); } } @@ -248,12 +250,14 @@ library HeapOrdering { uint96 _removedValue ) private { uint256 index = _heap.indexOf[_id]; - uint256 accountsLength = _heap.accounts.length; - - // Swap the last account and the account to remove, then pop it. - swap(_heap, index, accountsLength - 1); - if (_size == accountsLength) _heap.size = --_size; - _heap.accounts.pop(); + uint256 accountsLength = _heap.arrayLength; + + // Swap the last account and the account to remove, then "pop" it. + uint256 newArrayLength = accountsLength - 1; + swap(_heap, index, newArrayLength); + if (_size == accountsLength) _heap.size = uint128(--_size); + delete _heap.accounts[newArrayLength]; + _heap.arrayLength = uint128(newArrayLength); delete _heap.indexOf[_id]; // If the swapped account is in the heap, restore the invariant: its value can be smaller or larger than the removed value. @@ -269,7 +273,7 @@ library HeapOrdering { /// @param _heap The heap parameter. /// @return The length of the heap. function length(HeapArray storage _heap) internal view returns (uint256) { - return _heap.accounts.length; + return _heap.arrayLength; } /// @notice Returns the value of the account linked to `_id`. @@ -278,27 +282,27 @@ library HeapOrdering { /// @return The value of the account. function getValueOf(HeapArray storage _heap, address _id) internal view returns (uint256) { uint256 index = _heap.indexOf[_id]; - if (index >= _heap.accounts.length) return 0; + if (index >= _heap.arrayLength) return 0; Account memory account = _heap.accounts[index]; if (account.id != _id) return 0; - else return account.value; + return account.value; } /// @notice Returns the address at the head of the `_heap`. /// @param _heap The heap to get the head. /// @return The address of the head. function getHead(HeapArray storage _heap) internal view returns (address) { - if (_heap.accounts.length > 0) return _heap.accounts[ROOT].id; - else return address(0); + if (_heap.arrayLength > 0) return _heap.accounts[ROOT].id; + return address(0); } /// @notice Returns the address at the tail of unsorted portion of the `_heap`. /// @param _heap The heap to get the tail. /// @return The address of the tail. function getTail(HeapArray storage _heap) internal view returns (address) { - uint256 accountsLength = _heap.accounts.length; + uint256 accountsLength = _heap.arrayLength; if (accountsLength > 0) return _heap.accounts[accountsLength - 1].id; - else return address(0); + return address(0); } /// @notice Returns the address coming before `_id` in accounts. @@ -309,7 +313,7 @@ library HeapOrdering { function getPrev(HeapArray storage _heap, address _id) internal view returns (address) { uint256 index = _heap.indexOf[_id]; if (index > ROOT) return _heap.accounts[index - 1].id; - else return address(0); + return address(0); } /// @notice Returns the address coming after `_id` in accounts. @@ -319,8 +323,7 @@ library HeapOrdering { /// @return The address of the next account. function getNext(HeapArray storage _heap, address _id) internal view returns (address) { uint256 index = _heap.indexOf[_id]; - if (index + 1 >= _heap.accounts.length || _heap.accounts[index].id != _id) - return address(0); - else return _heap.accounts[index + 1].id; + if (index + 1 >= _heap.arrayLength || _heap.accounts[index].id != _id) return address(0); + return _heap.accounts[index + 1].id; } }