feat: add select for eaddress

zama-ai · Apr 5, 2024 · 869fbb4 · 869fbb4
1 parent cc9b1f6
commit 869fbb4
Show file tree

Hide file tree

Showing 4 changed files with 66 additions and 6 deletions.
diff --git a/fhevm/tfhe/tfhe_ciphertext.go b/fhevm/tfhe/tfhe_ciphertext.go
@@ -778,7 +778,8 @@ func (first *TfheCiphertext) executeTernaryCiphertextOperation(lhs *TfheCipherte
 	op8 func(first unsafe.Pointer, lhs unsafe.Pointer, rhs unsafe.Pointer) unsafe.Pointer,
 	op16 func(first unsafe.Pointer, lhs unsafe.Pointer, rhs unsafe.Pointer) unsafe.Pointer,
 	op32 func(first unsafe.Pointer, lhs unsafe.Pointer, rhs unsafe.Pointer) unsafe.Pointer,
-	op64 func(first unsafe.Pointer, lhs unsafe.Pointer, rhs unsafe.Pointer) unsafe.Pointer) (*TfheCiphertext, error) {
+	op64 func(first unsafe.Pointer, lhs unsafe.Pointer, rhs unsafe.Pointer) unsafe.Pointer,
+	op160 func(first unsafe.Pointer, lhs unsafe.Pointer, rhs unsafe.Pointer) unsafe.Pointer) (*TfheCiphertext, error) {
 	if lhs.FheUintType != rhs.FheUintType {
 		return nil, errors.New("ternary operations are only well-defined for identical types")
 	}
@@ -932,6 +933,35 @@ func (first *TfheCiphertext) executeTernaryCiphertextOperation(lhs *TfheCipherte
 		}
 		res.Serialization = C.GoBytes(unsafe.Pointer(res_ser.pointer), C.int(res_ser.length))
 		C.destroy_dynamic_buffer(res_ser)
+	case FheUint160:
+		lhs_ptr := C.deserialize_fhe_uint160(toDynamicBufferView((lhs.Serialization)))
+		if lhs_ptr == nil {
+			return nil, errors.New("160 bit binary op deserialization failed")
+		}
+		rhs_ptr := C.deserialize_fhe_uint160(toDynamicBufferView((rhs.Serialization)))
+		if rhs_ptr == nil {
+			C.destroy_fhe_uint160(lhs_ptr)
+			return nil, errors.New("160 bit binary op deserialization failed")
+		}
+		first_ptr := C.deserialize_fhe_bool(toDynamicBufferView((first.Serialization)))
+		if first_ptr == nil {
+			C.destroy_fhe_uint160(lhs_ptr)
+			C.destroy_fhe_uint160(rhs_ptr)
+			return nil, errors.New("Bool binary op deserialization failed")
+		}
+		res_ptr := op160(first_ptr, lhs_ptr, rhs_ptr)
+		C.destroy_fhe_uint160(lhs_ptr)
+		C.destroy_fhe_uint160(rhs_ptr)
+		if res_ptr == nil {
+			return nil, errors.New("160 bit binary op failed")
+		}
+		ret := C.serialize_fhe_uint160(res_ptr, res_ser)
+		C.destroy_fhe_uint160(res_ptr)
+		if ret != 0 {
+			return nil, errors.New("160 bit binary op serialization failed")
+		}
+		res.Serialization = C.GoBytes(unsafe.Pointer(res_ser.pointer), C.int(res_ser.length))
+		C.destroy_dynamic_buffer(res_ser)
 	default:
 		panic("ternary op unexpected ciphertext type")
 	}
@@ -1485,7 +1515,6 @@ func (lhs *TfheCiphertext) ScalarShr(rhs *big.Int) (*TfheCiphertext, error) {
 		fheUint160BinaryScalarNotSupportedOp, false)
 }
 
-
 func (lhs *TfheCiphertext) Rotl(rhs *TfheCiphertext) (*TfheCiphertext, error) {
 	return lhs.executeBinaryCiphertextOperation(rhs,
 		boolBinaryNotSupportedOp,
@@ -1984,6 +2013,9 @@ func (condition *TfheCiphertext) IfThenElse(lhs *TfheCiphertext, rhs *TfheCipher
 		},
 		func(condition unsafe.Pointer, lhs unsafe.Pointer, rhs unsafe.Pointer) unsafe.Pointer {
 			return C.if_then_else_fhe_uint64(condition, lhs, rhs, sks)
+		},
+		func(condition unsafe.Pointer, lhs unsafe.Pointer, rhs unsafe.Pointer) unsafe.Pointer {
+			return C.if_then_else_fhe_uint160(condition, lhs, rhs, sks)
 		})
 }
 

