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