From 2cd523cb53911277c727307a05ca8bbcaee65859 Mon Sep 17 00:00:00 2001 From: Chance Date: Tue, 27 Feb 2024 15:56:17 -0600 Subject: [PATCH 1/4] refactor: build outer levels before getting merkle elements --- .../contracts/internal/InternalLazyIMT.sol | 56 +++++++++++++------ 1 file changed, 39 insertions(+), 17 deletions(-) diff --git a/packages/imt.sol/contracts/internal/InternalLazyIMT.sol b/packages/imt.sol/contracts/internal/InternalLazyIMT.sol index 056356d9e..d95499cc6 100644 --- a/packages/imt.sol/contracts/internal/InternalLazyIMT.sol +++ b/packages/imt.sol/contracts/internal/InternalLazyIMT.sol @@ -171,9 +171,17 @@ library InternalLazyIMT { require(depth <= MAX_DEPTH, "LazyIMT: depth must be <= MAX_DEPTH"); // this should always short circuit if self.numberOfLeaves == 0 if (numberOfLeaves == 0) return _defaultZero(depth); + uint256[] memory levels = _levels(self, numberOfLeaves, depth); + return levels[depth]; + } + + function _levels(LazyIMTData storage self, uint40 numberOfLeaves, uint8 depth) internal view returns (uint256[] memory) { + require(depth <= MAX_DEPTH, "LazyIMT: depth must be <= MAX_DEPTH"); + require(numberOfLeaves > 0, "LazyIMT: number of leaves must be > 0"); + // this should always short circuit if self.numberOfLeaves == 0 uint40 index = numberOfLeaves - 1; - uint256[MAX_DEPTH + 1] memory levels; + uint256[] memory levels = new uint256[](MAX_DEPTH + 1); if (index & 1 == 0) { levels[0] = self.elements[_indexForElement(0, index)]; @@ -199,7 +207,7 @@ library InternalLazyIMT { i++; } } - return levels[depth]; + return levels; } function _merkleProofElements( @@ -209,32 +217,46 @@ library InternalLazyIMT { ) internal view returns (uint256[] memory) { uint40 numberOfLeaves = self.numberOfLeaves; require(index < numberOfLeaves, "LazyIMT: leaf must exist"); - require(depth > 0, "LazyIMT: depth must be > 0"); - require(depth <= MAX_DEPTH, "LazyIMT: depth must be < MAX_DEPTH"); - uint256[] memory proof = new uint256[](depth); - uint256 levelCount = self.numberOfLeaves; + // pass depth -1 because we don't need the root value + uint256[] memory outerLevels = _levels(self, numberOfLeaves, depth - 1); + uint256[] memory _elements = new uint256[](depth); - for (uint8 i = 0; i < depth; ) { - // Left leaf + // unroll the bottom entry of the tree because it will never need to + // be pulled from outerLevels + if (index & 1 == 0) { + if (index + 1 >= numberOfLeaves) { + _elements[0] = _defaultZero(0); + } else { + _elements[0] = self.elements[_indexForElement(0, index + 1)]; + } + } else { + _elements[0] = self.elements[_indexForElement(0, index - 1)]; + } + index >>= 1; + + for (uint8 i = 1; i < depth; ) { + uint256 currentLevelCount = numberOfLeaves >> i; if (index & 1 == 0) { - bool outsideLimits = ((index + 1) << i) % self.numberOfLeaves != 0; - if ((index + 1) < levelCount) { - proof[i] = self.elements[_indexForElement(i, index + 1)]; - } else if (((index + 1) == levelCount) && (outsideLimits)) { - proof[i] = _root(self, numberOfLeaves, i); + // pull from outerLevels if we need an element from + // a subtree that hasn't been built. e.g. a lazy segment of the tree + if (index + 1 >= currentLevelCount) { + if ((numberOfLeaves - 1 >> i) > index) { + _elements[i] = outerLevels[i]; + } else { + _elements[i] = _defaultZero(i); + } } else { - proof[i] = _defaultZero(i); + _elements[i] = self.elements[_indexForElement(i, index + 1)]; } } else { - proof[i] = self.elements[_indexForElement(i, index - 1)]; + _elements[i] = self.elements[_indexForElement(i, index - 1)]; } unchecked { index >>= 1; i++; } - levelCount = numberOfLeaves >> i; } - return proof; + return _elements; } } From 29afa0157eaf7e29823068a0d02823f9a4ec68ac Mon Sep 17 00:00:00 2001 From: Chance Date: Tue, 27 Feb 2024 16:58:07 -0600 Subject: [PATCH 2/4] chore: optimize chore: add depth assertion chore: revert package refactor: function comments --- .../contracts/internal/InternalLazyIMT.sol | 30 +++++++++---------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/packages/imt.sol/contracts/internal/InternalLazyIMT.sol b/packages/imt.sol/contracts/internal/InternalLazyIMT.sol index d95499cc6..d4127ecf3 100644 --- a/packages/imt.sol/contracts/internal/InternalLazyIMT.sol +++ b/packages/imt.sol/contracts/internal/InternalLazyIMT.sol @@ -171,18 +171,17 @@ library InternalLazyIMT { require(depth <= MAX_DEPTH, "LazyIMT: depth must be <= MAX_DEPTH"); // this should always short circuit if self.numberOfLeaves == 0 if (numberOfLeaves == 0) return _defaultZero(depth); - uint256[] memory levels = _levels(self, numberOfLeaves, depth); + uint256[] memory levels = new uint256[](depth + 1); + _levels(self, numberOfLeaves, depth, levels); return levels[depth]; } - function _levels(LazyIMTData storage self, uint40 numberOfLeaves, uint8 depth) internal view returns (uint256[] memory) { + function _levels(LazyIMTData storage self, uint40 numberOfLeaves, uint8 depth, uint256[] memory levels) internal view/* returns (uint256[] memory)*/ { require(depth <= MAX_DEPTH, "LazyIMT: depth must be <= MAX_DEPTH"); require(numberOfLeaves > 0, "LazyIMT: number of leaves must be > 0"); // this should always short circuit if self.numberOfLeaves == 0 uint40 index = numberOfLeaves - 1; - uint256[] memory levels = new uint256[](MAX_DEPTH + 1); - if (index & 1 == 0) { levels[0] = self.elements[_indexForElement(0, index)]; } else { @@ -207,7 +206,6 @@ library InternalLazyIMT { i++; } } - return levels; } function _merkleProofElements( @@ -218,12 +216,17 @@ library InternalLazyIMT { uint40 numberOfLeaves = self.numberOfLeaves; require(index < numberOfLeaves, "LazyIMT: leaf must exist"); + uint8 targetDepth = 1; + while (uint40(2) ** uint40(targetDepth) < numberOfLeaves) { + targetDepth++; + } + require(depth >= targetDepth, "LazyIMT: proof depth"); // pass depth -1 because we don't need the root value - uint256[] memory outerLevels = _levels(self, numberOfLeaves, depth - 1); uint256[] memory _elements = new uint256[](depth); + _levels(self, numberOfLeaves, targetDepth - 1, _elements); // unroll the bottom entry of the tree because it will never need to - // be pulled from outerLevels + // be pulled from _levels if (index & 1 == 0) { if (index + 1 >= numberOfLeaves) { _elements[0] = _defaultZero(0); @@ -238,14 +241,11 @@ library InternalLazyIMT { for (uint8 i = 1; i < depth; ) { uint256 currentLevelCount = numberOfLeaves >> i; if (index & 1 == 0) { - // pull from outerLevels if we need an element from - // a subtree that hasn't been built. e.g. a lazy segment of the tree - if (index + 1 >= currentLevelCount) { - if ((numberOfLeaves - 1 >> i) > index) { - _elements[i] = outerLevels[i]; - } else { - _elements[i] = _defaultZero(i); - } + // if the element is an uncomputed edge node we'll use the value set + // from _levels above + // otherwise set as usual below + if (index + 1 >= currentLevelCount && (numberOfLeaves - 1 >> i) <= index) { + _elements[i] = _defaultZero(i); } else { _elements[i] = self.elements[_indexForElement(i, index + 1)]; } From 8efa766ad79e5302fc0cea2a579d8e2d76f5bfbe Mon Sep 17 00:00:00 2001 From: Chance Date: Wed, 28 Feb 2024 02:01:50 -0600 Subject: [PATCH 3/4] fix: conditional logic --- packages/imt.sol/contracts/internal/InternalLazyIMT.sol | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/packages/imt.sol/contracts/internal/InternalLazyIMT.sol b/packages/imt.sol/contracts/internal/InternalLazyIMT.sol index d4127ecf3..c32bb67f8 100644 --- a/packages/imt.sol/contracts/internal/InternalLazyIMT.sol +++ b/packages/imt.sol/contracts/internal/InternalLazyIMT.sol @@ -244,10 +244,10 @@ library InternalLazyIMT { // if the element is an uncomputed edge node we'll use the value set // from _levels above // otherwise set as usual below - if (index + 1 >= currentLevelCount && (numberOfLeaves - 1 >> i) <= index) { - _elements[i] = _defaultZero(i); - } else { + if (index + 1 < currentLevelCount) { _elements[i] = self.elements[_indexForElement(i, index + 1)]; + } else if ((numberOfLeaves - 1 >> i) <= index) { + _elements[i] = _defaultZero(i); } } else { _elements[i] = self.elements[_indexForElement(i, index - 1)]; From b8db58cd886dc5a61048982e117a5285dda6eda4 Mon Sep 17 00:00:00 2001 From: Chance Date: Wed, 28 Feb 2024 15:12:40 -0600 Subject: [PATCH 4/4] chore: remove comment --- packages/imt.sol/contracts/internal/InternalLazyIMT.sol | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/imt.sol/contracts/internal/InternalLazyIMT.sol b/packages/imt.sol/contracts/internal/InternalLazyIMT.sol index c32bb67f8..8f0c69da2 100644 --- a/packages/imt.sol/contracts/internal/InternalLazyIMT.sol +++ b/packages/imt.sol/contracts/internal/InternalLazyIMT.sol @@ -176,7 +176,7 @@ library InternalLazyIMT { return levels[depth]; } - function _levels(LazyIMTData storage self, uint40 numberOfLeaves, uint8 depth, uint256[] memory levels) internal view/* returns (uint256[] memory)*/ { + function _levels(LazyIMTData storage self, uint40 numberOfLeaves, uint8 depth, uint256[] memory levels) internal view { require(depth <= MAX_DEPTH, "LazyIMT: depth must be <= MAX_DEPTH"); require(numberOfLeaves > 0, "LazyIMT: number of leaves must be > 0"); // this should always short circuit if self.numberOfLeaves == 0