Skip to content

Commit

Permalink
Merge pull request #61 from 1inch/feature/optimize-address-array
Browse files Browse the repository at this point in the history 
[SC-583] Optimize AddressArray and AddressSet libs
  • Loading branch information
@ZumZoom
ZumZoom authored Mar 22, 2024
2 parents 4664c8d + cdee296 commit 5ca7647
Show file tree
Hide file tree
Showing 6 changed files with 216 additions and 64 deletions.
185 changes: 133 additions & 52 deletions contracts/libraries/AddressArray.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,92 +8,173 @@ library AddressArray {
error PopFromEmptyArray();
error OutputArrayTooSmall();

uint256 internal constant _ZERO_ADDRESS = 0x8000000000000000000000000000000000000000000000000000000000000000; // Next tx gas optimization
uint256 internal constant _LENGTH_MASK = 0x0000000000000000ffffffff0000000000000000000000000000000000000000;
uint256 internal constant _ADDRESS_MASK = 0x000000000000000000000000ffffffffffffffffffffffffffffffffffffffff;
uint256 internal constant _ONE_LENGTH = 0x0000000000000000000000010000000000000000000000000000000000000000;
uint256 internal constant _LENGTH_OFFSET = 160;

/// @dev Data struct containing raw mapping.
struct Data {
mapping(uint256 => uint256) _raw;
uint256[1 << 32] _raw;
}

/// @dev Length of array.
function length(Data storage self) internal view returns (uint256) {
return self._raw[0] >> 160;
return (self._raw[0] & _LENGTH_MASK) >> _LENGTH_OFFSET;
}

/// @dev Returns data item from `self` storage at `i`.
function at(Data storage self, uint256 i) internal view returns (address) {
return address(uint160(self._raw[i]));
if (i >= 1 << 32) revert IndexOutOfBounds();
return address(uint160(self._raw[i] & _ADDRESS_MASK));
}

/// @dev Returns list of addresses from storage `self`.
function get(Data storage self) internal view returns (address[] memory arr) {
uint256 lengthAndFirst = self._raw[0];
arr = new address[](lengthAndFirst >> 160);
_get(self, arr, lengthAndFirst);
function get(Data storage self) internal view returns (address[] memory output) {
/// @solidity memory-safe-assembly
assembly { // solhint-disable-line no-inline-assembly
let lengthAndFirst := sload(self.slot)
let len := shr(_LENGTH_OFFSET, and(lengthAndFirst, _LENGTH_MASK))
let fst := and(lengthAndFirst, _ADDRESS_MASK)

// Allocate array
output := mload(0x40)
mstore(0x40, add(output, mul(0x20, add(1, len))))
mstore(output, len)

if len {
// Copy first element and then the rest in a loop
let ptr := add(output, 0x20)
mstore(ptr, fst)
for { let i := 1 } lt(i, len) { i:= add(i, 1) } {
let item := and(sload(add(self.slot, i)), _ADDRESS_MASK)
mstore(add(ptr, mul(0x20, i)), item)
}
}
}
}

/// @dev Puts list of addresses from `self` storage into `output` array.
function get(Data storage self, address[] memory output) internal view returns (address[] memory) {
return _get(self, output, self._raw[0]);
}
function get(Data storage self, address[] memory input) internal view returns (address[] memory output) {
output = input;
bytes4 err = OutputArrayTooSmall.selector;
/// @solidity memory-safe-assembly
assembly { // solhint-disable-line no-inline-assembly
let lengthAndFirst := sload(self.slot)
let len := shr(_LENGTH_OFFSET, and(lengthAndFirst, _LENGTH_MASK))
let fst := and(lengthAndFirst, _ADDRESS_MASK)

function _get(
Data storage self,
address[] memory output,
uint256 lengthAndFirst
) private view returns (address[] memory) {
uint256 len = lengthAndFirst >> 160;
if (len > output.length) revert OutputArrayTooSmall();
if (len > 0) {
output[0] = address(uint160(lengthAndFirst));
unchecked {
for (uint256 i = 1; i < len; i++) {
output[i] = address(uint160(self._raw[i]));
if gt(len, mload(input)) {
mstore(0, err)
revert(0, 4)
}
if len {
// Copy first element and then the rest in a loop
let ptr := add(output, 0x20)
mstore(ptr, fst)
for { let i := 1 } lt(i, len) { i:= add(i, 1) } {
let item := and(sload(add(self.slot, i)), _ADDRESS_MASK)
mstore(add(ptr, mul(0x20, i)), item)
}
}
}
return output;
}

/// @dev Array push back `account` operation on storage `self`.
function push(Data storage self, address account) internal returns (uint256) {
unchecked {
uint256 lengthAndFirst = self._raw[0];
uint256 len = lengthAndFirst >> 160;
if (len == 0) {
self._raw[0] = (1 << 160) + uint160(account);
} else {
self._raw[0] = lengthAndFirst + (1 << 160);
self._raw[len] = uint160(account);
function push(Data storage self, address account) internal returns (uint256 res) {
/// @solidity memory-safe-assembly
assembly { // solhint-disable-line no-inline-assembly
let lengthAndFirst := sload(self.slot)
let len := shr(_LENGTH_OFFSET, and(lengthAndFirst, _LENGTH_MASK))

switch len
case 0 {
sstore(self.slot, or(account, _ONE_LENGTH))
}
return len + 1;
default {
sstore(self.slot, add(lengthAndFirst, _ONE_LENGTH))
sstore(add(self.slot, len), or(account, _ZERO_ADDRESS))
}
res := add(len, 1)
}
}

/// @dev Array pop back operation for storage `self`.
function pop(Data storage self) internal {
unchecked {
uint256 lengthAndFirst = self._raw[0];
uint256 len = lengthAndFirst >> 160;
if (len == 0) revert PopFromEmptyArray();
self._raw[len - 1] = 0;
if (len > 1) {
self._raw[0] = lengthAndFirst - (1 << 160);
bytes4 err = PopFromEmptyArray.selector;
/// @solidity memory-safe-assembly
assembly { // solhint-disable-line no-inline-assembly
let lengthAndFirst := sload(self.slot)
let len := shr(_LENGTH_OFFSET, and(lengthAndFirst, _LENGTH_MASK))

switch len
case 0 {
mstore(0, err)
revert(0, 4)
}
case 1 {
sstore(self.slot, _ZERO_ADDRESS)
}
default {
sstore(self.slot, sub(lengthAndFirst, _ONE_LENGTH))
}
}
}

/// @dev Array pop back operation for storage `self` that returns popped element.
function popGet(Data storage self) internal returns(address res) {
bytes4 err = PopFromEmptyArray.selector;
/// @solidity memory-safe-assembly
assembly { // solhint-disable-line no-inline-assembly
let lengthAndFirst := sload(self.slot)
let len := shr(_LENGTH_OFFSET, and(lengthAndFirst, _LENGTH_MASK))

switch len
case 0 {
mstore(0, err)
revert(0, 4)
}
case 1 {
res := and(lengthAndFirst, _ADDRESS_MASK)
sstore(self.slot, _ZERO_ADDRESS)
}
default {
res := and(sload(add(self.slot, sub(len, 1))), _ADDRESS_MASK)
sstore(self.slot, sub(lengthAndFirst, _ONE_LENGTH))
}
}
}

/// @dev Set element for storage `self` at `index` to `account`.
function set(
Data storage self,
uint256 index,
address account
) internal {
uint256 len = length(self);
if (index >= len) revert IndexOutOfBounds();

if (index == 0) {
self._raw[0] = (len << 160) | uint160(account);
} else {
self._raw[index] = uint160(account);
function set(Data storage self, uint256 index, address account) internal {
bytes4 err = IndexOutOfBounds.selector;
/// @solidity memory-safe-assembly
assembly { // solhint-disable-line no-inline-assembly
let lengthAndFirst := sload(self.slot)
let len := shr(_LENGTH_OFFSET, and(lengthAndFirst, _LENGTH_MASK))
let fst := and(lengthAndFirst, _ADDRESS_MASK)

if iszero(lt(index, len)) {
mstore(0, err)
revert(0, 4)
}

switch index
case 0 {
sstore(self.slot, or(xor(lengthAndFirst, fst), account))
}
default {
sstore(add(self.slot, index), or(account, _ZERO_ADDRESS))
}
}
}

/// @dev Erase length of the array
function erase(Data storage self) internal {
/// @solidity memory-safe-assembly
assembly { // solhint-disable-line no-inline-assembly
sstore(self.slot, _ADDRESS_MASK)
}
}
}
46 changes: 36 additions & 10 deletions contracts/libraries/AddressSet.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,9 +13,9 @@ import "./AddressArray.sol";
library AddressSet {
using AddressArray for AddressArray.Data;

/** @dev Data struct from AddressArray.Data items
* and lookup mapping address => index in data array.
*/
uint256 internal constant _NULL_INDEX = type(uint256).max;

/// @dev Data struct from AddressArray.Data items and lookup mapping address => index in data array.
struct Data {
AddressArray.Data items;
mapping(address => uint256) lookup;
Expand All @@ -33,12 +33,24 @@ library AddressSet {

/// @dev Returns true if storage `s` has `item`.
function contains(Data storage s, address item) internal view returns (bool) {
return s.lookup[item] != 0;
uint256 index = s.lookup[item];
return index != 0 && index != _NULL_INDEX;
}

/// @dev Returns list of addresses from storage `s`.
function get(Data storage s) internal view returns (address[] memory) {
return s.items.get();
}

/// @dev Puts list of addresses from `s` storage into `output` array.
function get(Data storage s, address[] memory input) internal view returns (address[] memory) {
return s.items.get(input);
}

/// @dev Adds `item` into storage `s` and returns true if successful.
function add(Data storage s, address item) internal returns (bool) {
if (s.lookup[item] > 0) {
uint256 index = s.lookup[item];
if (index != 0 && index != _NULL_INDEX) {
return false;
}
s.lookup[item] = s.items.push(item);
Expand All @@ -48,18 +60,32 @@ library AddressSet {
/// @dev Removes `item` from storage `s` and returns true if successful.
function remove(Data storage s, address item) internal returns (bool) {
uint256 index = s.lookup[item];
if (index == 0) {
s.lookup[item] = _NULL_INDEX;
if (index == 0 || index == _NULL_INDEX) {
return false;
}
if (index < s.items.length()) {

address lastItem = s.items.popGet();
if (lastItem != item) {
unchecked {
address lastItem = s.items.at(s.items.length() - 1);
s.items.set(index - 1, lastItem);
s.lookup[lastItem] = index;
}
}
s.items.pop();
delete s.lookup[item];
return true;
}

/// @dev Erases set from storage `s` and returns all removed items
function erase(Data storage s) internal returns(address[] memory items) {
items = s.items.get();
uint256 len = items.length;
if (len > 0) {
s.items.erase();
unchecked {
for (uint256 i = 0; i < len; i++) {
s.lookup[items[i]] = _NULL_INDEX;
}
}
}
}
}
2 changes: 2 additions & 0 deletions contracts/tests/mocks/AddressArrayMock.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@ contract AddressArrayMock {

AddressArray.Data private _self;

error PopFromEmptyArray();

function length() external view returns (uint256) {
return AddressArray.length(_self);
}
Expand Down
4 changes: 4 additions & 0 deletions contracts/tests/mocks/AddressSetMock.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,10 @@ contract AddressSetMock {
return AddressSet.contains(_self, item);
}

function get() external view returns (address[] memory) {
return AddressSet.get(_self);
}

function add(address item) external returns (bool) {
return AddressSet.add(_self, item);
}
Expand Down
Loading

0 comments on commit 5ca7647

Please sign in to comment.