[SC-583] Optimize AddressArray and AddressSet libs

contracts/libraries/AddressArray.sol

@@ -8,92 +8,173 @@
     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)
         }
     }
 }

contracts/libraries/AddressSet.sol

@@ -13,9 +13,9 @@
 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 @@
 
     /// @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 @@
     /// @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;
+                }
+            }
+        }
+    }
 }

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);
     }

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);
     }