diff --git a/contracts/libraries/AddressArray.sol b/contracts/libraries/AddressArray.sol index dc21b3cf..9b39c871 100644 --- a/contracts/libraries/AddressArray.sol +++ b/contracts/libraries/AddressArray.sol @@ -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) } } } diff --git a/contracts/libraries/AddressSet.sol b/contracts/libraries/AddressSet.sol index dd0f29e5..9fa707ce 100644 --- a/contracts/libraries/AddressSet.sol +++ b/contracts/libraries/AddressSet.sol @@ -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; @@ -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); @@ -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; + } + } + } + } } diff --git a/contracts/tests/mocks/AddressArrayMock.sol b/contracts/tests/mocks/AddressArrayMock.sol index cb968e70..30a0c65e 100644 --- a/contracts/tests/mocks/AddressArrayMock.sol +++ b/contracts/tests/mocks/AddressArrayMock.sol @@ -9,6 +9,8 @@ contract AddressArrayMock { AddressArray.Data private _self; + error PopFromEmptyArray(); + function length() external view returns (uint256) { return AddressArray.length(_self); } diff --git a/contracts/tests/mocks/AddressSetMock.sol b/contracts/tests/mocks/AddressSetMock.sol index 2844a645..579c39a4 100644 --- a/contracts/tests/mocks/AddressSetMock.sol +++ b/contracts/tests/mocks/AddressSetMock.sol @@ -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); } diff --git a/test/contracts/AddressArray.test.ts b/test/contracts/AddressArray.test.ts index 785cc1ea..14633332 100644 --- a/test/contracts/AddressArray.test.ts +++ b/test/contracts/AddressArray.test.ts @@ -22,12 +22,14 @@ describe('AddressArray', function () { describe('length', function () { it('should calculate length 0', async function () { const { addressArrayMock } = await loadFixture(deployAddressArrayMock); + await signer1.sendTransaction(await addressArrayMock.length.populateTransaction()); expect(await addressArrayMock.length()).to.be.equal('0'); }); it('should calculate length 1', async function () { const { addressArrayMock } = await loadFixture(deployAddressArrayMock); await addressArrayMock.push(signer1); + await signer1.sendTransaction(await addressArrayMock.length.populateTransaction()); expect(await addressArrayMock.length()).to.be.equal('1'); }); }); @@ -58,12 +60,14 @@ describe('AddressArray', function () { describe('get', function () { it('should get empty array', async function () { const { addressArrayMock } = await loadFixture(deployAddressArrayMock); + await signer1.sendTransaction(await addressArrayMock.get.populateTransaction()); expect(await addressArrayMock.get()).to.be.deep.equal([]); }); it('should get array with 1 element', async function () { const { addressArrayMock } = await loadFixture(deployAddressArrayMock); await addressArrayMock.push(signer1); + await signer1.sendTransaction(await addressArrayMock.get.populateTransaction()); expect(await addressArrayMock.get()).to.be.deep.equal([signer1.address]); }); @@ -71,6 +75,7 @@ describe('AddressArray', function () { const { addressArrayMock } = await loadFixture(deployAddressArrayMock); await addressArrayMock.push(signer1); await addressArrayMock.push(signer2); + await signer1.sendTransaction(await addressArrayMock.get.populateTransaction()); expect(await addressArrayMock.get()).to.be.deep.equal([signer1.address, signer2.address]); }); @@ -79,6 +84,7 @@ describe('AddressArray', function () { await addressArrayMock.push(signer1); await addressArrayMock.push(signer2); await addressArrayMock.push(signer3); + await signer1.sendTransaction(await addressArrayMock.get.populateTransaction()); expect(await addressArrayMock.get()).to.be.deep.equal([signer1.address, signer2.address, signer3.address]); }); }); @@ -119,6 +125,7 @@ describe('AddressArray', function () { const { addressArrayMock } = await loadFixture(deployAddressArrayMock); await addressArrayMock.push(signer1); await addressArrayMock.pop(); + await signer1.sendTransaction(await addressArrayMock.get.populateTransaction()); expect(await addressArrayMock.get()).to.be.deep.equal([]); }); @@ -127,6 +134,7 @@ describe('AddressArray', function () { await addressArrayMock.push(signer1); await addressArrayMock.push(signer2); await addressArrayMock.pop(); + await signer1.sendTransaction(await addressArrayMock.get.populateTransaction()); expect(await addressArrayMock.get()).to.be.deep.equal([signer1.address]); }); @@ -136,6 +144,7 @@ describe('AddressArray', function () { await addressArrayMock.push(signer2); await addressArrayMock.push(signer3); await addressArrayMock.pop(); + await signer1.sendTransaction(await addressArrayMock.get.populateTransaction()); expect(await addressArrayMock.get()).to.be.deep.equal([signer1.address, signer2.address]); }); @@ -160,6 +169,7 @@ describe('AddressArray', function () { const { addressArrayMock } = await loadFixture(deployAddressArrayMock); await addressArrayMock.push(signer1); await addressArrayMock.set(0, signer2); + await signer1.sendTransaction(await addressArrayMock.get.populateTransaction()); expect(await addressArrayMock.get()).to.be.deep.equal([signer2.address]); }); @@ -168,6 +178,7 @@ describe('AddressArray', function () { await addressArrayMock.push(signer1); await addressArrayMock.push(signer2); await addressArrayMock.set(0, signer3); + await signer1.sendTransaction(await addressArrayMock.get.populateTransaction()); expect(await addressArrayMock.get()).to.be.deep.equal([signer3.address, signer2.address]); }); @@ -176,7 +187,22 @@ describe('AddressArray', function () { await addressArrayMock.push(signer1); await addressArrayMock.push(signer2); await addressArrayMock.set(1, signer3); + await signer1.sendTransaction(await addressArrayMock.get.populateTransaction()); expect(await addressArrayMock.get()).to.be.deep.equal([signer1.address, signer3.address]); }); }); + + describe('multiple add/remove', function () { + it('should add and remove multiple times', async function () { + const { addressArrayMock } = await loadFixture(deployAddressArrayMock); + await addressArrayMock.push(signer1); + await addressArrayMock.push(signer2); + await addressArrayMock.pop(); + await addressArrayMock.pop(); + await addressArrayMock.push(signer1); + await addressArrayMock.push(signer2); + await addressArrayMock.pop(); + await addressArrayMock.pop(); + }); + }); }); diff --git a/test/contracts/AddressSet.test.ts b/test/contracts/AddressSet.test.ts index 7376872b..f74bba29 100644 --- a/test/contracts/AddressSet.test.ts +++ b/test/contracts/AddressSet.test.ts @@ -142,8 +142,21 @@ describe('AddressSet', function () { const isRemoved = await addressSetMock.remove.staticCall(signer1); await addressSetMock.remove(signer1); expect(isRemoved).to.be.true; - expect(await addressSetMock.contains(signer1)).to.be.false; - expect(await addressSetMock.contains(signer2)).to.be.true; + expect(await addressSetMock.get()).to.be.deep.equal([signer2.address]); + }); + }); + + describe('multiple add/remove', function () { + it('should add and remove multiple times', async function () { + const { addressSetMock } = await loadFixture(deployAddressSetMock); + await addressSetMock.add(signer1); + await addressSetMock.add(signer2); + await addressSetMock.remove(signer2); + await addressSetMock.remove(signer1); + await addressSetMock.add(signer1); + await addressSetMock.add(signer2); + await addressSetMock.remove(signer2); + await addressSetMock.remove(signer1); }); }); });