From 305d714ca789f17890a71f81ddb0f947cd30d677 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nicol=C3=A1s=20Venturo?= <nicolas.venturo@gmail.com>
Date: Fri, 24 Jan 2020 12:21:33 -0300
Subject: [PATCH] Fix bug in remove, improve tests.

---
 contracts/utils/EnumerableSet.sol | 42 ++++++++++++++++++++-----------
 contracts/utils/README.adoc       |  2 ++
 test/utils/EnumerableSet.test.js  | 27 +++++++++++++++++---
 3 files changed, 53 insertions(+), 18 deletions(-)

diff --git a/contracts/utils/EnumerableSet.sol b/contracts/utils/EnumerableSet.sol
index 76eb8d55c..7b4b586ec 100644
--- a/contracts/utils/EnumerableSet.sol
+++ b/contracts/utils/EnumerableSet.sol
@@ -1,20 +1,22 @@
 pragma solidity ^0.5.0;
 
-
 /**
- * @title EnumerableSet
- * @dev Data structure - https://en.wikipedia.org/wiki/Set_(abstract_data_type)
+ * @dev Library for managing
+ * https://en.wikipedia.org/wiki/Set_(abstract_data_type)[sets] of primitive
+ * types.
  *
- * An EnumerableSet.AddressSet is a data structure containing a number of unique addresses.
+ * Sets have the following properties:
  *
- *  - It is possible to add and remove addresses in O(1).
- *  - It is also possible to query if the AddressSet contains an address in O(1).
- *  - It is possible to retrieve an array with all the addresses in the AddressSet using enumerate.
- *    This operation is O(N) where N is the number of addresses in the AddressSet. The order in
- *    which the addresses are retrieved is not guaranteed.
+ * - Elements are added, removed, and checked for existence in constant time
+ * (O(1)).
+ * - Elements are enumerated in O(n). The ordering in this list may change.
  *
- * Initialization of a set must include an empty array:
- * `EnumerableSet.AddressSet set = EnumerableSet.AddressSet({values: new address[](0)});`
+ * 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
  */
@@ -61,10 +63,22 @@ library EnumerableSet {
         returns (bool)
     {
         if (contains(set, value)){
-            // Replaced the value to remove with the last one in the array. O(1)
-            set.values[set.index[value] - 1] = set.values[set.values.length - 1];
-            set.values.pop();
+            uint256 toDeleteIndex = set.index[value] - 1;
+            uint256 lastIndex = set.values.length - 1;
+
+            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;
diff --git a/contracts/utils/README.adoc b/contracts/utils/README.adoc
index 229080211..cb66b4020 100644
--- a/contracts/utils/README.adoc
+++ b/contracts/utils/README.adoc
@@ -10,4 +10,6 @@ Miscellaneous contracts containing utility functions, often related to working w
 
 {{Arrays}}
 
+{{EnumerableSet}}
+
 {{ReentrancyGuard}}
diff --git a/test/utils/EnumerableSet.test.js b/test/utils/EnumerableSet.test.js
index b411cea6a..0162da5fa 100644
--- a/test/utils/EnumerableSet.test.js
+++ b/test/utils/EnumerableSet.test.js
@@ -64,26 +64,45 @@ describe('EnumerableSet', function () {
 
   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]
 
-    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);
+    await this.set.add(accountA);
+    await this.set.add(accountB);
 
-    expect(await this.set.enumerate()).to.have.same.members([ accountA, 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 ]);
   });
 });