@cofhe/hardhat-plugin 0.7.1 → 1.0.0-alpha.0

package/src/gas.ts

@@ -0,0 +1,227 @@
+import chalk from 'chalk';
+import { ethers } from 'ethers';
+import { MOCKS_TASK_MANAGER_ADDRESS } from '@cofhe/sdk';
+import type { HardhatRuntimeEnvironment } from 'hardhat/types';
+
+/// Topic of MockCoFHE's `MockGasConsumed(uint256)` event, emitted once per block of
+/// mock-only work (op replication, decrypt-task storage, log building) with the gas it
+/// consumed. Summing these per transaction gives the mock overhead of that transaction.
+export const MOCK_GAS_CONSUMED_TOPIC = ethers.id('MockGasConsumed(uint256)');
+
+/// Minimal receipt shape - compatible with ethers v6 TransactionReceipt and viem receipts.
+export type AdjustableGasReceipt = {
+  gasUsed: bigint | number | string;
+  logs?: ReadonlyArray<{ address: string; topics: ReadonlyArray<string>; data: string }>;
+};
+
+export type AdjustedGasBreakdown = {
+  /** Raw gas used by the transaction (includes mock overhead). */
+  gasUsed: bigint;
+  /** Gas consumed by mock-only work (FHE op replication, decrypt-task storage, mock logging). */
+  mockGas: bigint;
+  /** gasUsed - mockGas: an estimate of what the transaction would cost on a real CoFHE network. */
+  adjustedGasUsed: bigint;
+  /** Number of MockGasConsumed events - roughly the number of mock-only blocks executed. */
+  mockGasEvents: number;
+};
+
+/// Computes the gas breakdown of a transaction receipt by summing the MockGasConsumed
+/// events emitted by the mock task manager. Pure function of the receipt - no RPC calls.
+/// On a real network (no mock events in the logs) `adjustedGasUsed` equals `gasUsed`.
+/// Caveat: `mockGas` is measured pre-refund while `gasUsed` is post-refund, so transactions
+/// that earn large gas refunds (storage clearing) can slightly overstate the mock share;
+/// `adjustedGasUsed` is clamped at 0 for the pathological case.
+export const mock_getAdjustedGasBreakdown = (receipt: AdjustableGasReceipt): AdjustedGasBreakdown => {
+  let mockGas = 0n;
+  let mockGasEvents = 0;
+
+  for (const log of receipt.logs ?? []) {
+    if (
+      log.address.toLowerCase() === MOCKS_TASK_MANAGER_ADDRESS.toLowerCase() &&
+      log.topics[0] === MOCK_GAS_CONSUMED_TOPIC
+    ) {
+      mockGas += BigInt(log.data);
+      mockGasEvents += 1;
+    }
+  }
+
+  const gasUsed = BigInt(receipt.gasUsed);
+  return { gasUsed, mockGas, adjustedGasUsed: mockGas > gasUsed ? 0n : gasUsed - mockGas, mockGasEvents };
+};
+
+/// Returns the receipt's gas usage excluding mock overhead - an estimate of what the
+/// transaction would cost on a real CoFHE network. See mock_getAdjustedGasBreakdown.
+export const mock_getAdjustedGasUsed = (receipt: AdjustableGasReceipt): bigint =>
+  mock_getAdjustedGasBreakdown(receipt).adjustedGasUsed;
+
+// =====================
+//  END-OF-RUN SUMMARY
+// =====================
+
+type MethodAggregate = {
+  to: string;
+  selector: string;
+  calls: number;
+  totalGasUsed: bigint;
+  totalMockGas: bigint;
+};
+
+const shortAddress = (address: string) => `${address.slice(0, 6)}..${address.slice(-4)}`;
+
+/// Builds a selector -> "signature" map from every compiled artifact so summary rows can
+/// show method names instead of raw selectors.
+const buildSelectorNameMap = async (hre: HardhatRuntimeEnvironment): Promise<Map<string, string>> => {
+  const selectorNames = new Map<string, string>();
+  for (const fqn of await hre.artifacts.getAllFullyQualifiedNames()) {
+    try {
+      const artifact = await hre.artifacts.readArtifact(fqn);
+      const iface = new ethers.Interface(artifact.abi);
+      iface.forEachFunction((fn) => {
+        if (!selectorNames.has(fn.selector)) selectorNames.set(fn.selector, fn.format('sighash'));
+      });
+    } catch {
+      // Artifacts without a usable ABI are skipped
+    }
+  }
+  return selectorNames;
+};
+
+/// Resolves deployed addresses to contract names by matching on-chain code against
+/// artifact deployedBytecode. Best effort: contracts with immutables won't match and
+/// fall back to a shortened address.
+const buildAddressNameMap = async (
+  hre: HardhatRuntimeEnvironment,
+  addresses: string[]
+): Promise<Map<string, string>> => {
+  const codeToName = new Map<string, string>();
+  for (const fqn of await hre.artifacts.getAllFullyQualifiedNames()) {
+    try {
+      const artifact = await hre.artifacts.readArtifact(fqn);
+      if (artifact.deployedBytecode && artifact.deployedBytecode !== '0x') {
+        codeToName.set(artifact.deployedBytecode.toLowerCase(), artifact.contractName);
+      }
+    } catch {
+      // skip
+    }
+  }
+
+  const addressNames = new Map<string, string>();
+  for (const address of addresses) {
+    try {
+      const code: string = await hre.network.provider.send('eth_getCode', [address, 'latest']);
+      addressNames.set(address, codeToName.get(code.toLowerCase()) ?? shortAddress(address));
+    } catch {
+      addressNames.set(address, shortAddress(address));
+    }
+  }
+  return addressNames;
+};
+
+/// Prints a per-method gas summary for every transaction that performed mock FHE work,
+/// showing raw gas (as reported on the mock network) next to adjusted gas (an estimate of
+/// real-network cost). Enabled via `cofhe: { gasSummary: true }` in the hardhat config;
+/// runs after the test task completes.
+export const printMockGasSummary = async (hre: HardhatRuntimeEnvironment) => {
+  // Only meaningful against the in-process Hardhat network: elsewhere the mocks aren't
+  // auto-deployed and a genesis-to-latest eth_getLogs could hammer an external RPC.
+  if (hre.network.name !== 'hardhat') return;
+
+  let logs: Array<{ transactionHash: string; data: string }>;
+  try {
+    logs = await hre.network.provider.send('eth_getLogs', [
+      { fromBlock: '0x0', toBlock: 'latest', address: MOCKS_TASK_MANAGER_ADDRESS, topics: [MOCK_GAS_CONSUMED_TOPIC] },
+    ]);
+  } catch {
+    return;
+  }
+  if (!logs || logs.length === 0) return;
+
+  // Sum mock gas per transaction
+  const mockGasPerTx = new Map<string, bigint>();
+  for (const log of logs) {
+    mockGasPerTx.set(log.transactionHash, (mockGasPerTx.get(log.transactionHash) ?? 0n) + BigInt(log.data));
+  }
+
+  // Aggregate per called contract + method; deployments (FHE ops in constructors) are
+  // attributed to the created contract under a '(deployment)' method label.
+  const aggregates = new Map<string, MethodAggregate>();
+  for (const [txHash, mockGas] of mockGasPerTx) {
+    try {
+      const receipt = await hre.network.provider.send('eth_getTransactionReceipt', [txHash]);
+      const tx = await hre.network.provider.send('eth_getTransactionByHash', [txHash]);
+      if (!receipt) continue;
+      const target = tx?.to ?? receipt.contractAddress;
+      if (!target) continue;
+
+      const to = ethers.getAddress(target);
+      const selector = tx?.to ? (tx.input ?? '0x').slice(0, 10) : '(deployment)';
+      const key = `${to}:${selector}`;
+      const aggregate = aggregates.get(key) ?? { to, selector, calls: 0, totalGasUsed: 0n, totalMockGas: 0n };
+      aggregate.calls += 1;
+      aggregate.totalGasUsed += BigInt(receipt.gasUsed);
+      aggregate.totalMockGas += mockGas;
+      aggregates.set(key, aggregate);
+    } catch {
+      // Skip transactions that can't be resolved
+    }
+  }
+  if (aggregates.size === 0) return;
+
+  const selectorNames = await buildSelectorNameMap(hre);
+  const addressNames = await buildAddressNameMap(hre, [...new Set([...aggregates.values()].map((a) => a.to))]);
+
+  // Merge per-address aggregates by resolved contract name + method, so multiple deployed
+  // instances of the same contract (one per test, typically) share a row.
+  const merged = new Map<string, { contract: string; method: string } & Omit<MethodAggregate, 'to' | 'selector'>>();
+  for (const a of aggregates.values()) {
+    const contract = addressNames.get(a.to) ?? shortAddress(a.to);
+    const method = selectorNames.get(a.selector) ?? a.selector;
+    const key = `${contract}:${method}`;
+    const row = merged.get(key) ?? { contract, method, calls: 0, totalGasUsed: 0n, totalMockGas: 0n };
+    row.calls += a.calls;
+    row.totalGasUsed += a.totalGasUsed;
+    row.totalMockGas += a.totalMockGas;
+    merged.set(key, row);
+  }
+
+  const rows = [...merged.values()]
+    .map((r) => {
+      const avgGasUsed = r.totalGasUsed / BigInt(r.calls);
+      const avgMockGas = r.totalMockGas / BigInt(r.calls);
+      return {
+        contract: r.contract,
+        method: r.method,
+        calls: r.calls,
+        avgGasUsed,
+        avgAdjusted: avgMockGas > avgGasUsed ? 0n : avgGasUsed - avgMockGas,
+        overheadPct: avgGasUsed === 0n ? 0 : Math.min(100, Number((avgMockGas * 100n) / avgGasUsed)),
+      };
+    })
+    .sort((a, b) => a.contract.localeCompare(b.contract) || a.method.localeCompare(b.method));
+
+  const headers = ['Contract', 'Method', 'Calls', 'Avg gas (mocks)', 'Avg gas (adjusted)', 'Mock overhead'];
+  const cells = rows.map((r) => [
+    r.contract,
+    r.method,
+    String(r.calls),
+    r.avgGasUsed.toLocaleString(),
+    r.avgAdjusted.toLocaleString(),
+    `${r.overheadPct}%`,
+  ]);
+  const widths = headers.map((h, i) => Math.max(h.length, ...cells.map((row) => row[i].length)));
+  const line = (row: string[]) => `│ ${row.map((cell, i) => cell.padEnd(widths[i])).join(' │ ')} │`;
+  const divider = (l: string, m: string, r: string) => `${l}${widths.map((w) => '─'.repeat(w + 2)).join(m)}${r}`;
+
+  console.log('');
+  console.log(chalk.bold('[COFHE-MOCKS] Gas summary') + chalk.dim(' — adjusted ≈ cost excluding mock-only overhead'));
+  console.log(divider('┌', '┬', '┐'));
+  console.log(line(headers));
+  console.log(divider('├', '┼', '┤'));
+  for (const row of cells) console.log(line(row));
+  console.log(divider('└', '┴', '┘'));
+  console.log(
+    chalk.dim(
+      'Raw numbers include on-chain replication of off-chain CoFHE work; adjusted numbers estimate real-network cost. estimateGas remains unadjusted.'
+    )
+  );
+};

