diff --git a/contracts/TornadoTrees.sol b/contracts/TornadoTrees.sol index b76b256..077ea35 100644 --- a/contracts/TornadoTrees.sol +++ b/contracts/TornadoTrees.sol @@ -62,11 +62,19 @@ contract TornadoTrees is EnsResolve { _; } + struct SearchParams { + uint256 unprocessedDeposits; + uint256 unprocessedWithdrawals; + uint256 depositsPerDay; + uint256 withdrawalsPerDay; + } + constructor( address _governance, address _tornadoProxy, ITornadoTreesV1 _tornadoTreesV1, - IVerifier _treeUpdateVerifier + IVerifier _treeUpdateVerifier, + SearchParams memory _searchParams ) public { governance = _governance; tornadoProxy = _tornadoProxy; @@ -79,30 +87,63 @@ contract TornadoTrees is EnsResolve { uint256 depositLeaf = _tornadoTreesV1.lastProcessedDepositLeaf(); require(depositLeaf % CHUNK_SIZE == 0, "Incorrect TornadoTrees state"); lastProcessedDepositLeaf = depositLeaf; - depositsLength = depositV1Length = findArrayLength(_tornadoTreesV1, "deposits(uint256)", 3); // todo + depositsLength = depositV1Length = findArrayLength( + _tornadoTreesV1, + "deposits(uint256)", + _searchParams.unprocessedDeposits, + _searchParams.depositsPerDay + ); uint256 withdrawalLeaf = _tornadoTreesV1.lastProcessedWithdrawalLeaf(); require(withdrawalLeaf % CHUNK_SIZE == 0, "Incorrect TornadoTrees state"); lastProcessedWithdrawalLeaf = withdrawalLeaf; - withdrawalsLength = withdrawalsV1Length = findArrayLength(_tornadoTreesV1, "withdrawals(uint256)", 3); // todo + withdrawalsLength = withdrawalsV1Length = findArrayLength( + _tornadoTreesV1, + "withdrawals(uint256)", + _searchParams.unprocessedWithdrawals, + _searchParams.withdrawalsPerDay + ); } - // todo implement binary search function findArrayLength( ITornadoTreesV1 _tornadoTreesV1, - string memory _signature, - uint256 _from + string memory _type, + uint256 _from, + uint256 _step ) public view returns (uint256) { - bool success; - bytes memory data; + require(_from != 0 && _step != 0, "_from and _step should be > 0"); + require(elementExists(_tornadoTreesV1, _type, _from), "Inccorrect _from param"); - (success, data) = address(_tornadoTreesV1).staticcall{ gas: 3000 }(abi.encodeWithSignature(_signature, _from)); - while (success) { - _from++; - (success, data) = address(_tornadoTreesV1).staticcall{ gas: 3000 }(abi.encodeWithSignature(_signature, _from)); + uint256 index = _from + _step; + while (elementExists(_tornadoTreesV1, _type, index)) { + index = index + _step; } - return _from; + uint256 high = index; + uint256 low = index - _step; + uint256 mid = (low + high) / 2; + while (!elementExists(_tornadoTreesV1, _type, mid)) { + high = mid - 1; + 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; + } + + function elementExists( + ITornadoTreesV1 _tornadoTreesV1, + string memory _type, + uint256 index + ) public view returns (bool success) { + (success, ) = address(_tornadoTreesV1).staticcall{ gas: 2500 }(abi.encodeWithSignature(_type, index)); } function registerDeposit(address _instance, bytes32 _commitment) public onlyTornadoProxy { diff --git a/contracts/mocks/PublicArray.sol b/contracts/mocks/PublicArray.sol new file mode 100644 index 0000000..173ccb7 --- /dev/null +++ b/contracts/mocks/PublicArray.sol @@ -0,0 +1,28 @@ +// SPDX-License-Identifier: MIT + +pragma solidity ^0.6.0; + +contract PublicArray { + uint256[] public deposits; + uint256[] public withdrawals; + + function lastProcessedDepositLeaf() external view returns (uint256) {} + + function lastProcessedWithdrawalLeaf() external view returns (uint256) {} + + function depositRoot() external view returns (bytes32) {} + + function withdrawalRoot() external view returns (bytes32) {} + + function setDeposits(uint256[] memory _deposits) public { + for (uint256 i = 0; i < _deposits.length; i++) { + deposits.push(_deposits[i]); + } + } + + function setWithdrawals(uint256[] memory _withdrawals) public { + for (uint256 i = 0; i < _withdrawals.length; i++) { + withdrawals.push(_withdrawals[i]); + } + } +} diff --git a/contracts/mocks/TornadoTreesMock.sol b/contracts/mocks/TornadoTreesMock.sol index fb3176e..4a06428 100644 --- a/contracts/mocks/TornadoTreesMock.sol +++ b/contracts/mocks/TornadoTreesMock.sol @@ -14,8 +14,9 @@ contract TornadoTreesMock is TornadoTrees { address _governance, address _tornadoProxy, ITornadoTreesV1 _tornadoTreesV1, - IVerifier _treeUpdateVerifier - ) public TornadoTrees(_governance, _tornadoProxy, _tornadoTreesV1, _treeUpdateVerifier) {} + IVerifier _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))); diff --git a/hardhat.config.js b/hardhat.config.js index ea3d71f..7666668 100644 --- a/hardhat.config.js +++ b/hardhat.config.js @@ -21,7 +21,7 @@ const config = { solidity: '0.6.12', networks: { hardhat: { - blockGasLimit: 950000000, + blockGasLimit: 9500000, }, }, mocha: { diff --git a/test/binarySearch.test.js b/test/binarySearch.test.js new file mode 100644 index 0000000..f8bc5b3 --- /dev/null +++ b/test/binarySearch.test.js @@ -0,0 +1,78 @@ +/* global ethers */ + +const { expect } = require('chai') + +const depositsEven = [10, 11, 12, 13, 14, 15, 16, 17, 18] +const depositsOdd = [10, 11, 12, 13, 14, 15, 16, 17] + +describe('findArrayLength', () => { + let publicArray + let tornadoTrees + let PublicArray + + beforeEach(async function () { + const [operator, tornadoProxy] = await ethers.getSigners() + PublicArray = await ethers.getContractFactory('PublicArray') + publicArray = await PublicArray.deploy() + await publicArray.setDeposits(depositsEven) + await publicArray.setWithdrawals(depositsEven) + + const TornadoTrees = await ethers.getContractFactory('TornadoTreesMock') + tornadoTrees = await TornadoTrees.deploy( + operator.address, + tornadoProxy.address, + publicArray.address, + publicArray.address, + { + unprocessedDeposits: 3, + unprocessedWithdrawals: 3, + depositsPerDay: 2, + withdrawalsPerDay: 2, + }, + ) + }) + + it('should work for even array', async () => { + const depositsLength = await tornadoTrees.findArrayLength(publicArray.address, 'deposits(uint256)', 4, 2) + expect(depositsLength).to.be.equal(depositsEven.length) + }) + + 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) + 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) + 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) + 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) + expect(depositsLength).to.be.equal(depositsOdd.length) + }) + + it('should work for big array and big step', async () => { + const deposits = Array.from(Array(100).keys()) + publicArray = await PublicArray.deploy() + await publicArray.setDeposits(deposits) + const depositsLength = await tornadoTrees.findArrayLength( + publicArray.address, + 'deposits(uint256)', + 67, + 10, + ) + expect(depositsLength).to.be.equal(deposits.length) + }) +}) diff --git a/test/tornadoTrees.test.js b/test/tornadoTrees.test.js index 464e20d..5796037 100644 --- a/test/tornadoTrees.test.js +++ b/test/tornadoTrees.test.js @@ -77,6 +77,12 @@ describe('TornadoTrees', function () { tornadoProxy.address, tornadoTreesV1.address, verifier.address, + { + unprocessedDeposits: 1, + unprocessedWithdrawals: 1, + depositsPerDay: 2, + withdrawalsPerDay: 2, + }, ) depositDataEventFilter = tornadoTrees.filters.DepositData() }) @@ -159,6 +165,12 @@ describe('TornadoTrees', function () { tornadoProxy.address, tornadoTreesV1.address, verifier.address, + { + unprocessedDeposits: 1, + unprocessedWithdrawals: 1, + depositsPerDay: 2, + withdrawalsPerDay: 2, + }, ) let { input, args } = controller.batchTreeUpdate(tree, depositEvents)