@cofhe/mock-contracts 0.7.1 → 0.8.0

package/CHANGELOG.md

@@ -1,5 +1,17 @@
 # @cofhe/mock-contracts Changelog
 
+## 0.8.0
+
+### Minor Changes
+
+- 6f7a4c5: Resolve the Task Manager through the `CoFHEAddressBook` (`@fhenixprotocol/cofhe-contracts@1.0.0`). **Breaking:** `TASK_MANAGER_ADDRESS` is removed from `@cofhe/sdk`. Use `getTaskManagerAddress(publicClient)`, which reads `getTm(TASK_MANAGER_ID)` from the book at `COFHE_ADDRESS_BOOK_ADDRESS` (both exported) and caches the result per chain. If the chain has no book, or the id is unset, it throws `TASK_MANAGER_UNRESOLVED`. `verifyDecryptResult`, the ACL/ACP lookups and `useCofheEnabled` / `useCofheReadDecryptionResults` now resolve the Task Manager this way. `@cofhe/react` adds `useCofheTaskManagerAddress()`.
+
+  Mocks: a new `MockCoFHEAddressBook` is deployed at FHE.sol's `COFHE_ADDRESS_BOOK`, and `MockTaskManager` moves to `MOCKS_TASK_MANAGER_ADDRESS` (`0x…5000`) and is registered in the book. `MockACL` now stores its Task Manager (`setTaskManager` / `getTaskManagerAddress`) instead of reading a constant, so its `TASK_MANAGER_ADDRESS_()` getter is gone. All three plugins deploy and wire the book. Hardhat 2 exposes `hre.cofhe.mocks.getMockCoFHEAddressBook()`, Hardhat 3 exposes `conn.cofhe.mocks.MockCoFHEAddressBook`, and Foundry exposes `CofheTest.mockAddressBook`. Contracts must be built against `@fhenixprotocol/cofhe-contracts@1.0.0` or later (the hardhat plugin's peer range is now `>=1.0.0`). Each FHE op now makes one extra `getTm` call on the book, as it does on real networks, so expect slightly higher gas numbers.
+
+### Patch Changes
+
+- ccb7b68: Realistic gas reporting for mocks. Under Foundry, mock-only work is now excluded from gas metering by default (opt out with `mockTaskManager.setMockGasExcluded(false)`), so expect `forge snapshot` numbers to drop. Under Hardhat, mock transactions emit `MockGasConsumed` events (extra receipt logs), and the plugins add `getAdjustedGasUsed(receipt)` / `getAdjustedGasBreakdown(receipt)` plus an opt-in `cofhe.gasSummary` report.
+
 ## 0.7.1
 
 ### Patch Changes

package/README.md