diff --git a/fhevm/tfhe/tfhe_test.go b/fhevm/tfhe/tfhe_test.go
@@ -752,7 +752,6 @@ func TfheRotl(t *testing.T, fheUintType FheUintType) {
 	}
 }
 
-
 func TfheScalarRotl(t *testing.T, fheUintType FheUintType) {
 	var a, b big.Int
 	var expected uint64
@@ -1497,6 +1496,19 @@ func TfheIfThenElse(t *testing.T, fheUintType FheUintType) {
 	case FheUint64:
 		a.SetUint64(13333377777777777)
 		b.SetUint64(133337)
+	case FheUint160:
+		hexValue := "12345676876661323221435343"
+		hexValue2 := "12345676876661323221435344"
+		byteValue, err := hex.DecodeString(hexValue)
+		if err != nil {
+			log.Fatalf("Failed to decode hex string: %v", err)
+		}
+		byteValue2, err := hex.DecodeString(hexValue2)
+		if err != nil {
+			log.Fatalf("Failed to decode hex string: %v", err)
+		}
+		a.SetBytes(byteValue)
+		b.SetBytes(byteValue2)
 	}
 	ctCondition := new(TfheCiphertext)
 	ctCondition.Encrypt(condition, FheBool)
@@ -1510,10 +1522,10 @@ func TfheIfThenElse(t *testing.T, fheUintType FheUintType) {
 	ctRes2, _ := ctCondition2.IfThenElse(ctA, ctB)
 	res1, err1 := ctRes1.Decrypt()
 	res2, err2 := ctRes2.Decrypt()
-	if err1 != nil || res1.Uint64() != a.Uint64() {
+	if err1 != nil || res1.Cmp(&a) != 0 {
 		t.Fatalf("%d != %d", 0, res1.Uint64())
 	}
-	if err2 != nil || res2.Uint64() != b.Uint64() {
+	if err2 != nil || res2.Cmp(&b) != 0 {
 		t.Fatalf("%d != %d", 0, res2.Uint64())
 	}
 }
@@ -2068,7 +2080,6 @@ func TestTfheScalarShr64(t *testing.T) {
 	TfheScalarShr(t, FheUint64)
 }
 
-
 func TestTfheRotl4(t *testing.T) {
 	TfheRotl(t, FheUint4)
 }
@@ -2529,6 +2540,10 @@ func TestTfheIfThenElse64(t *testing.T) {
 	TfheIfThenElse(t, FheUint64)
 }
 
+func TestTfheIfThenElse160(t *testing.T) {
+	TfheIfThenElse(t, FheUint160)
+}
+
 func TestTfhe4Cast8(t *testing.T) {
 	TfheCast(t, FheUint4, FheUint8)
 }

diff --git a/fhevm/tfhe/tfhe_wrappers.c b/fhevm/tfhe/tfhe_wrappers.c
@@ -2593,6 +2593,17 @@ void* if_then_else_fhe_uint64(void* condition, void* ct1, void* ct2, void* sks)
 	return result;
 }
 
+void* if_then_else_fhe_uint160(void* condition, void* ct1, void* ct2, void* sks)
+{
+	FheUint160* result = NULL;
+
+	checked_set_server_key(sks);
+
+	const int r = fhe_uint160_if_then_else(condition, ct1, ct2, &result);
+	if(r != 0) return NULL;
+	return result;
+}
+
 int decrypt_fhe_bool(void* cks, void* ct, bool* res)
 {
 	*res = false;

diff --git a/fhevm/tfhe/tfhe_wrappers.h b/fhevm/tfhe/tfhe_wrappers.h
@@ -471,6 +471,8 @@ void* if_then_else_fhe_uint32(void* condition, void* ct1, void* ct2, void* sks);
 
 void* if_then_else_fhe_uint64(void* condition, void* ct1, void* ct2, void* sks);
 
+void* if_then_else_fhe_uint160(void* condition, void* ct1, void* ct2, void* sks);
+
 int decrypt_fhe_bool(void* cks, void* ct, bool* res);
 
 int decrypt_fhe_uint4(void* cks, void* ct, uint8_t* res);