Skip to content

Commit 834bbc4

Browse files
VectorizedVectorized
authored
✨ EnumerableSetLib indexOf (#1460)
1 parent 2718072 commit 834bbc4

File tree

4 files changed

+334
-1
lines changed

4 files changed

+334
-1
lines changed

‎docs/utils/enumerablesetlib.md‎

Lines changed: 66 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,16 @@ error ExceedsCapacity()
4545

4646
Cannot accommodate a new unique value with the capacity.
4747

48+
## Constants
49+
50+
### NOT_FOUND
51+
52+
```solidity
53+
uint256 internal constant NOT_FOUND = type(uint256).max
54+
```
55+
56+
The index to represent a value that does not exist.
57+
4858
## Structs
4959

5060
### AddressSet
@@ -538,4 +548,59 @@ function at(Uint8Set storage set, uint256 i)
538548
returns (uint8 result)
539549
```
540550

541-
Returns the element at index `i` in the set. Reverts if `i` is out-of-bounds.
551+
Returns the element at index `i` in the set. Reverts if `i` is out-of-bounds.
552+
553+
### indexOf(AddressSet,address)
554+
555+
```solidity
556+
function indexOf(AddressSet storage set, address value)
557+
internal
558+
view
559+
returns (uint256 result)
560+
```
561+
562+
Returns the index of `value`. Returns `NOT_FOUND` if the value does not exist.
563+
564+
### indexOf(Bytes32Set,bytes32)
565+
566+
```solidity
567+
function indexOf(Bytes32Set storage set, bytes32 value)
568+
internal
569+
view
570+
returns (uint256 result)
571+
```
572+
573+
Returns the index of `value`. Returns `NOT_FOUND` if the value does not exist.
574+
575+
### indexOf(Uint256Set,uint256)
576+
577+
```solidity
578+
function indexOf(Uint256Set storage set, uint256 i)
579+
internal
580+
view
581+
returns (uint256 result)
582+
```
583+
584+
Returns the index of `value`. Returns `NOT_FOUND` if the value does not exist.
585+
586+
### indexOf(Int256Set,int256)
587+
588+
```solidity
589+
function indexOf(Int256Set storage set, int256 i)
590+
internal
591+
view
592+
returns (uint256 result)
593+
```
594+
595+
Returns the index of `value`. Returns `NOT_FOUND` if the value does not exist.
596+
597+
### indexOf(Uint8Set,uint8)
598+
599+
```solidity
600+
function indexOf(Uint8Set storage set, uint8 value)
601+
internal
602+
view
603+
returns (uint256 result)
604+
```
605+
606+
Returns the index of `value`. Returns `NOT_FOUND` if the value does not exist.

‎src/utils/EnumerableSetLib.sol‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,9 @@ library EnumerableSetLib {
3232
/* CONSTANTS */
3333
/*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/
3434

35+
/// @dev The index to represent a value that does not exist.
36+
uint256 internal constant NOT_FOUND = type(uint256).max;
37+
3538
/// @dev A sentinel value to denote the zero value in storage.
3639
/// No elements can be equal to this value.
3740
/// `uint72(bytes9(keccak256(bytes("_ZERO_SENTINEL"))))`.
@@ -756,6 +759,107 @@ library EnumerableSetLib {
756759
}
757760
}
758761

762+
/// @dev Returns the index of `value`. Returns `NOT_FOUND` if the value does not exist.
763+
function indexOf(AddressSet storage set, address value)
764+
internal
765+
view
766+
returns (uint256 result)
767+
{
768+
result = NOT_FOUND;
769+
if (uint160(value) == _ZERO_SENTINEL) return result;
770+
bytes32 rootSlot = _rootSlot(set);
771+
/// @solidity memory-safe-assembly
772+
assembly {
773+
if iszero(value) { value := _ZERO_SENTINEL }
774+
result := not(0)
775+
let rootPacked := sload(rootSlot)
776+
for {} 1 {} {
777+
if iszero(shr(160, shl(160, rootPacked))) {
778+
if eq(shr(96, rootPacked), value) {
779+
result := 0
780+
break
781+
}
782+
if eq(shr(96, sload(add(rootSlot, 1))), value) {
783+
result := 1
784+
break
785+
}
786+
if eq(shr(96, sload(add(rootSlot, 2))), value) {
787+
result := 2
788+
break
789+
}
790+
break
791+
}
792+
mstore(0x20, rootSlot)
793+
mstore(0x00, value)
794+
result := sub(sload(keccak256(0x00, 0x40)), 1)
795+
break
796+
}
797+
}
798+
}
799+
800+
/// @dev Returns the index of `value`. Returns `NOT_FOUND` if the value does not exist.
801+
function indexOf(Bytes32Set storage set, bytes32 value)
802+
internal
803+
view
804+
returns (uint256 result)
805+
{
806+
result = NOT_FOUND;
807+
if (uint256(value) == _ZERO_SENTINEL) return result;
808+
bytes32 rootSlot = _rootSlot(set);
809+
/// @solidity memory-safe-assembly
810+
assembly {
811+
if iszero(value) { value := _ZERO_SENTINEL }
812+
for {} 1 {} {
813+
if iszero(sload(not(rootSlot))) {
814+
if eq(sload(rootSlot), value) {
815+
result := 0
816+
break
817+
}
818+
if eq(sload(add(rootSlot, 1)), value) {
819+
result := 1
820+
break
821+
}
822+
if eq(sload(add(rootSlot, 2)), value) {
823+
result := 2
824+
break
825+
}
826+
break
827+
}
828+
mstore(0x20, rootSlot)
829+
mstore(0x00, value)
830+
result := sub(sload(keccak256(0x00, 0x40)), 1)
831+
break
832+
}
833+
}
834+
}
835+
836+
/// @dev Returns the index of `value`. Returns `NOT_FOUND` if the value does not exist.
837+
function indexOf(Uint256Set storage set, uint256 i) internal view returns (uint256 result) {
838+
result = indexOf(_toBytes32Set(set), bytes32(i));
839+
}
840+
841+
/// @dev Returns the index of `value`. Returns `NOT_FOUND` if the value does not exist.
842+
function indexOf(Int256Set storage set, int256 i) internal view returns (uint256 result) {
843+
result = indexOf(_toBytes32Set(set), bytes32(uint256(i)));
844+
}
845+
846+
/// @dev Returns the index of `value`. Returns `NOT_FOUND` if the value does not exist.
847+
function indexOf(Uint8Set storage set, uint8 value) internal view returns (uint256 result) {
848+
/// @solidity memory-safe-assembly
849+
assembly {
850+
result := not(0)
851+
let packed := sload(set.slot)
852+
let m := shl(and(0xff, value), 1)
853+
if and(packed, m) {
854+
result := 0
855+
for { let p := and(packed, sub(m, 1)) } p {} {
856+
p := xor(p, and(p, add(1, not(p))))
857+
result := add(result, 1)
858+
}
859+
}
860+
}
861+
}
862+
759863
/*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
760864
/* PRIVATE HELPERS */
761865
/*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/

‎src/utils/g/EnumerableSetLib.sol‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,9 @@ library EnumerableSetLib {
6969
/* CONSTANTS */
7070
/*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/
7171

72+
/// @dev The index to represent a value that does not exist.
73+
uint256 internal constant NOT_FOUND = type(uint256).max;
74+
7275
/// @dev A sentinel value to denote the zero value in storage.
7376
/// No elements can be equal to this value.
7477
/// `uint72(bytes9(keccak256(bytes("_ZERO_SENTINEL"))))`.
@@ -764,6 +767,107 @@ library EnumerableSetLib {
764767
}
765768
}
766769

770+
/// @dev Returns the index of `value`. Returns `NOT_FOUND` if the value does not exist.
771+
function indexOf(AddressSet storage set, address value)
772+
internal
773+
view
774+
returns (uint256 result)
775+
{
776+
result = NOT_FOUND;
777+
if (uint160(value) == _ZERO_SENTINEL) return result;
778+
bytes32 rootSlot = _rootSlot(set);
779+
/// @solidity memory-safe-assembly
780+
assembly {
781+
if iszero(value) { value := _ZERO_SENTINEL }
782+
result := not(0)
783+
let rootPacked := sload(rootSlot)
784+
for {} 1 {} {
785+
if iszero(shr(160, shl(160, rootPacked))) {
786+
if eq(shr(96, rootPacked), value) {
787+
result := 0
788+
break
789+
}
790+
if eq(shr(96, sload(add(rootSlot, 1))), value) {
791+
result := 1
792+
break
793+
}
794+
if eq(shr(96, sload(add(rootSlot, 2))), value) {
795+
result := 2
796+
break
797+
}
798+
break
799+
}
800+
mstore(0x20, rootSlot)
801+
mstore(0x00, value)
802+
result := sub(sload(keccak256(0x00, 0x40)), 1)
803+
break
804+
}
805+
}
806+
}
807+
808+
/// @dev Returns the index of `value`. Returns `NOT_FOUND` if the value does not exist.
809+
function indexOf(Bytes32Set storage set, bytes32 value)
810+
internal
811+
view
812+
returns (uint256 result)
813+
{
814+
result = NOT_FOUND;
815+
if (uint256(value) == _ZERO_SENTINEL) return result;
816+
bytes32 rootSlot = _rootSlot(set);
817+
/// @solidity memory-safe-assembly
818+
assembly {
819+
if iszero(value) { value := _ZERO_SENTINEL }
820+
for {} 1 {} {
821+
if iszero(sload(not(rootSlot))) {
822+
if eq(sload(rootSlot), value) {
823+
result := 0
824+
break
825+
}
826+
if eq(sload(add(rootSlot, 1)), value) {
827+
result := 1
828+
break
829+
}
830+
if eq(sload(add(rootSlot, 2)), value) {
831+
result := 2
832+
break
833+
}
834+
break
835+
}
836+
mstore(0x20, rootSlot)
837+
mstore(0x00, value)
838+
result := sub(sload(keccak256(0x00, 0x40)), 1)
839+
break
840+
}
841+
}
842+
}
843+
844+
/// @dev Returns the index of `value`. Returns `NOT_FOUND` if the value does not exist.
845+
function indexOf(Uint256Set storage set, uint256 i) internal view returns (uint256 result) {
846+
result = indexOf(_toBytes32Set(set), bytes32(i));
847+
}
848+
849+
/// @dev Returns the index of `value`. Returns `NOT_FOUND` if the value does not exist.
850+
function indexOf(Int256Set storage set, int256 i) internal view returns (uint256 result) {
851+
result = indexOf(_toBytes32Set(set), bytes32(uint256(i)));
852+
}
853+
854+
/// @dev Returns the index of `value`. Returns `NOT_FOUND` if the value does not exist.
855+
function indexOf(Uint8Set storage set, uint8 value) internal view returns (uint256 result) {
856+
/// @solidity memory-safe-assembly
857+
assembly {
858+
result := not(0)
859+
let packed := sload(set.slot)
860+
let m := shl(and(0xff, value), 1)
861+
if and(packed, m) {
862+
result := 0
863+
for { let p := and(packed, sub(m, 1)) } p {} {
864+
p := xor(p, and(p, add(1, not(p))))
865+
result := add(result, 1)
866+
}
867+
}
868+
}
869+
}
870+
767871
/*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
768872
/* PRIVATE HELPERS */
769873
/*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/

‎test/EnumerableSetLib.t.sol‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -844,4 +844,64 @@ contract EnumerableSetLibTest is SoladyTest {
844844
function updateUint8Set(uint8 value, bool isAdd, uint256 cap) public returns (bool) {
845845
return uint8Set.update(value, isAdd, cap);
846846
}
847+
848+
function testAddressSetIndexOf(bytes32 r) public {
849+
LibPRNG.PRNG memory prng;
850+
prng.state = uint256(r);
851+
address[] memory a = new address[](_bound(_random(), 0, 16));
852+
853+
for (uint256 i; i != a.length; ++i) {
854+
address value = address(uint160(prng.next()));
855+
addressSet.add(value);
856+
a[i] = value;
857+
}
858+
859+
if (a.length != 0) {
860+
uint256 i = _random() % a.length;
861+
address value = a[i];
862+
assertEq(addressSet.indexOf(value), i);
863+
assertEq(addressSet.at(addressSet.indexOf(value)), value);
864+
if (_randomChance(2)) {
865+
value = address(bytes20(keccak256(abi.encode(value, value, "abcdef"))));
866+
assertEq(addressSet.indexOf(value), EnumerableSetLib.NOT_FOUND);
867+
}
868+
}
869+
}
870+
871+
function testBytes32SetIndexOf(bytes32 r) public {
872+
LibPRNG.PRNG memory prng;
873+
prng.state = uint256(r);
874+
bytes32[] memory a = new bytes32[](_bound(_random(), 0, 16));
875+
876+
for (uint256 i; i != a.length; ++i) {
877+
bytes32 value = bytes32(prng.next());
878+
bytes32Set.add(value);
879+
a[i] = value;
880+
}
881+
882+
if (a.length != 0) {
883+
uint256 i = _random() % a.length;
884+
bytes32 value = a[i];
885+
assertEq(bytes32Set.indexOf(value), i);
886+
assertEq(bytes32Set.at(bytes32Set.indexOf(value)), value);
887+
if (_randomChance(2)) {
888+
value = keccak256(abi.encode(value, value, "abcdef"));
889+
assertEq(bytes32Set.indexOf(value), EnumerableSetLib.NOT_FOUND);
890+
}
891+
}
892+
}
893+
894+
function testEnumerableUint8SetIndexOf(uint256 flags, uint8 value) public {
895+
uint8[] memory ordinals = _flagsToOrdinals(flags);
896+
unchecked {
897+
for (uint256 i; i != ordinals.length; ++i) {
898+
uint8Set.add(ordinals[i]);
899+
}
900+
if (uint8Set.contains(value)) {
901+
assertEq(uint8Set.at(uint8Set.indexOf(value)), value);
902+
} else {
903+
assertEq(uint8Set.indexOf(value), EnumerableSetLib.NOT_FOUND);
904+
}
905+
}
906+
}
847907
}

0 commit comments

Comments
 (0)