From f783b455593e6271744a12eb5108becf8998c113 Mon Sep 17 00:00:00 2001
From: poma <Semenov.Roman@mail.ru>
Date: Sun, 3 Nov 2019 12:06:58 +0300
Subject: [PATCH] refactor isKnownRoot, add test

---
 contracts/MerkleTreeWithHistory.sol | 24 ++++++++++-------------
 test/MerkleTreeWithHistory.test.js  | 30 +++++++++++++++++++++++++++--
 2 files changed, 38 insertions(+), 16 deletions(-)

diff --git a/contracts/MerkleTreeWithHistory.sol b/contracts/MerkleTreeWithHistory.sol
index 7043237..6c3b1d7 100644
--- a/contracts/MerkleTreeWithHistory.sol
+++ b/contracts/MerkleTreeWithHistory.sol
@@ -97,20 +97,16 @@ contract MerkleTreeWithHistory {
     if (_root == 0) {
       return false;
     }
-    // search most recent first
-    uint256 i;
-    for(i = currentRootIndex; i < 2**256 - 1; i--) {
-      if (_root == roots[i]) {
-        return true;
-      }
-    }
-
-    // process the rest of roots
-    for(i = ROOT_HISTORY_SIZE - 1; i > currentRootIndex; i--) {
-      if (_root == roots[i]) {
-        return true;
-      }
-    }
+    uint256 i = currentRootIndex;
+    do {
+       if (_root == roots[i]) {
+           return true;
+       }
+       if (i == 0) {
+           i = ROOT_HISTORY_SIZE;
+       }
+       i--;
+    } while (i != currentRootIndex);
     return false;
   }
 
diff --git a/test/MerkleTreeWithHistory.test.js b/test/MerkleTreeWithHistory.test.js
index 279dde1..b10630e 100644
--- a/test/MerkleTreeWithHistory.test.js
+++ b/test/MerkleTreeWithHistory.test.js
@@ -172,7 +172,7 @@ contract('MerkleTreeWithHistory', accounts => {
 
     it('should reject if tree is full', async () => {
       levels = 6
-      merkleTreeWithHistory = await MerkleTreeWithHistory.new(levels)
+      const merkleTreeWithHistory = await MerkleTreeWithHistory.new(levels)
 
       for (let i = 0; i < 2**levels; i++) {
         await merkleTreeWithHistory.insert(i+42).should.be.fulfilled
@@ -187,7 +187,7 @@ contract('MerkleTreeWithHistory', accounts => {
 
     it.skip('hasher gas', async () => {
       levels = 6
-      merkleTreeWithHistory = await MerkleTreeWithHistory.new(levels)
+      const merkleTreeWithHistory = await MerkleTreeWithHistory.new(levels)
       const zeroValue = await merkleTreeWithHistory.zeroValue()
 
       const gas = await merkleTreeWithHistory.hashLeftRight.estimateGas(zeroValue, zeroValue)
@@ -195,6 +195,32 @@ contract('MerkleTreeWithHistory', accounts => {
     })
   })
 
+  describe('#isKnownRoot', () => {
+    it('should work', async () => {
+      let path
+
+      for (let i = 1; i < 5; i++) {
+        await merkleTreeWithHistory.insert(i, { from: sender }).should.be.fulfilled
+        await tree.insert(i)
+        path = await tree.path(i - 1)
+        let isKnown = await merkleTreeWithHistory.isKnownRoot(path.root)
+        isKnown.should.be.equal(true)
+      }
+
+      await merkleTreeWithHistory.insert(42, { from: sender }).should.be.fulfilled
+      // check outdated root
+      let isKnown = await merkleTreeWithHistory.isKnownRoot(path.root)
+      isKnown.should.be.equal(true)
+    })
+
+    it('should not return uninitialized roots', async () => {
+      await merkleTreeWithHistory.insert(42, { from: sender }).should.be.fulfilled
+      let isKnown = await merkleTreeWithHistory.isKnownRoot(0)
+      isKnown.should.be.equal(false)
+    })
+  })
+
+
   afterEach(async () => {
     await revertSnapshot(snapshotId.result)
     // eslint-disable-next-line require-atomic-updates