amaters_core/compute/
operations.rs

1//! FHE operations module
2//!
3//! This module provides encrypted data types and operations for TFHE-based computations.
4//! All operations are performed on encrypted data without decryption.
5
6use 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/// Encrypted boolean value
15#[cfg(feature = "compute")]
16#[derive(Clone)]
17pub struct EncryptedBool {
18    inner: FheBool,
19}
20
21#[cfg(feature = "compute")]
22impl EncryptedBool {
23    /// Create from TFHE FheBool
24    pub fn from_fhe(value: FheBool) -> Self {
25        Self { inner: value }
26    }
27
28    /// Get inner FheBool
29    pub fn inner(&self) -> &FheBool {
30        &self.inner
31    }
32
33    /// Encrypt a boolean value
34    pub fn encrypt(value: bool, client_key: &tfhe::ClientKey) -> Self {
35        Self {
36            inner: FheBool::encrypt(value, client_key),
37        }
38    }
39
40    /// Decrypt to boolean
41    pub fn decrypt(&self, client_key: &tfhe::ClientKey) -> bool {
42        self.inner.decrypt(client_key)
43    }
44
45    /// Serialize to CipherBlob
46    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    /// Deserialize from CipherBlob
57    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    /// Logical AND operation
68    pub fn and(&self, other: &Self) -> Self {
69        Self {
70            inner: &self.inner & &other.inner,
71        }
72    }
73
74    /// Logical OR operation
75    pub fn or(&self, other: &Self) -> Self {
76        Self {
77            inner: &self.inner | &other.inner,
78        }
79    }
80
81    /// Logical XOR operation
82    pub fn xor(&self, other: &Self) -> Self {
83        Self {
84            inner: &self.inner ^ &other.inner,
85        }
86    }
87
88    /// Logical NOT operation
89    pub fn not(&self) -> Self {
90        Self {
91            inner: !&self.inner,
92        }
93    }
94}
95
96/// Encrypted unsigned 8-bit integer
97#[cfg(feature = "compute")]
98#[derive(Clone)]
99pub struct EncryptedU8 {
100    inner: FheUint8,
101}
102
103#[cfg(feature = "compute")]
104impl EncryptedU8 {
105    /// Create from TFHE FheUint8
106    pub fn from_fhe(value: FheUint8) -> Self {
107        Self { inner: value }
108    }
109
110    /// Get inner FheUint8
111    pub fn inner(&self) -> &FheUint8 {
112        &self.inner
113    }
114
115    /// Encrypt a u8 value
116    pub fn encrypt(value: u8, client_key: &tfhe::ClientKey) -> Self {
117        Self {
118            inner: FheUint8::encrypt(value, client_key),
119        }
120    }
121
122    /// Decrypt to u8
123    pub fn decrypt(&self, client_key: &tfhe::ClientKey) -> u8 {
124        self.inner.decrypt(client_key)
125    }
126
127    /// Serialize to CipherBlob
128    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    /// Deserialize from CipherBlob
139    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    /// Add two encrypted values
150    pub fn add(&self, other: &Self) -> Self {
151        Self {
152            inner: &self.inner + &other.inner,
153        }
154    }
155
156    /// Subtract two encrypted values
157    pub fn sub(&self, other: &Self) -> Self {
158        Self {
159            inner: &self.inner - &other.inner,
160        }
161    }
162
163    /// Multiply two encrypted values
164    pub fn mul(&self, other: &Self) -> Self {
165        Self {
166            inner: &self.inner * &other.inner,
167        }
168    }
169
170    /// Compare equality
171    pub fn eq(&self, other: &Self) -> EncryptedBool {
172        EncryptedBool {
173            inner: self.inner.eq(&other.inner),
174        }
175    }
176
177    /// Compare less than
178    pub fn lt(&self, other: &Self) -> EncryptedBool {
179        EncryptedBool {
180            inner: self.inner.lt(&other.inner),
181        }
182    }
183
184    /// Compare less than or equal
185    pub fn le(&self, other: &Self) -> EncryptedBool {
186        EncryptedBool {
187            inner: self.inner.le(&other.inner),
188        }
189    }
190
191    /// Compare greater than
192    pub fn gt(&self, other: &Self) -> EncryptedBool {
193        EncryptedBool {
194            inner: self.inner.gt(&other.inner),
195        }
196    }
197
198    /// Compare greater than or equal
199    pub fn ge(&self, other: &Self) -> EncryptedBool {
200        EncryptedBool {
201            inner: self.inner.ge(&other.inner),
202        }
203    }
204
205    /// Compare not equal
206    pub fn ne(&self, other: &Self) -> EncryptedBool {
207        EncryptedBool {
208            inner: self.inner.ne(&other.inner),
209        }
210    }
211}
212
213/// Encrypted unsigned 16-bit integer
214#[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/// Encrypted unsigned 32-bit integer
316#[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/// Encrypted unsigned 64-bit integer
418#[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/// Stub implementations when compute feature is disabled
520#[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        // Test AND
639        let result = a.and(&b);
640        assert!(!result.decrypt(keypair.client_key()));
641
642        // Test OR
643        let result = a.or(&b);
644        assert!(result.decrypt(keypair.client_key()));
645
646        // Test XOR
647        let result = a.xor(&b);
648        assert!(result.decrypt(keypair.client_key()));
649
650        // Test NOT
651        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        // Test addition
666        let result = a.add(&b);
667        assert_eq!(result.decrypt(keypair.client_key()), 8);
668
669        // Test subtraction
670        let result = a.sub(&b);
671        assert_eq!(result.decrypt(keypair.client_key()), 2);
672
673        // Test multiplication
674        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        // Test equality
689        let result = a.eq(&b);
690        assert!(!result.decrypt(keypair.client_key()));
691
692        // Test less than
693        let result = a.lt(&b);
694        assert!(!result.decrypt(keypair.client_key()));
695
696        // Test greater than
697        let result = a.gt(&b);
698        assert!(result.decrypt(keypair.client_key()));
699
700        // Test less than or equal
701        let result = a.le(&b);
702        assert!(!result.decrypt(keypair.client_key()));
703
704        // Test greater than or equal
705        let result = a.ge(&b);
706        assert!(result.decrypt(keypair.client_key()));
707
708        // Test not equal
709        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}