bump

contracts/TornadoTrees.sol

@@ -3,20 +3,17 @@
 pragma solidity ^0.6.0;
 pragma experimental ABIEncoderV2;
 
-import "torn-token/contracts/ENS.sol";
 import "./interfaces/ITornadoTreesV1.sol";
-import "./interfaces/IVerifier.sol";
+import "./interfaces/IBatchTreeUpdateVerifier.sol";
 
-import "hardhat/console.sol";
-
-contract TornadoTrees is EnsResolve {
+contract TornadoTrees {
   address public immutable governance;
   bytes32 public depositRoot;
   bytes32 public previousDepositRoot;
   bytes32 public withdrawalRoot;
   bytes32 public previousWithdrawalRoot;
   address public tornadoProxy;
-  IVerifier public treeUpdateVerifier;
+  IBatchTreeUpdateVerifier public treeUpdateVerifier;
   ITornadoTreesV1 public immutable tornadoTreesV1;
 
   // make sure CHUNK_TREE_HEIGHT has the same value in BatchTreeUpdate.circom
@@ -73,7 +70,7 @@ contract TornadoTrees is EnsResolve {
     address _governance,
     address _tornadoProxy,
     ITornadoTreesV1 _tornadoTreesV1,
-    IVerifier _treeUpdateVerifier,
+    IBatchTreeUpdateVerifier _treeUpdateVerifier,
     SearchParams memory _searchParams
   ) public {
     governance = _governance;
@@ -82,11 +79,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 +89,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 +101,37 @@ 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");
-
-    uint256 index = _from + _step;
-    while (elementExists(_tornadoTreesV1, _type, index)) {
-      index = index + _step;
+    if (_from == 0 && _step == 0) {
+      return 0; // for tests
     }
+    // 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 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));
   }
 
@@ -277,7 +274,7 @@ contract TornadoTrees is EnsResolve {
     tornadoProxy = _tornadoProxy;
   }
 
-  function setVerifierContract(IVerifier _treeUpdateVerifier) external onlyGovernance {
+  function setVerifierContract(IBatchTreeUpdateVerifier _treeUpdateVerifier) external onlyGovernance {
     treeUpdateVerifier = _treeUpdateVerifier;
   }
 

contracts/interfaces/IBatchTreeUpdateVerifier.sol

@@ -2,6 +2,6 @@
 
 pragma solidity ^0.6.0;
 
-interface IVerifier {
+interface IBatchTreeUpdateVerifier {
   function verifyProof(bytes calldata proof, uint256[1] calldata input) external view returns (bool);
 }

contracts/mocks/TornadoTreesMock.sol

@@ -5,7 +5,7 @@ pragma experimental ABIEncoderV2;
 
 import "../TornadoTrees.sol";
 import "../interfaces/ITornadoTreesV1.sol";
-import "../interfaces/IVerifier.sol";
+import "../interfaces/IBatchTreeUpdateVerifier.sol";
 
 contract TornadoTreesMock is TornadoTrees {
   uint256 public currentBlock;
@@ -14,14 +14,10 @@ contract TornadoTreesMock is TornadoTrees {
     address _governance,
     address _tornadoProxy,
     ITornadoTreesV1 _tornadoTreesV1,
-    IVerifier _treeUpdateVerifier,
+    IBatchTreeUpdateVerifier _treeUpdateVerifier,
     SearchParams memory _searchParams
   ) public TornadoTrees(_governance, _tornadoProxy, _tornadoTreesV1, _treeUpdateVerifier, _searchParams) {}
 
-  function resolve(bytes32 _addr) public view override returns (address) {
-    return address(uint160(uint256(_addr) >> (12 * 8)));
-  }
-
   function setBlockNumber(uint256 _blockNumber) public {
     currentBlock = _blockNumber;
   }
@@ -44,6 +40,11 @@ contract TornadoTreesMock is TornadoTrees {
     registerWithdrawal(_instance, _nullifier);
   }
 
+  function updateRoots(bytes32 _depositRoot, bytes32 _withdrawalRoot) public {
+    depositRoot = _depositRoot;
+    withdrawalRoot = _withdrawalRoot;
+  }
+
   function updateDepositTreeMock(
     bytes32 _oldRoot,
     bytes32 _newRoot,

package.json

@@ -1,6 +1,6 @@
 {
   "name": "tornado-trees",
-  "version": "0.0.1",
+  "version": "0.0.4",
   "main": "src/index.js",
   "repository": "https://github.com/tornadocash/tornado-trees.git",
   "author": "Tornadocash team <hello@tornado.cash>",

test/binarySearch.test.js

@@ -37,6 +37,13 @@ describe('findArrayLength', () => {
     expect(depositsLength).to.be.equal(depositsEven.length)
   })
 
+  it('should work for empty array', async () => {
+    publicArray = await PublicArray.deploy()
+    // will throw out of gas if you pass non zero params
+    const depositsLength = await tornadoTrees.findArrayLength(publicArray.address, 'deposits(uint256)', 0, 0)
+    expect(depositsLength).to.be.equal(0)
+  })
+
   it('should work for odd array', async () => {
     publicArray = await PublicArray.deploy()
     await publicArray.setDeposits(depositsOdd)
@@ -75,4 +82,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)
+    }
+  })
 })