make findArrayLength internal

contracts/TornadoTrees.sol

@@ -17,7 +17,6 @@ contract TornadoTrees is Initializable {
   IBatchTreeUpdateVerifier public treeUpdateVerifier;
   ITornadoTreesV1 public immutable tornadoTreesV1;
 
-  // make sure CHUNK_TREE_HEIGHT has the same value in BatchTreeUpdate.circom
   uint256 public constant CHUNK_TREE_HEIGHT = 8;
   uint256 public constant CHUNK_SIZE = 2**CHUNK_TREE_HEIGHT;
   uint256 public constant ITEM_SIZE = 32 + 20 + 4;
@@ -100,48 +99,6 @@ contract TornadoTrees is Initializable {
     withdrawalsLength = withdrawalsV1Length;
   }
 
-  // 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, // 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) {
-    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;
-
-    // Perform a binary search in this segment
-    while (low < mid) {
-      if (elementExists(_tornadoTreesV1, _type, mid)) {
-        low = mid;
-      } else {
-        high = mid;
-      }
-      mid = (low + high) / 2;
-    }
-    return mid + 1;
-  }
-
-  function elementExists(
-    ITornadoTreesV1 _tornadoTreesV1,
-    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));
-  }
-
   function registerDeposit(address _instance, bytes32 _commitment) public onlyTornadoProxy {
     uint256 _depositsLength = depositsLength;
     deposits[_depositsLength] = keccak256(abi.encode(_instance, _commitment, blockNumber()));
@@ -251,6 +208,48 @@ contract TornadoTrees is Initializable {
     require(_withdrawalRoot == withdrawalRoot || _withdrawalRoot == previousWithdrawalRoot, "Incorrect withdrawal tree root");
   }
 
+  /// @dev There is no array length getter for deposit and withdrawal arrays
+  /// in previous contract, so we have to find them length manually.
+  /// Used only during deployment
+  function findArrayLength(
+    ITornadoTreesV1 _tornadoTreesV1,
+    string memory _type,
+    uint256 _from, // most likely array length after the proposal has passed
+    uint256 _step // optimal step size to find first match, approximately equals dispersion
+  ) internal view returns (uint256) {
+    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;
+
+    // Perform a binary search in this segment
+    while (low < mid) {
+      if (elementExists(_tornadoTreesV1, _type, mid)) {
+        low = mid;
+      } else {
+        high = mid;
+      }
+      mid = (low + high) / 2;
+    }
+    return mid + 1;
+  }
+
+  function elementExists(
+    ITornadoTreesV1 _tornadoTreesV1,
+    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));
+  }
+
   function getRegisteredDeposits() external view returns (bytes32[] memory _deposits) {
     uint256 count = depositsLength - lastProcessedDepositLeaf;
     _deposits = new bytes32[](count);

contracts/mocks/TornadoTreesMock.sol

@@ -24,6 +24,15 @@ contract TornadoTreesMock is TornadoTrees {
     return currentBlock == 0 ? block.number : currentBlock;
   }
 
+  function findArrayLengthMock(
+    ITornadoTreesV1 _tornadoTreesV1,
+    string memory _type,
+    uint256 _from,
+    uint256 _step
+  ) public view returns (uint256) {
+    return findArrayLength(_tornadoTreesV1, _type, _from, _step);
+  }
+
   function register(
     address _instance,
     bytes32 _commitment,

test/binarySearch.test.js

@@ -28,40 +28,40 @@ describe('findArrayLength', () => {
   })
 
   it('should work for even array', async () => {
-    const depositsLength = await tornadoTrees.findArrayLength(publicArray.address, 'deposits(uint256)', 4, 2)
+    const depositsLength = await tornadoTrees.findArrayLengthMock(publicArray.address, 'deposits(uint256)', 4, 2)
     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)
+    const depositsLength = await tornadoTrees.findArrayLengthMock(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)
-    const depositsLength = await tornadoTrees.findArrayLength(publicArray.address, 'deposits(uint256)', 4, 2)
+    const depositsLength = await tornadoTrees.findArrayLengthMock(publicArray.address, 'deposits(uint256)', 4, 2)
     expect(depositsLength).to.be.equal(depositsOdd.length)
   })
 
   it('should work for even array and odd step', async () => {
-    const depositsLength = await tornadoTrees.findArrayLength(publicArray.address, 'deposits(uint256)', 4, 3)
+    const depositsLength = await tornadoTrees.findArrayLengthMock(publicArray.address, 'deposits(uint256)', 4, 3)
     expect(depositsLength).to.be.equal(depositsEven.length)
   })
 
   it('should work for odd array and odd step', async () => {
     publicArray = await PublicArray.deploy()
     await publicArray.setDeposits(depositsOdd)
-    const depositsLength = await tornadoTrees.findArrayLength(publicArray.address, 'deposits(uint256)', 4, 3)
+    const depositsLength = await tornadoTrees.findArrayLengthMock(publicArray.address, 'deposits(uint256)', 4, 3)
     expect(depositsLength).to.be.equal(depositsOdd.length)
   })
 
   it('should work for odd array and step 1', async () => {
     publicArray = await PublicArray.deploy()
     await publicArray.setDeposits(depositsOdd)
-    const depositsLength = await tornadoTrees.findArrayLength(publicArray.address, 'deposits(uint256)', 4, 1)
+    const depositsLength = await tornadoTrees.findArrayLengthMock(publicArray.address, 'deposits(uint256)', 4, 1)
     expect(depositsLength).to.be.equal(depositsOdd.length)
   })
 
@@ -69,7 +69,7 @@ describe('findArrayLength', () => {
     const deposits = Array.from(Array(100).keys())
     publicArray = await PublicArray.deploy()
     await publicArray.setDeposits(deposits)
-    const depositsLength = await tornadoTrees.findArrayLength(
+    const depositsLength = await tornadoTrees.findArrayLengthMock(
       publicArray.address,
       'deposits(uint256)',
       67,
@@ -82,7 +82,7 @@ describe('findArrayLength', () => {
     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)
+    const depositsLength = await tornadoTrees.findArrayLengthMock(publicArray.address, 'deposits(uint256)', 1, 50)
     expect(depositsLength).to.be.equal(deposits.length)
   })
 
@@ -100,7 +100,7 @@ describe('findArrayLength', () => {
       const deposits = Array.from(Array(len).keys())
       publicArray = await PublicArray.deploy()
       await publicArray.setDeposits(deposits)
-      const depositsLength = await tornadoTrees.findArrayLength(
+      const depositsLength = await tornadoTrees.findArrayLengthMock(
         publicArray.address,
         'deposits(uint256)',
         days * depositsPerDay,