rs_tfhe 0.1.1

A high-performance Rust implementation of TFHE (Torus Fully Homomorphic Encryption) with advanced programmable bootstrapping capabilities
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156

use crate::fft::FFT_PLAN;
use crate::params;
use crate::params::Torus;
use crate::params::TORUS_SIZE;
use crate::params::ZERO_TORUS;
use crate::tlwe;
use crate::trgsw;
use crate::trlwe;
use crate::utils;
use rand::Rng;

const TRGSWLV1_N: usize = params::trgsw_lv1::N;
const TRGSWLV1_IKS_T: usize = params::trgsw_lv1::IKS_T;
const TRGSWLV1_BASE: usize = 1 << params::trgsw_lv1::BASEBIT;

pub type SecretKeyLv0 = [Torus; params::tlwe_lv0::N];
pub type SecretKeyLv1 = [Torus; params::tlwe_lv1::N];
pub type KeySwitchingKey = Vec<tlwe::TLWELv0>;
pub type BootstrappingKey = Vec<trgsw::TRGSWLv1FFT>;

pub struct SecretKey {
  pub key_lv0: SecretKeyLv0,
  pub key_lv1: SecretKeyLv1,
}

impl Default for SecretKey {
    fn default() -> Self {
        Self::new()
    }
}

impl SecretKey {
  pub fn new() -> Self {
    let mut rng = rand::thread_rng();
    let mut key = SecretKey {
      key_lv0: [ZERO_TORUS; params::tlwe_lv0::N],
      key_lv1: [ZERO_TORUS; params::tlwe_lv1::N],
    };
    key
      .key_lv0
      .iter_mut()
      .for_each(|e| *e = rng.gen::<bool>() as Torus);
    key
      .key_lv1
      .iter_mut()
      .for_each(|e| *e = rng.gen::<bool>() as Torus);
    key
  }
}

pub struct CloudKey {
  pub decomposition_offset: Torus,
  pub blind_rotate_testvec: trlwe::TRLWELv1,
  pub key_switching_key: KeySwitchingKey,
  pub bootstrapping_key: BootstrappingKey,
}

impl CloudKey {
  pub fn new(secret_key: &SecretKey) -> Self {
    CloudKey {
      decomposition_offset: gen_decomposition_offset(),
      blind_rotate_testvec: gen_testvec(),
      key_switching_key: gen_key_switching_key(secret_key),
      bootstrapping_key: gen_bootstrapping_key(secret_key),
    }
  }

  pub fn new_no_ksk() -> Self {
    CloudKey {
      decomposition_offset: gen_decomposition_offset(),
      blind_rotate_testvec: gen_testvec(),
      key_switching_key: vec![tlwe::TLWELv0::new(); TRGSWLV1_BASE * TRGSWLV1_IKS_T * TRGSWLV1_N],
      bootstrapping_key: vec![trgsw::TRGSWLv1FFT::new_dummy(); params::tlwe_lv0::N],
    }
  }
}

pub fn gen_decomposition_offset() -> Torus {
  let mut offset: Torus = 0;

  for i in 0..(params::trgsw_lv1::L as Torus) {
    offset = offset.wrapping_add(
      params::trgsw_lv1::BG / 2
        * (1 << (TORUS_SIZE - (i + 1) as usize * params::trgsw_lv1::BGBIT as usize)),
    );
  }

  offset
}

pub fn gen_testvec() -> trlwe::TRLWELv1 {
  let mut testvec = trlwe::TRLWELv1::new();
  let b_torus = utils::f64_to_torus(0.125);
  for i in 0..params::trgsw_lv1::N {
    testvec.a[i] = 0;
    testvec.b[i] = b_torus;
  }

  testvec
}

pub fn gen_key_switching_key(secret_key: &SecretKey) -> KeySwitchingKey {
  const BASEBIT: usize = params::trgsw_lv1::BASEBIT;
  const IKS_T: usize = params::trgsw_lv1::IKS_T;
  let mut res = vec![tlwe::TLWELv0::new(); TRGSWLV1_BASE * TRGSWLV1_IKS_T * TRGSWLV1_N];

  for i in 0..params::trgsw_lv1::N {
    for j in 0..IKS_T {
      for k in 0..TRGSWLV1_BASE {
        if k == 0 {
          continue;
        }
        let p =
          ((k as u32 * secret_key.key_lv1[i]) as f64) / ((1 << ((j + 1) * BASEBIT)) as f64);
        let idx = (TRGSWLV1_BASE * TRGSWLV1_IKS_T * i) + (TRGSWLV1_BASE * j) + k;
        res[idx] = tlwe::TLWELv0::encrypt_f64(p, params::KSK_ALPHA, &secret_key.key_lv0);
      }
    }
  }

  res
}

pub fn gen_bootstrapping_key(secret_key: &SecretKey) -> BootstrappingKey {
  gen_bootstrapping_key_with_railgun(secret_key, crate::parallel::default_railgun())
}

pub fn gen_bootstrapping_key_with_railgun<R: crate::parallel::Railgun>(
  secret_key: &SecretKey,
  railgun: &R,
) -> BootstrappingKey {
  // OPTIMIZATION: Parallel bootstrap key generation
  // The bootstrap key has 500-700 independent TRGSW encryptions (one per LWE coefficient)
  // Each encryption is completely independent and can be parallelized
  // Expected speedup: 4-8x on multi-core systems
  //
  // Each worker thread uses its own thread-local FFT_PLAN for efficient caching
  //
  // NOTE: Threads need larger stack (8MB) due to deep call chains in TRGSW encryption
  let config = crate::parallel::ParallelConfig {
    stack_size: Some(8 * 1024 * 1024), // 8MB stack per thread
    num_threads: None,
  };

  railgun.with_config(config, || {
    railgun.par_map(&secret_key.key_lv0, |&kval| {
      FFT_PLAN.with(|plan| {
        let p = &mut plan.borrow_mut();
        trgsw::TRGSWLv1FFT::new(
          &trgsw::TRGSWLv1::encrypt_torus(kval, params::BSK_ALPHA, &secret_key.key_lv1, p),
          p,
        )
      })
    })
  })
}