findArrayLength fix

This commit is contained in:
Alexey 2021-02-09 12:27:12 +03:00
parent 73554b3ae3
commit 31d48c2b57
2 changed files with 38 additions and 12 deletions

View File

@ -105,6 +105,7 @@ contract TornadoTrees is EnsResolve {
); );
} }
// todo make things internal. Think of dedundant calls
function findArrayLength( function findArrayLength(
ITornadoTreesV1 _tornadoTreesV1, ITornadoTreesV1 _tornadoTreesV1,
string memory _type, string memory _type,
@ -112,7 +113,7 @@ contract TornadoTrees is EnsResolve {
uint256 _step uint256 _step
) public view returns (uint256) { ) public view returns (uint256) {
require(_from != 0 && _step != 0, "_from and _step should be > 0"); require(_from != 0 && _step != 0, "_from and _step should be > 0");
require(elementExists(_tornadoTreesV1, _type, _from), "Inccorrect _from param"); // require(elementExists(_tornadoTreesV1, _type, _from), "Inccorrect _from param");
uint256 index = _from + _step; uint256 index = _from + _step;
while (elementExists(_tornadoTreesV1, _type, index)) { while (elementExists(_tornadoTreesV1, _type, index)) {
@ -121,21 +122,38 @@ contract TornadoTrees is EnsResolve {
uint256 high = index; uint256 high = index;
uint256 low = index - _step; uint256 low = index - _step;
uint256 mid = (low + high) / 2; uint256 lastIndex = binarySearch(_tornadoTreesV1, _type, high, low);
while (!elementExists(_tornadoTreesV1, _type, mid)) {
high = mid - 1; return lastIndex + 1;
mid = (low + high) / 2; }
function binarySearch(
ITornadoTreesV1 _tornadoTreesV1,
string memory _type,
uint256 _high,
uint256 _low
) public view returns (uint256) {
require(_high >= _low, "Incorrect params");
uint256 mid = (_high + _low) / 2;
(bool isLast, bool exists) = isLastElement(_tornadoTreesV1, _type, mid);
if (isLast) {
return mid;
} }
high += 1; if (exists) {
low = mid + 1; return binarySearch(_tornadoTreesV1, _type, _high, mid + 1);
mid = (low + high) / 2; } else {
while (elementExists(_tornadoTreesV1, _type, mid)) { return binarySearch(_tornadoTreesV1, _type, mid - 1, _low);
low = mid + 1;
mid = (low + high) / 2;
} }
}
return high == low ? high : low; function isLastElement(
ITornadoTreesV1 _tornadoTreesV1,
string memory _type,
uint256 index
) public view returns (bool success, bool exists) {
exists = elementExists(_tornadoTreesV1, _type, index);
success = exists && !elementExists(_tornadoTreesV1, _type, index + 1);
} }
function elementExists( function elementExists(

View File

@ -75,4 +75,12 @@ describe('findArrayLength', () => {
) )
expect(depositsLength).to.be.equal(deposits.length) 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)
})
}) })