Merge pull request #116 from zama-ai/feature/test-codegen

Add binary operator tests in solidity

README.md

@@ -49,14 +49,30 @@ npm install
 
 Note: Solidity files are formatted with prettier.
 
-### Generate TFHE lib
+### Generate TFHE lib and tests
 
 ```
 npm run codegen
 ```
 
 WARNING: Use this command to generate Solidity code and prettier result automatically!
 
+Files that are generated now (can be seen inside `codegen/main.ts`)
+
+```
+lib/Common.sol
+lib/Precompiles.sol
+lib/Impl.sol
+lib/TFHE.sol
+contracts/tests/TFHETestSuiteX.sol
+test/tfheOperations/tfheOperations.ts
+```
+
+### Adding new operators
+
+Operators can be defined as data inside `codegen/common.ts` file and code automatically generates solidity overloads.
+Test for overloads must be added (or the build doesn't pass) inside `codegen/overloadsTests.ts` file.
+
 ### Test
 
 ```

codegen/common.ts

@@ -1,13 +1,18 @@
-import assert from 'assert';
+import { strict as assert } from 'node:assert';
 
 export type Operator = {
   name: string;
+  // express left scalar operation as different operation with arguments swapped
   leftScalarInvertOp?: string;
   precompileName: string;
   hasScalar: boolean;
   hasEncrypted: boolean;
   arguments: OperatorArguments;
   returnType: ReturnType;
+  // if true do trivial encryption for left scalar operand, this is workaround until tfhe-rs supports left scalar operands
+  leftScalarEncrypt?: boolean;
+  // disable left scalar operator
+  leftScalarDisable?: boolean;
 };
 
 export type Precompile = {
@@ -73,6 +78,7 @@ export const ALL_OPERATORS: Operator[] = [
     hasEncrypted: true,
     arguments: OperatorArguments.Binary,
     returnType: ReturnType.Uint,
+    leftScalarEncrypt: true,
   },
   {
     name: 'mul',
@@ -89,6 +95,7 @@ export const ALL_OPERATORS: Operator[] = [
     hasEncrypted: false,
     arguments: OperatorArguments.Binary,
     returnType: ReturnType.Uint,
+    leftScalarDisable: true,
   },
   {
     name: 'and',
@@ -121,6 +128,7 @@ export const ALL_OPERATORS: Operator[] = [
     hasEncrypted: true,
     arguments: OperatorArguments.Binary,
     returnType: ReturnType.Uint,
+    leftScalarEncrypt: true,
   },
   {
     name: 'shr',
@@ -129,6 +137,7 @@ export const ALL_OPERATORS: Operator[] = [
     hasEncrypted: true,
     arguments: OperatorArguments.Binary,
     returnType: ReturnType.Uint,
+    leftScalarEncrypt: true,
   },
   {
     name: 'eq',

codegen/main.ts

@@ -1,15 +1,23 @@
-import { writeFileSync } from 'fs';
+import { mkdirSync, writeFileSync } from 'fs';
 
 import { ALL_OPERATORS, ALL_PRECOMPILES, SUPPORTED_BITS, checks } from './common';
 import * as t from './templates';
+import * as testgen from './testgen';
 
 function generateAllFiles() {
   const operators = checks(ALL_OPERATORS);
 
-  writeFileSync(`lib/Common.sol`, t.commonSolHeader());
-  writeFileSync(`lib/Precompiles.sol`, t.precompiles(ALL_PRECOMPILES));
-  writeFileSync(`lib/Impl.sol`, t.implSol(operators));
-  writeFileSync(`lib/TFHE.sol`, t.tfheSol(operators, SUPPORTED_BITS));
+  const [tfheSolSource, overloads] = t.tfheSol(operators, SUPPORTED_BITS);
+  const ovShards = testgen.splitOverloadsToShards(overloads);
+  writeFileSync('lib/Common.sol', t.commonSolHeader());
+  writeFileSync('lib/Precompiles.sol', t.precompiles(ALL_PRECOMPILES));
+  writeFileSync('lib/Impl.sol', t.implSol(operators));
+  writeFileSync('lib/TFHE.sol', tfheSolSource);
+  mkdirSync('contracts/tests', { recursive: true });
+  ovShards.forEach((os) => {
+    writeFileSync(`contracts/tests/TFHETestSuite${os.shardNumber}.sol`, testgen.generateSmartContract(os));
+  });
+  writeFileSync('test/tfheOperations/tfheOperations.ts', testgen.generateTestCode(ovShards));
 }
 
 generateAllFiles();