diff --git a/packages/incremental-merkle-tree.sol/contracts/IncrementalBinaryTree.sol b/packages/incremental-merkle-tree.sol/contracts/IncrementalBinaryTree.sol index 0b9cc4f..3658c5d 100644 --- a/packages/incremental-merkle-tree.sol/contracts/IncrementalBinaryTree.sol +++ b/packages/incremental-merkle-tree.sol/contracts/IncrementalBinaryTree.sol @@ -31,8 +31,14 @@ library IncrementalBinaryTree { uint8 depth, uint256 zero ) public { - require(zero < SNARK_SCALAR_FIELD, "IncrementalBinaryTree: leaf must be < SNARK_SCALAR_FIELD"); - require(depth > 0 && depth <= MAX_DEPTH, "IncrementalBinaryTree: tree depth must be between 1 and 32"); + require( + zero < SNARK_SCALAR_FIELD, + "IncrementalBinaryTree: leaf must be < SNARK_SCALAR_FIELD" + ); + require( + depth > 0 && depth <= MAX_DEPTH, + "IncrementalBinaryTree: tree depth must be between 1 and 32" + ); self.depth = depth; @@ -48,8 +54,14 @@ library IncrementalBinaryTree { /// @param self: Tree data. /// @param leaf: Leaf to be inserted. function insert(IncrementalTreeData storage self, uint256 leaf) public { - require(leaf < SNARK_SCALAR_FIELD, "IncrementalBinaryTree: leaf must be < SNARK_SCALAR_FIELD"); - require(self.numberOfLeaves < 2**self.depth, "IncrementalBinaryTree: tree is full"); + require( + leaf < SNARK_SCALAR_FIELD, + "IncrementalBinaryTree: leaf must be < SNARK_SCALAR_FIELD" + ); + require( + self.numberOfLeaves < 2**self.depth, + "IncrementalBinaryTree: tree is full" + ); uint256 index = self.numberOfLeaves; uint256 hash = leaf; @@ -69,6 +81,38 @@ library IncrementalBinaryTree { self.numberOfLeaves += 1; } + /// @dev Updates a leaf in the tree. + /// @param self: Tree data. + /// @param leaf: Leaf to be updated. + /// @param proofSiblings: Array of the sibling nodes of the proof of membership. + /// @param proofPathIndices: Path of the proof of membership. + function update( + IncrementalTreeData storage self, + uint256 leaf, + uint256[] calldata proofSiblings, + uint8[] calldata proofPathIndices + ) public { + + uint256 hash = leaf; + for (uint8 i = 0; i < self.depth; i++) { + if (proofPathIndices[i] == 0) { + if (proofSiblings[i] == self.lastSubtrees[i][1]) { + self.lastSubtrees[i][0] = hash; + } + + hash = PoseidonT3.poseidon([hash, proofSiblings[i]]); + } else { + if (proofSiblings[i] == self.lastSubtrees[i][0]) { + self.lastSubtrees[i][1] = hash; + } + + hash = PoseidonT3.poseidon([proofSiblings[i], hash]); + } + } + + self.root = hash; + } + /// @dev Removes a leaf from the tree. /// @param self: Tree data. /// @param leaf: Leaf to be removed. @@ -118,9 +162,13 @@ library IncrementalBinaryTree { uint256[] calldata proofSiblings, uint8[] calldata proofPathIndices ) private view returns (bool) { - require(leaf < SNARK_SCALAR_FIELD, "IncrementalBinaryTree: leaf must be < SNARK_SCALAR_FIELD"); require( - proofPathIndices.length == self.depth && proofSiblings.length == self.depth, + leaf < SNARK_SCALAR_FIELD, + "IncrementalBinaryTree: leaf must be < SNARK_SCALAR_FIELD" + ); + require( + proofPathIndices.length == self.depth && + proofSiblings.length == self.depth, "IncrementalBinaryTree: length of path is not correct" ); diff --git a/packages/incremental-merkle-tree.sol/contracts/IncrementalQuinTree.sol b/packages/incremental-merkle-tree.sol/contracts/IncrementalQuinTree.sol index 1934237..84e50c9 100644 --- a/packages/incremental-merkle-tree.sol/contracts/IncrementalQuinTree.sol +++ b/packages/incremental-merkle-tree.sol/contracts/IncrementalQuinTree.sol @@ -31,8 +31,14 @@ library IncrementalQuinTree { uint8 depth, uint256 zero ) public { - require(zero < SNARK_SCALAR_FIELD, "IncrementalBinaryTree: leaf must be < SNARK_SCALAR_FIELD"); - require(depth > 0 && depth <= MAX_DEPTH, "IncrementalQuinTree: tree depth must be between 1 and 32"); + require( + zero < SNARK_SCALAR_FIELD, + "IncrementalBinaryTree: leaf must be < SNARK_SCALAR_FIELD" + ); + require( + depth > 0 && depth <= MAX_DEPTH, + "IncrementalQuinTree: tree depth must be between 1 and 32" + ); self.depth = depth; @@ -54,8 +60,14 @@ library IncrementalQuinTree { /// @param self: Tree data. /// @param leaf: Leaf to be inserted. function insert(IncrementalTreeData storage self, uint256 leaf) public { - require(leaf < SNARK_SCALAR_FIELD, "IncrementalQuinTree: leaf must be < SNARK_SCALAR_FIELD"); - require(self.numberOfLeaves < 5**self.depth, "IncrementalQuinTree: tree is full"); + require( + leaf < SNARK_SCALAR_FIELD, + "IncrementalQuinTree: leaf must be < SNARK_SCALAR_FIELD" + ); + require( + self.numberOfLeaves < 5**self.depth, + "IncrementalQuinTree: tree is full" + ); uint256 index = self.numberOfLeaves; uint256 hash = leaf; @@ -79,6 +91,50 @@ library IncrementalQuinTree { self.numberOfLeaves += 1; } + /// @dev Updates a leaf in the tree. + /// @param self: Tree data. + /// @param leaf: Leaf to be updated. + /// @param proofSiblings: Array of the sibling nodes of the proof of membership. + /// @param proofPathIndices: Path of the proof of membership. + function update( + IncrementalTreeData storage self, + uint256 leaf, + uint256[4][] calldata proofSiblings, + uint8[] calldata proofPathIndices + ) public { + require( + verify(self, leaf, proofSiblings, proofPathIndices), + "IncrementalQuinTree: leaf is not part of the tree" + ); + + uint256 hash = self.zeroes[0]; + + for (uint8 i = 0; i < self.depth; i++) { + uint256[5] memory nodes; + + for (uint8 j = 0; j < 5; j++) { + if (j < proofPathIndices[i]) { + nodes[j] = proofSiblings[i][j]; + } else if (j == proofPathIndices[i]) { + nodes[j] = hash; + } else { + nodes[j] = proofSiblings[i][j - 1]; + } + } + + if ( + nodes[0] == self.lastSubtrees[i][0] || + nodes[4] == self.lastSubtrees[i][4] + ) { + self.lastSubtrees[i][proofPathIndices[i]] = hash; + } + + hash = PoseidonT6.poseidon(nodes); + } + + self.root = hash; + } + /// @dev Removes a leaf from the tree. /// @param self: Tree data. /// @param leaf: Leaf to be removed. @@ -110,7 +166,10 @@ library IncrementalQuinTree { } } - if (nodes[0] == self.lastSubtrees[i][0] || nodes[4] == self.lastSubtrees[i][4]) { + if ( + nodes[0] == self.lastSubtrees[i][0] || + nodes[4] == self.lastSubtrees[i][4] + ) { self.lastSubtrees[i][proofPathIndices[i]] = hash; } @@ -132,9 +191,13 @@ library IncrementalQuinTree { uint256[4][] calldata proofSiblings, uint8[] calldata proofPathIndices ) private view returns (bool) { - require(leaf < SNARK_SCALAR_FIELD, "IncrementalQuinTree: leaf must be < SNARK_SCALAR_FIELD"); require( - proofPathIndices.length == self.depth && proofSiblings.length == self.depth, + leaf < SNARK_SCALAR_FIELD, + "IncrementalQuinTree: leaf must be < SNARK_SCALAR_FIELD" + ); + require( + proofPathIndices.length == self.depth && + proofSiblings.length == self.depth, "IncrementalQuinTree: length of path is not correct" ); diff --git a/packages/incremental-merkle-tree.sol/contracts/test/IncrementalBinaryTreeTest.sol b/packages/incremental-merkle-tree.sol/contracts/test/IncrementalBinaryTreeTest.sol index 1594e35..1f658d5 100644 --- a/packages/incremental-merkle-tree.sol/contracts/test/IncrementalBinaryTreeTest.sol +++ b/packages/incremental-merkle-tree.sol/contracts/test/IncrementalBinaryTreeTest.sol @@ -9,6 +9,7 @@ contract IncrementalBinaryTreeTest { event TreeCreated(bytes32 id, uint8 depth); event LeafInserted(bytes32 indexed treeId, uint256 leaf, uint256 root); + event LeafUpdated(bytes32 indexed treeId, uint256 leaf, uint256 root); event LeafRemoved(bytes32 indexed treeId, uint256 leaf, uint256 root); mapping(bytes32 => IncrementalTreeData) public trees; @@ -22,20 +23,42 @@ contract IncrementalBinaryTreeTest { } function insertLeaf(bytes32 _treeId, uint256 _leaf) external { - require(trees[_treeId].depth != 0, "BinaryTreeTest: tree does not exist"); + require( + trees[_treeId].depth != 0, + "BinaryTreeTest: tree does not exist" + ); trees[_treeId].insert(_leaf); emit LeafInserted(_treeId, _leaf, trees[_treeId].root); } + function updateLeaf( + bytes32 _treeId, + uint256 _leaf, + uint256[] calldata _proofSiblings, + uint8[] calldata _proofPathIndices + ) external { + require( + trees[_treeId].depth != 0, + "BinaryTreeTest: tree does not exist" + ); + + trees[_treeId].update(_leaf, _proofSiblings, _proofPathIndices); + + emit LeafUpdated(_treeId, _leaf, trees[_treeId].root); + } + function removeLeaf( bytes32 _treeId, uint256 _leaf, uint256[] calldata _proofSiblings, uint8[] calldata _proofPathIndices ) external { - require(trees[_treeId].depth != 0, "BinaryTreeTest: tree does not exist"); + require( + trees[_treeId].depth != 0, + "BinaryTreeTest: tree does not exist" + ); trees[_treeId].remove(_leaf, _proofSiblings, _proofPathIndices); diff --git a/packages/incremental-merkle-tree.sol/contracts/test/IncrementalQuinTreeTest.sol b/packages/incremental-merkle-tree.sol/contracts/test/IncrementalQuinTreeTest.sol index 62f7060..090af54 100644 --- a/packages/incremental-merkle-tree.sol/contracts/test/IncrementalQuinTreeTest.sol +++ b/packages/incremental-merkle-tree.sol/contracts/test/IncrementalQuinTreeTest.sol @@ -9,6 +9,7 @@ contract IncrementalQuinTreeTest { event TreeCreated(bytes32 id, uint8 depth); event LeafInserted(bytes32 indexed treeId, uint256 leaf, uint256 root); + event LeafUpdated(bytes32 indexed treeId, uint256 leaf, uint256 root); event LeafRemoved(bytes32 indexed treeId, uint256 leaf, uint256 root); mapping(bytes32 => IncrementalTreeData) public trees; @@ -29,6 +30,19 @@ contract IncrementalQuinTreeTest { emit LeafInserted(_treeId, _leaf, trees[_treeId].root); } + function updateLeaf( + bytes32 _treeId, + uint256 _leaf, + uint256[4][] calldata _proofSiblings, + uint8[] calldata _proofPathIndices + ) external { + require(trees[_treeId].depth != 0, "QuinTreeTest: tree does not exist"); + + trees[_treeId].update(_leaf, _proofSiblings, _proofPathIndices); + + emit LeafUpdated(_treeId, _leaf, trees[_treeId].root); + } + function removeLeaf( bytes32 _treeId, uint256 _leaf, diff --git a/yarn.lock b/yarn.lock index fb3fdf0..5412021 100644 --- a/yarn.lock +++ b/yarn.lock @@ -8834,7 +8834,7 @@ __metadata: dependencies: bn.js: ^4.11.8 ethereumjs-util: ^6.0.0 - checksum: 03127d09960e5f8a44167463faf25b2894db2f746376dbb8195b789ed11762f93db9c574eaa7c498c400063508e9dfc1c80de2edf5f0e1406b25c87d860ff2f1 + checksum: ae074be0bb012857ab5d3ae644d1163b908a48dd724b7d2567cfde309dc72222d460438f2411936a70dc949dc604ce1ef7118f7273bd525815579143c907e336 languageName: node linkType: hard