chore: add assertions to integration tests, fix linting issues, improve docs (#51)

.github/workflows/contracts.yml

@@ -112,7 +112,7 @@ jobs:
         run: yarn build
 
       - name: Run tests
-        run: yarn test:integration
+        run: yarn test:symbolic
 
   lint:
     name: Lint Commit Messages

packages/contracts/package.json

@@ -24,7 +24,7 @@
     "test": "forge test -vvv",
     "test:fuzz": "medusa fuzz",
     "test:integration": "forge test --match-contract Integration -vv --ffi",
-    "test:symbolic": "halmos",
+    "test:symbolic": "echo 'update this yarn script when symbolic testing is implemented'",
     "test:unit": "forge test --match-contract Unit -vvv",
     "test:unit:deep": "FOUNDRY_FUZZ_RUNS=5000 yarn test:unit"
   },

packages/contracts/script/Entrypoint.s.sol

@@ -4,9 +4,7 @@ pragma solidity 0.8.28;
 
 import {Entrypoint} from 'contracts/Entrypoint.sol';
 
-import {UnsafeUpgrades} from '@upgrades/Upgrades.sol';
 import {Script} from 'forge-std/Script.sol';
-import {console} from 'forge-std/console.sol';
 
 import {IERC20} from '@oz/interfaces/IERC20.sol';
 import {IPrivacyPool} from 'interfaces/IPrivacyPool.sol';
