1use crate::error::{AmateRSError, ErrorContext, Result};
7use crate::types::CipherBlob;
8
9#[cfg(feature = "compute")]
10use tfhe::prelude::*;
11#[cfg(feature = "compute")]
12use tfhe::{FheBool, FheUint8, FheUint16, FheUint32, FheUint64};
13
14#[cfg(feature = "compute")]
16#[derive(Clone)]
17pub struct EncryptedBool {
18 inner: FheBool,
19}
20
21#[cfg(feature = "compute")]
22impl EncryptedBool {
23 pub fn from_fhe(value: FheBool) -> Self {
25 Self { inner: value }
26 }
27
28 pub fn inner(&self) -> &FheBool {
30 &self.inner
31 }
32
33 pub fn encrypt(value: bool, client_key: &tfhe::ClientKey) -> Self {
35 Self {
36 inner: FheBool::encrypt(value, client_key),
37 }
38 }
39
40 pub fn decrypt(&self, client_key: &tfhe::ClientKey) -> bool {
42 self.inner.decrypt(client_key)
43 }
44
45 pub fn to_cipher_blob(&self) -> Result<CipherBlob> {
47 let bytes = bincode::serialize(&self.inner).map_err(|e| {
48 AmateRSError::Serialization(ErrorContext::new(format!(
49 "Failed to serialize EncryptedBool: {}",
50 e
51 )))
52 })?;
53 Ok(CipherBlob::new(bytes))
54 }
55
56 pub fn from_cipher_blob(blob: &CipherBlob) -> Result<Self> {
58 let inner: FheBool = bincode::deserialize(blob.as_bytes()).map_err(|e| {
59 AmateRSError::Deserialization(ErrorContext::new(format!(
60 "Failed to deserialize EncryptedBool: {}",
61 e
62 )))
63 })?;
64 Ok(Self { inner })
65 }
66
67 pub fn and(&self, other: &Self) -> Self {
69 Self {
70 inner: &self.inner & &other.inner,
71 }
72 }
73
74 pub fn or(&self, other: &Self) -> Self {
76 Self {
77 inner: &self.inner | &other.inner,
78 }
79 }
80
81 pub fn xor(&self, other: &Self) -> Self {
83 Self {
84 inner: &self.inner ^ &other.inner,
85 }
86 }
87
88 pub fn not(&self) -> Self {
90 Self {
91 inner: !&self.inner,
92 }
93 }
94}
95
96#[cfg(feature = "compute")]
98#[derive(Clone)]
99pub struct EncryptedU8 {
100 inner: FheUint8,
101}
102
103#[cfg(feature = "compute")]
104impl EncryptedU8 {
105 pub fn from_fhe(value: FheUint8) -> Self {
107 Self { inner: value }
108 }
109
110 pub fn inner(&self) -> &FheUint8 {
112 &self.inner
113 }
114
115 pub fn encrypt(value: u8, client_key: &tfhe::ClientKey) -> Self {
117 Self {
118 inner: FheUint8::encrypt(value, client_key),
119 }
120 }
121
122 pub fn decrypt(&self, client_key: &tfhe::ClientKey) -> u8 {
124 self.inner.decrypt(client_key)
125 }
126
127 pub fn to_cipher_blob(&self) -> Result<CipherBlob> {
129 let bytes = bincode::serialize(&self.inner).map_err(|e| {
130 AmateRSError::Serialization(ErrorContext::new(format!(
131 "Failed to serialize EncryptedU8: {}",
132 e
133 )))
134 })?;
135 Ok(CipherBlob::new(bytes))
136 }
137
138 pub fn from_cipher_blob(blob: &CipherBlob) -> Result<Self> {
140 let inner: FheUint8 = bincode::deserialize(blob.as_bytes()).map_err(|e| {
141 AmateRSError::Deserialization(ErrorContext::new(format!(
142 "Failed to deserialize EncryptedU8: {}",
143 e
144 )))
145 })?;
146 Ok(Self { inner })
147 }
148
149 pub fn add(&self, other: &Self) -> Self {
151 Self {
152 inner: &self.inner + &other.inner,
153 }
154 }
155
156 pub fn sub(&self, other: &Self) -> Self {
158 Self {
159 inner: &self.inner - &other.inner,
160 }
161 }
162
163 pub fn mul(&self, other: &Self) -> Self {
165 Self {
166 inner: &self.inner * &other.inner,
167 }
168 }
169
170 pub fn eq(&self, other: &Self) -> EncryptedBool {
172 EncryptedBool {
173 inner: self.inner.eq(&other.inner),
174 }
175 }
176
177 pub fn lt(&self, other: &Self) -> EncryptedBool {
179 EncryptedBool {
180 inner: self.inner.lt(&other.inner),
181 }
182 }
183
184 pub fn le(&self, other: &Self) -> EncryptedBool {
186 EncryptedBool {
187 inner: self.inner.le(&other.inner),
188 }
189 }
190
191 pub fn gt(&self, other: &Self) -> EncryptedBool {
193 EncryptedBool {
194 inner: self.inner.gt(&other.inner),
195 }
196 }
197
198 pub fn ge(&self, other: &Self) -> EncryptedBool {
200 EncryptedBool {
201 inner: self.inner.ge(&other.inner),
202 }
203 }
204
205 pub fn ne(&self, other: &Self) -> EncryptedBool {
207 EncryptedBool {
208 inner: self.inner.ne(&other.inner),
209 }
210 }
211}
212
213#[cfg(feature = "compute")]
215#[derive(Clone)]
216pub struct EncryptedU16 {
217 inner: FheUint16,
218}
219
220#[cfg(feature = "compute")]
221impl EncryptedU16 {
222 pub fn from_fhe(value: FheUint16) -> Self {
223 Self { inner: value }
224 }
225
226 pub fn inner(&self) -> &FheUint16 {
227 &self.inner
228 }
229
230 pub fn encrypt(value: u16, client_key: &tfhe::ClientKey) -> Self {
231 Self {
232 inner: FheUint16::encrypt(value, client_key),
233 }
234 }
235
236 pub fn decrypt(&self, client_key: &tfhe::ClientKey) -> u16 {
237 self.inner.decrypt(client_key)
238 }
239
240 pub fn to_cipher_blob(&self) -> Result<CipherBlob> {
241 let bytes = bincode::serialize(&self.inner).map_err(|e| {
242 AmateRSError::Serialization(ErrorContext::new(format!(
243 "Failed to serialize EncryptedU16: {}",
244 e
245 )))
246 })?;
247 Ok(CipherBlob::new(bytes))
248 }
249
250 pub fn from_cipher_blob(blob: &CipherBlob) -> Result<Self> {
251 let inner: FheUint16 = bincode::deserialize(blob.as_bytes()).map_err(|e| {
252 AmateRSError::Deserialization(ErrorContext::new(format!(
253 "Failed to deserialize EncryptedU16: {}",
254 e
255 )))
256 })?;
257 Ok(Self { inner })
258 }
259
260 pub fn add(&self, other: &Self) -> Self {
261 Self {
262 inner: &self.inner + &other.inner,
263 }
264 }
265
266 pub fn sub(&self, other: &Self) -> Self {
267 Self {
268 inner: &self.inner - &other.inner,
269 }
270 }
271
272 pub fn mul(&self, other: &Self) -> Self {
273 Self {
274 inner: &self.inner * &other.inner,
275 }
276 }
277
278 pub fn eq(&self, other: &Self) -> EncryptedBool {
279 EncryptedBool {
280 inner: self.inner.eq(&other.inner),
281 }
282 }
283
284 pub fn lt(&self, other: &Self) -> EncryptedBool {
285 EncryptedBool {
286 inner: self.inner.lt(&other.inner),
287 }
288 }
289
290 pub fn le(&self, other: &Self) -> EncryptedBool {
291 EncryptedBool {
292 inner: self.inner.le(&other.inner),
293 }
294 }
295
296 pub fn gt(&self, other: &Self) -> EncryptedBool {
297 EncryptedBool {
298 inner: self.inner.gt(&other.inner),
299 }
300 }
301
302 pub fn ge(&self, other: &Self) -> EncryptedBool {
303 EncryptedBool {
304 inner: self.inner.ge(&other.inner),
305 }
306 }
307
308 pub fn ne(&self, other: &Self) -> EncryptedBool {
309 EncryptedBool {
310 inner: self.inner.ne(&other.inner),
311 }
312 }
313}
314
315#[cfg(feature = "compute")]
317#[derive(Clone)]
318pub struct EncryptedU32 {
319 inner: FheUint32,
320}
321
322#[cfg(feature = "compute")]
323impl EncryptedU32 {
324 pub fn from_fhe(value: FheUint32) -> Self {
325 Self { inner: value }
326 }
327
328 pub fn inner(&self) -> &FheUint32 {
329 &self.inner
330 }
331
332 pub fn encrypt(value: u32, client_key: &tfhe::ClientKey) -> Self {
333 Self {
334 inner: FheUint32::encrypt(value, client_key),
335 }
336 }
337
338 pub fn decrypt(&self, client_key: &tfhe::ClientKey) -> u32 {
339 self.inner.decrypt(client_key)
340 }
341
342 pub fn to_cipher_blob(&self) -> Result<CipherBlob> {
343 let bytes = bincode::serialize(&self.inner).map_err(|e| {
344 AmateRSError::Serialization(ErrorContext::new(format!(
345 "Failed to serialize EncryptedU32: {}",
346 e
347 )))
348 })?;
349 Ok(CipherBlob::new(bytes))
350 }
351
352 pub fn from_cipher_blob(blob: &CipherBlob) -> Result<Self> {
353 let inner: FheUint32 = bincode::deserialize(blob.as_bytes()).map_err(|e| {
354 AmateRSError::Deserialization(ErrorContext::new(format!(
355 "Failed to deserialize EncryptedU32: {}",
356 e
357 )))
358 })?;
359 Ok(Self { inner })
360 }
361
362 pub fn add(&self, other: &Self) -> Self {
363 Self {
364 inner: &self.inner + &other.inner,
365 }
366 }
367
368 pub fn sub(&self, other: &Self) -> Self {
369 Self {
370 inner: &self.inner - &other.inner,
371 }
372 }
373
374 pub fn mul(&self, other: &Self) -> Self {
375 Self {
376 inner: &self.inner * &other.inner,
377 }
378 }
379
380 pub fn eq(&self, other: &Self) -> EncryptedBool {
381 EncryptedBool {
382 inner: self.inner.eq(&other.inner),
383 }
384 }
385
386 pub fn lt(&self, other: &Self) -> EncryptedBool {
387 EncryptedBool {
388 inner: self.inner.lt(&other.inner),
389 }
390 }
391
392 pub fn le(&self, other: &Self) -> EncryptedBool {
393 EncryptedBool {
394 inner: self.inner.le(&other.inner),
395 }
396 }
397
398 pub fn gt(&self, other: &Self) -> EncryptedBool {
399 EncryptedBool {
400 inner: self.inner.gt(&other.inner),
401 }
402 }
403
404 pub fn ge(&self, other: &Self) -> EncryptedBool {
405 EncryptedBool {
406 inner: self.inner.ge(&other.inner),
407 }
408 }
409
410 pub fn ne(&self, other: &Self) -> EncryptedBool {
411 EncryptedBool {
412 inner: self.inner.ne(&other.inner),
413 }
414 }
415}
416
417#[cfg(feature = "compute")]
419#[derive(Clone)]
420pub struct EncryptedU64 {
421 inner: FheUint64,
422}
423
424#[cfg(feature = "compute")]
425impl EncryptedU64 {
426 pub fn from_fhe(value: FheUint64) -> Self {
427 Self { inner: value }
428 }
429
430 pub fn inner(&self) -> &FheUint64 {
431 &self.inner
432 }
433
434 pub fn encrypt(value: u64, client_key: &tfhe::ClientKey) -> Self {
435 Self {
436 inner: FheUint64::encrypt(value, client_key),
437 }
438 }
439
440 pub fn decrypt(&self, client_key: &tfhe::ClientKey) -> u64 {
441 self.inner.decrypt(client_key)
442 }
443
444 pub fn to_cipher_blob(&self) -> Result<CipherBlob> {
445 let bytes = bincode::serialize(&self.inner).map_err(|e| {
446 AmateRSError::Serialization(ErrorContext::new(format!(
447 "Failed to serialize EncryptedU64: {}",
448 e
449 )))
450 })?;
451 Ok(CipherBlob::new(bytes))
452 }
453
454 pub fn from_cipher_blob(blob: &CipherBlob) -> Result<Self> {
455 let inner: FheUint64 = bincode::deserialize(blob.as_bytes()).map_err(|e| {
456 AmateRSError::Deserialization(ErrorContext::new(format!(
457 "Failed to deserialize EncryptedU64: {}",
458 e
459 )))
460 })?;
461 Ok(Self { inner })
462 }
463
464 pub fn add(&self, other: &Self) -> Self {
465 Self {
466 inner: &self.inner + &other.inner,
467 }
468 }
469
470 pub fn sub(&self, other: &Self) -> Self {
471 Self {
472 inner: &self.inner - &other.inner,
473 }
474 }
475
476 pub fn mul(&self, other: &Self) -> Self {
477 Self {
478 inner: &self.inner * &other.inner,
479 }
480 }
481
482 pub fn eq(&self, other: &Self) -> EncryptedBool {
483 EncryptedBool {
484 inner: self.inner.eq(&other.inner),
485 }
486 }
487
488 pub fn lt(&self, other: &Self) -> EncryptedBool {
489 EncryptedBool {
490 inner: self.inner.lt(&other.inner),
491 }
492 }
493
494 pub fn le(&self, other: &Self) -> EncryptedBool {
495 EncryptedBool {
496 inner: self.inner.le(&other.inner),
497 }
498 }
499
500 pub fn gt(&self, other: &Self) -> EncryptedBool {
501 EncryptedBool {
502 inner: self.inner.gt(&other.inner),
503 }
504 }
505
506 pub fn ge(&self, other: &Self) -> EncryptedBool {
507 EncryptedBool {
508 inner: self.inner.ge(&other.inner),
509 }
510 }
511
512 pub fn ne(&self, other: &Self) -> EncryptedBool {
513 EncryptedBool {
514 inner: self.inner.ne(&other.inner),
515 }
516 }
517}
518
519#[cfg(not(feature = "compute"))]
521#[derive(Clone, Debug)]
522pub struct EncryptedBool {
523 _phantom: std::marker::PhantomData<()>,
524}
525
526#[cfg(not(feature = "compute"))]
527impl EncryptedBool {
528 pub fn to_cipher_blob(&self) -> Result<CipherBlob> {
529 Err(AmateRSError::FeatureNotEnabled(ErrorContext::new(
530 "FHE compute feature is not enabled".to_string(),
531 )))
532 }
533
534 pub fn from_cipher_blob(_blob: &CipherBlob) -> Result<Self> {
535 Err(AmateRSError::FeatureNotEnabled(ErrorContext::new(
536 "FHE compute feature is not enabled".to_string(),
537 )))
538 }
539}
540
541#[cfg(not(feature = "compute"))]
542#[derive(Clone, Debug)]
543pub struct EncryptedU8 {
544 _phantom: std::marker::PhantomData<()>,
545}
546
547#[cfg(not(feature = "compute"))]
548impl EncryptedU8 {
549 pub fn to_cipher_blob(&self) -> Result<CipherBlob> {
550 Err(AmateRSError::FeatureNotEnabled(ErrorContext::new(
551 "FHE compute feature is not enabled".to_string(),
552 )))
553 }
554
555 pub fn from_cipher_blob(_blob: &CipherBlob) -> Result<Self> {
556 Err(AmateRSError::FeatureNotEnabled(ErrorContext::new(
557 "FHE compute feature is not enabled".to_string(),
558 )))
559 }
560}
561
562#[cfg(not(feature = "compute"))]
563#[derive(Clone, Debug)]
564pub struct EncryptedU16 {
565 _phantom: std::marker::PhantomData<()>,
566}
567
568#[cfg(not(feature = "compute"))]
569impl EncryptedU16 {
570 pub fn to_cipher_blob(&self) -> Result<CipherBlob> {
571 Err(AmateRSError::FeatureNotEnabled(ErrorContext::new(
572 "FHE compute feature is not enabled".to_string(),
573 )))
574 }
575
576 pub fn from_cipher_blob(_blob: &CipherBlob) -> Result<Self> {
577 Err(AmateRSError::FeatureNotEnabled(ErrorContext::new(
578 "FHE compute feature is not enabled".to_string(),
579 )))
580 }
581}
582
583#[cfg(not(feature = "compute"))]
584#[derive(Clone, Debug)]
585pub struct EncryptedU32 {
586 _phantom: std::marker::PhantomData<()>,
587}
588
589#[cfg(not(feature = "compute"))]
590impl EncryptedU32 {
591 pub fn to_cipher_blob(&self) -> Result<CipherBlob> {
592 Err(AmateRSError::FeatureNotEnabled(ErrorContext::new(
593 "FHE compute feature is not enabled".to_string(),
594 )))
595 }
596
597 pub fn from_cipher_blob(_blob: &CipherBlob) -> Result<Self> {
598 Err(AmateRSError::FeatureNotEnabled(ErrorContext::new(
599 "FHE compute feature is not enabled".to_string(),
600 )))
601 }
602}
603
604#[cfg(not(feature = "compute"))]
605#[derive(Clone, Debug)]
606pub struct EncryptedU64 {
607 _phantom: std::marker::PhantomData<()>,
608}
609
610#[cfg(not(feature = "compute"))]
611impl EncryptedU64 {
612 pub fn to_cipher_blob(&self) -> Result<CipherBlob> {
613 Err(AmateRSError::FeatureNotEnabled(ErrorContext::new(
614 "FHE compute feature is not enabled".to_string(),
615 )))
616 }
617
618 pub fn from_cipher_blob(_blob: &CipherBlob) -> Result<Self> {
619 Err(AmateRSError::FeatureNotEnabled(ErrorContext::new(
620 "FHE compute feature is not enabled".to_string(),
621 )))
622 }
623}
624
625#[cfg(all(test, feature = "compute"))]
626mod tests {
627 use super::*;
628 use crate::compute::keys::FheKeyPair;
629
630 #[test]
631 fn test_encrypted_bool_operations() -> Result<()> {
632 let keypair = FheKeyPair::generate()?;
633 keypair.set_as_global_server_key();
634
635 let a = EncryptedBool::encrypt(true, keypair.client_key());
636 let b = EncryptedBool::encrypt(false, keypair.client_key());
637
638 let result = a.and(&b);
640 assert!(!result.decrypt(keypair.client_key()));
641
642 let result = a.or(&b);
644 assert!(result.decrypt(keypair.client_key()));
645
646 let result = a.xor(&b);
648 assert!(result.decrypt(keypair.client_key()));
649
650 let result = a.not();
652 assert!(!result.decrypt(keypair.client_key()));
653
654 Ok(())
655 }
656
657 #[test]
658 fn test_encrypted_u8_arithmetic() -> Result<()> {
659 let keypair = FheKeyPair::generate()?;
660 keypair.set_as_global_server_key();
661
662 let a = EncryptedU8::encrypt(5, keypair.client_key());
663 let b = EncryptedU8::encrypt(3, keypair.client_key());
664
665 let result = a.add(&b);
667 assert_eq!(result.decrypt(keypair.client_key()), 8);
668
669 let result = a.sub(&b);
671 assert_eq!(result.decrypt(keypair.client_key()), 2);
672
673 let result = a.mul(&b);
675 assert_eq!(result.decrypt(keypair.client_key()), 15);
676
677 Ok(())
678 }
679
680 #[test]
681 fn test_encrypted_u8_comparison() -> Result<()> {
682 let keypair = FheKeyPair::generate()?;
683 keypair.set_as_global_server_key();
684
685 let a = EncryptedU8::encrypt(5, keypair.client_key());
686 let b = EncryptedU8::encrypt(3, keypair.client_key());
687
688 let result = a.eq(&b);
690 assert!(!result.decrypt(keypair.client_key()));
691
692 let result = a.lt(&b);
694 assert!(!result.decrypt(keypair.client_key()));
695
696 let result = a.gt(&b);
698 assert!(result.decrypt(keypair.client_key()));
699
700 let result = a.le(&b);
702 assert!(!result.decrypt(keypair.client_key()));
703
704 let result = a.ge(&b);
706 assert!(result.decrypt(keypair.client_key()));
707
708 let result = a.ne(&b);
710 assert!(result.decrypt(keypair.client_key()));
711
712 Ok(())
713 }
714
715 #[test]
716 fn test_encrypted_value_serialization() -> Result<()> {
717 let keypair = FheKeyPair::generate()?;
718 keypair.set_as_global_server_key();
719
720 let value = EncryptedU8::encrypt(42, keypair.client_key());
721 let blob = value.to_cipher_blob()?;
722
723 assert!(!blob.is_empty());
724
725 let restored = EncryptedU8::from_cipher_blob(&blob)?;
726 assert_eq!(restored.decrypt(keypair.client_key()), 42);
727
728 Ok(())
729 }
730}