@@ -10,6 +10,7 @@ A mock smart contract library for testing CoFHE (Confidential Computing Framewor
 ## Features
 
 - Mock implementations of core CoFHE contracts:
+  - MockCoFHEAddressBook (deployed at FHE.sol's `COFHE_ADDRESS_BOOK`, resolves the TaskManager)
   - MockTaskManager
   - MockThresholdNetwork
   - MockZkVerifier
@@ -137,6 +138,15 @@ When interacting with CoFHE this request is routed to the Threshold Network, whi
 
 When working with the mocks, the `cofheClient` instead queries the `MockThresholdNetwork` contract, which verifies the request `acp` and returns the decrypted result.
 
+### Gas
+
+Because the mocks replicate the off-chain FHE computation on-chain (plaintext math, mock storage writes, logging), transactions consume more gas than they would on a real CoFHE network. The mocks account for this in two ways:
+
+- **Foundry**: when gas exclusion is enabled (`setMockGasExcluded(true)` — the foundry plugin's `deployMocks()` does this automatically), the mock-only work runs between `pauseGasMetering` / `resumeGasMetering` cheatcodes, so forge's reported gas approximates real-network cost. Requires cheatcode access (`vm.allowCheatcodes(taskManagerAddress)`) and forge ≥ 1.0; without either, execution silently stays metered.
+- **Everywhere else** (Hardhat, anvil): `MockTaskManager` measures each block of mock-only work with `gasleft()` and emits a `MockGasConsumed(uint256)` event. Summing these events from a receipt and subtracting from `gasUsed` yields the adjusted cost — the hardhat plugins expose this as `getAdjustedGasUsed(receipt)` / `getAdjustedGasBreakdown(receipt)` and an opt-in `gasSummary` table.
+
+`eth_estimateGas` is never adjusted: the mock work really executes, so transactions still need the raw gas limit.
+
 ### Using Foundry
 
 Use [`@cofhe/foundry-plugin`](../foundry-plugin/README.md), which builds on these mocks. Inherit its

package/contracts/MockACL.sol

@@ -3,7 +3,6 @@ pragma solidity >=0.8.19 <0.9.0;
 
 import { Strings } from '@openzeppelin/contracts/utils/Strings.sol';
 import { MockPermissioned, ACP, SCOPE_GLOBAL, SCOPE_CONTRACT, SCOPE_HANDLES } from './Permissioned.sol';
-import { TASK_MANAGER_ADDRESS } from '@fhenixprotocol/cofhe-contracts/FHE.sol';
 
 /**
  * @title  ACL
@@ -32,6 +31,9 @@ contract MockACL is MockPermissioned {
   /// @param receiver Address the share would have been directed at.
   error NotShared(uint256 handle, address receiver);
 
+  /// @notice Returned when the TaskManager address being set is zero.
+  error InvalidTaskManagerAddress();
+
   /// @notice          Returned when a share is pending but was written by a different party.
   /// @param expected  Sharer the receiver named.
   /// @param actual    Sharer recorded in the slot.
@@ -53,6 +55,9 @@ contract MockACL is MockPermissioned {
   /// @notice Emitted when the share registry address is updated (zero = unset).
   event ShareRegistryUpdated(address oldAddress, address newAddress);
 
+  /// @notice Emitted when the TaskManager address is updated.
+  event TaskManagerUpdated(address oldAddress, address newAddress);
+
   /// @custom:storage-location erc7201:cofhe.storage.ACL
   struct ACLStorage {
     mapping(uint256 handle => bool isGlobal) globalHandles;
@@ -62,6 +67,7 @@ contract MockACL is MockPermissioned {
     /// @dev ACP infrastructure addresses served to SDKs (appended fields — do not reorder)
     address defaultRevokerContract;
     address shareRegistry;
+    address taskManager;
   }
 
   /// @notice Name of the contract.
@@ -76,9 +82,6 @@ contract MockACL is MockPermissioned {
   /// @notice Patch version of the contract.
   uint256 private constant PATCH_VERSION = 0;
 
-  /// @notice TaskManagerAddress address.
-  address public constant TASK_MANAGER_ADDRESS_ = TASK_MANAGER_ADDRESS;
-
   /// @dev keccak256(abi.encode(uint256(keccak256("cofhe.storage.ACL")) - 1)) & ~bytes32(uint256(0xff))
   bytes32 private constant ACL_SLOT =
     keccak256(abi.encode(uint256(keccak256('cofhe.storage.ACL')) - 1)) & ~bytes32(uint256(0xff));
@@ -103,7 +106,7 @@ contract MockACL is MockPermissioned {
    * @param requester     Address of the account giving the permissions.
    */
   function allow(uint256 handle, address account, address requester) public virtual {
-    if (msg.sender != TASK_MANAGER_ADDRESS_) {
+    if (msg.sender != _taskManager()) {
       revert DirectAllowForbidden(msg.sender);
     }
 
@@ -122,7 +125,7 @@ contract MockACL is MockPermissioned {
    * @param requester     Address of the account giving the permissions.
    */
   function allowGlobal(uint256 handle, address requester) public virtual {
-    if (msg.sender != TASK_MANAGER_ADDRESS_) {
+    if (msg.sender != _taskManager()) {
       revert DirectAllowForbidden(msg.sender);
     }
 
@@ -140,7 +143,7 @@ contract MockACL is MockPermissioned {
    * @param handlesList   List of handles.
    */
   function allowForDecryption(uint256[] memory handlesList, address requester) public virtual {
-    if (msg.sender != TASK_MANAGER_ADDRESS_) {
+    if (msg.sender != _taskManager()) {
       revert DirectAllowForbidden(msg.sender);
     }
 
@@ -167,11 +170,11 @@ contract MockACL is MockPermissioned {
    * @param requester     Address of the requester.
    */
   function allowTransient(uint256 handle, address account, address requester) public virtual {
-    if (msg.sender != TASK_MANAGER_ADDRESS_) {
+    if (msg.sender != _taskManager()) {
       revert DirectAllowForbidden(msg.sender);
     }
 
-    if (!isAllowed(handle, requester) && requester != TASK_MANAGER_ADDRESS_) {
+    if (!isAllowed(handle, requester) && requester != _taskManager()) {
       revert SenderNotAllowed(requester);
     }
 
@@ -217,14 +220,14 @@ contract MockACL is MockPermissioned {
    * @notice          Grants `receiver` transient access to `handle` and records `sharer` as the
    *                  party that handed it over, for the duration of this transaction.
    * @dev             The caller must be the Task Manager contract.
-   * @dev             Stricter than allowTransient: there is no TASK_MANAGER_ADDRESS bypass, since
+   * @dev             Stricter than allowTransient: there is no TaskManager bypass, since
    *                  nothing shares on the Task Manager's own behalf.
    * @param handle    Handle.
    * @param sharer    Address handing the handle over.
    * @param receiver  Address the handle is being handed to.
    */
   function shareCtHash(uint256 handle, address sharer, address receiver) public virtual {
-    if (msg.sender != TASK_MANAGER_ADDRESS_) {
+    if (msg.sender != _taskManager()) {
       revert DirectAllowForbidden(msg.sender);
     }
 
@@ -247,7 +250,7 @@ contract MockACL is MockPermissioned {
    * @param receiver        Address consuming the share.
    */
   function receiveCtHash(uint256 handle, address expectedSharer, address receiver) public virtual {
-    if (msg.sender != TASK_MANAGER_ADDRESS_) {
+    if (msg.sender != _taskManager()) {
       revert DirectAllowForbidden(msg.sender);
     }
 
@@ -270,7 +273,7 @@ contract MockACL is MockPermissioned {
    * @param delegateeContract Delegatee contract.
    */
   function delegateAccount(address delegatee, address delegateeContract) public virtual {
-    if (msg.sender != TASK_MANAGER_ADDRESS_) {
+    if (msg.sender != _taskManager()) {
       revert DirectAllowForbidden(msg.sender);
     }
     if (delegateeContract == msg.sender) {
@@ -328,7 +331,7 @@ contract MockACL is MockPermissioned {
    * @return taskManagerAddress  Address of the TaskManager.
    */
   function getTaskManagerAddress() public view virtual returns (address) {
-    return TASK_MANAGER_ADDRESS_;
+    return _taskManager();
   }
 
   /**
@@ -378,7 +381,7 @@ contract MockACL is MockPermissioned {
    *      Account Abstraction when bundling several UserOps calling the TaskManagerCoprocessor.
    */
   function cleanTransientStorage() external virtual {
-    if (msg.sender != TASK_MANAGER_ADDRESS_) {
+    if (msg.sender != _taskManager()) {
       revert DirectAllowForbidden(msg.sender);
     }
 
@@ -458,6 +461,21 @@ contract MockACL is MockPermissioned {
     $.shareRegistry = newAddress;
   }
 
+  /// @notice Sets the TaskManager allowed to drive this ACL (unrestricted in the mock;
+  ///         the production setter is onlyRole(DEFAULT_ADMIN_ROLE)).
+  function setTaskManager(address newAddress) external {
+    if (newAddress == address(0)) {
+      revert InvalidTaskManagerAddress();
+    }
+    ACLStorage storage $ = _getACLStorage();
+    emit TaskManagerUpdated($.taskManager, newAddress);
+    $.taskManager = newAddress;
+  }
+
+  function _taskManager() private view returns (address) {
+    return _getACLStorage().taskManager;
+  }
+
   /// @notice V3 (ACP) access check — the scope table.
   ///
   ///         | condition                                            | result |

package/contracts/MockCoFHE.sol

@@ -6,6 +6,11 @@ import { FHE } from '@fhenixprotocol/cofhe-contracts/FHE.sol';
 import { FunctionId, Utils } from '@fhenixprotocol/cofhe-contracts/ICofhe.sol';
 import { console } from 'hardhat/console.sol';
 
+/// @dev Foundry/hevm cheatcode address: address(uint160(uint256(keccak256('hevm cheat code')))).
+///      Only reachable when running under forge with cheatcode access granted (vm.allowCheatcodes);
+///      on Hardhat this address has no code and the shim is never enabled.
+address constant CHEATCODE_ADDRESS = 0x7109709ECfa91a80626fF3989D68f67F5b1DD12D;
+
 address constant ZK_VERIFIER_SIGNER_ADDRESS = 0x6E12D8C87503D4287c294f2Fdef96ACd9DFf6bd2;
 uint256 constant ZK_VERIFIER_SIGNER_PRIVATE_KEY = 49099792800763675079532137679706322989817545357788440619111868498148356080914;
 
@@ -33,9 +38,27 @@ abstract contract MockCoFHE {
 
   bool public logOps = true;
 
+  /// @dev When true, mock-only work (plaintext replication + logging) runs between
+  ///      pauseGasMetering/resumeGasMetering cheatcodes so it is excluded from forge's
+  ///      gas accounting, giving gas numbers closer to the real CoFHE task manager.
+  ///      Only enable under forge AFTER granting this contract cheatcode access
+  ///      (vm.allowCheatcodes(address(taskManager))); must stay false on Hardhat.
+  bool public mockGasExcluded = false;
+
   mapping(uint256 => uint256) public mockStorage;
   mapping(uint256 => bool) public inMockStorage;
 
+  /// @dev Emitted for each block of mock-only work (op replication, decrypt-task storage,
+  ///      log building) with the gas it consumed, whenever that work runs metered (i.e. not
+  ///      under forge's paused gas metering). Tooling sums these per transaction to compute
+  ///      gas usage excluding mock overhead - see the hardhat plugin's getAdjustedGasUsed.
+  event MockGasConsumed(uint256 gas);
+
+  /// @dev Gas spent emitting MockGasConsumed itself (LOG1 + 32 data bytes + bookkeeping).
+  ///      It lands after the closing gasleft() read, so it is added back as a constant to
+  ///      keep the reported figure covering the full mock-only cost.
+  uint256 internal constant MOCK_GAS_EVENT_COST = 1030;
+
   error InputNotInMockStorage(uint256 ctHash);
 
   // Used internally to check if we missed any operations in the mocks
@@ -49,6 +72,59 @@ abstract contract MockCoFHE {
     logOps = _logOps;
   }
 
+  function setMockGasExcluded(bool _mockGasExcluded) public {
+    mockGasExcluded = _mockGasExcluded;
+  }
+
+  /// @dev How a mock-only block is being accounted for.
+  ///      MEASURE           - metering is live: measure with gasleft() and emit MockGasConsumed.
+  ///      PAUSED_BY_MOCK    - we paused forge's metering and must resume it afterwards.
+  ///      ALREADY_UNMETERED - metering was already paused by the caller's own test; leave it
+  ///                          untouched (the mock work is unmetered either way) and emit nothing.
+  enum MockGasMode {
+    MEASURE,
+    PAUSED_BY_MOCK,
+    ALREADY_UNMETERED
+  }
+
+  /// @dev True while forge's gas metering is paused: gasleft() is frozen, so two consecutive
+  ///      reads are equal. While metered they differ by the cost of the second GAS opcode.
+  ///      (Foundry has no query cheatcode for this state.)
+  function _isGasMeteringPaused() internal view returns (bool) {
+    uint256 a = gasleft();
+    uint256 b = gasleft();
+    return a == b;
+  }
+
+  /// @dev Opens a mock-only block. Under forge (detected by code at the cheatcode address -
+  ///      forge plants a marker byte there, other environments have none) with exclusion
+  ///      enabled, pauses gas metering unless the caller's test already paused it - never
+  ///      steal a pause we don't own, or `_mockGasTrackEnd` would resume metering out from
+  ///      under the user's own measurement region. Everywhere else, records gasleft() so the
+  ///      block's cost can be reported via MockGasConsumed.
+  ///      Only wrap code that cannot revert - a revert inside the block would skip
+  ///      `_mockGasTrackEnd` and leak paused metering (see sendEventCreated's trampoline
+  ///      for the revert-safe variant).
+  function _mockGasTrackStart() internal returns (MockGasMode mode, uint256 startGas) {
+    if (mockGasExcluded && CHEATCODE_ADDRESS.code.length > 0) {
+      if (_isGasMeteringPaused()) return (MockGasMode.ALREADY_UNMETERED, 0);
+      (bool ok, ) = CHEATCODE_ADDRESS.call(abi.encodeWithSignature('pauseGasMetering()'));
+      if (ok) return (MockGasMode.PAUSED_BY_MOCK, 0);
+    }
+    return (MockGasMode.MEASURE, gasleft());
+  }
+
+  /// @dev Closes a mock-only block: resumes metering (if we paused it) or emits the
+  ///      measured cost. A caller-owned pause is left exactly as we found it.
+  function _mockGasTrackEnd(MockGasMode mode, uint256 startGas) internal {
+    if (mode == MockGasMode.PAUSED_BY_MOCK) {
+      (bool ok, ) = CHEATCODE_ADDRESS.call(abi.encodeWithSignature('resumeGasMetering()'));
+      ok;
+    } else if (mode == MockGasMode.MEASURE) {
+      emit MockGasConsumed(startGas - gasleft() + MOCK_GAS_EVENT_COST);
+    }
+  }
+
   // Utils
 
   function getUintTypeFromHash(uint256 hash) internal pure returns (uint8) {
@@ -224,76 +300,87 @@ abstract contract MockCoFHE {
     logAllow(operation, ctHash, account);
   }
 
+  /// @dev Same as MOCK_logAllow, but the log-string building runs excluded from forge gas
+  ///      metering / reported via MockGasConsumed, like the rest of the mock-only work.
+  ///      With logging disabled the whole thing is skipped - the residual cost is negligible.
+  ///      Safe without the revert trampoline: logAllow only formats bounded strings.
+  function MOCK_trackedLogAllow(string memory operation, uint256 ctHash, address account) internal {
+    if (!logOps) return;
+    (MockGasMode mode, uint256 startGas) = _mockGasTrackStart();
+    logAllow(operation, ctHash, account);
+    _mockGasTrackEnd(mode, startGas);
+  }
+
   // Mock functions
 
   function MOCK_verifyKeyInStorage(uint256 ctHash) internal view {
     if (!inMockStorage[ctHash]) revert InputNotInMockStorage(ctHash);
   }
 
-  function MOCK_unaryOperation(uint256 ctHash, string memory operation, uint256 input) internal {
-    if (opIs(operation, FunctionId.random)) {
+  function MOCK_unaryOperation(uint256 ctHash, FunctionId funcId, uint256 input) internal {
+    if (funcId == FunctionId.random) {
       _set(ctHash, uint256(blockhash(block.number - 1)));
-      logOperation('FHE.random', '', logCtHash(ctHash));
+      if (logOps) logOperation('FHE.random', '', logCtHash(ctHash));
       return;
     }
-    if (opIs(operation, FunctionId.cast)) {
+    if (funcId == FunctionId.cast) {
       _set(ctHash, _get(input));
-      logOperation('FHE.cast', logCtHash(input), logCtHash(ctHash));
+      if (logOps) logOperation('FHE.cast', logCtHash(input), logCtHash(ctHash));
       return;
     }
-    if (opIs(operation, FunctionId.not)) {
+    if (funcId == FunctionId.not) {
       bool inputIsTruthy = _get(input) == 1;
       _set(ctHash, !inputIsTruthy);
-      logOperation('FHE.not', logCtHash(input), logCtHash(ctHash));
+      if (logOps) logOperation('FHE.not', logCtHash(input), logCtHash(ctHash));
       return;
     }
-    if (opIs(operation, FunctionId.square)) {
+    if (funcId == FunctionId.square) {
       unchecked {
         _set(ctHash, _get(input) * _get(input));
       }
-      logOperation('FHE.square', string.concat(logCtHash(input), ' * ', logCtHash(input)), logCtHash(ctHash));
+      if (logOps) logOperation('FHE.square', string.concat(logCtHash(input), ' * ', logCtHash(input)), logCtHash(ctHash));
       return;
     }
-    revert InvalidUnaryOperation(operation);
+    revert InvalidUnaryOperation(Utils.functionIdToString(funcId));
   }
 
-  function MOCK_twoInputOperation(uint256 ctHash, string memory operation, uint256 input1, uint256 input2) internal {
-    if (opIs(operation, FunctionId.sub)) {
+  function MOCK_twoInputOperation(uint256 ctHash, FunctionId funcId, uint256 input1, uint256 input2) internal {
+    if (funcId == FunctionId.sub) {
       unchecked {
         _set(ctHash, _get(input1) - _get(input2));
       }
-      logOperation('FHE.sub', string.concat(logCtHash(input1), ' - ', logCtHash(input2)), logCtHash(ctHash));
+      if (logOps) logOperation('FHE.sub', string.concat(logCtHash(input1), ' - ', logCtHash(input2)), logCtHash(ctHash));
       return;
     }
-    if (opIs(operation, FunctionId.add)) {
+    if (funcId == FunctionId.add) {
       unchecked {
         _set(ctHash, _get(input1) + _get(input2));
       }
-      logOperation('FHE.add', string.concat(logCtHash(input1), ' + ', logCtHash(input2)), logCtHash(ctHash));
+      if (logOps) logOperation('FHE.add', string.concat(logCtHash(input1), ' + ', logCtHash(input2)), logCtHash(ctHash));
       return;
     }
-    if (opIs(operation, FunctionId.xor)) {
+    if (funcId == FunctionId.xor) {
       unchecked {
         _set(ctHash, _get(input1) ^ _get(input2));
       }
-      logOperation('FHE.xor', string.concat(logCtHash(input1), ' ^ ', logCtHash(input2)), logCtHash(ctHash));
+      if (logOps) logOperation('FHE.xor', string.concat(logCtHash(input1), ' ^ ', logCtHash(input2)), logCtHash(ctHash));
       return;
     }
-    if (opIs(operation, FunctionId.and)) {
+    if (funcId == FunctionId.and) {
       unchecked {
         _set(ctHash, _get(input1) & _get(input2));
       }
-      logOperation('FHE.and', string.concat(logCtHash(input1), ' & ', logCtHash(input2)), logCtHash(ctHash));
+      if (logOps) logOperation('FHE.and', string.concat(logCtHash(input1), ' & ', logCtHash(input2)), logCtHash(ctHash));
       return;
     }
-    if (opIs(operation, FunctionId.or)) {
+    if (funcId == FunctionId.or) {
       unchecked {
         _set(ctHash, _get(input1) | _get(input2));
       }
-      logOperation('FHE.or', string.concat(logCtHash(input1), ' | ', logCtHash(input2)), logCtHash(ctHash));
+      if (logOps) logOperation('FHE.or', string.concat(logCtHash(input1), ' | ', logCtHash(input2)), logCtHash(ctHash));
       return;
     }
-    if (opIs(operation, FunctionId.div)) {
+    if (funcId == FunctionId.div) {
       uint256 cleartext2 = _get(input2);
       if (cleartext2 == 0) {
         _set(ctHash, type(uint256).max);
@@ -302,143 +389,147 @@ abstract contract MockCoFHE {
           _set(ctHash, _get(input1) / cleartext2);
         }
       }
-      logOperation('FHE.div', string.concat(logCtHash(input1), ' / ', logCtHash(input2)), logCtHash(ctHash));
+      if (logOps) logOperation('FHE.div', string.concat(logCtHash(input1), ' / ', logCtHash(input2)), logCtHash(ctHash));
       return;
     }
-    if (opIs(operation, FunctionId.rem)) {
+    if (funcId == FunctionId.rem) {
       unchecked {
         _set(ctHash, _get(input1) % _get(input2));
       }
-      logOperation('FHE.rem', string.concat(logCtHash(input1), ' % ', logCtHash(input2)), logCtHash(ctHash));
+      if (logOps) logOperation('FHE.rem', string.concat(logCtHash(input1), ' % ', logCtHash(input2)), logCtHash(ctHash));
       return;
     }
-    if (opIs(operation, FunctionId.mul)) {
+    if (funcId == FunctionId.mul) {
       unchecked {
         _set(ctHash, _get(input1) * _get(input2));
       }
-      logOperation('FHE.mul', string.concat(logCtHash(input1), ' * ', logCtHash(input2)), logCtHash(ctHash));
+      if (logOps) logOperation('FHE.mul', string.concat(logCtHash(input1), ' * ', logCtHash(input2)), logCtHash(ctHash));
       return;
     }
-    if (opIs(operation, FunctionId.shl)) {
+    if (funcId == FunctionId.shl) {
       unchecked {
         _set(ctHash, _get(input1) << _get(input2));
       }
-      logOperation('FHE.shl', string.concat(logCtHash(input1), ' << ', logCtHash(input2)), logCtHash(ctHash));
+      if (logOps) logOperation('FHE.shl', string.concat(logCtHash(input1), ' << ', logCtHash(input2)), logCtHash(ctHash));
       return;
     }
-    if (opIs(operation, FunctionId.shr)) {
+    if (funcId == FunctionId.shr) {
       unchecked {
         _set(ctHash, _get(input1) >> _get(input2));
       }
-      logOperation('FHE.shr', string.concat(logCtHash(input1), ' >> ', logCtHash(input2)), logCtHash(ctHash));
+      if (logOps) logOperation('FHE.shr', string.concat(logCtHash(input1), ' >> ', logCtHash(input2)), logCtHash(ctHash));
       return;
     }
-    if (opIs(operation, FunctionId.gte)) {
+    if (funcId == FunctionId.gte) {
       _set(ctHash, _get(input1) >= _get(input2));
-      logOperation('FHE.gte', string.concat(logCtHash(input1), ' >= ', logCtHash(input2)), logCtHash(ctHash));
+      if (logOps) logOperation('FHE.gte', string.concat(logCtHash(input1), ' >= ', logCtHash(input2)), logCtHash(ctHash));
       return;
     }
-    if (opIs(operation, FunctionId.lte)) {
+    if (funcId == FunctionId.lte) {
       _set(ctHash, _get(input1) <= _get(input2));
-      logOperation('FHE.lte', string.concat(logCtHash(input1), ' <= ', logCtHash(input2)), logCtHash(ctHash));
+      if (logOps) logOperation('FHE.lte', string.concat(logCtHash(input1), ' <= ', logCtHash(input2)), logCtHash(ctHash));
       return;
     }
-    if (opIs(operation, FunctionId.lt)) {
+    if (funcId == FunctionId.lt) {
       _set(ctHash, _get(input1) < _get(input2));
-      logOperation('FHE.lt', string.concat(logCtHash(input1), ' < ', logCtHash(input2)), logCtHash(ctHash));
+      if (logOps) logOperation('FHE.lt', string.concat(logCtHash(input1), ' < ', logCtHash(input2)), logCtHash(ctHash));
       return;
     }
-    if (opIs(operation, FunctionId.gt)) {
+    if (funcId == FunctionId.gt) {
       _set(ctHash, _get(input1) > _get(input2));
-      logOperation('FHE.gt', string.concat(logCtHash(input1), ' > ', logCtHash(input2)), logCtHash(ctHash));
+      if (logOps) logOperation('FHE.gt', string.concat(logCtHash(input1), ' > ', logCtHash(input2)), logCtHash(ctHash));
       return;
     }
-    if (opIs(operation, FunctionId.min)) {
+    if (funcId == FunctionId.min) {
       uint256 min;
       unchecked {
         min = _get(input1) < _get(input2) ? _get(input1) : _get(input2);
       }
       _set(ctHash, min);
 
-      logOperation(
-        'FHE.min',
-        string.concat('min(', logCtHash(input1), ', ', logCtHash(input2), ')'),
-        logCtHash(ctHash)
-      );
+      if (logOps)
+        logOperation(
+          'FHE.min',
+          string.concat('min(', logCtHash(input1), ', ', logCtHash(input2), ')'),
+          logCtHash(ctHash)
+        );
       return;
     }
-    if (opIs(operation, FunctionId.max)) {
+    if (funcId == FunctionId.max) {
       uint256 max;
       unchecked {
         max = _get(input1) > _get(input2) ? _get(input1) : _get(input2);
       }
       _set(ctHash, max);
 
-      logOperation(
-        'FHE.max',
-        string.concat('max(', logCtHash(input1), ', ', logCtHash(input2), ')'),
-        logCtHash(ctHash)
-      );
+      if (logOps)
+        logOperation(
+          'FHE.max',
+          string.concat('max(', logCtHash(input1), ', ', logCtHash(input2), ')'),
+          logCtHash(ctHash)
+        );
       return;
     }
-    if (opIs(operation, FunctionId.eq)) {
+    if (funcId == FunctionId.eq) {
       _set(ctHash, _get(input1) == _get(input2));
 
-      logOperation('FHE.eq', string.concat(logCtHash(input1), ' == ', logCtHash(input2)), logCtHash(ctHash));
+      if (logOps) logOperation('FHE.eq', string.concat(logCtHash(input1), ' == ', logCtHash(input2)), logCtHash(ctHash));
       return;
     }
-    if (opIs(operation, FunctionId.ne)) {
+    if (funcId == FunctionId.ne) {
       _set(ctHash, _get(input1) != _get(input2));
 
-      logOperation('FHE.ne', string.concat(logCtHash(input1), ' != ', logCtHash(input2)), logCtHash(ctHash));
+      if (logOps) logOperation('FHE.ne', string.concat(logCtHash(input1), ' != ', logCtHash(input2)), logCtHash(ctHash));
       return;
     }
-    if (opIs(operation, FunctionId.rol)) {
+    if (funcId == FunctionId.rol) {
       unchecked {
         _set(ctHash, _get(input1) << _get(input2));
       }
 
-      logOperation('FHE.rol', string.concat(logCtHash(input1), ' << ', logCtHash(input2)), logCtHash(ctHash));
+      if (logOps) logOperation('FHE.rol', string.concat(logCtHash(input1), ' << ', logCtHash(input2)), logCtHash(ctHash));
       return;
     }
-    if (opIs(operation, FunctionId.ror)) {
+    if (funcId == FunctionId.ror) {
       unchecked {
         _set(ctHash, _get(input1) >> _get(input2));
       }
 
-      logOperation('FHE.ror', string.concat(logCtHash(input1), ' >> ', logCtHash(input2)), logCtHash(ctHash));
+      if (logOps) logOperation('FHE.ror', string.concat(logCtHash(input1), ' >> ', logCtHash(input2)), logCtHash(ctHash));
       return;
     }
-    revert InvalidTwoInputOperation(operation);
+    revert InvalidTwoInputOperation(Utils.functionIdToString(funcId));
   }
 
   function MOCK_threeInputOperation(
     uint256 ctHash,
-    string memory operation,
+    FunctionId funcId,
     uint256 input1,
     uint256 input2,
     uint256 input3
   ) internal {
-    if (opIs(operation, FunctionId.trivialEncrypt)) {
+    if (funcId == FunctionId.trivialEncrypt) {
       _set(ctHash, input1);
 
-      logOperation(
-        string.concat('FHE.asE', removeFirstLetter(getUtypeStringFromHash(ctHash))),
-        string.concat(removeFirstLetter(getUtypeStringFromHash(ctHash)), '(', Strings.toString(input1), ')'),
-        logCtHash(ctHash)
-      );
+      if (logOps)
+        logOperation(
+          string.concat('FHE.asE', removeFirstLetter(getUtypeStringFromHash(ctHash))),
+          string.concat(removeFirstLetter(getUtypeStringFromHash(ctHash)), '(', Strings.toString(input1), ')'),
+          logCtHash(ctHash)
+        );
       return;
     }
-    if (opIs(operation, FunctionId.select)) {
+    if (funcId == FunctionId.select) {
       _set(ctHash, _get(input1) == 1 ? _get(input2) : _get(input3));
 
-      logOperation(
-        'FHE.select',
-        string.concat(logCtHash(input1), ' ? ', logCtHash(input2), ' : ', logCtHash(input3)),
-        logCtHash(ctHash)
-      );
+      if (logOps)
+        logOperation(
+          'FHE.select',
+          string.concat(logCtHash(input1), ' ? ', logCtHash(input2), ' : ', logCtHash(input3)),
+          logCtHash(ctHash)
+        );
       return;
     }
-    revert InvalidThreeInputOperation(operation);
+    revert InvalidThreeInputOperation(Utils.functionIdToString(funcId));
   }
 }