Implementation of an address Enumerable Set (OpenZeppelin#2061)

* Drafted Enumerable.sol.

* Drafted test framework.

* Tweaked the tests to follow oz structure.

* Coded EnumerableSet.

* Moved EnumerableSet to `utils`.

* Fixed linting.

* Improved comments.

* Tweaked contract description.

* Renamed struct to AddressSet.

* Relaxed version pragma to 0.5.0

* Removed events.

* Revert on useless operations.

* Small comment.

* Created AddressSet factory method.

* Failed transactions return false.

* Transactions now return false on failure.

* Remove comments from mock

* Rename mock functions

* Adapt tests to code style, use test-helpers

* Fix bug in remove, improve tests.

* Add changelog entry

* Add entry on Utils doc

* Add optimization for removal of last slot

* Update docs

* Fix headings of utilities documentation

Co-authored-by: Nicolás Venturo <[email protected]>

Author: 2 people

CHANGELOG.md

@@ -5,6 +5,7 @@
 ### New features
  * `SafeCast.toUintXX`: new library for integer downcasting, which allows for safe operation on smaller types (e.g. `uint32`) when combined with `SafeMath`. ([#1926](https://github.com/OpenZeppelin/openzeppelin-solidity/pull/1926))
  * `ERC721Metadata`: added `baseURI`, which can be used for dramatic gas savings when all token URIs share a prefix (e.g. `http://api.myapp.com/tokens/<id>`). ([#1970](https://github.com/OpenZeppelin/openzeppelin-solidity/pull/1970))
+ * `EnumerableSet`: new library for storing enumerable sets of values. Only `AddressSet` is supported in this release. ([#2061](https://github.com/OpenZeppelin/openzeppelin-solidity/pull/2061))
  * `Create2`: simple library to make usage of the `CREATE2` opcode easier. ([#1744](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/1744))
 
 ### Improvements

contracts/mocks/EnumerableSetMock.sol

@@ -0,0 +1,33 @@
+pragma solidity ^0.5.0;
+
+import "../utils/EnumerableSet.sol";
+
+contract EnumerableSetMock{
+    using EnumerableSet for EnumerableSet.AddressSet;
+
+    event TransactionResult(bool result);
+
+    EnumerableSet.AddressSet private set;
+
+    constructor() public {
+        set = EnumerableSet.newAddressSet();
+    }
+
+    function contains(address value) public view returns (bool) {
+        return EnumerableSet.contains(set, value);
+    }
+
+    function add(address value) public {
+        bool result = EnumerableSet.add(set, value);
+        emit TransactionResult(result);
+    }
+
+    function remove(address value) public {
+        bool result = EnumerableSet.remove(set, value);
+        emit TransactionResult(result);
+    }
+
+    function enumerate() public view returns (address[] memory) {
+        return EnumerableSet.enumerate(set);
+    }
+}

contracts/utils/EnumerableSet.sol

@@ -0,0 +1,120 @@
+pragma solidity ^0.5.0;
+
+/**
+ * @dev Library for managing
+ * https://en.wikipedia.org/wiki/Set_(abstract_data_type)[sets] of primitive
+ * types.
+ *
+ * Sets have the following properties:
+ *
+ * - Elements are added, removed, and checked for existence in constant time
+ * (O(1)).
+ * - Elements are enumerated in O(n). No guarantees are made on the ordering.
+ *
+ * As of v2.5.0, only `address` sets are supported.
+ *
+ * Include with `using EnumerableSet for EnumerableSet.AddressSet;`, and use
+ * {newAddressSet} to create a new `AddressSet`.
+ *
+ * _Available since v2.5.0._
+ *
+ * @author Alberto Cuesta Cañada
+ */
+library EnumerableSet {
+
+    struct AddressSet {
+        // Position of the value in the `values` array, plus 1 because index 0
+        // means a value is not in the set.
+        mapping (address => uint256) index;
+        address[] values;
+    }
+
+    /**
+     * @dev Creates a new empty address set.
+     */
+    function newAddressSet()
+        internal
+        pure
+        returns (AddressSet memory)
+    {
+        return AddressSet({values: new address[](0)});
+    }
+
+    /**
+     * @dev Add a value to a set. O(1).
+     * Returns false if the value was already in the set.
+     */
+    function add(AddressSet storage set, address value)
+        internal
+        returns (bool)
+    {
+        if (!contains(set, value)){
+            set.index[value] = set.values.push(value);
+            return true;
+        } else {
+            return false;
+        }
+    }
+
+    /**
+     * @dev Removes a value from a set. O(1).
+     * Returns false if the value was not present in the set.
+     */
+    function remove(AddressSet storage set, address value)
+        internal
+        returns (bool)
+    {
+        if (contains(set, value)){
+            uint256 toDeleteIndex = set.index[value] - 1;
+            uint256 lastIndex = set.values.length - 1;
+
+            // If the element we're deleting is the last one, we can just remove it without doing a swap
+            if (lastIndex != toDeleteIndex) {
+                address lastValue = set.values[lastIndex];
+
+                // Move the last value to the index where the deleted value is
+                set.values[toDeleteIndex] = lastValue;
+                // Update the index for the moved value
+                set.index[lastValue] = toDeleteIndex + 1; // All indexes are 1-based
+            }
+
+            // Delete the index entry for the deleted value
+            delete set.index[value];
+
+            // Delete the old entry for the moved value
+            set.values.pop();
+
+            return true;
+        } else {
+            return false;
+        }
+    }
+
+    /**
+     * @dev Returns true if the value is in the set. O(1).
+     */
+    function contains(AddressSet storage set, address value)
+        internal
+        view
+        returns (bool)
+    {
+        return set.index[value] != 0;
+    }
+
+    /**
+     * @dev Returns an array with all values in the set. O(N).
+     * Note that there are no guarantees on the ordering of values inside the
+     * array, and it may change when more values are added or removed.
+     */
+    function enumerate(AddressSet storage set)
+        internal
+        view
+        returns (address[] memory)
+    {
+        address[] memory output = new address[](set.values.length);
+        for (uint256 i; i < set.values.length; i++){
+            output[i] = set.values[i];
+        }
+        return output;
+    }
+}

contracts/utils/README.adoc

@@ -10,6 +10,8 @@ Miscellaneous contracts containing utility functions, often related to working w
 
 {{Arrays}}
 
+{{EnumerableSet}}
+
 {{Create2}}
 
 {{ReentrancyGuard}}

docs/modules/ROOT/pages/utilities.adoc

@@ -83,13 +83,18 @@ Easy!
 
 Want to split some payments between multiple people? Maybe you have an app that sends 30% of art purchases to the original creator and 70% of the profits to the current owner; you can build that with xref:api:payment.adoc#PaymentSplitter[`PaymentSplitter`]!
 
-In solidity, there are some security concerns with blindly sending money to accounts, since it allows them to execute arbitrary code. You can read up on these security concerns in the https://consensys.github.io/smart-contract-best-practices/[Ethereum Smart Contract Best Practices] website. One of the ways to fix reentrancy and stalling problems is, instead of immediately sending Ether to accounts that need it, you can use xref:api:payment.adoc#PullPayment[`PullPayment`], which offers an xref:api:payment.adoc#PullPayment-_asyncTransfer-address-uint256-[`_asyncTransfer`] function for sending money to something and requesting that they xref:api:payment.adoc#PullPayment-withdrawPayments-address-payable-[`withdrawPayments()`] it later.
+In Solidity, there are some security concerns with blindly sending money to accounts, since it allows them to execute arbitrary code. You can read up on these security concerns in the https://consensys.github.io/smart-contract-best-practices/[Ethereum Smart Contract Best Practices] website. One of the ways to fix reentrancy and stalling problems is, instead of immediately sending Ether to accounts that need it, you can use xref:api:payment.adoc#PullPayment[`PullPayment`], which offers an xref:api:payment.adoc#PullPayment-_asyncTransfer-address-uint256-[`_asyncTransfer`] function for sending money to something and requesting that they xref:api:payment.adoc#PullPayment-withdrawPayments-address-payable-[`withdrawPayments()`] it later.
 
 If you want to Escrow some funds, check out xref:api:payment.adoc#Escrow[`Escrow`] and xref:api:payment.adoc#ConditionalEscrow[`ConditionalEscrow`] for governing the release of some escrowed Ether.
 
+[[collections]]
+== Collections
+
+If you need support for more powerful collections than Solidity's native arrays and mappings, take a look at xref:api:utils.adoc#EnumerableSet[`EnumerableSet`]. It is similar to a mapping in that it stores and removes elements in constant time and doesn't allow for repeated entries, but it also supports _enumeration_, which means you can easily query all elements of the set both on and off-chain.
+
 [[misc]]
-=== Misc
+== Misc
 
 Want to check if an address is a contract? Use xref:api:utils.adoc#Address[`Address`] and xref:api:utils.adoc#Address-isContract-address-[`Address.isContract()`].
 
-Want to keep track of some numbers that increment by 1 every time you want another one? Check out xref:api:drafts.adoc#Counter[`Counter`]. This is especially useful for creating incremental ERC721 `tokenId` s like we did in the last section.
+Want to keep track of some numbers that increment by 1 every time you want another one? Check out xref:api:drafts.adoc#Counter[`Counter`]. This is useful for lots of things, like creating incremental identifiers, as shown on the xref:721.adoc[ERC721 guide].

test/utils/EnumerableSet.test.js

@@ -0,0 +1,108 @@
+const { accounts, contract } = require('@openzeppelin/test-environment');
+const { expectEvent } = require('@openzeppelin/test-helpers');
+const { expect } = require('chai');
+
+const EnumerableSetMock = contract.fromArtifact('EnumerableSetMock');
+
+describe('EnumerableSet', function () {
+  const [ accountA, accountB, accountC ] = accounts;
+
+  beforeEach(async function () {
+    this.set = await EnumerableSetMock.new();
+  });
+
+  it('starts empty', async function () {
+    expect(await this.set.contains(accountA)).to.equal(false);
+    expect(await this.set.enumerate()).to.have.same.members([]);
+  });
+
+  it('adds a value', async function () {
+    const receipt = await this.set.add(accountA);
+    expectEvent(receipt, 'TransactionResult', { result: true });
+
+    expect(await this.set.contains(accountA)).to.equal(true);
+    expect(await this.set.enumerate()).to.have.same.members([ accountA ]);
+  });
+
+  it('adds several values', async function () {
+    await this.set.add(accountA);
+    await this.set.add(accountB);
+
+    expect(await this.set.contains(accountA)).to.equal(true);
+    expect(await this.set.contains(accountB)).to.equal(true);
+
+    expect(await this.set.contains(accountC)).to.equal(false);
+
+    expect(await this.set.enumerate()).to.have.same.members([ accountA, accountB ]);
+  });
+
+  it('returns false when adding elements already in the set', async function () {
+    await this.set.add(accountA);
+
+    const receipt = (await this.set.add(accountA));
+    expectEvent(receipt, 'TransactionResult', { result: false });
+
+    expect(await this.set.enumerate()).to.have.same.members([ accountA ]);
+  });
+
+  it('removes added values', async function () {
+    await this.set.add(accountA);
+
+    const receipt = await this.set.remove(accountA);
+    expectEvent(receipt, 'TransactionResult', { result: true });
+
+    expect(await this.set.contains(accountA)).to.equal(false);
+    expect(await this.set.enumerate()).to.have.same.members([]);
+  });
+
+  it('returns false when removing elements not in the set', async function () {
+    const receipt = await this.set.remove(accountA);
+    expectEvent(receipt, 'TransactionResult', { result: false });
+
+    expect(await this.set.contains(accountA)).to.equal(false);
+  });
+
+  it('adds and removes multiple values', async function () {
+    // []
+
+    await this.set.add(accountA);
+    await this.set.add(accountC);
+
+    // [A, C]
+
+    await this.set.remove(accountA);
+    await this.set.remove(accountB);
+
+    // [C]
+
+    await this.set.add(accountB);
+
+    // [C, B]
+
+    await this.set.add(accountA);
+    await this.set.remove(accountC);
+
+    // [A, B]
+
+    await this.set.add(accountA);
+    await this.set.add(accountB);
+
+    // [A, B]
+
+    await this.set.add(accountC);
+    await this.set.remove(accountA);
+
+    // [B, C]
+
+    await this.set.add(accountA);
+    await this.set.remove(accountB);
+
+    // [A, C]
+
+    expect(await this.set.contains(accountA)).to.equal(true);
+    expect(await this.set.contains(accountB)).to.equal(false);
+    expect(await this.set.contains(accountC)).to.equal(true);
+
+    expect(await this.set.enumerate()).to.have.same.members([ accountA, accountC ]);
+  });
+});