Remove unbounded loop in size() implementation

gnosis · Apr 8, 2020 · 8ada5a1 · 8ada5a1
1 parent 065b0f8
commit 8ada5a1
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 56 deletions.
diff --git a/contracts/libraries/IterableAppendOnlySet.sol b/contracts/libraries/IterableAppendOnlySet.sol
@@ -5,6 +5,7 @@ library IterableAppendOnlySet {
     struct Data {
         mapping(address => address) nextMap;
         address last;
+        uint96 size; // width is chosen to align struct size to full words
     }
 
     function insert(Data storage self, address value) public returns (bool) {
@@ -13,6 +14,7 @@ library IterableAppendOnlySet {
         }
         self.nextMap[self.last] = value;
         self.last = value;
+        self.size += 1;
         return true;
     }
 
@@ -31,26 +33,4 @@ library IterableAppendOnlySet {
         require(value != self.last, "Trying to get next of last element");
         return self.nextMap[value];
     }
-
-    function size(Data storage self) public view returns (uint256) {
-        if (self.last == address(0)) {
-            return 0;
-        }
-        uint256 count = 1;
-        address current = first(self);
-        while (current != self.last) {
-            current = next(self, current);
-            count++;
-        }
-        return count;
-    }
-
-    function atIndex(Data storage self, uint256 index) public view returns (address) {
-        require(index < size(self), "requested index too large");
-        address res = first(self);
-        for (uint256 i = 0; i < index; i++) {
-            res = next(self, res);
-        }
-        return res;
-    }
 }
diff --git a/contracts/libraries/wrappers/IterableAppendOnlySetWrapper.sol b/contracts/libraries/wrappers/IterableAppendOnlySetWrapper.sol
@@ -23,15 +23,11 @@ contract IterableAppendOnlySetWrapper {
         return data.last;
     }
 
-    function size() public view returns (uint256) {
-        return data.size();
+    function size() public view returns (uint96) {
+        return data.size;
     }
 
     function next(address value) public view returns (address) {
         return data.next(value);
     }
-
-    function atIndex(uint256 index) public view returns (address) {
-        return data.atIndex(index);
-    }
 }
diff --git a/test/libraries/iterable_append_only_set.js b/test/libraries/iterable_append_only_set.js
@@ -1,5 +1,7 @@
 const IterableSetLib = artifacts.require(".libraries/IterableAppendOnlySet.sol")
-const IterableSet = artifacts.require(".libraries/wrappers/IterableAppendOnlySetWrapper.sol")
+const IterableSet = artifacts.require(
+  ".libraries/wrappers/IterableAppendOnlySetWrapper.sol"
+)
 
 const truffleAssert = require("truffle-assertions")
 
@@ -19,7 +21,7 @@ async function getSetContent(set) {
   return result
 }
 
-contract("IterableSet", function (accounts) {
+contract("IterableSet", function(accounts) {
   beforeEach(async () => {
     const lib = await IterableSetLib.new()
     await IterableSet.link(IterableSetLib, lib.address)
@@ -29,10 +31,18 @@ contract("IterableSet", function (accounts) {
     const set = await IterableSet.new()
     assert.deepEqual(await getSetContent(set), [])
 
-    assert.equal(await set.contains(accounts[0]), false, "The element should not be there")
+    assert.equal(
+      await set.contains(accounts[0]),
+      false,
+      "The element should not be there"
+    )
     await set.insert(accounts[0])
 
-    assert.equal(await set.contains(accounts[0]), true, "The element should be there")
+    assert.equal(
+      await set.contains(accounts[0]),
+      true,
+      "The element should be there"
+    )
     assert.deepEqual(await getSetContent(set), accounts.slice(0, 1))
   })
 
@@ -54,11 +64,19 @@ contract("IterableSet", function (accounts) {
   it("should insert the same value only once", async () => {
     const set = await IterableSet.new()
 
-    assert.equal(await set.insert.call(accounts[0]), true, "First insert should insert")
+    assert.equal(
+      await set.insert.call(accounts[0]),
+      true,
+      "First insert should insert"
+    )
     await set.insert(accounts[0])
     assert.deepEqual(await getSetContent(set), accounts.slice(0, 1))
 
-    assert.equal(await set.insert.call(accounts[0]), false, "Second insert should do nothing")
+    assert.equal(
+      await set.insert.call(accounts[0]),
+      false,
+      "Second insert should do nothing"
+    )
     await set.insert(accounts[0])
     assert.deepEqual(await getSetContent(set), accounts.slice(0, 1))
   })
@@ -114,25 +132,4 @@ contract("IterableSet", function (accounts) {
     await set.insert(accounts[0])
     await truffleAssert.reverts(set.next(accounts[0]))
   })
-
-  it("fetches elements at index (by insertion order)", async () => {
-    const set = await IterableSet.new()
-
-    await set.insert(accounts[0])
-    await set.insert(accounts[1])
-
-    assert.equal(await set.atIndex(0), accounts[0])
-    assert.equal(await set.atIndex(1), accounts[1])
-  })
-  it("fails to fetches at index when requested index is too large", async () => {
-    const set = await IterableSet.new()
-
-    // Nothing in the set yet!
-    await truffleAssert.reverts(set.atIndex(0))
-
-    await set.insert(accounts[0])
-    await truffleAssert.reverts(set.atIndex(1))
-  })
-
-
-})
+})