Skip to content

Commit

Permalink
Merge pull request #246 from zama-ai/feature/new-cmux
Browse files Browse the repository at this point in the history 
feat() use new cmux
  • Loading branch information
@immortal-tofu
immortal-tofu authored Jan 8, 2024
2 parents ddc6ad2 + cec03b8 commit bd6eb1c
Show file tree
Hide file tree
Showing 6 changed files with 21 additions and 38 deletions.
4 changes: 2 additions & 2 deletions codegen/main.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,9 +13,9 @@ function generateAllFiles() {
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));
writeFileSync('mocks/Impl.sol', t.implSolMock(context, operators));
const [tfheSolSourceMock, _] = t.tfheSol(context, operators, SUPPORTED_BITS, true);
writeFileSync('lib_mock/TFHE.sol', tfheSolSourceMock);
writeFileSync('mocks/TFHE.sol', tfheSolSourceMock);
mkdirSync('examples/tests', { recursive: true });
ovShards.forEach((os) => {
writeFileSync(`examples/tests/TFHETestSuite${os.shardNumber}.sol`, testgen.generateSmartContract(os));
Expand Down
18 changes: 3 additions & 15 deletions codegen/templates.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -133,6 +133,7 @@ function fheLibCustomInterfaceFunctions(): string {
function trivialEncrypt(uint256 ct, bytes1 toType) external pure returns (uint256 result);
function decrypt(uint256 ct) external view returns (uint256 result);
function fheRand(bytes1 inp) external view returns (uint256 result);
function fheIfThenElse(uint256 control, uint256 ifTrue, uint256 ifFalse) external pure returns (uint256 result);
`;
}

Expand Down Expand Up @@ -491,22 +492,12 @@ function tfheShiftOperators(inputBits: number, operator: Operator, signatures: O
}

function tfheCmux(inputBits: number): string {
if (inputBits == 8) {
return `
return `
// 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 cmux(ebool control, euint${inputBits} a, euint${inputBits} b) internal pure returns (euint${inputBits}) {
return euint${inputBits}.wrap(Impl.cmux(ebool.unwrap(control), euint${inputBits}.unwrap(a), euint${inputBits}.unwrap(b)));
}`;
}

return `
// 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 cmux(ebool control, euint${inputBits} a, euint${inputBits} b) internal pure returns (euint${inputBits}) {
euint${inputBits} ctrl = asEuint${inputBits}(asEuint8(control));
return euint${inputBits}.wrap(Impl.cmux(euint${inputBits}.unwrap(ctrl), euint${inputBits}.unwrap(a), euint${inputBits}.unwrap(b)));
}`;
}

function tfheAsEboolCustomCast(inputBits: number, outputBits: number): string {
Expand Down Expand Up @@ -734,10 +725,7 @@ function implCustomMethods(ctx: CodegenContext): string {
// If 'control's value is 'true', the result has the same value as 'ifTrue'.
// If 'control's value is 'false', the result has the same value as 'ifFalse'.
function cmux(uint256 control, uint256 ifTrue, uint256 ifFalse) internal pure returns (uint256 result) {
// result = (ifTrue - ifFalse) * control + ifFalse
uint256 subOutput = FhevmLib(address(EXT_TFHE_LIBRARY)).fheSub(ifTrue, ifFalse, bytes1(0x00));
uint256 mulOutput = FhevmLib(address(EXT_TFHE_LIBRARY)).fheMul(control, subOutput, bytes1(0x00));
result = FhevmLib(address(EXT_TFHE_LIBRARY)).fheAdd(mulOutput, ifFalse, bytes1(0x00));
result = FhevmLib(address(EXT_TFHE_LIBRARY)).fheIfThenElse(control, ifTrue, ifFalse);
}
// We do assembly here because ordinary call will emit extcodesize check which is zero for our precompiles
Expand Down
7 changes: 3 additions & 4 deletions lib/Impl.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,8 @@ interface FhevmLib {
function decrypt(uint256 ct) external view returns (uint256 result);

function fheRand(bytes1 inp) external view returns (uint256 result);

function fheIfThenElse(uint256 control, uint256 ifTrue, uint256 ifFalse) external pure returns (uint256 result);
}

address constant EXT_TFHE_LIBRARY = address(0x000000000000000000000000000000000000005d);
Expand Down Expand Up @@ -234,10 +236,7 @@ library Impl {
// If 'control's value is 'true', the result has the same value as 'ifTrue'.
// If 'control's value is 'false', the result has the same value as 'ifFalse'.
function cmux(uint256 control, uint256 ifTrue, uint256 ifFalse) internal pure returns (uint256 result) {
// result = (ifTrue - ifFalse) * control + ifFalse
uint256 subOutput = FhevmLib(address(EXT_TFHE_LIBRARY)).fheSub(ifTrue, ifFalse, bytes1(0x00));
uint256 mulOutput = FhevmLib(address(EXT_TFHE_LIBRARY)).fheMul(control, subOutput, bytes1(0x00));
result = FhevmLib(address(EXT_TFHE_LIBRARY)).fheAdd(mulOutput, ifFalse, bytes1(0x00));
result = FhevmLib(address(EXT_TFHE_LIBRARY)).fheIfThenElse(control, ifTrue, ifFalse);
}

// We do assembly here because ordinary call will emit extcodesize check which is zero for our precompiles
Expand Down
14 changes: 6 additions & 8 deletions lib/TFHE.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2122,18 +2122,16 @@ library TFHE {
return euint8.wrap(Impl.cmux(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'.
// 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 cmux(ebool control, euint16 a, euint16 b) internal pure returns (euint16) {
euint16 ctrl = asEuint16(asEuint8(control));
return euint16.wrap(Impl.cmux(euint16.unwrap(ctrl), euint16.unwrap(a), euint16.unwrap(b)));
return euint16.wrap(Impl.cmux(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'.
// 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 cmux(ebool control, euint32 a, euint32 b) internal pure returns (euint32) {
euint32 ctrl = asEuint32(asEuint8(control));
return euint32.wrap(Impl.cmux(euint32.unwrap(ctrl), euint32.unwrap(a), euint32.unwrap(b)));
return euint32.wrap(Impl.cmux(ebool.unwrap(control), euint32.unwrap(a), euint32.unwrap(b)));
}

// Cast an encrypted integer from euint16 to euint8.
Expand Down
14 changes: 6 additions & 8 deletions mocks/TFHE.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2122,18 +2122,16 @@ library TFHE {
return euint8.wrap(Impl.cmux(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'.
// 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 cmux(ebool control, euint16 a, euint16 b) internal pure returns (euint16) {
euint16 ctrl = asEuint16(asEuint8(control));
return euint16.wrap(Impl.cmux(euint16.unwrap(ctrl), euint16.unwrap(a), euint16.unwrap(b)));
return euint16.wrap(Impl.cmux(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'.
// 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 cmux(ebool control, euint32 a, euint32 b) internal pure returns (euint32) {
euint32 ctrl = asEuint32(asEuint8(control));
return euint32.wrap(Impl.cmux(euint32.unwrap(ctrl), euint32.unwrap(a), euint32.unwrap(b)));
return euint32.wrap(Impl.cmux(ebool.unwrap(control), euint32.unwrap(a), euint32.unwrap(b)));
}

// Cast an encrypted integer from euint16 to euint8.
Expand Down
2 changes: 1 addition & 1 deletion package.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@
"codegen": "npx ts-node codegen/main.ts && npm run prettier",
"task:getEthereumAddress": "hardhat task:getEthereumAddress",
"task:deployIdentity": "hardhat task:deployIdentity",
"fhevm:start": "docker run -i -p 8545:8545 --rm --name fhevm ghcr.io/zama-ai/evmos-dev-node:v0.1.10",
"fhevm:start": "docker run -i -p 8545:8545 --rm --name fhevm ghcr.io/zama-ai/ethermint-dev-node:v0.2.2",
"fhevm:stop": "docker rm -f fhevm",
"fhevm:restart": "fhevm:stop && fhevm:start",
"fhevm:faucet": "npm run fhevm:faucet:alice && sleep 5 && npm run fhevm:faucet:bob && sleep 5 && npm run fhevm:faucet:carol && sleep 5 && npm run fhevm:faucet:dave && sleep 5 && npm run fhevm:faucet:eve",
Expand Down

0 comments on commit bd6eb1c

Please sign in to comment.