libraries: merkle: MerkleTree: Add insert methods into Merkle tree

src/Darkpool.sol

@@ -8,12 +8,11 @@ import { VerificationKeys } from "./libraries/darkpool/VerificationKeys.sol";
 import { IHasher } from "./libraries/poseidon2/IHasher.sol";
 import { IVerifier } from "./libraries/verifier/IVerifier.sol";
 import { ValidWalletCreateStatement, StatementSerializer } from "./libraries/darkpool/PublicInputs.sol";
-import { MerkleTypes } from "./libraries/merkle/MerkleTypes.sol";
-import { MerkleTreeLib } from "./libraries/merkle/MerkleTreeLib.sol";
+import { MerkleTreeLib } from "./libraries/merkle/MerkleTree.sol";
 
 // Use the StatementSerializer for all statements
 using StatementSerializer for ValidWalletCreateStatement;
-using MerkleTreeLib for MerkleTypes.MerkleTree;
+using MerkleTreeLib for MerkleTreeLib.MerkleTree;
 
 contract Darkpool {
     /// @notice The hasher for the darkpool
@@ -22,14 +21,15 @@ contract Darkpool {
     IVerifier public verifier;
 
     /// @notice The Merkle tree for wallet commitments
-    MerkleTypes.MerkleTree public walletTree;
+    MerkleTreeLib.MerkleTree public walletTree;
 
     /// @notice The constructor for the darkpool
     /// @param hasher_ The hasher for the darkpool
     /// @param verifier_ The verifier for the darkpool
     constructor(IHasher hasher_, IVerifier verifier_) {
         hasher = hasher_;
         verifier = verifier_;
+        walletTree.initialize();
     }
 
     /// @notice Create a wallet in the darkpool
@@ -48,25 +48,7 @@ contract Darkpool {
         uint256 walletCommitment = hasher.spongeHash(hashInputs);
 
         // 3. Insert the wallet commitment into the Merkle tree
-        require(walletTree.isInitialized, "Merkle tree not initialized");
-        bytes32 leaf = bytes32(walletCommitment);
-        walletTree.insertLeaf(hasher, leaf);
-    }
-
-    /// @notice Initialize the Merkle tree for wallet commitments
-    /// @param depth The depth of the Merkle tree
-    function initializeMerkleTree(uint8 depth) public {
-        require(!walletTree.isInitialized, "Merkle tree already initialized");
-
-        walletTree.depth = depth;
-        walletTree.isInitialized = true;
-        walletTree.nextLeafIndex = 0;
-        walletTree.maxLeaves = 1 << depth; // 2^depth
-        walletTree.rootHistorySize = 0;
-
-        // Initialize the root to a default value (implementation-specific)
-        // This would typically involve computing the root of an empty tree
-        // based on your specific hashing implementation
-        walletTree.root = bytes32(0);
+        BN254.ScalarField walletCommitmentScalar = BN254.ScalarField.wrap(walletCommitment);
+        walletTree.insertLeaf(walletCommitmentScalar, hasher);
     }
 }

src/libraries/darkpool/Constants.sol

@@ -0,0 +1,11 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.0;
+
+/// @title Darkpool Constants
+/// @notice Constants used in the darkpool
+library DarkpoolConstants {
+    /// @notice The number of shares in a wallet
+    uint256 constant N_WALLET_SHARES = 70;
+    /// @notice The depth of the Merkle tree
+    uint256 constant MERKLE_DEPTH = 32;
+}

src/libraries/darkpool/PublicInputs.sol

@@ -5,9 +5,6 @@ import { BN254 } from "solidity-bn254/BN254.sol";
 
 // This file represents the public inputs (statements) for various proofs used by the darkpool
 
-/// @dev The number of public shares in a wallet
-uint256 constant N_WALLET_SHARES = 70;
-
 // -------------------
 // | Statement Types |
 // -------------------

src/libraries/merkle/MerkleTree.sol

@@ -0,0 +1,86 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.0;
+
+import { BN254 } from "solidity-bn254/BN254.sol";
+import { IHasher } from "../poseidon2/IHasher.sol";
+import { DarkpoolConstants } from "../darkpool/Constants.sol";
+
+/// @title MerkleTreeLib
+/// @notice Library for Merkle tree operations
+
+library MerkleTreeLib {
+    /// @notice Structure containing Merkle tree state
+    struct MerkleTree {
+        /// @notice The next available leaf index
+        uint64 nextIndex;
+        /// @notice The current root of the tree
+        BN254.ScalarField root;
+        /// @notice The current path of siblings for the next leaf to be inserted.
+        BN254.ScalarField[] siblingPath;
+        /// @notice The root history, mapping from historic roots to a boolean
+        mapping(BN254.ScalarField => bool) rootHistory;
+    }
+
+    /// @notice Initialize the Merkle tree
+    /// @param tree The tree to initialize
+    function initialize(MerkleTree storage tree) internal {
+        tree.nextIndex = 0;
+        tree.root = BN254.ScalarField.wrap(0);
+
+        // Initialize the sibling path
+        BN254.ScalarField zero = BN254.ScalarField.wrap(0);
+        tree.siblingPath = new BN254.ScalarField[](DarkpoolConstants.MERKLE_DEPTH);
+        for (uint256 i = 0; i < DarkpoolConstants.MERKLE_DEPTH; i++) {
+            tree.siblingPath[i] = zero;
+        }
+    }
+
+    /// @notice Returns the root of the tree
+    /// @param tree The tree to get the root of
+    /// @return The root of the tree
+    function root(MerkleTree storage tree) internal view returns (BN254.ScalarField) {
+        return tree.root;
+    }
+
+    /// @notice Returns whether the given root is in the history of the tree
+    /// @param tree The tree to check the root history of
+    /// @param root The root to check
+    /// @return Whether the root is in the history of the tree
+    function rootInHistory(MerkleTree storage tree, BN254.ScalarField root) internal view returns (bool) {
+        return tree.rootHistory[root];
+    }
+
+    /// @notice Insert a leaf into the tree
+    /// @param tree The tree to insert the leaf into
+    /// @param leaf The leaf to insert
+    function insertLeaf(MerkleTree storage tree, BN254.ScalarField leaf, IHasher hasher) internal {
+        // Compute the hash of the leaf into the tree
+        uint256 idx = tree.nextIndex;
+        uint256 leafUint = BN254.ScalarField.unwrap(leaf);
+        uint256[] memory sisterLeaves = new uint256[](tree.siblingPath.length);
+        for (uint256 i = 0; i < tree.siblingPath.length; i++) {
+            sisterLeaves[i] = BN254.ScalarField.unwrap(tree.siblingPath[i]);
+        }
+        uint256[] memory hashes = hasher.merkleHash(idx, leafUint, sisterLeaves);
+
+        // Update the tree
+        tree.nextIndex++;
+        BN254.ScalarField newRoot = BN254.ScalarField.wrap(hashes[hashes.length - 1]);
+        tree.root = newRoot;
+        tree.rootHistory[newRoot] = true;
+
+        // Update the sibling paths, switching between left and right nodes as appropriate
+        for (uint256 i = 0; i < DarkpoolConstants.MERKLE_DEPTH; i++) {
+            uint256 idxBit = (idx >> i) & 1;
+            if (idxBit == 0) {
+                // Left node, the new sibling is the intermediate hash computed in the merkle insertion
+                tree.siblingPath[i] = BN254.ScalarField.wrap(hashes[i]);
+            } else {
+                // Right node, the new sibling is in a new sub-tree, and is the zero value
+                // for this depth in the tree
+                // TODO: Use depth-dependent zero values
+                tree.siblingPath[i] = BN254.ScalarField.wrap(0);
+            }
+        }
+    }
+}

test/Darkpool.t.sol

@@ -5,6 +5,7 @@ import { BN254 } from "solidity-bn254/BN254.sol";
 import { Test } from "forge-std/Test.sol";
 import { TestUtils } from "./utils/TestUtils.sol";
 import { HuffDeployer } from "foundry-huff/HuffDeployer.sol";
+import { console2 } from "forge-std/console2.sol";
 
 import { PlonkProof } from "../src/libraries/verifier/Types.sol";
 import { Darkpool } from "../src/Darkpool.sol";
@@ -44,6 +45,9 @@ contract DarkpoolTest is TestUtils {
         ValidWalletCreateStatement memory statement =
             ValidWalletCreateStatement({ privateShareCommitment: privateShareCommitment, publicShares: publicShares });
 
+        uint256 gasStart = gasleft();
         darkpool.createWallet(statement, proof);
+        uint256 gasEnd = gasleft();
+        console2.log("Gas used:", gasStart - gasEnd);
     }
 }

test/utils/TestUtils.sol

@@ -3,7 +3,7 @@ pragma solidity ^0.8.0;
 
 import { Test } from "forge-std/Test.sol";
 import { BN254 } from "lib/solidity-bn254/src/BN254.sol";
-import { N_WALLET_SHARES } from "src/libraries/darkpool/PublicInputs.sol";
+import { DarkpoolConstants } from "src/libraries/darkpool/Constants.sol";
 
 contract TestUtils is Test {
     /// @dev The BN254 field modulus from roundUtils.huff
@@ -45,8 +45,8 @@ contract TestUtils is Test {
 
     /// @dev Generate a random set of wallet shares
     function randomWalletShares() internal returns (BN254.ScalarField[] memory) {
-        BN254.ScalarField[] memory shares = new BN254.ScalarField[](N_WALLET_SHARES);
-        for (uint256 i = 0; i < N_WALLET_SHARES; i++) {
+        BN254.ScalarField[] memory shares = new BN254.ScalarField[](DarkpoolConstants.N_WALLET_SHARES);
+        for (uint256 i = 0; i < DarkpoolConstants.N_WALLET_SHARES; i++) {
             shares[i] = BN254.ScalarField.wrap(randomFelt());
         }
         return shares;