@@ -19,7 +17,7 @@ contract RegisterPool is Script {
   uint256 internal _minimumDepositAmount;
   uint256 internal _vettingFeeBPS;
 
-  address internal ETH = 0xEeeeeEeeeEeEeeEeEeEeeEEEeeeeEeeeeeeeEEeE;
+  address internal constant _ETH = 0xEeeeeEeeeEeEeeEeEeEeeEEEeeeeEeeeeeeeEEeE;
 
   function setUp() public {
     entrypoint = Entrypoint(payable(vm.envAddress('ENTRYPOINT_ADDRESS')));
@@ -27,7 +25,7 @@ contract RegisterPool is Script {
     try vm.parseAddress(vm.prompt('Enter asset address')) returns (address _assetAddress) {
       _asset = IERC20(_assetAddress);
     } catch {
-      _asset = IERC20(ETH);
+      _asset = IERC20(_ETH);
     }
 
     _pool = IPrivacyPool(vm.parseAddress(vm.prompt('Enter pool address')));

packages/contracts/src/contracts/PrivacyPool.sol

@@ -44,7 +44,7 @@ abstract contract PrivacyPool is State, IPrivacyPool {
     if (msg.sender != _w.processooor) revert InvalidProcesooor();
 
     // Check the context matches the proof's public signal to ensure its integrity
-    if (_p.context() != uint256(keccak256(abi.encode(_w, SCOPE))) % SNARK_SCALAR_FIELD) revert ContextMismatch();
+    if (_p.context() != uint256(keccak256(abi.encode(_w, SCOPE))) % _SNARK_SCALAR_FIELD) revert ContextMismatch();
 
     // Check the state root is known
     if (!_isKnownRoot(_p.stateRoot())) revert UnknownStateRoot();
@@ -69,7 +69,7 @@ abstract contract PrivacyPool is State, IPrivacyPool {
     // Store asset address
     ASSET = _asset;
     // Compute SCOPE
-    SCOPE = uint256(keccak256(abi.encodePacked(address(this), block.chainid, _asset))) % SNARK_SCALAR_FIELD;
+    SCOPE = uint256(keccak256(abi.encodePacked(address(this), block.chainid, _asset))) % _SNARK_SCALAR_FIELD;
   }
 
   /*///////////////////////////////////////////////////////////////
@@ -86,7 +86,7 @@ abstract contract PrivacyPool is State, IPrivacyPool {
     if (dead) revert PoolIsDead();
 
     // Compute label
-    uint256 _label = uint256(keccak256(abi.encodePacked(SCOPE, ++nonce))) % SNARK_SCALAR_FIELD;
+    uint256 _label = uint256(keccak256(abi.encodePacked(SCOPE, ++nonce))) % _SNARK_SCALAR_FIELD;
     // Store depositor and ragequit cooldown
     deposits[_label] = Deposit(_depositor, _value, block.timestamp + 1 weeks);
 

packages/contracts/src/contracts/State.sol

@@ -30,7 +30,8 @@ import {IVerifier} from 'interfaces/IVerifier.sol';
 abstract contract State is IState {
   using InternalLeanIMT for LeanIMTData;
 
-  uint256 internal constant SNARK_SCALAR_FIELD =
+  // @notice The SNARK scalar field
+  uint256 internal constant _SNARK_SCALAR_FIELD =
     21_888_242_871_839_275_222_246_405_745_257_275_088_548_364_400_416_034_343_698_204_186_575_808_495_617;
 
   /// @inheritdoc IState
@@ -109,7 +110,7 @@ abstract contract State is IState {
     currentRootIndex = newRootIndex;
 
     // Store the new root at the new index
-    roots[newRootIndex] = _updatedRoot % SNARK_SCALAR_FIELD;
+    roots[newRootIndex] = _updatedRoot % _SNARK_SCALAR_FIELD;
 
     emit LeafInserted(_merkleTree.size, _leaf, _updatedRoot);
   }

packages/contracts/src/interfaces/IState.sol

@@ -10,6 +10,21 @@ import {IVerifier} from 'interfaces/IVerifier.sol';
  * @notice Interface for the State contract
  */
 interface IState {
+  /*///////////////////////////////////////////////////////////////
+                               ENUMS
+  //////////////////////////////////////////////////////////////*/
+
+  /**
+   * @notice Enum representing statuses of a nullifier
+   */
+  enum NullifierStatus {
+    NONE,
+    SPENT, // Nullifier is spent
+    RAGEQUIT_PENDING, // Nullifier is being ragequitted
+    RAGEQUIT_FINALIZED // Nullifier has been ragequitted
+
+  }
+
   /*///////////////////////////////////////////////////////////////
                               STRUCTS
   //////////////////////////////////////////////////////////////*/
@@ -26,21 +41,6 @@ interface IState {
     uint256 whenRagequitteable;
   }
 
-  /*///////////////////////////////////////////////////////////////
-                               ENUMS
-  //////////////////////////////////////////////////////////////*/
-
-  /**
-   * @notice Enum representing statuses of a nullifier
-   */
-  enum NullifierStatus {
-    NONE,
-    SPENT, // Nullifier is spent
-    RAGEQUIT_PENDING, // Nullifier is being ragequitted
-    RAGEQUIT_FINALIZED // Nullifier has been ragequitted
-
-  }
-
   /*///////////////////////////////////////////////////////////////
                               EVENTS
   //////////////////////////////////////////////////////////////*/

packages/contracts/src/interfaces/IVerifier.sol

@@ -1,18 +1,36 @@
 // SPDX-License-Identifier: MIT
 pragma solidity 0.8.28;
 
+/**
+ * @title IVerifier
+ * @notice Interface the Groth16 verifier contracts
+ */
 interface IVerifier {
+  /**
+   * @notice Verifies a Withdrawal Proof
+   * @param _pA First elliptic curve point (π_A) of the Groth16 proof, encoded as two field elements
+   * @param _pB Second elliptic curve point (π_B) of the Groth16 proof, encoded as 2x2 matrix of field elements
+   * @param _pC Third elliptic curve point (π_C) of the Groth16 proof, encoded as two field elements
+   * @return _valid The boolean indicating if the proof is valid
+   */
   function verifyProof(
-    uint256[2] memory pA,
-    uint256[2][2] memory pB,
-    uint256[2] memory pC,
-    uint256[8] memory pubSignals
-  ) external returns (bool);
+    uint256[2] memory _pA,
+    uint256[2][2] memory _pB,
+    uint256[2] memory _pC,
+    uint256[8] memory _pubSignals
+  ) external returns (bool _valid);
 
+  /**
+   * @notice Verifies a Ragequit Proof
+   * @param _pA First elliptic curve point (π_A) of the Groth16 proof, encoded as two field elements
+   * @param _pB Second elliptic curve point (π_B) of the Groth16 proof, encoded as 2x2 matrix of field elements
+   * @param _pC Third elliptic curve point (π_C) of the Groth16 proof, encoded as two field elements
+   * @return _valid The boolean indicating if the proof is valid
+   */
   function verifyProof(
-    uint256[2] memory pA,
-    uint256[2][2] memory pB,
-    uint256[2] memory pC,
-    uint256[5] memory pubSignals
-  ) external returns (bool);
+    uint256[2] memory _pA,
+    uint256[2][2] memory _pB,
+    uint256[2] memory _pC,
+    uint256[5] memory _pubSignals
+  ) external returns (bool _valid);
 }

packages/contracts/test/helper/MerkleProofGenerator.mjs

@@ -2,12 +2,18 @@ import { ethers } from "ethers";
 import { generateMerkleProof } from "@privacy-pool-core/sdk";
 
 async function main() {
+  // Fetch script arguments
   const args = process.argv.slice(2);
+
+  // Leaf we want to prove
   const leaf = BigInt(args[0]);
+  // All leaves in tree
   const leaves = args.slice(1).map(BigInt);
 
+  // Generate merkle proof using the SDK method
   const proof = generateMerkleProof(leaves, leaf);
 
+  // Fix the issue with index being NaN for depth 0 (LeanIMT issue)
   proof.index = Object.is(proof.index, NaN) ? 0 : proof.index;
 
   // Convert proof to ABI-encoded bytes

packages/contracts/test/helper/WithdrawalProofGenerator.mjs

@@ -57,6 +57,7 @@ async function main() {
   const paddedStateSiblings = padSiblings(stateMerkleProof[2], 32);
   const paddedAspSiblings = padSiblings(aspMerkleProof[2], 32);
 
+  // Generate withdrawal proof
   const { proof, publicSignals } = await sdk.proveWithdrawal(commitment, {
     context,
     withdrawalAmount: withdrawnValue,
@@ -80,6 +81,7 @@ async function main() {
     newNullifier,
   });
 
+  // Format the proof to match the Solidity struct
   const withdrawalProof = {
     _pA: [BigInt(proof.pi_a[0]), BigInt(proof.pi_a[1])],
     _pB: [
@@ -99,6 +101,7 @@ async function main() {
     ].map((x) => BigInt(x)),
   };
 
+  // ABI encode the proof
   const encodedProof = encodeAbiParameters(
     [
       {

packages/contracts/test/integration/IntegrationBase.sol

@@ -4,8 +4,8 @@ pragma solidity 0.8.28;
 import {Entrypoint, IEntrypoint} from 'contracts/Entrypoint.sol';
 import {IPrivacyPool} from 'contracts/PrivacyPool.sol';
 
-import {IPrivacyPoolComplex, PrivacyPoolComplex} from 'contracts/implementations/PrivacyPoolComplex.sol';
-import {IPrivacyPoolSimple, PrivacyPoolSimple} from 'contracts/implementations/PrivacyPoolSimple.sol';
+import {PrivacyPoolComplex} from 'contracts/implementations/PrivacyPoolComplex.sol';
+import {PrivacyPoolSimple} from 'contracts/implementations/PrivacyPoolSimple.sol';
 
 import {CommitmentVerifier} from 'contracts/verifiers/CommitmentVerifier.sol';
 import {WithdrawalVerifier} from 'contracts/verifiers/WithdrawalVerifier.sol';
@@ -96,7 +96,7 @@ contract IntegrationBase is Test {
   uint256[] internal _aspLeaves;
 
   // Snark Scalar Field
-  uint256 internal constant SNARK_SCALAR_FIELD =
+  uint256 public constant SNARK_SCALAR_FIELD =
     21_888_242_871_839_275_222_246_405_745_257_275_088_548_364_400_416_034_343_698_204_186_575_808_495_617;
 
   // Pranked addresses
@@ -116,7 +116,7 @@ contract IntegrationBase is Test {
   uint256 internal constant _DEFAULT_ASP_ROOT = uint256(keccak256('ASP_ROOT')) % Constants.SNARK_SCALAR_FIELD;
   uint256 internal constant _DEFAULT_NEW_COMMITMENT_HASH =
     uint256(keccak256('NEW_COMMITMENT_HASH')) % Constants.SNARK_SCALAR_FIELD;
-  bytes4 internal constant NONE = 0xb4dc0dee;
+  bytes4 public constant NONE = 0xb4dc0dee;
 
   /*///////////////////////////////////////////////////////////////
                               SETUP
@@ -169,6 +169,7 @@ contract IntegrationBase is Test {
     // Deal the asset to the depositor
     _deal(_params.depositor, _params.asset, _params.amount);
 
+    // If not ETH, approve Entrypoint to deposit funds
     if (_params.asset != _ETH) {
       vm.prank(_params.depositor);
       _params.asset.approve(address(_entrypoint), _params.amount);
@@ -177,6 +178,7 @@ contract IntegrationBase is Test {
     // Define pool to deposit to
     IPrivacyPool _pool = _params.asset == _ETH ? _ethPool : _daiPool;
 
+    // Fetch current nonce
     uint256 _currentNonce = _pool.nonce();
 
     // Compute deposit parameters
@@ -240,23 +242,29 @@ contract IntegrationBase is Test {
   //////////////////////////////////////////////////////////////*/
 
   function _selfWithdraw(WithdrawalParams memory _params) internal returns (Commitment memory _commitment) {
+    // Define pool to deposit to
     IPrivacyPool _pool = _params.commitment.asset == _ETH ? _ethPool : _daiPool;
 
+    // Build `Withdrawal` object for direct withdrawal
     IPrivacyPool.Withdrawal memory _withdrawal =
       IPrivacyPool.Withdrawal({processooor: _params.recipient, scope: _pool.SCOPE(), data: ''});
 
+    // Withdraw
     _commitment = _withdraw(_params.recipient, _pool, _withdrawal, _params, true);
   }
 
   function _withdrawThroughRelayer(WithdrawalParams memory _params) internal returns (Commitment memory _commitment) {
+    // Define pool to deposit to
     IPrivacyPool _pool = _params.commitment.asset == _ETH ? _ethPool : _daiPool;
 
+    // Build `Withdrawal` object for relayed withdrawal
     IPrivacyPool.Withdrawal memory _withdrawal = IPrivacyPool.Withdrawal({
       processooor: address(_entrypoint),
       scope: _pool.SCOPE(),
       data: abi.encode(_params.recipient, _RELAYER, _VETTING_FEE_BPS)
     });
 
+    // Withdraw
     _commitment = _withdraw(_RELAYER, _pool, _withdrawal, _params, false);
   }
 
@@ -267,9 +275,15 @@ contract IntegrationBase is Test {
     WithdrawalParams memory _params,
     bool _direct
   ) private returns (Commitment memory _commitment) {
+    // Fetch balances before withdrawal
+    uint256 _recipientInitialBalance = _balance(_params.recipient, _params.commitment.asset);
+    uint256 _entrypointInitialBalance = _balance(address(_entrypoint), _params.commitment.asset);
+    uint256 _poolInitialBalance = _balance(address(_pool), _params.commitment.asset);
+
     // Compute context hash
     uint256 _context = uint256(keccak256(abi.encode(_withdrawal, _pool.SCOPE()))) % SNARK_SCALAR_FIELD;
 
+    // Compute new commitment properties
     _commitment.value = _params.commitment.value - _params.withdrawnAmount;
     _commitment.label = _params.commitment.label;
     _commitment.nullifier = _genSecretBySeed(_params.newNullifier);
@@ -302,7 +316,34 @@ contract IntegrationBase is Test {
       _entrypoint.relay(_withdrawal, _proof);
     }
 
-    _insertIntoShadowMerkleTree(_commitment.hash);
+    if (_params.revertReason == NONE) {
+      // Check nullifier hash has been spent
+      assertTrue(_pool.nullifierHashes(_proof.pubSignals[1]), 'Existing nullifier hash must be spent');
+
+      // Insert new commitment in mirrored state tree
+      _insertIntoShadowMerkleTree(_commitment.hash);
+
+      // Discount fees if applicable
+      uint256 _withdrawnAmountAfterFees =
+        _direct ? _params.withdrawnAmount : _deductFee(_params.withdrawnAmount, _VETTING_FEE_BPS);
+
+      // Check balance changes
+      assertEq(
+        _balance(_params.recipient, _params.commitment.asset),
+        _recipientInitialBalance + _withdrawnAmountAfterFees,
+        'User balance mismatch'
+      );
+      assertEq(
+        _balance(address(_entrypoint), _params.commitment.asset),
+        _entrypointInitialBalance,
+        "Entrypoint balance shouldn't change"
+      );
+      assertEq(
+        _balance(address(_pool), _params.commitment.asset),
+        _poolInitialBalance - _params.withdrawnAmount,
+        'Pool balance mismatch'
+      );
+    }
   }
 
   /*///////////////////////////////////////////////////////////////
@@ -310,15 +351,36 @@ contract IntegrationBase is Test {
   //////////////////////////////////////////////////////////////*/
 
   function _ragequit(address _depositor, Commitment memory _commitment) internal {
+    // Define pool to ragequit from
     IPrivacyPool _pool = _commitment.asset == _ETH ? _ethPool : _daiPool;
 
+    uint256 _depositorInitialBalance = _balance(_depositor, _commitment.asset);
+    uint256 _entrypointInitialBalance = _balance(address(_entrypoint), _commitment.asset);
+    uint256 _poolInitialBalance = _balance(address(_pool), _commitment.asset);
+
     // Generate ragequit proof
-    ProofLib.RagequitProof memory _ragequitProof =
+    ProofLib.RagequitProof memory _proof =
       _generateRagequitProof(_commitment.value, _commitment.label, _commitment.nullifier, _commitment.secret);
 
-    // Initiate Ragequit
+    // Ragequit
     vm.prank(_depositor);
-    _pool.ragequit(_ragequitProof);
+    _pool.ragequit(_proof);
+
+    // Insert new commitment in mirrored state tree
+    assertTrue(_pool.nullifierHashes(_proof.pubSignals[2]), 'Existing nullifier hash must be spent');
+
+    // Check balance changes
+    assertEq(
+      _balance(_depositor, _commitment.asset), _depositorInitialBalance + _commitment.value, 'User balance mismatch'
+    );
+    assertEq(
+      _balance(address(_entrypoint), _commitment.asset),
+      _entrypointInitialBalance,
+      "Entrypoint balance shouldn't change"
+    );
+    assertEq(
+      _balance(address(_pool), _commitment.asset), _poolInitialBalance - _commitment.value, 'Pool balance mismatch'
+    );
   }
 
   /*///////////////////////////////////////////////////////////////
@@ -366,7 +428,9 @@ contract IntegrationBase is Test {
     private
     returns (ProofLib.WithdrawProof memory _proof)
   {
+    // Generate state merkle proof
     bytes memory _stateMerkleProof = _generateMerkleProof(_merkleLeaves, _params.existingCommitment);
+    // Generate ASP merkle proof
     bytes memory _aspMerkleProof = _generateMerkleProof(_aspLeaves, _params.label);
 
     string[] memory _inputs = new string[](12);

packages/contracts/test/unit/core/PrivacyPool.t.sol

@@ -4,8 +4,6 @@ pragma solidity 0.8.28;
 import {Test} from 'forge-std/Test.sol';
 import {InternalLeanIMT, LeafAlreadyExists, LeanIMTData} from 'lean-imt/InternalLeanIMT.sol';
 
-import {PoseidonT2} from 'poseidon/PoseidonT2.sol';
-import {PoseidonT3} from 'poseidon/PoseidonT3.sol';
 import {PoseidonT4} from 'poseidon/PoseidonT4.sol';
 
 import {IPrivacyPool, PrivacyPool} from 'contracts/PrivacyPool.sol';
@@ -14,7 +12,6 @@ import {ProofLib} from 'contracts/lib/ProofLib.sol';
 import {Constants} from 'test/helper/Constants.sol';
 
 import {IState} from 'interfaces/IState.sol';
-import {IVerifier} from 'interfaces/IVerifier.sol';
 
 /**
  * @notice Test contract for the PrivacyPool