From feb8a98ab921494b0261f5346bc52da48d5931f8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cl=C3=A9ment=20=27birdy=27=20Danjou?= <clement@danjou.io>
Date: Thu, 8 Feb 2024 00:48:23 +0100
Subject: [PATCH] feat: add encrypt64 for 64bits

---
 src/sdk/encrypt.test.ts | 14 +++++++++++++-
 src/sdk/encrypt.ts      | 13 +++++++++++++
 src/sdk/index.ts        |  9 ++++++++-
 3 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/src/sdk/encrypt.test.ts b/src/sdk/encrypt.test.ts
index 09e8f3a..8463b87 100644
--- a/src/sdk/encrypt.test.ts
+++ b/src/sdk/encrypt.test.ts
@@ -2,14 +2,16 @@ import {
   FheUint8,
   FheUint16,
   FheUint32,
+  FheUint64,
   CompactFheUint8List,
   CompactFheUint16List,
   CompactFheUint32List,
+  CompactFheUint64List,
   TfheCompactPublicKey,
   TfheClientKey,
 } from 'node-tfhe';
 import { createTfheKeypair } from '../tfhe';
-import { encrypt8, encrypt16, encrypt32 } from './encrypt';
+import { encrypt8, encrypt16, encrypt32, encrypt64 } from './encrypt';
 
 describe('encrypt8', () => {
   let clientKey: TfheClientKey;
@@ -80,4 +82,14 @@ describe('encrypt8', () => {
       expect(decrypted).toBe(30210);
     });
   });
+
+  it('encrypt/decrypt 32bits', async () => {
+    const buffer = encrypt64(3021094839202949, publicKey);
+    const compactList = CompactFheUint64List.deserialize(buffer);
+    let encryptedList = compactList.expand();
+    encryptedList.forEach((v: FheUint64) => {
+      const decrypted = v.decrypt(clientKey);
+      expect(decrypted.toString()).toBe('3021094839202949');
+    });
+  });
 });
diff --git a/src/sdk/encrypt.ts b/src/sdk/encrypt.ts
index 1a76316..3d3b594 100644
--- a/src/sdk/encrypt.ts
+++ b/src/sdk/encrypt.ts
@@ -3,6 +3,7 @@ import {
   CompactFheUint8List,
   CompactFheUint16List,
   CompactFheUint32List,
+  CompactFheUint64List,
 } from 'node-tfhe';
 
 export const encrypt8 = (
@@ -40,3 +41,15 @@ export const encrypt32 = (
   );
   return encrypted.serialize();
 };
+
+export const encrypt64 = (
+  value: number,
+  publicKey: TfheCompactPublicKey,
+): Uint8Array => {
+  const uint64Array = new BigUint64Array([BigInt(value)]);
+  const encrypted = CompactFheUint64List.encrypt_with_compact_public_key(
+    uint64Array,
+    publicKey,
+  );
+  return encrypted.serialize();
+};
diff --git a/src/sdk/index.ts b/src/sdk/index.ts
index aabf6ef..bad78af 100644
--- a/src/sdk/index.ts
+++ b/src/sdk/index.ts
@@ -1,6 +1,6 @@
 import { TfheCompactPublicKey } from 'node-tfhe';
 import sodium from 'libsodium-wrappers';
-import { encrypt8, encrypt16, encrypt32 } from './encrypt';
+import { encrypt8, encrypt16, encrypt32, encrypt64 } from './encrypt';
 import {
   EIP712,
   GeneratePublicKeyParams,
@@ -14,6 +14,7 @@ export type FhevmInstance = {
   encrypt8: (value: number) => Uint8Array;
   encrypt16: (value: number) => Uint8Array;
   encrypt32: (value: number) => Uint8Array;
+  encrypt64: (value: number) => Uint8Array;
   generateToken: (
     options: GeneratePublicKeyParams & {
       force?: boolean;
@@ -120,6 +121,12 @@ export const createInstance = async (
       return encrypt32(value, tfheCompactPublicKey);
     },
 
+    encrypt64(value) {
+      if (value == null) throw new Error('Missing value');
+      if (typeof value !== 'number') throw new Error('Value must be a number');
+      return encrypt64(value, tfheCompactPublicKey);
+    },
+
     /**
      * @deprecated Since version 0.3.0. Will be deleted in version 0.4.0. Use generatePublicKey instead.
      */