package/src/index.ts

@@ -22,12 +22,21 @@ import { deployMocks, type DeployMocksArgs, type LogMocksDeploy } from './deploy
 import { mock_setLoggingEnabled, mock_withLogs } from './logging.js';
 import { getFixedMockContract, mock_expectPlaintext } from './utils.js';
 import { mock_getPlaintext } from './utils.js';
+import {
+  mock_getAdjustedGasBreakdown,
+  mock_getAdjustedGasUsed,
+  printMockGasSummary,
+  type AdjustableGasReceipt,
+  type AdjustedGasBreakdown,
+} from './gas.js';
 import { hardhat } from '@cofhe/sdk/chains';
 import {
   MockThresholdNetworkArtifact,
+  MockCoFHEAddressBookArtifact,
   MockTaskManagerArtifact,
   MockZkVerifierArtifact,
   type MockACL,
+  type MockCoFHEAddressBook,
   type MockTaskManager,
   type MockThresholdNetwork,
   type MockZkVerifier,
@@ -44,6 +53,14 @@ declare module 'hardhat/types/config' {
       logMocks?: boolean;
       /** Whether to show gas usage warnings for mock operations (default: true) */
       gasWarning?: boolean;
+      /**
+       * Print a per-method gas summary after `hardhat test`, showing raw gas next to
+       * adjusted gas (mock-only overhead excluded - an estimate of real-network cost).
+       * Reconstructed from chain history: transactions rolled back by loadFixture /
+       * evm_revert won't appear - use getAdjustedGasUsed(receipt) for exact numbers.
+       * (default: false)
+       */
+      gasSummary?: boolean;
       /**
        * Controls deploy-mocks console output.
        * - `''`   — silent, no output
@@ -60,6 +77,8 @@ declare module 'hardhat/types/config' {
       logMocks: boolean;
       /** Whether to show gas usage warnings for mock operations (default: true) */
       gasWarning: boolean;
+      /** Print a per-method adjusted-gas summary after `hardhat test` (default: false) */
+      gasSummary: boolean;
       mocksDeployVerbosity: LogMocksDeploy;
     };
   }
