tornado-packages/tornado-trees
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

make findArrayLength internal

Browse Source
This commit is contained in:
poma 2021-03-06 15:05:54 +03:00
parent 76543d068b
commit 272c7a94e0
No known key found for this signature in database
GPG Key ID: BA20CB01FE165657
3 changed files with 60 additions and 52 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

85
contracts/TornadoTrees.sol
View File

@ -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);

9
contracts/mocks/TornadoTreesMock.sol
View File

@ -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,

18
test/binarySearch.test.js
View File

@ -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,