Merge pull request #242 from zama-ai/fhelib-mock

Mocked version of TFHE.sol for mocked tests and coverage

codegen/main.ts

@@ -9,10 +9,13 @@ function generateAllFiles() {
 
   const network = Network[(process.env.TARGET_NETWORK as keyof typeof Network) || 'Evmos'];
   const context = networkCodegenContext(network);
-  const [tfheSolSource, overloads] = t.tfheSol(context, operators, SUPPORTED_BITS);
+  const [tfheSolSource, overloads] = t.tfheSol(context, operators, SUPPORTED_BITS, false);
   const ovShards = testgen.splitOverloadsToShards(overloads);
   writeFileSync('lib/Impl.sol', t.implSol(context, operators));
   writeFileSync('lib/TFHE.sol', tfheSolSource);
+  writeFileSync('lib_mock/Impl.sol', t.implSolMock(context, operators));
+  const [tfheSolSourceMock, _] = t.tfheSol(context, operators, SUPPORTED_BITS, true);
+  writeFileSync('lib_mock/TFHE.sol', tfheSolSourceMock);
   mkdirSync('examples/tests', { recursive: true });
   ovShards.forEach((os) => {
     writeFileSync(`examples/tests/TFHETestSuite${os.shardNumber}.sol`, testgen.generateSmartContract(os));

codegen/templates.ts

@@ -140,6 +140,7 @@ export function tfheSol(
   ctx: CodegenContext,
   operators: Operator[],
   supportedBits: number[],
+  mocked: boolean,
 ): [string, OverloadSignature[]] {
   const signatures: OverloadSignature[] = [];
   const res: string[] = [];
@@ -159,14 +160,25 @@ library TFHE {
     res.push(`     euint${b} constant NIL${b} = euint${b}.wrap(0);
 `);
   });
-  supportedBits.forEach((b) => {
-    res.push(`
-    // Return true if the enrypted integer is initialized and false otherwise.
-    function isInitialized(euint${b} v) internal pure returns (bool) {
-        return euint${b}.unwrap(v) != 0;
-    }
-`);
-  });
+  if (mocked) {
+    supportedBits.forEach((b) => {
+      res.push(`
+      // Return true if the enrypted integer is initialized and false otherwise.
+      function isInitialized(euint${b} v) internal pure returns (bool) {
+          return true;
+      }
+    `);
+    });
+  } else {
+    supportedBits.forEach((b) => {
+      res.push(`
+      // Return true if the enrypted integer is initialized and false otherwise.
+      function isInitialized(euint${b} v) internal pure returns (bool) {
+          return euint${b}.unwrap(v) != 0;
+      }
+    `);
+    });
+  }
 
   supportedBits.forEach((lhsBits) => {
     supportedBits.forEach((rhsBits) => {
@@ -783,3 +795,175 @@ function implCustomMethods(ctx: CodegenContext): string {
     }
     `;
 }
+
+export function implSolMock(ctx: CodegenContext, operators: Operator[]): string {
+  const res: string[] = [];
+
+  res.push(`
+// SPDX-License-Identifier: BSD-3-Clause-Clear
+
+pragma solidity 0.8.19;
+
+library Impl {
+  function add(uint256 lhs, uint256 rhs, bool scalar) internal pure returns (uint256 result) {
+    unchecked {
+        result = lhs + rhs;
+    }
+  }
+
+  function sub(uint256 lhs, uint256 rhs, bool scalar) internal pure returns (uint256 result) {
+      unchecked {
+          result = lhs - rhs;
+      }
+  }
+
+  function mul(uint256 lhs, uint256 rhs, bool scalar) internal pure returns (uint256 result) {
+      unchecked {
+          result = lhs * rhs;
+      }
+  }
+
+  function div(uint256 lhs, uint256 rhs) internal pure returns (uint256 result) {
+      result = lhs / rhs; // unchecked does not change behaviour even when dividing by 0
+  }
+
+  function rem(uint256 lhs, uint256 rhs) internal pure returns (uint256 result) {
+      result = lhs % rhs;
+  }
+
+  function and(uint256 lhs, uint256 rhs) internal pure returns (uint256 result) {
+      result = lhs & rhs;
+  }
+
+  function or(uint256 lhs, uint256 rhs) internal pure returns (uint256 result) {
+      result = lhs | rhs;
+  }
+
+  function xor(uint256 lhs, uint256 rhs) internal pure returns (uint256 result) {
+      result = lhs ^ rhs;
+  }
+
+  function shl(uint256 lhs, uint256 rhs, bool scalar) internal pure returns (uint256 result) {
+      result = lhs << rhs;
+  }
+
+  function shr(uint256 lhs, uint256 rhs, bool scalar) internal pure returns (uint256 result) {
+      result = lhs >> rhs;
+  }
+
+  function eq(uint256 lhs, uint256 rhs, bool scalar) internal pure returns (uint256 result) {
+      result = (lhs == rhs) ? 1 : 0;
+  }
+
+  function ne(uint256 lhs, uint256 rhs, bool scalar) internal pure returns (uint256 result) {
+      result = (lhs != rhs) ? 1 : 0;
+  }
+
+  function ge(uint256 lhs, uint256 rhs, bool scalar) internal pure returns (uint256 result) {
+      result = (lhs >= rhs) ? 1 : 0;
+  }
+
+  function gt(uint256 lhs, uint256 rhs, bool scalar) internal pure returns (uint256 result) {
+      result = (lhs > rhs) ? 1 : 0;
+  }
+
+  function le(uint256 lhs, uint256 rhs, bool scalar) internal pure returns (uint256 result) {
+      result = (lhs <= rhs) ? 1 : 0;
+  }
+
+  function lt(uint256 lhs, uint256 rhs, bool scalar) internal pure returns (uint256 result) {
+      result = (lhs < rhs) ? 1 : 0;
+  }
+
+  function min(uint256 lhs, uint256 rhs, bool scalar) internal pure returns (uint256 result) {
+      result = (lhs < rhs) ? lhs : rhs;
+  }
+
+  function max(uint256 lhs, uint256 rhs, bool scalar) internal pure returns (uint256 result) {
+      result = (lhs > rhs) ? lhs : rhs;
+  }
+
+  function neg(uint256 ct) internal pure returns (uint256 result) {
+      uint256 y;
+      assembly {
+          y := not(ct)
+      }
+      unchecked {
+          return y + 1;
+      }
+  }
+
+  function not(uint256 ct) internal pure returns (uint256 result) {
+      uint256 y;
+      assembly {
+          y := not(ct)
+      }
+      return y;
+  }
+
+  function cmux(uint256 control, uint256 ifTrue, uint256 ifFalse) internal pure returns (uint256 result) {
+      result = (control == 1) ? ifTrue : ifFalse;
+  }
+
+  function optReq(uint256 ciphertext) internal view {
+      require(ciphertext == 1, "transaction execution reverted");
+  }
+
+  function reencrypt(uint256 ciphertext, bytes32 publicKey) internal view returns (bytes memory reencrypted) {
+      reencrypted = new bytes(32);
+      assembly {
+          mstore(add(reencrypted, 32), ciphertext)
+      }
+      return reencrypted;
+  }
+
+  function fhePubKey() internal view returns (bytes memory key) {
+      key = hex"0123456789ABCDEF";
+  }
+
+  function verify(bytes memory _ciphertextBytes, uint8 _toType) internal pure returns (uint256 result) {
+      uint256 x;
+      assembly {
+          switch gt(mload(_ciphertextBytes), 31)
+          case 1 {
+              x := mload(add(_ciphertextBytes, add(32, sub(mload(_ciphertextBytes), 32))))
+          }
+          default {
+              x := mload(add(_ciphertextBytes, 32))
+          }
+      }
+      if (_ciphertextBytes.length < 32) {
+          x = x >> ((32 - _ciphertextBytes.length) * 8);
+      }
+      return x;
+  }
+
+  function cast(uint256 ciphertext, uint8 toType) internal pure returns (uint256 result) {
+      if (toType == 0) {
+          result = uint256(uint8(ciphertext));
+      }
+      if (toType == 1) {
+          result = uint256(uint16(ciphertext));
+      }
+      if (toType == 2) {
+          result = uint256(uint32(ciphertext));
+      }
+  }
+
+  function trivialEncrypt(uint256 value, uint8 toType) internal pure returns (uint256 result) {
+      result = value;
+  }
+
+  function decrypt(uint256 ciphertext) internal view returns (uint256 result) {
+      result = ciphertext;
+  }
+
+  function rand(uint8 randType) internal view returns (uint256 result) {
+      result = uint256(keccak256(abi.encodePacked(block.number, gasleft(), msg.sender))); // assuming no duplicated tx by same sender in a single block
+  }
+`);
+
+  res.push('}\n');
+
+  return res.join('');
+}

docs/howto/write_contract/hardhat.md

@@ -5,3 +5,36 @@
 The best way to start writing smart contracts with fhEVM is to use our [Hardhat template](https://github.com/zama-ai/fhevm-hardhat-template).
 
 It allows you to start a fhEVM docker image and run your smart contract on it. Read the [README](https://github.com/zama-ai/fhevm-hardhat-template/blob/main/README.md) for more information.
+
+For faster testing iterations, instead of launching all the tests on the local fhEVM node via `pnpm test`or `npx hardhat test` which could last several minutes, you could use instead a mocked version of the `TFHE.sol` library.
+The same tests should pass, as is (almost always), without any modification, neither the javascript files neither the solidity files need to be changed between the mocked and the real version.
+
+To run the mocked tests use either:
+
+```
+pnpm test:mock
+```
+
+Or equivalently:
+
+```
+HARDHAT_NETWORK=hardhat npx hardhat test --network hardhat
+```
+
+In mocked mode, all tests should pass in few seconds instead of few minutes, allowing a better developer experience.
+
+Furthermore, getting the coverage of tests is only possible in mocked mode. Just use the following command:
+
+```
+pnpm coverage:mock
+```
+
+Or equivalently:
+
+```
+HARDHAT_NETWORK=hardhat npx hardhat coverage-mock --network hardhat
+```
+
+then open the file `coverage/index.html`. This will allow increased security by pointing out missing branches not covered yet by the current test suite.
+
+Notice that, due to intrinsic limitations of the original EVM, the mocked version differ in few corner cases from the real fhevm, the most important change is the TFHE.isInitialized method which will always return true in the mocked version. This means that before deploying to production, developpers still need to run the tests with the original fhevm node, as a final check in non-mocked mode, with `pnpm test` or `npx hardhat test`.

docs/howto/write_contract/others.md

@@ -7,3 +7,5 @@ Our library compiles seamlessly with the traditional Solidity compiler and is ge
 ## Foundry
 
 The fhEVM does not work with Foundry as Foundry employs its own EVM, preventing us from incorporating a mock for our precompiled contract. An [ongoing discussion](https://github.com/foundry-rs/foundry/issues/5576) is exploring the possibility of incorporating a plugin system for precompiles, which could potentially pave the way for the utilization of Foundry at a later stage.
+
+However, you could still use Foundry with the mocked version of the fhEVM, but please be aware that this approach is **NOT** recommended, since the mocked version is not fully equivalent to the real fhEVM node's implementation (see explanation in [hardhat](hardhat.md)). In order to do this, you will need to rename your `TFHE.sol` imports from `../lib/TFHE.sol` to `../mocks/TFHE.sol` in your solidity sources and test files.

hardhat.config.ts

@@ -1,16 +1,59 @@
 import '@nomicfoundation/hardhat-toolbox';
 import { config as dotenvConfig } from 'dotenv';
+import * as fs from 'fs';
 import 'hardhat-deploy';
+import 'hardhat-preprocessor';
+import { TASK_PREPROCESS } from 'hardhat-preprocessor';
 import type { HardhatUserConfig } from 'hardhat/config';
+import { task } from 'hardhat/config';
 import type { NetworkUserConfig } from 'hardhat/types';
 import { resolve } from 'path';
+import * as path from 'path';
 
 import './tasks/accounts';
 import './tasks/getEthereumAddress';
 import './tasks/mint';
 import './tasks/taskDeploy';
 import './tasks/taskIdentity';
 
+// Function to recursively get all .sol files in a folder
+function getAllSolidityFiles(dir: string, fileList: string[] = []): string[] {
+  fs.readdirSync(dir).forEach((file) => {
+    const filePath = path.join(dir, file);
+    if (fs.statSync(filePath).isDirectory()) {
+      getAllSolidityFiles(filePath, fileList);
+    } else if (filePath.endsWith('.sol')) {
+      fileList.push(filePath);
+    }
+  });
+  return fileList;
+}
+
+task('coverage-mock', 'Run coverage after running pre-process task').setAction(async function (args, env) {
+  // Get all .sol files in the examples/ folder
+  const examplesPath = path.join(env.config.paths.root, 'examples/');
+  const solidityFiles = getAllSolidityFiles(examplesPath);
+
+  // Backup original files
+  const originalContents: Record<string, string> = {};
+  solidityFiles.forEach((filePath) => {
+    originalContents[filePath] = fs.readFileSync(filePath, { encoding: 'utf8' });
+  });
+
+  try {
+    // Run pre-process task
+    await env.run(TASK_PREPROCESS);
+
+    // Run coverage task
+    await env.run('coverage');
+  } finally {
+    // Restore original files
+    for (const filePath in originalContents) {
+      fs.writeFileSync(filePath, originalContents[filePath], { encoding: 'utf8' });
+    }
+  }
+});
+
 const dotenvConfigPath: string = process.env.DOTENV_CONFIG_PATH || './.env';
 dotenvConfig({ path: resolve(__dirname, dotenvConfigPath) });
 
@@ -20,6 +63,16 @@ if (!mnemonic) {
   throw new Error('Please set your MNEMONIC in a .env file');
 }
 
+const network = process.env.HARDHAT_NETWORK;
+
+function getRemappings() {
+  return fs
+    .readFileSync('remappings.txt', 'utf8')
+    .split('\n')
+    .filter(Boolean) // remove empty lines
+    .map((line: string) => line.trim().split('='));
+}
+
 const chainIds = {
   zama: 8009,
   local: 9000,
@@ -51,6 +104,25 @@ function getChainConfig(chain: keyof typeof chainIds): NetworkUserConfig {
 }
 
 const config: HardhatUserConfig = {
+  preprocess: {
+    eachLine: (hre) => ({
+      transform: (line: string) => {
+        if (network === 'hardhat') {
+          // checks if HARDHAT_NETWORK env variable is set to "hardhat" to use the remapping for the mocked version of TFHE.sol
+          if (line.match(/".*.sol";$/)) {
+            // match all lines with `"<any-import-path>.sol";`
+            for (const [from, to] of getRemappings()) {
+              if (line.includes(from)) {
+                line = line.replace(from, to);
+                break;
+              }
+            }
+          }
+        }
+        return line;
+      },
+    }),
+  },
   defaultNetwork: 'local',
   namedAccounts: {
     deployer: 0,