@cofhe/mock-contracts 0.7.1 → 1.0.0-alpha.0

package/contracts/MockACL.sol

@@ -2,8 +2,8 @@
 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';
+import { ACP, SCOPE_GLOBAL, SCOPE_CONTRACT, SCOPE_HANDLES } from './Permissioned.sol';
+import { MockPermissioned } from './MockPermissioned.sol';
 
 /**
  * @title  ACL
@@ -32,6 +32,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 +56,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 +68,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 +83,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 +107,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 +126,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 +144,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 +171,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 +221,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 +251,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 +274,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 +332,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 +382,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 +462,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,92 @@ 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)) {
-      bool inputIsTruthy = _get(input) == 1;
-      _set(ctHash, !inputIsTruthy);
-      logOperation('FHE.not', logCtHash(input), logCtHash(ctHash));
+    if (funcId == FunctionId.not) {
+      if (getIsBoolTypeFromHash(ctHash)) {
+        bool inputIsTruthy = _get(input) == 1;
+        _set(ctHash, !inputIsTruthy);
+      } else {
+        // bitwise NOT, masked to the type width by _set
+        _set(ctHash, ~_get(input));
+      }
+      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 +394,166 @@ 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)) {
-      unchecked {
-        _set(ctHash, _get(input1) % _get(input2));
+    if (funcId == FunctionId.rem) {
+      // Like TFHE-rs, the remainder by zero returns the first input unchanged
+      uint256 cleartext2 = _get(input2);
+      if (cleartext2 == 0) {
+        _set(ctHash, _get(input1));
+      } else {
+        unchecked {
+          _set(ctHash, _get(input1) % cleartext2);
+        }
       }
-      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)) {
-      unchecked {
-        _set(ctHash, _get(input1) << _get(input2));
-      }
-
-      logOperation('FHE.rol', string.concat(logCtHash(input1), ' << ', logCtHash(input2)), logCtHash(ctHash));
+    if (funcId == FunctionId.rol) {
+      // rotation within the type width, amount modulo the width (TFHE-rs semantics); _set masks the result
+      uint256 bits = getUtypeBits(ctHash);
+      uint256 shift = _get(input2) % bits;
+      uint256 value = _get(input1);
+      _set(ctHash, shift == 0 ? value : (value << shift) | (value >> (bits - shift)));
+
+      if (logOps)
+        logOperation(
+          'FHE.rol',
+          string.concat('rol(', logCtHash(input1), ', ', logCtHash(input2), ')'),
+          logCtHash(ctHash)
+        );
       return;
     }
-    if (opIs(operation, FunctionId.ror)) {
-      unchecked {
-        _set(ctHash, _get(input1) >> _get(input2));
-      }
-
-      logOperation('FHE.ror', string.concat(logCtHash(input1), ' >> ', logCtHash(input2)), logCtHash(ctHash));
+    if (funcId == FunctionId.ror) {
+      uint256 bits = getUtypeBits(ctHash);
+      uint256 shift = _get(input2) % bits;
+      uint256 value = _get(input1);
+      _set(ctHash, shift == 0 ? value : (value >> shift) | (value << (bits - shift)));
+
+      if (logOps)
+        logOperation(
+          'FHE.ror',
+          string.concat('ror(', 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));
   }
 }