test: Darkpool.t.sol: Add update wallet tests

src/Darkpool.sol

@@ -14,10 +14,10 @@ import {
     StatementSerializer
 } from "./libraries/darkpool/PublicInputs.sol";
 import { MerkleTreeLib } from "./libraries/merkle/MerkleTree.sol";
+import { NullifierLib } from "./libraries/darkpool/NullifierSet.sol";
 
-// Use the StatementSerializer for all statements
-using StatementSerializer for ValidWalletCreateStatement;
 using MerkleTreeLib for MerkleTreeLib.MerkleTree;
+using NullifierLib for NullifierLib.NullifierSet;
 
 contract Darkpool {
     /// @notice The hasher for the darkpool
@@ -26,17 +26,46 @@ contract Darkpool {
     IVerifier public verifier;
 
     /// @notice The Merkle tree for wallet commitments
-    MerkleTreeLib.MerkleTree public walletTree;
+    MerkleTreeLib.MerkleTree private merkleTree;
+    /// @notice The nullifier set for the darkpool
+    /// @dev Each time a wallet is updated (placing an order, settling a match, depositing, etc) a nullifier is spent.
+    /// @dev This ensures that a pre-update wallet cannot create two separate post-update wallets in the Merkle state
+    /// @dev The nullifier is computed deterministically from the shares of the pre-update wallet
+    NullifierLib.NullifierSet private nullifierSet;
 
     /// @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();
+        merkleTree.initialize();
     }
 
+    // --- State Getters --- //
+
+    /// @notice Get the current Merkle root
+    /// @return The current Merkle root
+    function getMerkleRoot() public view returns (BN254.ScalarField) {
+        return merkleTree.root;
+    }
+
+    /// @notice Check whether a root is in the Merkle root history
+    /// @param root The root to check
+    /// @return Whether the root is in the history
+    function rootInHistory(BN254.ScalarField root) public view returns (bool) {
+        return merkleTree.rootHistory[root];
+    }
+
+    /// @notice Check whether a nullifier has been spent
+    /// @param nullifier The nullifier to check
+    /// @return Whether the nullifier has been spent
+    function nullifierSpent(BN254.ScalarField nullifier) public view returns (bool) {
+        return nullifierSet.isSpent(nullifier);
+    }
+
+    // --- Core Wallet Methods --- //
+
     /// @notice Create a wallet in the darkpool
     /// @param statement The statement to verify
     /// @param proof The proof of `VALID WALLET CREATE`
@@ -47,22 +76,29 @@ contract Darkpool {
         // 2. Compute a commitment to the wallet shares, and insert into the Merkle tree
         BN254.ScalarField walletCommitment =
             computeWalletCommitment(statement.publicShares, statement.privateShareCommitment);
-        walletTree.insertLeaf(walletCommitment, hasher);
+        merkleTree.insertLeaf(walletCommitment, hasher);
     }
 
     /// @notice Update a wallet in the darkpool
     /// @param statement The statement to verify
     /// @param proof The proof of `VALID WALLET UPDATE`
     function updateWallet(ValidWalletUpdateStatement memory statement, PlonkProof memory proof) public {
-        // 1. Verify the proof
+        // 1. Verify the Merkle root to which the pre-update wallet's inclusion proof opens,
+        // and check that the nullifier has not been spent
+        require(merkleTree.rootInHistory(statement.merkleRoot), "Invalid Merkle root");
+        nullifierSet.spend(statement.previousNullifier);
+
+        // 2. Verify the proof
         verifier.verifyValidWalletUpdate(statement, proof);
 
         // 2. Compute a commitment to the wallet shares, and insert into the Merkle tree
         BN254.ScalarField walletCommitment =
             computeWalletCommitment(statement.newPublicShares, statement.newPrivateShareCommitment);
-        walletTree.insertLeaf(walletCommitment, hasher);
+        merkleTree.insertLeaf(walletCommitment, hasher);
     }
 
+    // --- Helper Methods --- //
+
     /// @dev Compute a commitment to a wallet's shares
     function computeWalletCommitment(
         BN254.ScalarField[] memory publicShares,

src/libraries/darkpool/NullifierSet.sol

@@ -6,8 +6,8 @@ import { BN254 } from "solidity-bn254/BN254.sol";
 /// @title NullifierSet
 /// @notice Tracks the set of spent nullifiers in the darkpool, ensuring that a pre-update wallet
 /// @notice cannot create two separate post-update wallets
-library Nullifiers {
-    using Nullifiers for Nullifiers.NullifierSet;
+library NullifierLib {
+    using NullifierLib for NullifierLib.NullifierSet;
 
     /// @notice The nullifiers in the set
     struct NullifierSet {

src/libraries/darkpool/PublicInputs.sol

@@ -45,6 +45,11 @@ library StatementSerializer {
     using StatementSerializer for ExternalTransfer;
     using StatementSerializer for PublicRootKey;
 
+    /// @notice The number of scalar field elements in a ValidWalletCreateStatement
+    uint256 constant VALID_WALLET_CREATE_SCALAR_SIZE = 71;
+    /// @notice The number of scalar field elements in a ValidWalletUpdateStatement
+    uint256 constant VALID_WALLET_UPDATE_SCALAR_SIZE = 81;
+
     // --- Valid Wallet Create --- //
 
     /// @notice Serializes a ValidWalletCreateStatement into an array of scalar field elements
@@ -56,7 +61,7 @@ library StatementSerializer {
         returns (BN254.ScalarField[] memory)
     {
         // Create array with size = 1 (for privateShareCommitment) + publicShares.length
-        BN254.ScalarField[] memory serialized = new BN254.ScalarField[](1 + self.publicShares.length);
+        BN254.ScalarField[] memory serialized = new BN254.ScalarField[](VALID_WALLET_CREATE_SCALAR_SIZE);
 
         // Add the wallet commitment
         serialized[0] = self.privateShareCommitment;
@@ -79,7 +84,7 @@ library StatementSerializer {
         pure
         returns (BN254.ScalarField[] memory)
     {
-        BN254.ScalarField[] memory serialized = new BN254.ScalarField[](1 + 1 + self.newPublicShares.length);
+        BN254.ScalarField[] memory serialized = new BN254.ScalarField[](VALID_WALLET_UPDATE_SCALAR_SIZE);
         serialized[0] = self.previousNullifier;
         serialized[1] = self.newPrivateShareCommitment;
 
@@ -114,7 +119,7 @@ library StatementSerializer {
         serialized[0] = BN254.ScalarField.wrap(uint256(uint160(self.account)));
         serialized[1] = BN254.ScalarField.wrap(uint256(uint160(self.mint)));
         serialized[2] = BN254.ScalarField.wrap(self.amount);
-        serialized[3] = BN254.ScalarField.wrap(self.timestamp);
+        serialized[3] = BN254.ScalarField.wrap(uint256(self.transferType));
 
         return serialized;
     }

src/libraries/darkpool/Types.sol

@@ -17,8 +17,6 @@ struct ExternalTransfer {
     address mint;
     /// @dev The amount of the transfer
     uint256 amount;
-    /// @dev The timestamp of the transfer
-    uint256 timestamp;
     /// @dev Indicates if it's a deposit or withdrawal
     TransferType transferType;
 }

test/Darkpool.t.sol

@@ -4,22 +4,23 @@ pragma solidity ^0.8.0;
 import { BN254 } from "solidity-bn254/BN254.sol";
 import { Test } from "forge-std/Test.sol";
 import { TestUtils } from "./utils/TestUtils.sol";
+import { CalldataUtils } from "./utils/CalldataUtils.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";
-import { Nullifiers } from "../src/libraries/darkpool/NullifierSet.sol";
+import { NullifierLib } from "../src/libraries/darkpool/NullifierSet.sol";
 import { IHasher } from "../src/libraries/poseidon2/IHasher.sol";
 import { IVerifier } from "../src/libraries/verifier/IVerifier.sol";
 import { TestVerifier } from "./test-contracts/TestVerifier.sol";
-import { ValidWalletCreateStatement } from "../src/libraries/darkpool/PublicInputs.sol";
+import { ValidWalletCreateStatement, ValidWalletUpdateStatement } from "../src/libraries/darkpool/PublicInputs.sol";
 
-contract DarkpoolTest is TestUtils {
-    using Nullifiers for Nullifiers.NullifierSet;
+contract DarkpoolTest is CalldataUtils {
+    using NullifierLib for NullifierLib.NullifierSet;
 
     Darkpool public darkpool;
-    Nullifiers.NullifierSet private testNullifierSet;
+    NullifierLib.NullifierSet private testNullifierSet;
 
     function setUp() public {
         IHasher hasher = IHasher(HuffDeployer.deploy("libraries/poseidon2/poseidonHasher"));
@@ -46,30 +47,50 @@ contract DarkpoolTest is TestUtils {
 
     /// @notice Test creating a wallet
     function test_createWallet() public {
-        BN254.ScalarField dummyScalar = BN254.ScalarField.wrap(1);
-        BN254.G1Point memory dummyPoint = BN254.P1();
-        BN254.ScalarField[] memory publicInputs = new BN254.ScalarField[](1);
-        publicInputs[0] = dummyScalar;
-
-        PlonkProof memory proof = PlonkProof({
-            wire_comms: [dummyPoint, dummyPoint, dummyPoint, dummyPoint, dummyPoint],
-            z_comm: dummyPoint,
-            quotient_comms: [dummyPoint, dummyPoint, dummyPoint, dummyPoint, dummyPoint],
-            w_zeta: dummyPoint,
-            w_zeta_omega: dummyPoint,
-            wire_evals: [dummyScalar, dummyScalar, dummyScalar, dummyScalar, dummyScalar],
-            sigma_evals: [dummyScalar, dummyScalar, dummyScalar, dummyScalar],
-            z_bar: dummyScalar
-        });
-
-        BN254.ScalarField privateShareCommitment = BN254.ScalarField.wrap(randomFelt());
-        BN254.ScalarField[] memory publicShares = randomWalletShares();
-        ValidWalletCreateStatement memory statement =
-            ValidWalletCreateStatement({ privateShareCommitment: privateShareCommitment, publicShares: publicShares });
-
-        uint256 gasStart = gasleft();
+        (ValidWalletCreateStatement memory statement, PlonkProof memory proof) = createWalletCalldata();
         darkpool.createWallet(statement, proof);
-        uint256 gasEnd = gasleft();
-        console2.log("Gas used:", gasStart - gasEnd);
+    }
+
+    /// @notice Test updating a wallet
+    function test_updateWallet_validUpdate() public {
+        // Setup calldata
+        (ValidWalletUpdateStatement memory statement, PlonkProof memory proof) = updateWalletCalldata();
+
+        // Modify the merkle root to be valid
+        BN254.ScalarField currRoot = darkpool.getMerkleRoot();
+        statement.merkleRoot = currRoot;
+
+        // Update the wallet
+        darkpool.updateWallet(statement, proof);
+    }
+
+    /// @notice Test updating a wallet with an invalid Merkle root
+    function test_updateWallet_invalidMerkleRoot() public {
+        // Setup calldata
+        (ValidWalletUpdateStatement memory statement, PlonkProof memory proof) = updateWalletCalldata();
+
+        // Modify the merkle root to be invalid
+        statement.merkleRoot = randomScalar();
+
+        // Should fail
+        vm.expectRevert("Invalid Merkle root");
+        darkpool.updateWallet(statement, proof);
+    }
+
+    /// @notice Test updating a wallet with a spent nullifier
+    function test_updateWallet_spentNullifier() public {
+        // Setup calldata
+        (ValidWalletUpdateStatement memory statement, PlonkProof memory proof) = updateWalletCalldata();
+
+        // Modify the merkle root to be valid
+        BN254.ScalarField currRoot = darkpool.getMerkleRoot();
+        statement.merkleRoot = currRoot;
+
+        // First update should succeed
+        darkpool.updateWallet(statement, proof);
+
+        // Second update with same nullifier should fail
+        vm.expectRevert("Nullifier already spent");
+        darkpool.updateWallet(statement, proof);
     }
 }

test/utils/CalldataUtils.sol

@@ -0,0 +1,87 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.20;
+
+import { BN254 } from "solidity-bn254/BN254.sol";
+import { Vm } from "forge-std/Vm.sol";
+import { TestUtils, uintToScalarWords } from "./TestUtils.sol";
+import { PlonkProof } from "../../src/libraries/verifier/Types.sol";
+import { ExternalTransfer, PublicRootKey, TransferType } from "../../src/libraries/darkpool/Types.sol";
+import { ValidWalletCreateStatement, ValidWalletUpdateStatement } from "../../src/libraries/darkpool/PublicInputs.sol";
+
+// Utilities for generating darkpool calldata
+
+contract CalldataUtils is TestUtils {
+    /// @dev The first testing address
+    address public constant DUMMY_ADDRESS = address(0x1);
+    /// @dev A dummy wallet address
+    address public constant DUMMY_WALLET_ADDRESS = address(0x2);
+
+    // ---------------------
+    // | Darkpool Calldata |
+    // ---------------------
+
+    /// @notice Generate calldata for creating a wallet
+    function createWalletCalldata()
+        internal
+        returns (ValidWalletCreateStatement memory statement, PlonkProof memory proof)
+    {
+        statement = ValidWalletCreateStatement({
+            privateShareCommitment: BN254.ScalarField.wrap(randomFelt()),
+            publicShares: randomWalletShares()
+        });
+        proof = dummyPlonkProof();
+    }
+
+    /// @notice Generate calldata for updating a wallet
+    function updateWalletCalldata()
+        internal
+        returns (ValidWalletUpdateStatement memory statement, PlonkProof memory proof)
+    {
+        statement = ValidWalletUpdateStatement({
+            previousNullifier: randomScalar(),
+            newPublicShares: randomWalletShares(),
+            newPrivateShareCommitment: randomScalar(),
+            merkleRoot: randomScalar(),
+            externalTransfer: emptyExternalTransfer(),
+            oldPkRoot: randomRootKey()
+        });
+        proof = dummyPlonkProof();
+    }
+
+    // ------------------
+    // | Calldata Types |
+    // ------------------
+
+    /// @notice Generate an empty external transfer
+    function emptyExternalTransfer() internal pure returns (ExternalTransfer memory transfer) {
+        transfer =
+            ExternalTransfer({ account: address(0), mint: address(0), amount: 0, transferType: TransferType.Deposit });
+    }
+
+    /// @notice Generate a random root key
+    function randomRootKey() internal returns (PublicRootKey memory rootKey) {
+        Vm.Wallet memory wallet = randomEthereumWallet();
+        (BN254.ScalarField xLow, BN254.ScalarField xHigh) = uintToScalarWords(wallet.publicKeyX);
+        (BN254.ScalarField yLow, BN254.ScalarField yHigh) = uintToScalarWords(wallet.publicKeyY);
+        rootKey = PublicRootKey({ x: [xLow, xHigh], y: [yLow, yHigh] });
+    }
+
+    /// @notice Generates a dummy PlonK proof
+    function dummyPlonkProof() internal pure returns (PlonkProof memory proof) {
+        BN254.ScalarField dummyScalar = BN254.ScalarField.wrap(1);
+        BN254.G1Point memory dummyPoint = BN254.P1();
+        BN254.ScalarField[] memory publicInputs = new BN254.ScalarField[](1);
+        publicInputs[0] = dummyScalar;
+
+        proof = PlonkProof({
+            wire_comms: [dummyPoint, dummyPoint, dummyPoint, dummyPoint, dummyPoint],
+            z_comm: dummyPoint,
+            quotient_comms: [dummyPoint, dummyPoint, dummyPoint, dummyPoint, dummyPoint],
+            w_zeta: dummyPoint,
+            w_zeta_omega: dummyPoint,
+            wire_evals: [dummyScalar, dummyScalar, dummyScalar, dummyScalar, dummyScalar],
+            sigma_evals: [dummyScalar, dummyScalar, dummyScalar, dummyScalar],
+            z_bar: dummyScalar
+        });
+    }
+}

test/utils/TestUtils.sol

@@ -2,12 +2,21 @@
 pragma solidity ^0.8.0;
 
 import { Test } from "forge-std/Test.sol";
+import { Vm } from "forge-std/Vm.sol";
 import { BN254 } from "lib/solidity-bn254/src/BN254.sol";
 import { DarkpoolConstants } from "src/libraries/darkpool/Constants.sol";
 
+/// @dev Split a uint256 into a pair of BN254.ScalarField words, in little-endian order
+function uintToScalarWords(uint256 value) pure returns (BN254.ScalarField low, BN254.ScalarField high) {
+    low = BN254.ScalarField.wrap(value % BN254.R_MOD);
+    high = BN254.ScalarField.wrap(value / BN254.R_MOD);
+}
+
 contract TestUtils is Test {
     /// @dev The BN254 field modulus from roundUtils.huff
     uint256 constant PRIME = BN254.R_MOD;
+    /// @dev The scalar field modulus for K256
+    uint256 constant K256_SCALAR_MOD = 0xfffffffffffffffffffffffffffffffebaaedce6af48a03bbfd25e8cd0364141;
 
     // --- Fuzzing Helpers --- //
 
@@ -18,6 +27,11 @@ contract TestUtils is Test {
         return vm.randomUint() % PRIME;
     }
 
+    /// @dev Generate a random BN254 scalar field element
+    function randomScalar() internal returns (BN254.ScalarField) {
+        return BN254.ScalarField.wrap(randomFelt());
+    }
+
     /// @dev Generate a random G1 point
     function randomG1Point() internal returns (BN254.G1Point memory) {
         BN254.ScalarField scalar = BN254.ScalarField.wrap(randomFelt());
@@ -52,6 +66,12 @@ contract TestUtils is Test {
         return shares;
     }
 
+    /// @dev Generate a random ethereum wallet
+    function randomEthereumWallet() internal returns (Vm.Wallet memory) {
+        uint256 seed = vm.randomUint() % K256_SCALAR_MOD;
+        return vm.createWallet(seed);
+    }
+
     // --- FFI Helpers --- //
 
     /// @dev Helper to compile a Rust binary