From 31d48c2b57b1c8b89ec1e229313a1b26d4229816 Mon Sep 17 00:00:00 2001
From: Alexey <peretsmobil@gmail.com>
Date: Tue, 9 Feb 2021 12:27:12 +0300
Subject: [PATCH] findArrayLength fix

---
 contracts/TornadoTrees.sol | 42 +++++++++++++++++++++++++++-----------
 test/binarySearch.test.js  |  8 ++++++++
 2 files changed, 38 insertions(+), 12 deletions(-)

diff --git a/contracts/TornadoTrees.sol b/contracts/TornadoTrees.sol
index 077ea35..ce0aeeb 100644
--- a/contracts/TornadoTrees.sol
+++ b/contracts/TornadoTrees.sol
@@ -105,6 +105,7 @@ contract TornadoTrees is EnsResolve {
     );
   }
 
+  // todo make things internal. Think of dedundant calls
   function findArrayLength(
     ITornadoTreesV1 _tornadoTreesV1,
     string memory _type,
@@ -112,7 +113,7 @@ contract TornadoTrees is EnsResolve {
     uint256 _step
   ) public view returns (uint256) {
     require(_from != 0 && _step != 0, "_from and _step should be > 0");
-    require(elementExists(_tornadoTreesV1, _type, _from), "Inccorrect _from param");
+    // require(elementExists(_tornadoTreesV1, _type, _from), "Inccorrect _from param");
 
     uint256 index = _from + _step;
     while (elementExists(_tornadoTreesV1, _type, index)) {
@@ -121,21 +122,38 @@ contract TornadoTrees is EnsResolve {
 
     uint256 high = index;
     uint256 low = index - _step;
-    uint256 mid = (low + high) / 2;
-    while (!elementExists(_tornadoTreesV1, _type, mid)) {
-      high = mid - 1;
-      mid = (low + high) / 2;
+    uint256 lastIndex = binarySearch(_tornadoTreesV1, _type, high, low);
+
+    return lastIndex + 1;
+  }
+
+  function binarySearch(
+    ITornadoTreesV1 _tornadoTreesV1,
+    string memory _type,
+    uint256 _high,
+    uint256 _low
+  ) public view returns (uint256) {
+    require(_high >= _low, "Incorrect params");
+    uint256 mid = (_high + _low) / 2;
+    (bool isLast, bool exists) = isLastElement(_tornadoTreesV1, _type, mid);
+    if (isLast) {
+      return mid;
     }
 
-    high += 1;
-    low = mid + 1;
-    mid = (low + high) / 2;
-    while (elementExists(_tornadoTreesV1, _type, mid)) {
-      low = mid + 1;
-      mid = (low + high) / 2;
+    if (exists) {
+      return binarySearch(_tornadoTreesV1, _type, _high, mid + 1);
+    } else {
+      return binarySearch(_tornadoTreesV1, _type, mid - 1, _low);
     }
+  }
 
-    return high == low ? high : low;
+  function isLastElement(
+    ITornadoTreesV1 _tornadoTreesV1,
+    string memory _type,
+    uint256 index
+  ) public view returns (bool success, bool exists) {
+    exists = elementExists(_tornadoTreesV1, _type, index);
+    success = exists && !elementExists(_tornadoTreesV1, _type, index + 1);
   }
 
   function elementExists(
diff --git a/test/binarySearch.test.js b/test/binarySearch.test.js
index f8bc5b3..f3633e9 100644
--- a/test/binarySearch.test.js
+++ b/test/binarySearch.test.js
@@ -75,4 +75,12 @@ describe('findArrayLength', () => {
     )
     expect(depositsLength).to.be.equal(deposits.length)
   })
+
+  it('should work for an array and big big step', async () => {
+    const deposits = Array.from(Array(30).keys())
+    publicArray = await PublicArray.deploy()
+    await publicArray.setDeposits(deposits)
+    const depositsLength = await tornadoTrees.findArrayLength(publicArray.address, 'deposits(uint256)', 1, 50)
+    expect(depositsLength).to.be.equal(deposits.length)
+  })
 })