phantom-zone/src/shortint/ops.rs
src/shortint/ops.rsRaw
1use std::mem::MaybeUninit;
2
3use itertools::{izip, Itertools};
4use num_traits::PrimInt;
5
6use crate::{
7 backend::ModularOpsU64,
8 bool::{
9 evaluator::{BoolEvaluator, BooleanGates, ClientKey, ServerKeyEvaluationDomain},
10 parameters::CiphertextModulus,
11 },
12 ntt::NttBackendU64,
13 random::DefaultSecureRng,
14 Decryptor,
15};
16
17pub(super) fn half_adder<E: BooleanGates>(
18 evaluator: &mut E,
19 a: &mut E::Ciphertext,
20 b: &E::Ciphertext,
21 key: &E::Key,
22) -> E::Ciphertext {
23 let carry = evaluator.and(a, b, key);
24 evaluator.xor_inplace(a, b, key);
25 carry
26}
27
28pub(super) fn full_adder_plain_carry_in<E: BooleanGates>(
29 evaluator: &mut E,
30 a: &mut E::Ciphertext,
31 b: &E::Ciphertext,
32 carry_in: bool,
33 key: &E::Key,
34) -> E::Ciphertext {
35 let mut a_and_b = evaluator.and(a, b, key);
36 evaluator.xor_inplace(a, b, key); //a = a ^ b
37 if carry_in {
38 // a_and_b = A & B | ((A^B) & C_in={True})
39 evaluator.or_inplace(&mut a_and_b, &a, key);
40 } else {
41 // a_and_b = A & B | ((A^B) & C_in={False})
42 // a_and_b = A & B
43 // noop
44 }
45
46 // In xor if a input is 0, output equals the firt variable. If input is 1 then
47 // output equals !(first variable)
48 if carry_in {
49 // (A^B)^1 = !(A^B)
50 evaluator.not_inplace(a);
51 } else {
52 // (A^B)^0
53 // no-op
54 }
55 a_and_b
56}
57
58pub(super) fn full_adder<E: BooleanGates>(
59 evaluator: &mut E,
60 a: &mut E::Ciphertext,
61 b: &E::Ciphertext,
62 carry_in: &E::Ciphertext,
63 key: &E::Key,
64) -> E::Ciphertext {
65 let mut a_and_b = evaluator.and(a, b, key);
66 evaluator.xor_inplace(a, b, key); //a = a ^ b
67 let a_xor_b_and_c = evaluator.and(&a, carry_in, key);
68 evaluator.or_inplace(&mut a_and_b, &a_xor_b_and_c, key); // a_and_b = A & B | ((A^B) & C_in)
69 evaluator.xor_inplace(a, &carry_in, key);
70 a_and_b
71}
72
73pub(super) fn arbitrary_bit_adder<E: BooleanGates>(
74 evaluator: &mut E,
75 a: &mut [E::Ciphertext],
76 b: &[E::Ciphertext],
77 carry_in: bool,
78 key: &E::Key,
79) -> (E::Ciphertext, E::Ciphertext)
80where
81 E::Ciphertext: Clone,
82{
83 assert!(a.len() == b.len());
84 let n = a.len();
85
86 let mut carry = if !carry_in {
87 half_adder(evaluator, &mut a[0], &b[0], key)
88 } else {
89 full_adder_plain_carry_in(evaluator, &mut a[0], &b[0], true, key)
90 };
91
92 izip!(a.iter_mut(), b.iter())
93 .skip(1)
94 .take(n - 3)
95 .for_each(|(a_bit, b_bit)| {
96 carry = full_adder(evaluator, a_bit, b_bit, &carry, key);
97 });
98
99 let carry_last_last = full_adder(evaluator, &mut a[n - 2], &b[n - 2], &carry, key);
100 let carry_last = full_adder(evaluator, &mut a[n - 1], &b[n - 1], &carry_last_last, key);
101
102 (carry_last, carry_last_last)
103}
104
105pub(super) fn arbitrary_bit_subtractor<E: BooleanGates>(
106 evaluator: &mut E,
107 a: &[E::Ciphertext],
108 b: &[E::Ciphertext],
109 key: &E::Key,
110) -> (Vec<E::Ciphertext>, E::Ciphertext, E::Ciphertext)
111where
112 E::Ciphertext: Clone,
113{
114 let mut neg_b: Vec<E::Ciphertext> = b.iter().map(|v| evaluator.not(v)).collect();
115 let (carry_last, carry_last_last) = arbitrary_bit_adder(evaluator, &mut neg_b, &a, true, key);
116 return (neg_b, carry_last, carry_last_last);
117}
118
119pub(super) fn bit_mux<E: BooleanGates>(
120 evaluator: &mut E,
121 selector: E::Ciphertext,
122 if_true: &E::Ciphertext,
123 if_false: &E::Ciphertext,
124 key: &E::Key,
125) -> E::Ciphertext {
126 // (s&a) | ((1-s)^b)
127 let not_selector = evaluator.not(&selector);
128
129 let s_and_a = evaluator.and(&selector, if_true, key);
130 let s_and_b = evaluator.and(¬_selector, if_false, key);
131 evaluator.or(&s_and_a, &s_and_b, key)
132}
133
134pub(super) fn arbitrary_bit_mux<E: BooleanGates>(
135 evaluator: &mut E,
136 selector: &E::Ciphertext,
137 if_true: &[E::Ciphertext],
138 if_false: &[E::Ciphertext],
139 key: &E::Key,
140) -> Vec<E::Ciphertext> {
141 // (s&a) | ((1-s)^b)
142 let not_selector = evaluator.not(&selector);
143
144 izip!(if_true.iter(), if_false.iter())
145 .map(|(a, b)| {
146 let s_and_a = evaluator.and(&selector, a, key);
147 let s_and_b = evaluator.and(¬_selector, b, key);
148 evaluator.or(&s_and_a, &s_and_b, key)
149 })
150 .collect()
151}
152
153pub(super) fn eight_bit_mul<E: BooleanGates>(
154 evaluator: &mut E,
155 a: &[E::Ciphertext],
156 b: &[E::Ciphertext],
157 key: &E::Key,
158) -> Vec<E::Ciphertext> {
159 assert!(a.len() == 8);
160 assert!(b.len() == 8);
161 let mut carries = Vec::with_capacity(7);
162 let mut out = Vec::with_capacity(8);
163
164 for i in (0..8) {
165 if i == 0 {
166 let s = evaluator.and(&a[0], &b[0], key);
167 out.push(s);
168 } else if i == 1 {
169 let mut tmp0 = evaluator.and(&a[1], &b[0], key);
170 let tmp1 = evaluator.and(&a[0], &b[1], key);
171 let carry = half_adder(evaluator, &mut tmp0, &tmp1, key);
172 carries.push(carry);
173 out.push(tmp0);
174 } else {
175 let mut sum = {
176 let mut sum = evaluator.and(&a[i], &b[0], key);
177 let tmp = evaluator.and(&a[i - 1], &b[1], key);
178 carries[0] = full_adder(evaluator, &mut sum, &tmp, &carries[0], key);
179 sum
180 };
181
182 for j in 2..i {
183 let tmp = evaluator.and(&a[i - j], &b[j], key);
184 carries[j - 1] = full_adder(evaluator, &mut sum, &tmp, &carries[j - 1], key);
185 }
186
187 let tmp = evaluator.and(&a[0], &b[i], key);
188 let carry = half_adder(evaluator, &mut sum, &tmp, key);
189 carries.push(carry);
190
191 out.push(sum)
192 }
193 debug_assert!(carries.len() <= 7);
194 }
195
196 out
197}
198
199pub(super) fn arbitrary_bit_division_for_quotient_and_rem<E: BooleanGates>(
200 evaluator: &mut E,
201 a: &[E::Ciphertext],
202 b: &[E::Ciphertext],
203 key: &E::Key,
204) -> (Vec<E::Ciphertext>, Vec<E::Ciphertext>)
205where
206 E::Ciphertext: Clone,
207{
208 let n = a.len();
209 let neg_b = b.iter().map(|v| evaluator.not(v)).collect_vec();
210
211 // Both remainder and quotient are initially stored in Big-endian in contract to
212 // the usual little endian we use. This is more friendly to vec pushes in
213 // division. After computing remainder and quotient, we simply reverse the
214 // vectors.
215 let mut remainder = vec![];
216 let mut quotient = vec![];
217 for i in 0..n {
218 // left shift
219 remainder.push(a[n - 1 - i].clone());
220
221 let mut subtract = remainder.clone();
222
223 // subtraction
224 // At i^th iteration remainder is only filled with i bits and the rest of the
225 // bits are zero. For example, at i = 1
226 // 0 0 0 0 0 0 X X => remainder
227 // - Y Y Y Y Y Y Y Y => divisor .
228 // --------------- .
229 // Z Z Z Z Z Z Z Z => result
230 // For the next iteration we only care about result if divisor is <= remainder
231 // (which implies result <= remainder). Otherwise we care about remainder
232 // (recall re-storing division). Hence we optimise subtraction and
233 // ignore full adders for places where remainder bits are known to be false
234 // bits. We instead use `ANDs` to compute the carry overs, since the
235 // last carry over indicates whether the value has overflown (i.e. divisor <=
236 // remainder). Last carry out is `true` if value has not overflown, otherwise
237 // false.
238 let mut carry =
239 full_adder_plain_carry_in(evaluator, &mut subtract[i], &neg_b[0], true, key);
240 for j in 1..i + 1 {
241 carry = full_adder(evaluator, &mut subtract[i - j], &neg_b[j], &carry, key);
242 }
243 for j in i + 1..n {
244 // All I care about are the carries
245 evaluator.and_inplace(&mut carry, &neg_b[j], key);
246 }
247
248 let not_carry = evaluator.not(&carry);
249 // Choose `remainder` if subtraction has overflown (i.e. carry = false).
250 // Otherwise choose `subtractor`.
251 //
252 // mux k^a | !(k)^b, where k is the selector.
253 izip!(remainder.iter_mut(), subtract.iter_mut()).for_each(|(r, s)| {
254 // choose `s` when carry is true, otherwise choose r
255 evaluator.and_inplace(s, &carry, key);
256 evaluator.and_inplace(r, ¬_carry, key);
257 evaluator.or_inplace(r, s, key);
258 });
259
260 // Set i^th MSB of quotient to 1 if carry = true, otherwise set it to 0.
261 // X&1 | X&0 => X&1 => X
262 quotient.push(carry);
263 }
264
265 remainder.reverse();
266 quotient.reverse();
267
268 (quotient, remainder)
269}
270
271fn is_zero<E: BooleanGates>(evaluator: &mut E, a: &[E::Ciphertext], key: &E::Key) -> E::Ciphertext {
272 let mut a = a.iter().map(|v| evaluator.not(v)).collect_vec();
273 let (out, rest_a) = a.split_at_mut(1);
274 rest_a.iter().for_each(|c| {
275 evaluator.and_inplace(&mut out[0], c, key);
276 });
277 return a.remove(0);
278}
279
280fn arbitrary_bit_equality<E: BooleanGates>(
281 evaluator: &mut E,
282 a: &[E::Ciphertext],
283 b: &[E::Ciphertext],
284 key: &E::Key,
285) -> E::Ciphertext {
286 assert!(a.len() == b.len());
287 let mut out = evaluator.and(&a[0], &b[0], key);
288 izip!(a.iter(), b.iter()).skip(1).for_each(|(abit, bbit)| {
289 let e = evaluator.xnor(abit, bbit, key);
290 evaluator.and(&mut out, &e, key);
291 });
292 return out;
293}
294
295/// Comaprator handle computes comparator result 2ns MSB onwards. It is
296/// separated because comparator subroutine for signed and unsgind integers
297/// differs only for 1st MSB and is common second MSB onwards
298fn _comparator_handler_from_second_msb<E: BooleanGates>(
299 evaluator: &mut E,
300 a: &[E::Ciphertext],
301 b: &[E::Ciphertext],
302 mut comp: E::Ciphertext,
303 mut casc: E::Ciphertext,
304 key: &E::Key,
305) -> E::Ciphertext {
306 let n = a.len();
307
308 // handle MSB - 1
309 let mut tmp = evaluator.not(&b[n - 2]);
310 evaluator.and(&mut tmp, &a[n - 2], key);
311 evaluator.and(&mut tmp, &casc, key);
312 evaluator.or(&mut comp, &tmp, key);
313
314 for i in 2..n {
315 // calculate cascading bit
316 let tmp_casc = evaluator.xnor(&a[n - 2 - i], &b[n - 2 - i], key);
317 evaluator.and(&mut casc, &tmp_casc, key);
318
319 // calculate computate bit
320 let mut tmp = evaluator.not(&b[n - 1 - i]);
321 evaluator.and(&mut tmp, &a[n - 1 - i], key);
322 evaluator.and(&mut tmp, &casc, key);
323 evaluator.or(&mut comp, &tmp, key);
324 }
325
326 return comp;
327}
328
329/// Signed integer comparison is same as unsigned integer with MSB flipped.
330fn arbitrary_signed_bit_comparator<E: BooleanGates>(
331 evaluator: &mut E,
332 a: &[E::Ciphertext],
333 b: &[E::Ciphertext],
334 key: &E::Key,
335) -> E::Ciphertext {
336 assert!(a.len() == b.len());
337 let n = a.len();
338
339 // handle MSB
340 let mut comp = evaluator.not(&a[n - 1]);
341 evaluator.and(&mut comp, &b[n - 1], key); // comp
342 let casc = evaluator.xnor(&a[n - 1], &b[n - 1], key); // casc
343
344 return _comparator_handler_from_second_msb(evaluator, a, b, comp, casc, key);
345}
346
347fn arbitrary_bit_comparator<E: BooleanGates>(
348 evaluator: &mut E,
349 a: &[E::Ciphertext],
350 b: &[E::Ciphertext],
351 key: &E::Key,
352) -> E::Ciphertext {
353 assert!(a.len() == b.len());
354 let n = a.len();
355
356 // handle MSB
357 let mut comp = evaluator.not(&b[n - 1]);
358 evaluator.and(&mut comp, &a[n - 1], key);
359 let casc = evaluator.xnor(&a[n - 1], &b[n - 1], key);
360
361 return _comparator_handler_from_second_msb(evaluator, a, b, comp, casc, key);
362}