fix: allow for ebytes256 + reencrypt tests

chore: removed console logs
zama-ai · Jul 9, 2024 · a1c02aa · a1c02aa
1 parent 29c05c3
commit a1c02aa
Show file tree

Hide file tree

Showing 8 changed files with 271 additions and 101 deletions.
diff --git a/codegen/main.ts b/codegen/main.ts
@@ -15,9 +15,6 @@ function generateAllFiles() {
   writeFileSync('lib/TFHE.sol', tfheSolSource);
   writeFileSync('lib/FhevmLib.sol', t.fhevmLibSol(operators));
   writeFileSync('lib/TFHEExecutor.sol', t.tfheExecutorSol(context, operators));
-  writeFileSync('mocks/Impl.sol', t.implSolMock(context, operators));
-  const [tfheSolSourceMock, _] = t.tfheSol(context, operators, SUPPORTED_BITS, true);
-  writeFileSync('mocks/TFHE.sol', tfheSolSourceMock);
   // mkdirSync('examples/tests', { recursive: true });
   // ovShards.forEach((os) => {
   //   writeFileSync(`examples/tests/TFHETestSuite${os.shardNumber}.sol`, testgen.generateSmartContract(os));

diff --git a/codegen/templates.ts b/codegen/templates.ts
@@ -876,6 +876,10 @@ function tfheAclMethods(supportedBits: number[]): string {
       return Impl.isAllowed(eaddress.unwrap(value), account);
     }
 
+    function isAllowed(ebytes256 value, address account) internal view returns (bool) {
+      return Impl.isAllowed(ebytes256.unwrap(value), account);
+    }
+
     function isSenderAllowed(ebool value) internal view returns (bool) {
       return Impl.isAllowed(ebool.unwrap(value), msg.sender);
     }
@@ -897,6 +901,10 @@ function tfheAclMethods(supportedBits: number[]): string {
     function isSenderAllowed(eaddress value) internal view returns(bool) {
       return Impl.isAllowed(eaddress.unwrap(value), msg.sender);
     }
+
+    function isSenderAllowed(ebytes256 value) internal view returns(bool) {
+      return Impl.isAllowed(ebytes256.unwrap(value), msg.sender);
+    }
     `,
   );
 
@@ -923,6 +931,10 @@ function tfheAclMethods(supportedBits: number[]): string {
     function allow(eaddress value, address account) internal {
       Impl.allow(eaddress.unwrap(value), account);
     }
+
+    function allow(ebytes256 value, address account) internal {
+      Impl.allow(ebytes256.unwrap(value), account);
+    }
     `,
   );
 
@@ -949,6 +961,10 @@ function tfheAclMethods(supportedBits: number[]): string {
     function allowTransient(eaddress value, address account) internal {
       Impl.allowTransient(eaddress.unwrap(value), account);
     }
+
+    function allowTransient(ebytes256 value, address account) internal {
+      Impl.allowTransient(ebytes256.unwrap(value), account);
+    }
     `,
   );
 

diff --git a/examples/Reencrypt.sol b/examples/Reencrypt.sol
@@ -0,0 +1,52 @@
+// SPDX-License-Identifier: BSD-3-Clause-Clear
+
+pragma solidity ^0.8.24;
+
+import "../lib/TFHE.sol";
+
+contract Reencrypt {
+    ebool public xBool;
+    euint4 public xUint4;
+    euint8 public xUint8;
+    euint16 public xUint16;
+    euint32 public xUint32;
+    euint64 public xUint64;
+    eaddress public xAddress;
+    ebytes256 public yBytes256;
+
+    constructor() {
+        xBool = TFHE.asEbool(true);
+        TFHE.allow(xBool, address(this));
+        TFHE.allow(xBool, msg.sender);
+
+        xUint4 = TFHE.asEuint4(4);
+        TFHE.allow(xUint4, address(this));
+        TFHE.allow(xUint4, msg.sender);
+
+        xUint8 = TFHE.asEuint8(42);
+        TFHE.allow(xUint8, address(this));
+        TFHE.allow(xUint8, msg.sender);
+
+        xUint16 = TFHE.asEuint16(16);
+        TFHE.allow(xUint16, address(this));
+        TFHE.allow(xUint16, msg.sender);
+
+        xUint32 = TFHE.asEuint32(32);
+        TFHE.allow(xUint32, address(this));
+        TFHE.allow(xUint32, msg.sender);
+
+        xUint64 = TFHE.asEuint64(18446744073709551600);
+        TFHE.allow(xUint64, address(this));
+        TFHE.allow(xUint64, msg.sender);
+
+        xAddress = TFHE.asEaddress(0x8ba1f109551bD432803012645Ac136ddd64DBA72);
+        TFHE.allow(xAddress, address(this));
+        TFHE.allow(xAddress, msg.sender);
+    }
+
+    function setEBytes256(einput inputHandleEBytes256, bytes memory inputProofEBytes256) external {
+        yBytes256 = TFHE.asEbytes256(inputHandleEBytes256, inputProofEBytes256);
+        TFHE.allow(yBytes256, address(this));
+        TFHE.allow(yBytes256, msg.sender);
+    }
+}
diff --git a/lib/FhevmLib.sol b/lib/FhevmLib.sol
@@ -4,68 +4,40 @@ pragma solidity ^0.8.24;
 
 interface FhevmLib {
     function fheAdd(uint256 lhs, uint256 rhs, bytes1 scalarByte) external pure returns (uint256 result);
-
     function fheSub(uint256 lhs, uint256 rhs, bytes1 scalarByte) external pure returns (uint256 result);
-
     function fheMul(uint256 lhs, uint256 rhs, bytes1 scalarByte) external pure returns (uint256 result);
-
     function fheDiv(uint256 lhs, uint256 rhs, bytes1 scalarByte) external pure returns (uint256 result);
-
     function fheRem(uint256 lhs, uint256 rhs, bytes1 scalarByte) external pure returns (uint256 result);
-
     function fheBitAnd(uint256 lhs, uint256 rhs, bytes1 scalarByte) external pure returns (uint256 result);
-
     function fheBitOr(uint256 lhs, uint256 rhs, bytes1 scalarByte) external pure returns (uint256 result);
-
     function fheBitXor(uint256 lhs, uint256 rhs, bytes1 scalarByte) external pure returns (uint256 result);
-
     function fheShl(uint256 lhs, uint256 rhs, bytes1 scalarByte) external pure returns (uint256 result);
-
     function fheShr(uint256 lhs, uint256 rhs, bytes1 scalarByte) external pure returns (uint256 result);
-
     function fheRotl(uint256 lhs, uint256 rhs, bytes1 scalarByte) external pure returns (uint256 result);
-
     function fheRotr(uint256 lhs, uint256 rhs, bytes1 scalarByte) external pure returns (uint256 result);
-
     function fheEq(uint256 lhs, uint256 rhs, bytes1 scalarByte) external pure returns (uint256 result);
-
     function fheNe(uint256 lhs, uint256 rhs, bytes1 scalarByte) external pure returns (uint256 result);
-
     function fheGe(uint256 lhs, uint256 rhs, bytes1 scalarByte) external pure returns (uint256 result);
-
     function fheGt(uint256 lhs, uint256 rhs, bytes1 scalarByte) external pure returns (uint256 result);
-
     function fheLe(uint256 lhs, uint256 rhs, bytes1 scalarByte) external pure returns (uint256 result);
-
     function fheLt(uint256 lhs, uint256 rhs, bytes1 scalarByte) external pure returns (uint256 result);
-
     function fheMin(uint256 lhs, uint256 rhs, bytes1 scalarByte) external pure returns (uint256 result);
-
     function fheMax(uint256 lhs, uint256 rhs, bytes1 scalarByte) external pure returns (uint256 result);
-
     function fheNeg(uint256 ct) external pure returns (uint256 result);
-
     function fheNot(uint256 ct) external pure returns (uint256 result);
 
     function fhePubKey(bytes1 fromLib) external view returns (bytes memory result);
-
     function verifyCiphertext(
         bytes32 inputHandle,
         address callerAddress,
         address contractAddress,
         bytes memory inputProof,
         bytes1 inputType
     ) external pure returns (uint256 result);
-
     function cast(uint256 ct, bytes1 toType) external pure returns (uint256 result);
-
     function trivialEncrypt(uint256 ct, bytes1 toType) external pure returns (uint256 result);
-
     function fheIfThenElse(uint256 control, uint256 ifTrue, uint256 ifFalse) external pure returns (uint256 result);
-
     function fheArrayEq(uint256[] memory lhs, uint256[] memory rhs) external pure returns (uint256 result);
-
     function fheRand(bytes1 randType, uint256 seed) external view returns (uint256 result);
-
     function fheRandBounded(uint256 upperBound, bytes1 randType, uint256 seed) external view returns (uint256 result);
 }
diff --git a/lib/Impl.sol b/lib/Impl.sol
@@ -8,76 +8,45 @@ import "./ACLAddress.sol";
 
 interface IFHEVMCoprocessor {
     function fheAdd(uint256 lhs, uint256 rhs, bytes1 scalarByte) external returns (uint256 result);
-
     function fheSub(uint256 lhs, uint256 rhs, bytes1 scalarByte) external returns (uint256 result);
-
     function fheMul(uint256 lhs, uint256 rhs, bytes1 scalarByte) external returns (uint256 result);
-
     function fheDiv(uint256 lhs, uint256 rhs, bytes1 scalarByte) external returns (uint256 result);
-
     function fheRem(uint256 lhs, uint256 rhs, bytes1 scalarByte) external returns (uint256 result);
-
     function fheBitAnd(uint256 lhs, uint256 rhs, bytes1 scalarByte) external returns (uint256 result);
-
     function fheBitOr(uint256 lhs, uint256 rhs, bytes1 scalarByte) external returns (uint256 result);
-
     function fheBitXor(uint256 lhs, uint256 rhs, bytes1 scalarByte) external returns (uint256 result);
-
     function fheShl(uint256 lhs, uint256 rhs, bytes1 scalarByte) external returns (uint256 result);
-
     function fheShr(uint256 lhs, uint256 rhs, bytes1 scalarByte) external returns (uint256 result);
-
     function fheRotl(uint256 lhs, uint256 rhs, bytes1 scalarByte) external returns (uint256 result);
-
     function fheRotr(uint256 lhs, uint256 rhs, bytes1 scalarByte) external returns (uint256 result);
-
     function fheEq(uint256 lhs, uint256 rhs, bytes1 scalarByte) external returns (uint256 result);
-
     function fheNe(uint256 lhs, uint256 rhs, bytes1 scalarByte) external returns (uint256 result);
-
     function fheGe(uint256 lhs, uint256 rhs, bytes1 scalarByte) external returns (uint256 result);
-
     function fheGt(uint256 lhs, uint256 rhs, bytes1 scalarByte) external returns (uint256 result);
-
     function fheLe(uint256 lhs, uint256 rhs, bytes1 scalarByte) external returns (uint256 result);
-
     function fheLt(uint256 lhs, uint256 rhs, bytes1 scalarByte) external returns (uint256 result);
-
     function fheMin(uint256 lhs, uint256 rhs, bytes1 scalarByte) external returns (uint256 result);
-
     function fheMax(uint256 lhs, uint256 rhs, bytes1 scalarByte) external returns (uint256 result);
-
     function fheNeg(uint256 ct) external returns (uint256 result);
-
     function fheNot(uint256 ct) external returns (uint256 result);
-
     function verifyCiphertext(
         bytes32 inputHandle,
         address callerAddress,
         bytes memory inputProof,
         bytes1 inputType
     ) external returns (uint256 result);
-
     function cast(uint256 ct, bytes1 toType) external returns (uint256 result);
-
     function trivialEncrypt(uint256 ct, bytes1 toType) external returns (uint256 result);
-
     function fheIfThenElse(uint256 control, uint256 ifTrue, uint256 ifFalse) external returns (uint256 result);
-
     function fheRand(bytes1 randType) external returns (uint256 result);
-
     function fheRandBounded(uint256 upperBound, bytes1 randType) external returns (uint256 result);
-
     function cleanTransientStorage() external;
 }
 
 interface IACL {
     function allowTransient(uint256 ciphertext, address account) external;
-
     function allow(uint256 handle, address account) external;
-
     function cleanTransientStorage() external;
-
     function isAllowed(uint256 handle, address account) external view returns (bool);
 }
 

diff --git a/lib/TFHE.sol b/lib/TFHE.sol
@@ -5223,31 +5223,26 @@ library TFHE {
     function select(ebool control, euint4 a, euint4 b) internal returns (euint4) {
         return euint4.wrap(Impl.select(ebool.unwrap(control), euint4.unwrap(a), euint4.unwrap(b)));
     }
-
     // If 'control''s value is 'true', the result has the same value as 'a'.
     // If 'control''s value is 'false', the result has the same value as 'b'.
     function select(ebool control, euint8 a, euint8 b) internal returns (euint8) {
         return euint8.wrap(Impl.select(ebool.unwrap(control), euint8.unwrap(a), euint8.unwrap(b)));
     }
-
     // If 'control''s value is 'true', the result has the same value as 'a'.
     // If 'control''s value is 'false', the result has the same value as 'b'.
     function select(ebool control, euint16 a, euint16 b) internal returns (euint16) {
         return euint16.wrap(Impl.select(ebool.unwrap(control), euint16.unwrap(a), euint16.unwrap(b)));
     }
-
     // If 'control''s value is 'true', the result has the same value as 'a'.
     // If 'control''s value is 'false', the result has the same value as 'b'.
     function select(ebool control, euint32 a, euint32 b) internal returns (euint32) {
         return euint32.wrap(Impl.select(ebool.unwrap(control), euint32.unwrap(a), euint32.unwrap(b)));
     }
-
     // If 'control''s value is 'true', the result has the same value as 'a'.
     // If 'control''s value is 'false', the result has the same value as 'b'.
     function select(ebool control, euint64 a, euint64 b) internal returns (euint64) {
         return euint64.wrap(Impl.select(ebool.unwrap(control), euint64.unwrap(a), euint64.unwrap(b)));
     }
-
     // Cast an encrypted integer from euint8 to euint4.
     function asEuint4(euint8 value) internal returns (euint4) {
         return euint4.wrap(Impl.cast(euint8.unwrap(value), Common.euint4_t));
@@ -5690,31 +5685,29 @@ library TFHE {
     function isAllowed(ebool value, address account) internal view returns (bool) {
         return Impl.isAllowed(ebool.unwrap(value), account);
     }
-
     function isAllowed(euint4 value, address account) internal view returns (bool) {
         return Impl.isAllowed(euint4.unwrap(value), account);
     }
-
     function isAllowed(euint8 value, address account) internal view returns (bool) {
         return Impl.isAllowed(euint8.unwrap(value), account);
     }
-
     function isAllowed(euint16 value, address account) internal view returns (bool) {
         return Impl.isAllowed(euint16.unwrap(value), account);
     }
-
     function isAllowed(euint32 value, address account) internal view returns (bool) {
         return Impl.isAllowed(euint32.unwrap(value), account);
     }
-
     function isAllowed(euint64 value, address account) internal view returns (bool) {
         return Impl.isAllowed(euint64.unwrap(value), account);
     }
-
     function isAllowed(eaddress value, address account) internal view returns (bool) {
         return Impl.isAllowed(eaddress.unwrap(value), account);
     }
 
+    function isAllowed(ebytes256 value, address account) internal view returns (bool) {
+        return Impl.isAllowed(ebytes256.unwrap(value), account);
+    }
+
     function isSenderAllowed(ebool value) internal view returns (bool) {
         return Impl.isAllowed(ebool.unwrap(value), msg.sender);
     }
@@ -5743,6 +5736,10 @@ library TFHE {
         return Impl.isAllowed(eaddress.unwrap(value), msg.sender);
     }
 
+    function isSenderAllowed(ebytes256 value) internal view returns (bool) {
+        return Impl.isAllowed(ebytes256.unwrap(value), msg.sender);
+    }
+
     function allow(ebool value, address account) internal {
         Impl.allow(ebool.unwrap(value), account);
     }
@@ -5771,6 +5768,10 @@ library TFHE {
         Impl.allow(eaddress.unwrap(value), account);
     }
 
+    function allow(ebytes256 value, address account) internal {
+        Impl.allow(ebytes256.unwrap(value), account);
+    }
+
     function allowTransient(ebool value, address account) internal {
         Impl.allowTransient(ebool.unwrap(value), account);
     }
@@ -5798,4 +5799,8 @@ library TFHE {
     function allowTransient(eaddress value, address account) internal {
         Impl.allowTransient(eaddress.unwrap(value), account);
     }
+
+    function allowTransient(ebytes256 value, address account) internal {
+        Impl.allowTransient(ebytes256.unwrap(value), account);
+    }
 }