minor

contracts/TornadoTrees.sol

@@ -82,11 +82,9 @@ contract TornadoTrees is EnsResolve {
     tornadoTreesV1 = _tornadoTreesV1;
 
     depositRoot = _tornadoTreesV1.depositRoot();
-    withdrawalRoot = _tornadoTreesV1.withdrawalRoot();
-
-    uint256 depositLeaf = _tornadoTreesV1.lastProcessedDepositLeaf();
-    require(depositLeaf % CHUNK_SIZE == 0, "Incorrect TornadoTrees state");
-    lastProcessedDepositLeaf = depositLeaf;
+    uint256 lastDepositLeaf = _tornadoTreesV1.lastProcessedDepositLeaf();
+    require(lastDepositLeaf % CHUNK_SIZE == 0, "Incorrect TornadoTrees state");
+    lastProcessedDepositLeaf = lastDepositLeaf;
     depositsLength = depositV1Length = findArrayLength(
       _tornadoTreesV1,
       "deposits(uint256)",
@@ -94,9 +92,10 @@ contract TornadoTrees is EnsResolve {
       _searchParams.depositsPerDay
     );
 
-    uint256 withdrawalLeaf = _tornadoTreesV1.lastProcessedWithdrawalLeaf();
-    require(withdrawalLeaf % CHUNK_SIZE == 0, "Incorrect TornadoTrees state");
-    lastProcessedWithdrawalLeaf = withdrawalLeaf;
+    withdrawalRoot = _tornadoTreesV1.withdrawalRoot();
+    uint256 lastWithdrawalLeaf = _tornadoTreesV1.lastProcessedWithdrawalLeaf();
+    require(lastWithdrawalLeaf % CHUNK_SIZE == 0, "Incorrect TornadoTrees state");
+    lastProcessedWithdrawalLeaf = lastWithdrawalLeaf;
     withdrawalsLength = withdrawalsV1Length = findArrayLength(
       _tornadoTreesV1,
       "withdrawals(uint256)",
@@ -105,37 +104,34 @@ contract TornadoTrees is EnsResolve {
     );
   }
 
+  // todo make things internal
+  /// @dev There is no array length getter for deposit and withdrawal arrays
+  /// in previous contract, so we have to find them length manually
   function findArrayLength(
     ITornadoTreesV1 _tornadoTreesV1,
     string memory _type,
-    uint256 _from,
-    uint256 _step
+    uint256 _from, // most likely array length after the proposal has passed
+    uint256 _step // optimal step size to find first match, approximately equals dispersion
   ) public view returns (uint256) {
-    require(_from != 0 && _step != 0, "_from and _step should be > 0");
-    require(elementExists(_tornadoTreesV1, _type, _from), "Inccorrect _from param");
+    // Find the segment with correct array length
+    bool direction = elementExists(_tornadoTreesV1, _type, _from);
+    do {
+      _from = direction ? _from + _step : _from - _step;
+    } while (direction == elementExists(_tornadoTreesV1, _type, _from));
+    uint256 high = direction ? _from : _from + _step;
+    uint256 low = direction ? _from - _step : _from;
+    uint256 mid = (high + low) / 2;
 
-    uint256 index = _from + _step;
-    while (elementExists(_tornadoTreesV1, _type, index)) {
-      index = index + _step;
-    }
-
-    uint256 high = index;
-    uint256 low = index - _step;
-    uint256 mid = (low + high) / 2;
-    while (!elementExists(_tornadoTreesV1, _type, mid)) {
-      high = mid - 1;
+    // Perform a binary search in this segment
+    while (low < mid) {
+      if (elementExists(_tornadoTreesV1, _type, mid)) {
+        low = mid;
+      } else {
+        high = mid;
+      }
       mid = (low + high) / 2;
     }
-
-    high += 1;
-    low = mid + 1;
-    mid = (low + high) / 2;
-    while (elementExists(_tornadoTreesV1, _type, mid)) {
-      low = mid + 1;
-      mid = (low + high) / 2;
-    }
-
-    return high == low ? high : low;
+    return mid + 1;
   }
 
   function elementExists(
@@ -143,6 +139,7 @@ contract TornadoTrees is EnsResolve {
     string memory _type,
     uint256 index
   ) public view returns (bool success) {
+    // Try to get the element. If it succeeds the array length is higher, it it reverts the length is equal or lower
     (success, ) = address(_tornadoTreesV1).staticcall{ gas: 2500 }(abi.encodeWithSignature(_type, index));
   }
 

test/binarySearch.test.js

@@ -75,4 +75,36 @@ 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)
+  })
+
+  it('should pass stress test', async () => {
+    const iterations = 30
+    const days = 10
+    const depositsPerDay = 10
+    const dispersion = 5
+
+    for (let i = 0; i < iterations; i++) {
+      let len = 0
+      for (let j = 0; j < days; j++) {
+        len += depositsPerDay + Math.round((Math.random() - 0.5) * 2 * dispersion)
+      }
+      const deposits = Array.from(Array(len).keys())
+      publicArray = await PublicArray.deploy()
+      await publicArray.setDeposits(deposits)
+      const depositsLength = await tornadoTrees.findArrayLength(
+        publicArray.address,
+        'deposits(uint256)',
+        days * depositsPerDay,
+        dispersion * 2,
+      )
+      expect(depositsLength).to.be.equal(deposits.length)
+    }
+  })
 })