@@ -132,27 +151,26 @@ subtask(TASK_COMPILE_SOLIDITY_GET_SOURCE_PATHS).setAction(async (taskArgs: { sou
 });
 
 extendConfig((config, userConfig) => {
-  // Allow users to override the localcofhe network config
-  if (userConfig.networks && userConfig.networks.localcofhe) {
-    return;
+  // Inject localcofhe preset only when the user hasn't defined it.
+  // Must not early-return: sepolia presets and config.cofhe still need to run
+  // (Hardhat 3 plugin already scopes the skip this way).
+  if (!userConfig.networks?.localcofhe) {
+    config.networks.localcofhe = {
+      gas: 'auto',
+      gasMultiplier: 1.2,
+      gasPrice: 'auto',
+      timeout: 10_000,
+      httpHeaders: {},
+      url: 'http://127.0.0.1:42069',
+      accounts: [
+        '0xac0974bec39a17e36ba4a6b4d238ff944bacb478cbed5efcae784d7bf4f2ff80',
+        '0x59c6995e998f97a5a0044966f0945389dc9e86dae88c7a8412f4603b6b78690d',
+        '0x5de4111afa1a4b94908f83103eb1f1706367c2e68ca870fc3fb9a804cdab365a',
+        '0x7c852118294e51e653712a81e05800f419141751be58f605c371e15141b007a6',
+      ],
+    };
   }
 
-  // Default config
-  config.networks.localcofhe = {
-    gas: 'auto',
-    gasMultiplier: 1.2,
-    gasPrice: 'auto',
-    timeout: 10_000,
-    httpHeaders: {},
-    url: 'http://127.0.0.1:42069',
-    accounts: [
-      '0xac0974bec39a17e36ba4a6b4d238ff944bacb478cbed5efcae784d7bf4f2ff80',
-      '0x59c6995e998f97a5a0044966f0945389dc9e86dae88c7a8412f4603b6b78690d',
-      '0x5de4111afa1a4b94908f83103eb1f1706367c2e68ca870fc3fb9a804cdab365a',
-      '0x7c852118294e51e653712a81e05800f419141751be58f605c371e15141b007a6',
-    ],
-  };
-
   // Only add Sepolia config if user hasn't defined it
   if (!userConfig.networks?.['eth-sepolia']) {
     config.networks['eth-sepolia'] = {
@@ -185,6 +203,7 @@ extendConfig((config, userConfig) => {
   config.cofhe = {
     logMocks: userConfig.cofhe?.logMocks ?? false,
     gasWarning: userConfig.cofhe?.gasWarning ?? false,
+    gasSummary: userConfig.cofhe?.gasSummary ?? false,
     mocksDeployVerbosity: (userConfig.cofhe?.mocksDeployVerbosity ?? 'v') as LogMocksDeploy,
   };
 });
@@ -243,7 +262,12 @@ task(TASK_TEST, 'Deploy mock contracts on hardhat').setAction(async ({}, hre, ru
       mocksDeployVerbosity: hre.config.cofhe.mocksDeployVerbosity,
     });
   }
-  return runSuper();
+
+  const result = await runSuper();
+  if (hre.config.cofhe.gasSummary) {
+    await printMockGasSummary(hre);
+  }
+  return result;
 });
 
 task(TASK_NODE, 'Deploy mock contracts on hardhat').setAction(async ({}, hre, runSuper) => {
@@ -277,6 +301,7 @@ export * from './utils.js';
 export * from './fund.js';
 export * from './logging.js';
 export * from './deploy.js';
+export * from './gas.js';
 
 /**
  * Runtime environment extensions for the CoFHE Hardhat plugin.
@@ -321,6 +346,26 @@ declare module 'hardhat/types/runtime' {
        */
       createClientWithBatteries: (signer?: HardhatEthersSigner) => Promise<CofheClient>;
 
+      /**
+       * Returns a transaction's gas usage excluding mock-only overhead (FHE op replication,
+       * decrypt-task storage, mock logging) - an estimate of what the transaction would cost
+       * on a real CoFHE network. Pure function of the receipt (sums the mock task manager's
+       * MockGasConsumed events from `receipt.logs`); on a real network the receipt carries no
+       * such events and the raw `gasUsed` is returned unchanged.
+       * @param {AdjustableGasReceipt} receipt - The transaction receipt (ethers or viem shape)
+       * @returns {bigint} gasUsed minus mock overhead
+       */
+      getAdjustedGasUsed: (receipt: AdjustableGasReceipt) => bigint;
+
+      /**
+       * Full gas breakdown of a transaction receipt: raw `gasUsed`, the mock-only `mockGas`,
+       * the `adjustedGasUsed` (raw minus mock), and the number of mock gas events.
+       * See getAdjustedGasUsed.
+       * @param {AdjustableGasReceipt} receipt - The transaction receipt (ethers or viem shape)
+       * @returns {AdjustedGasBreakdown} The breakdown
+       */
+      getAdjustedGasBreakdown: (receipt: AdjustableGasReceipt) => AdjustedGasBreakdown;
+
       mocks: {
         /**
          * **[MOCKS ONLY]**
@@ -393,6 +438,12 @@ declare module 'hardhat/types/runtime' {
          */
         expectPlaintext: (ctHash: bigint | string, expectedValue: bigint) => Promise<void>;
 
+        /**
+         * Get the MockCoFHEAddressBook contract (typed via typechain)
+         * @returns {Promise<MockCoFHEAddressBook>} The MockCoFHEAddressBook contract
+         */
+        getMockCoFHEAddressBook: () => Promise<MockCoFHEAddressBook>;
+
         /**
          * Get the MockTaskManager contract (typed via typechain)
          * @returns {Promise<MockTaskManager>} The MockTaskManager contract
@@ -493,6 +544,8 @@ extendEnvironment((hre) => {
       // Return client
       return client;
     },
+    getAdjustedGasUsed: mock_getAdjustedGasUsed,
+    getAdjustedGasBreakdown: mock_getAdjustedGasBreakdown,
     mocks: {
       withLogs: async (closureName: string, closure: () => Promise<void>) => {
         return mock_withLogs(hre, closureName, closure);
@@ -514,6 +567,8 @@ extendEnvironment((hre) => {
         const [signer] = await hre.ethers.getSigners();
         return mock_expectPlaintext(signer.provider, ctHash, expectedValue);
       },
+      getMockCoFHEAddressBook: async () =>
+        getFixedMockContract(hre, MockCoFHEAddressBookArtifact) as unknown as Promise<MockCoFHEAddressBook>,
       getMockTaskManager: async () =>
         getFixedMockContract(hre, MockTaskManagerArtifact) as unknown as Promise<MockTaskManager>,
       getMockACL: async () => {

package/src/utils.ts

@@ -1,4 +1,4 @@
-import { TASK_MANAGER_ADDRESS, MOCKS_ZK_VERIFIER_ADDRESS } from '@cofhe/sdk';
+import { MOCKS_TASK_MANAGER_ADDRESS, MOCKS_ZK_VERIFIER_ADDRESS } from '@cofhe/sdk';
 import { expect } from 'chai';
 import { Contract, ethers } from 'ethers';
 import { type HardhatEthersProvider } from '@nomicfoundation/hardhat-ethers/internal/hardhat-ethers-provider';
@@ -12,7 +12,11 @@ import type { HardhatRuntimeEnvironment } from 'hardhat/types';
 /// so that Hardhat can decode reverts and trace calls against these contracts.
 /// If the mock contract should be deployed to a fixed address, `hardhat_setCode` is used.
 /// Otherwise, ethers.js deploys it to a normal address.
-export const deployMockContractFromArtifact = async (hre: HardhatRuntimeEnvironment, artifact: MockArtifact) => {
+export const deployMockContractFromArtifact = async (
+  hre: HardhatRuntimeEnvironment,
+  artifact: MockArtifact,
+  constructorArgs: readonly unknown[] = []
+) => {
   const hardhatArtifact = await hre.artifacts.readArtifact(artifact.contractName);
 
   if (artifact.isFixed) {
@@ -22,7 +26,7 @@ export const deployMockContractFromArtifact = async (hre: HardhatRuntimeEnvironm
 
   const [signer] = await hre.ethers.getSigners();
   const factory = new hre.ethers.ContractFactory(hardhatArtifact.abi, hardhatArtifact.bytecode, signer);
-  const contract = await factory.deploy();
+  const contract = await factory.deploy(...constructorArgs);
   await contract.waitForDeployment();
   return contract as Contract;
 };
@@ -43,8 +47,10 @@ const mock_checkIsTestnet = async (fnName: string, provider: HardhatEthersProvid
   // Get bytecode at ZK_VERIFIER_ADDRESS
   const bytecode = await provider.getCode(MOCKS_ZK_VERIFIER_ADDRESS);
 
-  // If bytecode is empty, we are on a testnet
-  const isTestnet = bytecode.length === 0;
+  // provider.getCode() (eth_getCode) returns the string "0x" -- never a zero-length string --
+  // when there is no contract deployed at the address, so a real "no bytecode" check must
+  // compare against that, not against string length 0.
+  const isTestnet = !bytecode || bytecode.length <= 2;
 
   // Log if we are on a testnet
   if (isTestnet) {
@@ -63,7 +69,7 @@ export const mock_getPlaintext = async (
 
   // Connect to MockTaskManager
   const taskManager = new ethers.Contract(
-    TASK_MANAGER_ADDRESS,
+    MOCKS_TASK_MANAGER_ADDRESS,
     ['function mockStorage(uint256) view returns (uint256)'],
     provider
   );
@@ -83,7 +89,7 @@ export const mock_getPlaintextExists = async (
 
   // Connect to MockTaskManager
   const taskManager = new ethers.Contract(
-    TASK_MANAGER_ADDRESS,
+    MOCKS_TASK_MANAGER_ADDRESS,
     ['function inMockStorage(uint256) view returns (bool)'],
     provider
   );