|
| 1 | +// SPDX-License-Identifier: MIT |
| 2 | +pragma solidity ^0.8.28; |
| 3 | + |
| 4 | +library AssociatedArrayLib { |
| 5 | + using AssociatedArrayLib for *; |
| 6 | + |
| 7 | + struct Array { |
| 8 | + uint256 _spacer; |
| 9 | + } |
| 10 | + |
| 11 | + struct Bytes32Array { |
| 12 | + Array _inner; |
| 13 | + } |
| 14 | + |
| 15 | + struct AddressArray { |
| 16 | + Array _inner; |
| 17 | + } |
| 18 | + |
| 19 | + struct UintArray { |
| 20 | + Array _inner; |
| 21 | + } |
| 22 | + |
| 23 | + error AssociatedArray_OutOfBounds(uint256 index); |
| 24 | + |
| 25 | + function add(Bytes32Array storage s, address account, bytes32 value) internal { |
| 26 | + if (!_contains(s._inner, account, value)) { |
| 27 | + _push(s._inner, account, value); |
| 28 | + } |
| 29 | + } |
| 30 | + |
| 31 | + function set(Bytes32Array storage s, address account, uint256 index, bytes32 value) internal { |
| 32 | + _set(s._inner, account, index, value); |
| 33 | + } |
| 34 | + |
| 35 | + function push(Bytes32Array storage s, address account, bytes32 value) internal { |
| 36 | + _push(s._inner, account, value); |
| 37 | + } |
| 38 | + |
| 39 | + function pop(Bytes32Array storage s, address account) internal { |
| 40 | + _pop(s._inner, account); |
| 41 | + } |
| 42 | + |
| 43 | + function remove(Bytes32Array storage s, address account, uint256 index) internal { |
| 44 | + _remove(s._inner, account, index); |
| 45 | + } |
| 46 | + |
| 47 | + function add(UintArray storage s, address account, uint256 value) internal { |
| 48 | + if (!_contains(s._inner, account, bytes32(value))) { |
| 49 | + _push(s._inner, account, bytes32(value)); |
| 50 | + } |
| 51 | + } |
| 52 | + |
| 53 | + function set(UintArray storage s, address account, uint256 index, uint256 value) internal { |
| 54 | + _set(s._inner, account, index, bytes32(value)); |
| 55 | + } |
| 56 | + |
| 57 | + function push(UintArray storage s, address account, uint256 value) internal { |
| 58 | + _push(s._inner, account, bytes32(value)); |
| 59 | + } |
| 60 | + |
| 61 | + function pop(UintArray storage s, address account) internal { |
| 62 | + _pop(s._inner, account); |
| 63 | + } |
| 64 | + |
| 65 | + function remove(UintArray storage s, address account, uint256 index) internal { |
| 66 | + _remove(s._inner, account, index); |
| 67 | + } |
| 68 | + |
| 69 | + function add(AddressArray storage s, address account, address value) internal { |
| 70 | + if (!_contains(s._inner, account, bytes32(uint256(uint160(value))))) { |
| 71 | + _push(s._inner, account, bytes32(uint256(uint160(value)))); |
| 72 | + } |
| 73 | + } |
| 74 | + |
| 75 | + function set(AddressArray storage s, address account, uint256 index, address value) internal { |
| 76 | + _set(s._inner, account, index, bytes32(uint256(uint160(value)))); |
| 77 | + } |
| 78 | + |
| 79 | + function push(AddressArray storage s, address account, address value) internal { |
| 80 | + _push(s._inner, account, bytes32(uint256(uint160(value)))); |
| 81 | + } |
| 82 | + |
| 83 | + function pop(AddressArray storage s, address account) internal { |
| 84 | + _pop(s._inner, account); |
| 85 | + } |
| 86 | + |
| 87 | + function remove(AddressArray storage s, address account, uint256 index) internal { |
| 88 | + _remove(s._inner, account, index); |
| 89 | + } |
| 90 | + |
| 91 | + function length(Bytes32Array storage s, address account) internal view returns (uint256) { |
| 92 | + return _length(s._inner, account); |
| 93 | + } |
| 94 | + |
| 95 | + function get(Bytes32Array storage s, address account, uint256 index) internal view returns (bytes32) { |
| 96 | + return _get(s._inner, account, index); |
| 97 | + } |
| 98 | + |
| 99 | + function getAll(Bytes32Array storage s, address account) internal view returns (bytes32[] memory) { |
| 100 | + return _getAll(s._inner, account); |
| 101 | + } |
| 102 | + |
| 103 | + function contains(Bytes32Array storage s, address account, bytes32 value) internal view returns (bool) { |
| 104 | + return _contains(s._inner, account, value); |
| 105 | + } |
| 106 | + |
| 107 | + function length(AddressArray storage s, address account) internal view returns (uint256) { |
| 108 | + return _length(s._inner, account); |
| 109 | + } |
| 110 | + |
| 111 | + function get(AddressArray storage s, address account, uint256 index) internal view returns (address) { |
| 112 | + return address(uint160(uint256(_get(s._inner, account, index)))); |
| 113 | + } |
| 114 | + |
| 115 | + function getAll(AddressArray storage s, address account) internal view returns (address[] memory) { |
| 116 | + bytes32[] memory bytes32Array = _getAll(s._inner, account); |
| 117 | + address[] memory addressArray; |
| 118 | + |
| 119 | + /// @solidity memory-safe-assembly |
| 120 | + assembly { |
| 121 | + addressArray := bytes32Array |
| 122 | + } |
| 123 | + return addressArray; |
| 124 | + } |
| 125 | + |
| 126 | + function contains(AddressArray storage s, address account, address value) internal view returns (bool) { |
| 127 | + return _contains(s._inner, account, bytes32(uint256(uint160(value)))); |
| 128 | + } |
| 129 | + |
| 130 | + function length(UintArray storage s, address account) internal view returns (uint256) { |
| 131 | + return _length(s._inner, account); |
| 132 | + } |
| 133 | + |
| 134 | + function get(UintArray storage s, address account, uint256 index) internal view returns (uint256) { |
| 135 | + return uint256(_get(s._inner, account, index)); |
| 136 | + } |
| 137 | + |
| 138 | + function getAll(UintArray storage s, address account) internal view returns (uint256[] memory) { |
| 139 | + bytes32[] memory bytes32Array = _getAll(s._inner, account); |
| 140 | + uint256[] memory uintArray; |
| 141 | + |
| 142 | + /// @solidity memory-safe-assembly |
| 143 | + assembly { |
| 144 | + uintArray := bytes32Array |
| 145 | + } |
| 146 | + return uintArray; |
| 147 | + } |
| 148 | + |
| 149 | + function contains(UintArray storage s, address account, uint256 value) internal view returns (bool) { |
| 150 | + return _contains(s._inner, account, bytes32(value)); |
| 151 | + } |
| 152 | + |
| 153 | + function _set(Array storage s, address account, uint256 index, bytes32 value) private { |
| 154 | + _set(_slot(s, account), index, value); |
| 155 | + } |
| 156 | + |
| 157 | + function _set(bytes32 slot, uint256 index, bytes32 value) private { |
| 158 | + assembly { |
| 159 | + //if (index >= _length(s, account)) revert AssociatedArray_OutOfBounds(index); |
| 160 | + if iszero(lt(index, sload(slot))) { |
| 161 | + mstore(0, 0x8277484f) // `AssociatedArray_OutOfBounds(uint256)` |
| 162 | + mstore(0x20, index) |
| 163 | + revert(0x1c, 0x24) |
| 164 | + } |
| 165 | + sstore(add(slot, mul(0x20, add(index, 1))), value) |
| 166 | + } |
| 167 | + } |
| 168 | + |
| 169 | + function _push(Array storage s, address account, bytes32 value) private { |
| 170 | + bytes32 slot = _slot(s, account); |
| 171 | + assembly { |
| 172 | + // load length (stored @ slot), add 1 to it => index. |
| 173 | + // mul index by 0x20 and add it to orig slot to get the next free slot |
| 174 | + let index := add(sload(slot), 1) |
| 175 | + sstore(add(slot, mul(0x20, index)), value) |
| 176 | + sstore(slot, index) //increment length by 1 |
| 177 | + } |
| 178 | + } |
| 179 | + |
| 180 | + function _pop(Array storage s, address account) private { |
| 181 | + bytes32 slot = _slot(s, account); |
| 182 | + uint256 __length; |
| 183 | + assembly { |
| 184 | + __length := sload(slot) |
| 185 | + } |
| 186 | + if (__length == 0) return; |
| 187 | + _set(slot, __length - 1, 0); |
| 188 | + assembly { |
| 189 | + sstore(slot, sub(__length, 1)) |
| 190 | + } |
| 191 | + } |
| 192 | + |
| 193 | + function _remove(Array storage s, address account, uint256 index) private { |
| 194 | + bytes32 slot = _slot(s, account); |
| 195 | + uint256 __length; |
| 196 | + assembly { |
| 197 | + __length := sload(slot) |
| 198 | + if iszero(lt(index, __length)) { |
| 199 | + mstore(0, 0x8277484f) // `AssociatedArray_OutOfBounds(uint256)` |
| 200 | + mstore(0x20, index) |
| 201 | + revert(0x1c, 0x24) |
| 202 | + } |
| 203 | + } |
| 204 | + _set(slot, index, _get(s, account, __length - 1)); |
| 205 | + |
| 206 | + assembly { |
| 207 | + // clear the last slot |
| 208 | + // this is the 'unchecked' version of _set(slot, __length - 1, 0) |
| 209 | + // as we use length-1 as index, so the check is excessive. |
| 210 | + // also removes extra -1 and +1 operations |
| 211 | + sstore(add(slot, mul(0x20, __length)), 0) |
| 212 | + // store new length |
| 213 | + sstore(slot, sub(__length, 1)) |
| 214 | + } |
| 215 | + } |
| 216 | + |
| 217 | + function _length(Array storage s, address account) private view returns (uint256 __length) { |
| 218 | + bytes32 slot = _slot(s, account); |
| 219 | + assembly { |
| 220 | + __length := sload(slot) |
| 221 | + } |
| 222 | + } |
| 223 | + |
| 224 | + function _get(Array storage s, address account, uint256 index) private view returns (bytes32 value) { |
| 225 | + return _get(_slot(s, account), index); |
| 226 | + } |
| 227 | + |
| 228 | + function _get(bytes32 slot, uint256 index) private view returns (bytes32 value) { |
| 229 | + assembly { |
| 230 | + //if (index >= _length(s, account)) revert AssociatedArray_OutOfBounds(index); |
| 231 | + if iszero(lt(index, sload(slot))) { |
| 232 | + mstore(0, 0x8277484f) // `AssociatedArray_OutOfBounds(uint256)` |
| 233 | + mstore(0x20, index) |
| 234 | + revert(0x1c, 0x24) |
| 235 | + } |
| 236 | + value := sload(add(slot, mul(0x20, add(index, 1)))) |
| 237 | + } |
| 238 | + } |
| 239 | + |
| 240 | + function _getAll(Array storage s, address account) private view returns (bytes32[] memory values) { |
| 241 | + bytes32 slot = _slot(s, account); |
| 242 | + uint256 __length; |
| 243 | + assembly { |
| 244 | + __length := sload(slot) |
| 245 | + } |
| 246 | + values = new bytes32[](__length); |
| 247 | + for (uint256 i; i < __length; i++) { |
| 248 | + values[i] = _get(slot, i); |
| 249 | + } |
| 250 | + } |
| 251 | + |
| 252 | + // inefficient. complexity = O(n) |
| 253 | + // use with caution |
| 254 | + // in case of large arrays, consider using EnumerableSet4337 instead |
| 255 | + function _contains(Array storage s, address account, bytes32 value) private view returns (bool) { |
| 256 | + bytes32 slot = _slot(s, account); |
| 257 | + uint256 __length; |
| 258 | + assembly { |
| 259 | + __length := sload(slot) |
| 260 | + } |
| 261 | + for (uint256 i; i < __length; i++) { |
| 262 | + if (_get(slot, i) == value) { |
| 263 | + return true; |
| 264 | + } |
| 265 | + } |
| 266 | + return false; |
| 267 | + } |
| 268 | + |
| 269 | + function _slot(Array storage s, address account) private pure returns (bytes32 __slot) { |
| 270 | + assembly { |
| 271 | + mstore(0x00, account) |
| 272 | + mstore(0x20, s.slot) |
| 273 | + __slot := keccak256(0x00, 0x40) |
| 274 | + } |
| 275 | + } |
| 276 | +} |
0 commit comments