From 83b9e00e478f352589bc7b28d9c5b8a6d8cf633c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cl=C3=A9ment=20=27birdy=27=20Danjou?= <clement@danjou.io>
Date: Thu, 29 Jun 2023 16:23:15 +0200
Subject: [PATCH] feat() add default value for reencrypt

---
 codegen/codegen.py          |  8 +++++++
 examples/BlindAuction.sol   |  6 +----
 examples/EncryptedERC20.sol | 12 ++--------
 lib/TFHE.sol                | 48 +++++++++++++++++++++++++++++++++++++
 4 files changed, 59 insertions(+), 15 deletions(-)

diff --git a/codegen/codegen.py b/codegen/codegen.py
index add94fea..dbc7beb5 100644
--- a/codegen/codegen.py
+++ b/codegen/codegen.py
@@ -906,6 +906,14 @@
         return Impl.reencrypt(euint{i}.unwrap(ciphertext), publicKey);
     }}
 
+    function reencrypt(euint{i} ciphertext, bytes32 publicKey, uint{i} defaultValue) internal view returns (bytes memory reencrypted) {{
+        if (euint{i}.unwrap(ciphertext) != 0) {{
+            return Impl.reencrypt(euint{i}.unwrap(ciphertext), publicKey);
+        }} else {{
+            return Impl.reencrypt(euint{i}.unwrap(asEuint{i}(defaultValue)), publicKey);
+        }}
+    }}
+
     function req(euint{i} ciphertext) internal view {{
         Impl.req(euint{i}.unwrap(ciphertext));
     }}
diff --git a/examples/BlindAuction.sol b/examples/BlindAuction.sol
index b67edcec..47eaa020 100644
--- a/examples/BlindAuction.sol
+++ b/examples/BlindAuction.sol
@@ -98,11 +98,7 @@ contract BlindAuction is EIP712WithModifier {
         onlySignedPublicKey(publicKey, signature)
         returns (bytes memory)
     {
-        if (TFHE.isInitialized(bids[msg.sender])) {
-            return TFHE.reencrypt(bids[msg.sender], publicKey);
-        } else {
-            return TFHE.reencrypt(TFHE.asEuint32(0), publicKey);
-        }
+        return TFHE.reencrypt(bids[msg.sender], publicKey, 0);
     }
 
     // Returns the user bid
diff --git a/examples/EncryptedERC20.sol b/examples/EncryptedERC20.sol
index 6784009f..455e64fe 100644
--- a/examples/EncryptedERC20.sol
+++ b/examples/EncryptedERC20.sol
@@ -54,11 +54,7 @@ contract EncryptedERC20 is EIP712WithModifier {
         onlySignedPublicKey(publicKey, signature)
         returns (bytes memory)
     {
-        if (TFHE.isInitialized(totalSupply)) {
-            return TFHE.reencrypt(totalSupply, publicKey);
-        } else {
-            return TFHE.reencrypt(TFHE.asEuint32(0), publicKey);
-        }
+        return TFHE.reencrypt(totalSupply, publicKey, 0);
     }
 
     // Returns the balance of the caller encrypted under the provided public key.
@@ -71,11 +67,7 @@ contract EncryptedERC20 is EIP712WithModifier {
         onlySignedPublicKey(publicKey, signature)
         returns (bytes memory)
     {
-        if (TFHE.isInitialized(balances[msg.sender])) {
-            return TFHE.reencrypt(balances[msg.sender], publicKey);
-        } else {
-            return TFHE.reencrypt(TFHE.asEuint32(0), publicKey);
-        }
+        return TFHE.reencrypt(balances[msg.sender], publicKey, 0);
     }
 
     // Sets the `encryptedAmount` as the allowance of `spender` over the caller's tokens.
diff --git a/lib/TFHE.sol b/lib/TFHE.sol
index 8e7ed47f..cb2c5b87 100644
--- a/lib/TFHE.sol
+++ b/lib/TFHE.sol
@@ -2170,6 +2170,22 @@ library TFHE {
         return Impl.reencrypt(euint8.unwrap(ciphertext), publicKey);
     }
 
+    function reencrypt(
+        euint8 ciphertext,
+        bytes32 publicKey,
+        uint8 defaultValue
+    ) internal view returns (bytes memory reencrypted) {
+        if (euint8.unwrap(ciphertext) != 0) {
+            return Impl.reencrypt(euint8.unwrap(ciphertext), publicKey);
+        } else {
+            return
+                Impl.reencrypt(
+                    euint8.unwrap(asEuint8(defaultValue)),
+                    publicKey
+                );
+        }
+    }
+
     function req(euint8 ciphertext) internal view {
         Impl.req(euint8.unwrap(ciphertext));
     }
@@ -2195,6 +2211,22 @@ library TFHE {
         return Impl.reencrypt(euint16.unwrap(ciphertext), publicKey);
     }
 
+    function reencrypt(
+        euint16 ciphertext,
+        bytes32 publicKey,
+        uint16 defaultValue
+    ) internal view returns (bytes memory reencrypted) {
+        if (euint16.unwrap(ciphertext) != 0) {
+            return Impl.reencrypt(euint16.unwrap(ciphertext), publicKey);
+        } else {
+            return
+                Impl.reencrypt(
+                    euint16.unwrap(asEuint16(defaultValue)),
+                    publicKey
+                );
+        }
+    }
+
     function req(euint16 ciphertext) internal view {
         Impl.req(euint16.unwrap(ciphertext));
     }
@@ -2220,6 +2252,22 @@ library TFHE {
         return Impl.reencrypt(euint32.unwrap(ciphertext), publicKey);
     }
 
+    function reencrypt(
+        euint32 ciphertext,
+        bytes32 publicKey,
+        uint32 defaultValue
+    ) internal view returns (bytes memory reencrypted) {
+        if (euint32.unwrap(ciphertext) != 0) {
+            return Impl.reencrypt(euint32.unwrap(ciphertext), publicKey);
+        } else {
+            return
+                Impl.reencrypt(
+                    euint32.unwrap(asEuint32(defaultValue)),
+                    publicKey
+                );
+        }
+    }
+
     function req(euint32 ciphertext) internal view {
         Impl.req(euint32.unwrap(ciphertext));
     }