diff --git a/package.json b/package.json index 1fc7a61..c2a0257 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,7 @@ { "name": "zk-kit", "description": "A monorepo of reusable JS libraries for zero-knowledge technologies.", + "version": "0.1.1", "license": "MIT", "repository": "git@github.com:privacy-scaling-explorations/zk-kit.git", "homepage": "https://github.com/privacy-scaling-explorations/zk-kit", diff --git a/packages/incremental-merkle-tree.sol/contracts/IncrementalBinaryTree.sol b/packages/incremental-merkle-tree.sol/contracts/IncrementalBinaryTree.sol index 0b9cc4f..098baac 100644 --- a/packages/incremental-merkle-tree.sol/contracts/IncrementalBinaryTree.sol +++ b/packages/incremental-merkle-tree.sol/contracts/IncrementalBinaryTree.sol @@ -69,6 +69,45 @@ library IncrementalBinaryTree { self.numberOfLeaves += 1; } + /// @dev Updates a leaf in the tree. + /// @param self: Tree data. + /// @param leaf: Leaf to be updated. + /// @param newLeaf: New leaf. + /// @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 newLeaf, + uint256[] calldata proofSiblings, + uint8[] calldata proofPathIndices + ) public { + require( + verify(self, leaf, proofSiblings, proofPathIndices), + "IncrementalBinaryTree: leaf is not part of the tree" + ); + require(newLeaf < SNARK_SCALAR_FIELD, "IncrementalBinaryTree: leaf must be < SNARK_SCALAR_FIELD"); + + uint256 hash = newLeaf; + 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. diff --git a/packages/incremental-merkle-tree.sol/contracts/IncrementalQuinTree.sol b/packages/incremental-merkle-tree.sol/contracts/IncrementalQuinTree.sol index 1934237..c294cb3 100644 --- a/packages/incremental-merkle-tree.sol/contracts/IncrementalQuinTree.sol +++ b/packages/incremental-merkle-tree.sol/contracts/IncrementalQuinTree.sol @@ -79,6 +79,50 @@ library IncrementalQuinTree { self.numberOfLeaves += 1; } + /// @dev Updates a leaf in the tree. + /// @param self: Tree data. + /// @param leaf: Leaf to be updated. + /// @param newLeaf: New leaf. + /// @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 newLeaf, + uint256[4][] calldata proofSiblings, + uint8[] calldata proofPathIndices + ) public { + require( + verify(self, leaf, proofSiblings, proofPathIndices), + "IncrementalQuinTree: leaf is not part of the tree" + ); + require(newLeaf < SNARK_SCALAR_FIELD, "IncrementalQuinTree: leaf must be < SNARK_SCALAR_FIELD"); + + uint256 hash = newLeaf; + + 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. diff --git a/packages/incremental-merkle-tree.sol/contracts/README.md b/packages/incremental-merkle-tree.sol/contracts/README.md index 17d46a2..03d769a 100644 --- a/packages/incremental-merkle-tree.sol/contracts/README.md +++ b/packages/incremental-merkle-tree.sol/contracts/README.md @@ -39,7 +39,7 @@ ✔️ [IncrementalBinaryTree](https://github.com/privacy-scaling-explorations/zk-kit/blob/main/packages/incremental-merkle-tree.sol/contracts/IncrementalBinaryTree.sol) (Poseidon)\ ✔️ [IncrementalQuinTree](https://github.com/privacy-scaling-explorations/zk-kit/blob/main/packages/incremental-merkle-tree.sol/contracts/IncrementalQuinTree.sol) (Poseidon) -> The methods of each library are always the same (i.e `insert`, `remove`, `verify`). +> The methods of each library are always the same (i.e `insert`, `update`, `remove`, `verify`). --- @@ -74,6 +74,7 @@ contract Example { event TreeCreated(bytes32 id, uint8 depth); event LeafInserted(bytes32 indexed treeId, uint256 leaf, uint256 root); + event LeadUpdated(bytes32 indexed treeId, uint256 leaf, uint256 root); event LeafRemoved(bytes32 indexed treeId, uint256 leaf, uint256 root); mapping(bytes32 => IncrementalTreeData) public trees; @@ -94,6 +95,19 @@ contract Example { 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, "Example: 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/packages/incremental-merkle-tree.sol/contracts/test/IncrementalBinaryTreeTest.sol b/packages/incremental-merkle-tree.sol/contracts/test/IncrementalBinaryTreeTest.sol index 1594e35..95e863a 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; @@ -29,6 +30,20 @@ contract IncrementalBinaryTreeTest { emit LeafInserted(_treeId, _leaf, trees[_treeId].root); } + function updateLeaf( + bytes32 _treeId, + uint256 _leaf, + uint256 _newLeaf, + uint256[] calldata _proofSiblings, + uint8[] calldata _proofPathIndices + ) external { + require(trees[_treeId].depth != 0, "BinaryTreeTest: tree does not exist"); + + trees[_treeId].update(_leaf, _newLeaf, _proofSiblings, _proofPathIndices); + + emit LeafUpdated(_treeId, _newLeaf, trees[_treeId].root); + } + function removeLeaf( bytes32 _treeId, uint256 _leaf, diff --git a/packages/incremental-merkle-tree.sol/contracts/test/IncrementalQuinTreeTest.sol b/packages/incremental-merkle-tree.sol/contracts/test/IncrementalQuinTreeTest.sol index 62f7060..be65cea 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,20 @@ contract IncrementalQuinTreeTest { emit LeafInserted(_treeId, _leaf, trees[_treeId].root); } + function updateLeaf( + bytes32 _treeId, + uint256 _leaf, + uint256 _newLeaf, + uint256[4][] calldata _proofSiblings, + uint8[] calldata _proofPathIndices + ) external { + require(trees[_treeId].depth != 0, "QuinTreeTest: tree does not exist"); + + trees[_treeId].update(_leaf, _newLeaf, _proofSiblings, _proofPathIndices); + + emit LeafUpdated(_treeId, _newLeaf, trees[_treeId].root); + } + function removeLeaf( bytes32 _treeId, uint256 _leaf, diff --git a/packages/incremental-merkle-tree.sol/test/IncrementalBinaryTreeTest.ts b/packages/incremental-merkle-tree.sol/test/IncrementalBinaryTreeTest.ts index 2d338d7..a96e1fc 100644 --- a/packages/incremental-merkle-tree.sol/test/IncrementalBinaryTreeTest.ts +++ b/packages/incremental-merkle-tree.sol/test/IncrementalBinaryTreeTest.ts @@ -83,6 +83,60 @@ describe("IncrementalBinaryTreeTest", () => { await expect(transaction).to.be.revertedWith("IncrementalBinaryTree: tree is full") }) + it("Should not update a leaf if the tree does not exist", async () => { + const treeId = ethers.utils.formatBytes32String("none") + + const transaction = contract.updateLeaf(treeId, leaf, leaf, [0, 1], [0, 1]) + + await expect(transaction).to.be.revertedWith("BinaryTreeTest: tree does not exist") + }) + + it("Should not update a leaf if its value is > SNARK_SCALAR_FIELD", async () => { + const leaf = BigInt("21888242871839275222246405745257275088548364400416034343698204186575808495618") + + const transaction = contract.updateLeaf(treeId, leaf, leaf, [0, 1], [0, 1]) + + await expect(transaction).to.be.revertedWith("IncrementalBinaryTree: leaf must be < SNARK_SCALAR_FIELD") + }) + + it("Should not update a leaf if wrong current leaf is given", async () => { + const treeId = ethers.utils.formatBytes32String("tree2") + const tree = createTree(depth, 0) + for (let i = 0; i < 4; i += 1) tree.insert(BigInt(i + 1)) + + const leaf = BigInt(1337) + tree.update(2, leaf) + const { pathIndices, siblings } = tree.createProof(2) + const transaction = contract.updateLeaf( + treeId, + leaf, + leaf, + siblings.map((s) => s[0]), + pathIndices + ) + + await expect(transaction).to.be.revertedWith("IncrementalBinaryTree: leaf is not part of the tree") + }) + + it("Should update a leaf", async () => { + const treeId = ethers.utils.formatBytes32String("tree2") + const tree = createTree(depth, 0) + for (let i = 0; i < 4; i += 1) tree.insert(BigInt(i + 1)) + + const leaf = BigInt(1337) + tree.update(2, leaf) + const { root, pathIndices, siblings } = tree.createProof(2) + const transaction = contract.updateLeaf( + treeId, + BigInt(3), + leaf, + siblings.map((s) => s[0]), + pathIndices + ) + + await expect(transaction).to.emit(contract, "LeafUpdated").withArgs(treeId, leaf, root) + }) + it("Should not remove a leaf if the tree does not exist", async () => { const treeId = ethers.utils.formatBytes32String("none") @@ -111,7 +165,6 @@ describe("IncrementalBinaryTreeTest", () => { await contract.insertLeaf(treeId, BigInt(3)) const { siblings, pathIndices, root } = tree.createProof(0) - const transaction = contract.removeLeaf( treeId, BigInt(1), diff --git a/packages/incremental-merkle-tree.sol/test/IncrementalQuinTreeTest.ts b/packages/incremental-merkle-tree.sol/test/IncrementalQuinTreeTest.ts index 950ef55..32556d4 100644 --- a/packages/incremental-merkle-tree.sol/test/IncrementalQuinTreeTest.ts +++ b/packages/incremental-merkle-tree.sol/test/IncrementalQuinTreeTest.ts @@ -86,6 +86,48 @@ describe("IncrementalQuinTreeTest", () => { await expect(transaction).to.be.revertedWith("IncrementalQuinTree: tree is full") }) + it("Should not update a leaf if the tree does not exist", async () => { + const treeId = ethers.utils.formatBytes32String("none") + + const transaction = contract.updateLeaf(treeId, leaf, leaf, [[0, 1, 2, 3]], [0]) + + await expect(transaction).to.be.revertedWith("QuinTreeTest: tree does not exist") + }) + + it("Should not update a leaf if its value is > SNARK_SCALAR_FIELD", async () => { + const leaf = BigInt("21888242871839275222246405745257275088548364400416034343698204186575808495618") + + const transaction = contract.updateLeaf(treeId, leaf, leaf, [[0, 1, 2, 3]], [0]) + + await expect(transaction).to.be.revertedWith("IncrementalQuinTree: leaf must be < SNARK_SCALAR_FIELD") + }) + + it("Should not update a leaf if wrong current leaf is given", async () => { + const treeId = ethers.utils.formatBytes32String("tree2") + const tree = createTree(depth, 0, 5) + for (let i = 0; i < 6; i += 1) tree.insert(BigInt(i + 1)) + + const leaf = BigInt(1337) + tree.update(2, leaf) + const { pathIndices, siblings } = tree.createProof(2) + const transaction = contract.updateLeaf(treeId, leaf, leaf, siblings, pathIndices) + + await expect(transaction).to.be.revertedWith("IncrementalQuinTree: leaf is not part of the tree") + }) + + it("Should update a leaf", async () => { + const treeId = ethers.utils.formatBytes32String("tree2") + const tree = createTree(depth, 0, 5) + for (let i = 0; i < 6; i += 1) tree.insert(BigInt(i + 1)) + + const leaf = BigInt(1337) + tree.update(2, leaf) + const { pathIndices, siblings, root } = tree.createProof(2) + const transaction = contract.updateLeaf(treeId, BigInt(3), leaf, siblings, pathIndices) + + await expect(transaction).to.emit(contract, "LeafUpdated").withArgs(treeId, leaf, root) + }) + it("Should not remove a leaf if the tree does not exist", async () => { const treeId = ethers.utils.formatBytes32String("none") @@ -114,7 +156,6 @@ describe("IncrementalQuinTreeTest", () => { await contract.insertLeaf(treeId, BigInt(3)) const { siblings, pathIndices, root } = tree.createProof(0) - const transaction = contract.removeLeaf(treeId, BigInt(1), siblings, pathIndices) await expect(transaction).to.emit(contract, "LeafRemoved").withArgs(treeId, BigInt(1), root) 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