fix: subtree strategy

This commit is contained in:
rymnc 2024-04-13 01:45:40 +03:00
parent d3d53bbd9f
commit 055942a482
No known key found for this signature in database
GPG Key ID: AAA088D5C68ECD34
3 changed files with 29 additions and 21 deletions

View File

@ -89,25 +89,17 @@ library BinaryIMTMemory {
revert("IncrementalBinaryTree: defaultZero bad index");
}
/// @dev Computes the root of the tree given the leaves.
/// @param self: Tree data.
/// @param leaves: Leaves in the tree
function calcRoot(
BinaryIMTMemoryData memory self,
uint256 depth,
uint256[] memory leaves
)
public
pure
returns (uint256)
{
function calcSubtreeRoot(BinaryIMTMemoryData memory self, uint256[] memory leaves) public pure returns (uint256) {
// instead of the depth being 0..19, we go from 10..19.
// this way, we preserve the subtree root
uint8 depth = 10;
uint256[2][] memory lastSubtrees = new uint256[2][](depth);
for (uint256 j = 0; j < leaves.length;) {
uint256 index = self.numberOfLeaves;
uint256 hash = leaves[j];
for (uint8 i = 0; i < depth;) {
if (index & 1 == 0) {
lastSubtrees[i] = [hash, defaultZero(i)];
lastSubtrees[i] = [hash, defaultZero(i + depth)];
} else {
lastSubtrees[i][1] = hash;
}

View File

@ -308,18 +308,22 @@ abstract contract RlnBase {
function root(uint256 shardIndex) public view returns (uint256) {
BinaryIMTMemoryData memory imtData;
uint256 idCommitmentIndex = leavesSet[shardIndex];
uint256[] memory calcLeaves = new uint256[](idCommitmentIndex);
uint256[] memory calcLeaves = new uint256[](1024);
if (idCommitmentIndex == 0) {
return BinaryIMTMemory.calcSubtreeRoot(imtData, calcLeaves);
}
for (uint256 i = 0; i < idCommitmentIndex; i++) {
uint256 idCommitment = leavesIndex[shardIndex][i];
uint256 userMessageLimit = userMessageLimits[idCommitment];
calcLeaves[i] = PoseidonT3.hash([idCommitment, userMessageLimit]);
}
return BinaryIMTMemory.calcRoot(imtData, 10, calcLeaves);
return BinaryIMTMemory.calcSubtreeRoot(imtData, calcLeaves);
}
function full_root(uint256[] calldata roots) public pure returns (uint256) {
function fullRoot(uint256[] calldata roots) public pure returns (uint256) {
BinaryIMTMemoryData memory imtData;
return BinaryIMTMemory.calcRoot(imtData, 10, roots);
return BinaryIMTMemory.calcSubtreeRoot(imtData, roots);
}
function getLeafAtShard(uint256 shardIndex, uint256 index) public view returns (uint256) {

View File

@ -263,6 +263,11 @@ contract RlnTest is Test {
18_217_334_211_520_937_958_971_536_517_166_530_749_184_547_628_672_204_353_760_850_739_130_586_503_124;
vm.pauseGasMetering();
// before we insert into the tree, the kats must pass
assertEq(
rln.root(0),
12_413_880_268_183_407_374_852_357_075_976_609_371_175_688_755_676_981_206_018_884_971_008_854_919_922
);
for (uint256 i = 0; i < idCommitments.length; i++) {
// default 1 message limit
rln.register{ value: MEMBERSHIP_DEPOSIT }(idCommitments[i], 1);
@ -271,7 +276,14 @@ contract RlnTest is Test {
assertEq(
rln.root(0),
1_009_702_141_963_982_084_971_194_921_450_594_432_549_814_353_062_162_436_132_398_913_014_207_593_991
2_435_591_050_584_562_785_273_516_085_683_091_349_247_653_357_522_297_485_884_387_709_399_190_797_862
);
uint256[] memory roots = new uint256[](1024);
roots[0] = 2_435_591_050_584_562_785_273_516_085_683_091_349_247_653_357_522_297_485_884_387_709_399_190_797_862;
assertEq(
rln.fullRoot(roots),
3_881_711_806_223_377_894_370_509_479_419_308_415_025_371_832_767_296_831_624_691_168_726_496_442_721
);
}
@ -291,15 +303,15 @@ contract RlnTest is Test {
vm.resumeGasMetering();
assertEq(
rln.root(0),
15_191_547_950_571_875_204_938_158_154_881_229_414_576_184_216_597_114_679_496_514_970_950_611_126_017
10_424_638_537_885_815_866_594_133_227_550_095_274_095_476_771_249_698_162_779_016_221_356_823_894_456
);
}
function test__full_root() public {
function test__fullRoot() public {
uint256[] memory roots = new uint256[](1024);
assertEq(
// repeat 0 1024 times
rln.full_root(roots),
rln.fullRoot(roots),
12_413_880_268_183_407_374_852_357_075_976_609_371_175_688_755_676_981_206_018_884_971_008_854_919_922
);
}