Improve mock gas metering

.changeset/mock-gas-reporting.md

@@ -0,0 +1,8 @@
+---
+'@cofhe/mock-contracts': minor
+'@cofhe/foundry-plugin': minor
+'@cofhe/hardhat-plugin': minor
+'@cofhe/hardhat-3-plugin': minor
+---
+
+Realistic gas reporting for mocks. Under Foundry, `CofheTest.deployMocks()` now excludes mock-only work (FHE op replication, decrypt-task storage, logging) from gas metering by default, so reported gas approximates real-network cost — expect existing `forge snapshot` numbers to drop; opt out with `mockTaskManager.setMockGasExcluded(false)` (forge ≥ 1.0 recommended). A gas-metering pause owned by your own test is detected and left untouched. The mocks are also cheaper outright (enum op dispatch, log strings only built when logging is enabled). On Hardhat, the mock task manager emits a `MockGasConsumed(uint256)` event per block of mock-only work — note that FHE transaction receipts therefore carry additional logs, so tests using positional log access (`receipt.logs[i]`, `logs.length`) may need updating. Both hardhat plugins expose `getAdjustedGasUsed(receipt)` / `getAdjustedGasBreakdown(receipt)` (on `hre.cofhe` / `conn.cofhe`) and an opt-in `cofhe.gasSummary` config that prints a per-method raw-vs-adjusted table after `hardhat test` (reconstructed from chain history: snapshot-reverted transactions won't appear). `eth_estimateGas` remains unadjusted.

packages/foundry-plugin/README.md

@@ -66,17 +66,27 @@ contract MyTest is CofheTest {
 }
 ```
 
+## Gas reporting
+
+Mock FHE operations replicate the off-chain CoFHE work on-chain, which would normally inflate reported gas 2–3×. Under forge, `deployMocks()` excludes that mock-only work from gas metering by default (via `pauseGasMetering` cheatcodes), so `forge test --gas-report`, `forge snapshot`, and `gasleft()` measurements report numbers close to real-network costs.
+
+- Requires forge ≥ 1.0 (older versions had unreliable `pauseGasMetering` behavior). On unsupported setups the shim silently falls back to normal metering — correctness is never affected.
+- Opt out with `mockTaskManager.setMockGasExcluded(false)` to see raw mock gas.
+- When exclusion is off (or on Hardhat), the mock task manager instead emits a `MockGasConsumed(uint256)` event per block of mock-only work, so the overhead remains measurable from receipts.
+
+Note: upgrading to a version with this feature will lower existing `forge snapshot` numbers — the drop is the mock overhead disappearing from the report, not a change in your contracts.
+
 ## API
 
 ### `CofheTest` (abstract base)
 
-| Function                         | Description                                                  |
-| -------------------------------- | ------------------------------------------------------------ |
-| `deployMocks()`                  | Deploys all mock contracts and wires them together           |
-| `createCofheClient()`            | Returns a new unconnected `CofheClient`                      |
-| `enableLogs()` / `disableLogs()` | Toggle plaintext operation logging                           |
-| `getPlaintext(ctHash)`           | Returns the stored plaintext for a ciphertext handle         |
-| `expectPlaintext(handle, value)` | Asserts the plaintext of an encrypted handle matches `value` |
+| Function                         | Description                                                                                        |
+| -------------------------------- | -------------------------------------------------------------------------------------------------- |
+| `deployMocks()`                  | Deploys all mock contracts and wires them together; enables mock-gas exclusion (see Gas reporting) |
+| `createCofheClient()`            | Returns a new unconnected `CofheClient`                                                            |
+| `enableLogs()` / `disableLogs()` | Toggle plaintext operation logging                                                                 |
+| `getPlaintext(ctHash)`           | Returns the stored plaintext for a ciphertext handle                                               |
+| `expectPlaintext(handle, value)` | Asserts the plaintext of an encrypted handle matches `value`                                       |
 
 `getPlaintext` and `expectPlaintext` have typed overloads for `ebool`, `euint8`, `euint16`, `euint32`, `euint64`, `euint128`, and `eaddress`.
 

packages/foundry-plugin/contracts/CofheTest.sol

@@ -47,6 +47,11 @@ abstract contract CofheTest is Test {
     mockTaskManager = MockTaskManager(TASK_MANAGER_ADDRESS);
     mockTaskManager.initialize(TM_ADMIN);
     mockTaskManager.setLogOps(false);
+    // Exclude mock-only plaintext replication from gas metering so reported gas is closer
+    // to the real task manager. Requires cheatcode access, which etched contracts don't
+    // inherit from the test contract.
+    vm.allowCheatcodes(TASK_MANAGER_ADDRESS);
+    mockTaskManager.setMockGasExcluded(true);
     vm.label(address(mockTaskManager), 'MockTaskManager');
 
     // 2. ACL (non-fixed deploy so constructor runs and EIP712 domain is set)

packages/foundry-plugin/test/GasMetering.t.sol

@@ -0,0 +1,252 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.25;
+
+import { Test, Vm } from 'forge-std/Test.sol';
+import '@fhenixprotocol/cofhe-contracts/FHE.sol';
+import { MockTaskManager } from '@cofhe/mock-contracts/contracts/MockTaskManager.sol';
+import { CofheTest } from '../contracts/CofheTest.sol';
+
+/// @dev Minimal consumer contract exercising a typical FHE flow.
+contract GasConsumer {
+  euint32 public counter;
+
+  function init(uint32 start) public {
+    counter = FHE.asEuint32(start);
+    FHE.allowThis(counter);
+  }
+
+  function addToCounter(uint32 amount) public {
+    counter = FHE.add(counter, FHE.asEuint32(amount));
+    FHE.allowThis(counter);
+  }
+
+  function mixedOps(uint32 amount) public {
+    euint32 a = FHE.asEuint32(amount);
+    euint32 b = FHE.mul(a, FHE.asEuint32(amount + 1));
+    ebool gt = FHE.gt(b, counter);
+    counter = FHE.select(gt, b, counter);
+    FHE.allowThis(counter);
+  }
+
+  function requestDecrypt() public {
+    // createDecryptTask is a public task-manager entry point (not routed through the FHE
+    // library in this version); called directly to exercise its mock-gas tracking.
+    MockTaskManager(TASK_MANAGER_ADDRESS).createDecryptTask(uint256(euint32.unwrap(counter)), address(this));
+  }
+}
+
+contract GasMeteringTest is CofheTest {
+  bytes32 constant MOCK_GAS_CONSUMED_TOPIC = keccak256('MockGasConsumed(uint256)');
+
+  function setUp() public {
+    deployMocks();
+  }
+
+  /// @dev Deploys and warms up a consumer. Distinct `salt` values keep ctHashes (and thus
+  ///      storage-slot temperature) independent between consumers within one test.
+  function _newWarmConsumer(bool excluded, uint32 salt) internal returns (GasConsumer consumer) {
+    mockTaskManager.setMockGasExcluded(excluded);
+    consumer = new GasConsumer();
+    consumer.init(salt);
+    consumer.addToCounter(salt + 1); // warm up
+    consumer.mixedOps(salt + 2);
+  }
+
+  /// @dev Measures one addToCounter + one mixedOps + one decrypt on a warmed consumer.
+  function _measureFlows(
+    GasConsumer consumer,
+    uint32 salt
+  ) internal returns (uint256 gasAdd, uint256 gasMixed, uint256 gasDecrypt) {
+    uint256 g0 = gasleft();
+    consumer.addToCounter(salt + 3);
+    gasAdd = g0 - gasleft();
+
+    g0 = gasleft();
+    consumer.mixedOps(salt + 4);
+    gasMixed = g0 - gasleft();
+
+    g0 = gasleft();
+    consumer.requestDecrypt();
+    gasDecrypt = g0 - gasleft();
+  }
+
+  function _sumMockGasEvents(Vm.Log[] memory logs) internal view returns (uint256 total, uint256 count) {
+    for (uint256 i = 0; i < logs.length; i++) {
+      if (logs[i].emitter == address(mockTaskManager) && logs[i].topics[0] == MOCK_GAS_CONSUMED_TOPIC) {
+        total += abi.decode(logs[i].data, (uint256));
+        count++;
+      }
+    }
+  }
+
+  /// @dev The core promise of the exclusion feature: with it enabled, reported gas drops
+  ///      substantially. Guards against silent degradation (e.g. allowCheatcodes removed
+  ///      from deployMocks) - the shim is designed to fall back to metered execution
+  ///      without reverting, so only an assertion like this catches it.
+  function test_exclusionReducesReportedGas() public {
+    GasConsumer a = _newWarmConsumer(true, 10);
+    (uint256 exAdd, uint256 exMixed, uint256 exDecrypt) = _measureFlows(a, 10);
+    GasConsumer b = _newWarmConsumer(false, 20);
+    (uint256 mAdd, uint256 mMixed, uint256 mDecrypt) = _measureFlows(b, 20);
+
+    emit log_named_uint('excluded  addToCounter', exAdd);
+    emit log_named_uint('metered   addToCounter', mAdd);
+    emit log_named_uint('excluded  mixedOps    ', exMixed);
+    emit log_named_uint('metered   mixedOps    ', mMixed);
+    emit log_named_uint('excluded  decrypt     ', exDecrypt);
+    emit log_named_uint('metered   decrypt     ', mDecrypt);
+
+    assertLt(exAdd, mAdd, 'exclusion should reduce addToCounter gas');
+    assertGt(mAdd - exAdd, 50_000, 'addToCounter exclusion delta too small');
+    assertLt(exMixed, mMixed, 'exclusion should reduce mixedOps gas');
+    assertGt(mMixed - exMixed, 100_000, 'mixedOps exclusion delta too small');
+    assertLt(exDecrypt, mDecrypt, 'exclusion should reduce decrypt gas');
+    assertGt(mDecrypt - exDecrypt, 30_000, 'decrypt exclusion delta too small');
+  }
+
+  /// @dev Cross-checks the two mechanisms against each other: the forge-excluded gas and
+  ///      the metered gas minus the sum of MockGasConsumed events must agree. If either
+  ///      the pause bracketing or the measurement bracketing drifts (mock work added
+  ///      outside a bracket, miscalibrated event cost), the two diverge and this fails.
+  function _crossCheck(uint32 saltA, uint32 saltB) internal {
+    GasConsumer a = _newWarmConsumer(true, saltA);
+    (uint256 exAdd, uint256 exMixed, uint256 exDecrypt) = _measureFlows(a, saltA);
+    uint256 excludedTotal = exAdd + exMixed + exDecrypt;
+
+    GasConsumer b = _newWarmConsumer(false, saltB);
+    vm.recordLogs();
+    (uint256 mAdd, uint256 mMixed, uint256 mDecrypt) = _measureFlows(b, saltB);
+    (uint256 mockGas, uint256 eventCount) = _sumMockGasEvents(vm.getRecordedLogs());
+    uint256 adjustedTotal = mAdd + mMixed + mDecrypt - mockGas;
+
+    emit log_named_uint('excluded total ', excludedTotal);
+    emit log_named_uint('adjusted total ', adjustedTotal);
+    emit log_named_uint('mock gas       ', mockGas);
+    emit log_named_uint('mock events    ', eventCount);
+
+    assertGt(eventCount, 0, 'expected MockGasConsumed events on the metered path');
+    uint256 diff = adjustedTotal > excludedTotal ? adjustedTotal - excludedTotal : excludedTotal - adjustedTotal;
+    // Small asymmetries are expected (cheatcode call overhead on the excluded path,
+    // event-cost calibration), bounded per mock block.
+    assertLt(diff, eventCount * 1_500, 'adjusted gas diverges from excluded gas');
+  }
+
+  function test_adjustedMatchesExcluded_logsOff() public {
+    _crossCheck(30, 40);
+  }
+
+  function test_adjustedMatchesExcluded_logsOn() public {
+    enableLogs();
+    _crossCheck(50, 60);
+  }
+
+  /// @dev No MockGasConsumed events while gas metering is paused (forge exclusion active) -
+  ///      emitting there would double-signal overhead that was never counted.
+  function test_noEventsWhenExcluded() public {
+    GasConsumer a = _newWarmConsumer(true, 70);
+    vm.recordLogs();
+    _measureFlows(a, 70);
+    (, uint256 eventCount) = _sumMockGasEvents(vm.getRecordedLogs());
+    assertEq(eventCount, 0, 'no MockGasConsumed events expected while metering is paused');
+  }
+
+  /// @dev A pause owned by the caller's test must survive FHE ops: the shim detects the
+  ///      pre-existing pause and neither re-pauses nor resumes (and emits no events).
+  function test_userOwnedPauseNotClobbered() public {
+    mockTaskManager.setMockGasExcluded(true);
+    GasConsumer consumer = new GasConsumer();
+    consumer.init(0);
+    consumer.addToCounter(1); // warm
+
+    vm.recordLogs();
+    vm.pauseGasMetering();
+
+    consumer.addToCounter(5); // FHE ops inside the user's paused region
+
+    // If the shim resumed metering, this is now metered.
+    uint256 g0 = gasleft();
+    consumer.addToCounter(7);
+    uint256 gasInsidePause = g0 - gasleft();
+    vm.resumeGasMetering();
+
+    assertEq(gasInsidePause, 0, 'mock shim resumed metering inside user-owned paused region');
+    (, uint256 eventCount) = _sumMockGasEvents(vm.getRecordedLogs());
+    assertEq(eventCount, 0, 'no MockGasConsumed events expected while unmetered');
+    expectPlaintext(euint32.unwrap(consumer.counter()), 13);
+  }
+
+  /// @dev Correctness must be identical with the shim enabled.
+  function test_resultsUnchangedWithExclusion() public {
+    mockTaskManager.setMockGasExcluded(true);
+    GasConsumer consumer = new GasConsumer();
+    consumer.init(0);
+    consumer.addToCounter(5);
+    consumer.addToCounter(7);
+    expectPlaintext(euint32.unwrap(consumer.counter()), 12);
+
+    consumer.mixedOps(99); // 99 * 100 > 12 -> counter = 9900
+    expectPlaintext(euint32.unwrap(consumer.counter()), 9900);
+  }
+
+  /// @dev Logging path also works while metering is paused.
+  function test_worksWithLogsEnabled() public {
+    mockTaskManager.setMockGasExcluded(true);
+    enableLogs();
+    GasConsumer consumer = new GasConsumer();
+    consumer.init(0);
+    consumer.addToCounter(3);
+    expectPlaintext(euint32.unwrap(consumer.counter()), 3);
+  }
+}
+
+/// @dev Triggers a genuine revert inside the mock replication: trivially-encrypted
+///      hashes skip the ACL check in createTask, but have no plaintext in mock storage,
+///      so MOCK_twoInputOperation reverts with InputNotInMockStorage while metering is paused.
+contract GasMeteringRevertTest is CofheTest {
+  GasConsumer consumer;
+
+  function setUp() public {
+    deployMocks();
+    consumer = new GasConsumer();
+    consumer.init(0);
+  }
+
+  function _fabricatedTrivialHash(bytes32 seed) internal pure returns (uint256) {
+    // keccak-derived hash with metadata: [type|trivial bit] byte + securityZone byte
+    return (uint256(keccak256(abi.encode(seed))) & ~uint256(0xFFFF)) | (uint256(0x80 | Utils.EUINT32_TFHE) << 8);
+  }
+
+  function test_meteringResumesAfterMockRevert() public {
+    mockTaskManager.setMockGasExcluded(true);
+    consumer.addToCounter(1); // warm up
+
+    uint256[] memory hashes = new uint256[](2);
+    hashes[0] = _fabricatedTrivialHash('a');
+    hashes[1] = _fabricatedTrivialHash('b');
+    uint256[] memory extra = new uint256[](0);
+
+    vm.expectRevert();
+    mockTaskManager.createTask(Utils.EUINT32_TFHE, FunctionId.add, hashes, extra);
+
+    uint256 g0 = gasleft();
+    consumer.addToCounter(5);
+    uint256 gasAdd = g0 - gasleft();
+    emit log_named_uint('post-revert addToCounter', gasAdd);
+    // Same op measures ~100k when metering is healthy; near-zero means the paused state leaked
+    assertGt(gasAdd, 50000, 'gas metering stayed paused after mock revert');
+  }
+
+  /// @dev Same leak-check for createDecryptTask's tracked block: the plaintext read that can
+  ///      revert (InputNotInMockStorage) happens before metering is paused.
+  function test_meteringResumesAfterDecryptRevert() public {
+    mockTaskManager.setMockGasExcluded(true);
+    consumer.addToCounter(1); // warm up
+
+    vm.expectRevert();
+    mockTaskManager.createDecryptTask(_fabricatedTrivialHash('c'), address(this));
+
+    uint256 g0 = gasleft();
+    consumer.addToCounter(5);
+    assertGt(g0 - gasleft(), 50000, 'gas metering stayed paused after decrypt revert');
+  }
+}

packages/hardhat-3-plugin/README.md

@@ -36,18 +36,20 @@ export default defineConfig({
 
   // Optional cofhe config (all values shown are their defaults)
   cofhe: {
-    gasWarning: true, // warn when mock ops report higher gas than real FHE
-    logMocks: true, // enable mock contract event logging
+    gasWarning: false, // warn when mock ops report higher gas than real FHE
+    gasSummary: false, // print a per-method adjusted-gas summary after `hardhat test`
+    logMocks: false, // enable mock contract event logging
   },
 });
 ```
 
 ### Config options
 
-| Option       | Type      | Default | Description                                                                              |
-| ------------ | --------- | ------- | ---------------------------------------------------------------------------------------- |
-| `gasWarning` | `boolean` | `true`  | Print a warning after mock deployment reminding that mock gas costs differ from live FHE |
-| `logMocks`   | `boolean` | `true`  | Enable event-based logging inside mock contracts                                         |
+| Option       | Type      | Default | Description                                                                                       |
+| ------------ | --------- | ------- | ------------------------------------------------------------------------------------------------- |
+| `gasWarning` | `boolean` | `false` | Print a warning after mock deployment reminding that mock gas costs differ from live FHE          |
+| `gasSummary` | `boolean` | `false` | Print a per-method gas summary after `hardhat test`, with mock-only overhead excluded (see below) |
+| `logMocks`   | `boolean` | `false` | Enable event-based logging inside mock contracts                                                  |
 
 ---
 
@@ -127,6 +129,40 @@ const client = await cofhe.createClientWithBatteries(walletClient);
 
 ---
 
+### `cofhe.getAdjustedGasUsed(receipt)` / `cofhe.getAdjustedGasBreakdown(receipt)`
+
+Mock transactions consume more gas than they would on a real CoFHE network, because the off-chain FHE work is replicated on-chain. The mock task manager measures that overhead and emits a `MockGasConsumed(uint256)` event per block of mock-only work, letting the plugin report corrected numbers:
+
+```typescript
+const hash = await myContract.write.doFheThings();
+const receipt = await publicClient.waitForTransactionReceipt({ hash });
+
+// Gas usage excluding mock overhead — an estimate of real-network cost.
+const adjusted = cofhe.getAdjustedGasUsed(receipt);
+
+// Or the full breakdown:
+const { gasUsed, mockGas, adjustedGasUsed, mockGasEvents } = cofhe.getAdjustedGasBreakdown(receipt);
+```
+
+Both are pure functions of the receipt (no RPC calls). On a real network the receipt carries no mock events, so `adjustedGasUsed` equals `gasUsed` — the same code works everywhere.
+
+With `cofhe: { gasSummary: true }` in the config, a per-method table (raw vs adjusted gas) is printed after `hardhat test`:
+
+```
+[COFHE-MOCKS] Gas summary — adjusted ≈ cost excluding mock-only overhead
+┌──────────────────┬──────────────────────────┬───────┬─────────────────┬────────────────────┬───────────────┐
+│ Contract         │ Method                   │ Calls │ Avg gas (mocks) │ Avg gas (adjusted) │ Mock overhead │
+├──────────────────┼──────────────────────────┼───────┼─────────────────┼────────────────────┼───────────────┤
+│ MyFHEContract    │ doFheThings()            │ 9     │ 146,199         │ 109,595            │ 25%           │
+└──────────────────┴──────────────────────────┴───────┴─────────────────┴────────────────────┴───────────────┘
+```
+
+Note: `eth_estimateGas` is not adjusted — the mock work really does execute, so transactions still need the raw gas limit. Use a testnet for estimate-sensitive flows.
+
+The summary table is collected per test worker at process exit and reconstructed from chain history: transactions rolled back by snapshots won't appear, and it is best-effort after forced exits (`process.exit`, Ctrl+C) — use `getAdjustedGasUsed(receipt)` inside tests for exact per-transaction numbers.
+
+---
+
 ### `cofhe.mocks`
 
 #### Contract descriptors