phantom-zone/src/shortint/ops.rs
src/shortint/ops.rsRaw
1use itertools::{izip, Itertools};
2
3use crate::bool::evaluator::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 s_and_a = evaluator.and(&selector, if_true, key);
118 let s_and_b = evaluator.and(¬_selector, if_false, key);
119 evaluator.or(&s_and_a, &s_and_b, key)
120}
121
122pub(super) fn arbitrary_bit_mux<E: BooleanGates>(
123 evaluator: &mut E,
124 selector: &E::Ciphertext,
125 if_true: &[E::Ciphertext],
126 if_false: &[E::Ciphertext],
127 key: &E::Key,
128) -> Vec<E::Ciphertext> {
129 // (s&a) | ((1-s)^b)
130 let not_selector = evaluator.not(&selector);
131
132 izip!(if_true.iter(), if_false.iter())
133 .map(|(a, b)| {
134 let s_and_a = evaluator.and(&selector, a, key);
135 let s_and_b = evaluator.and(¬_selector, b, key);
136 evaluator.or(&s_and_a, &s_and_b, key)
137 })
138 .collect()
139}
140
141pub(super) fn eight_bit_mul<E: BooleanGates>(
142 evaluator: &mut E,
143 a: &[E::Ciphertext],
144 b: &[E::Ciphertext],
145 key: &E::Key,
146) -> Vec<E::Ciphertext> {
147 assert!(a.len() == 8);
148 assert!(b.len() == 8);
149 let mut carries = Vec::with_capacity(7);
150 let mut out = Vec::with_capacity(8);
151
152 for i in (0..8) {
153 if i == 0 {
154 let s = evaluator.and(&a[0], &b[0], key);
155 out.push(s);
156 } else if i == 1 {
157 let mut tmp0 = evaluator.and(&a[1], &b[0], key);
158 let tmp1 = evaluator.and(&a[0], &b[1], key);
159 let carry = half_adder(evaluator, &mut tmp0, &tmp1, key);
160 carries.push(carry);
161 out.push(tmp0);
162 } else {
163 let mut sum = {
164 let mut sum = evaluator.and(&a[i], &b[0], key);
165 let tmp = evaluator.and(&a[i - 1], &b[1], key);
166 carries[0] = full_adder(evaluator, &mut sum, &tmp, &carries[0], key);
167 sum
168 };
169
170 for j in 2..i {
171 let tmp = evaluator.and(&a[i - j], &b[j], key);
172 carries[j - 1] = full_adder(evaluator, &mut sum, &tmp, &carries[j - 1], key);
173 }
174
175 let tmp = evaluator.and(&a[0], &b[i], key);
176 let carry = half_adder(evaluator, &mut sum, &tmp, key);
177 carries.push(carry);
178
179 out.push(sum)
180 }
181 debug_assert!(carries.len() <= 7);
182 }
183
184 out
185}
186
187pub(super) fn arbitrary_bit_division_for_quotient_and_rem<E: BooleanGates>(
188 evaluator: &mut E,
189 a: &[E::Ciphertext],
190 b: &[E::Ciphertext],
191 key: &E::Key,
192) -> (Vec<E::Ciphertext>, Vec<E::Ciphertext>)
193where
194 E::Ciphertext: Clone,
195{
196 let n = a.len();
197 let neg_b = b.iter().map(|v| evaluator.not(v)).collect_vec();
198
199 // Both remainder and quotient are initially stored in Big-endian in contrast to
200 // the usual little endian we use. This is more friendly to vec pushes in
201 // division. After computing remainder and quotient, we simply reverse the
202 // vectors.
203 let mut remainder = vec![];
204 let mut quotient = vec![];
205 for i in 0..n {
206 // left shift
207 remainder.push(a[n - 1 - i].clone());
208
209 let mut subtract = remainder.clone();
210
211 // subtraction
212 // At i^th iteration remainder is only filled with i bits and the rest of the
213 // bits are zero. For example, at i = 1
214 // 0 0 0 0 0 0 X X => remainder
215 // - Y Y Y Y Y Y Y Y => divisor .
216 // --------------- .
217 // Z Z Z Z Z Z Z Z => result
218 // For the next iteration we only care about result if divisor is <= remainder
219 // (which implies result <= remainder). Otherwise we care about remainder
220 // (recall re-storing division). Hence we optimise subtraction and
221 // ignore full adders for places where remainder bits are known to be false
222 // bits. We instead use `ANDs` to compute the carry overs, since the
223 // last carry over indicates whether the value has overflown (i.e. divisor <=
224 // remainder). Last carry out is `true` if value has not overflown, otherwise
225 // false.
226 let mut carry =
227 full_adder_plain_carry_in(evaluator, &mut subtract[i], &neg_b[0], true, key);
228 for j in 1..i + 1 {
229 carry = full_adder(evaluator, &mut subtract[i - j], &neg_b[j], &carry, key);
230 }
231 for j in i + 1..n {
232 // All I care about are the carries
233 evaluator.and_inplace(&mut carry, &neg_b[j], key);
234 }
235
236 let not_carry = evaluator.not(&carry);
237 // Choose `remainder` if subtraction has overflown (i.e. carry = false).
238 // Otherwise choose `subtractor`.
239 //
240 // mux k^a | !(k)^b, where k is the selector.
241 izip!(remainder.iter_mut(), subtract.iter_mut()).for_each(|(r, s)| {
242 // choose `s` when carry is true, otherwise choose r
243 evaluator.and_inplace(s, &carry, key);
244 evaluator.and_inplace(r, ¬_carry, key);
245 evaluator.or_inplace(r, s, key);
246 });
247
248 // Set i^th MSB of quotient to 1 if carry = true, otherwise set it to 0.
249 // X&1 | X&0 => X&1 => X
250 quotient.push(carry);
251 }
252
253 remainder.reverse();
254 quotient.reverse();
255
256 (quotient, remainder)
257}
258
259fn is_zero<E: BooleanGates>(evaluator: &mut E, a: &[E::Ciphertext], key: &E::Key) -> E::Ciphertext {
260 let mut a = a.iter().map(|v| evaluator.not(v)).collect_vec();
261 let (out, rest_a) = a.split_at_mut(1);
262 rest_a.iter().for_each(|c| {
263 evaluator.and_inplace(&mut out[0], c, key);
264 });
265 return a.remove(0);
266}
267
268pub(super) fn arbitrary_bit_equality<E: BooleanGates>(
269 evaluator: &mut E,
270 a: &[E::Ciphertext],
271 b: &[E::Ciphertext],
272 key: &E::Key,
273) -> E::Ciphertext {
274 assert!(a.len() == b.len());
275 let mut out = evaluator.xnor(&a[0], &b[0], key);
276 izip!(a.iter(), b.iter()).skip(1).for_each(|(abit, bbit)| {
277 let e = evaluator.xnor(abit, bbit, key);
278 evaluator.and_inplace(&mut out, &e, key);
279 });
280 return out;
281}
282
283/// Comparator handle computes comparator result 2ns MSB onwards. It is
284/// separated because comparator subroutine for signed and unsgind integers
285/// differs only for 1st MSB and is common second MSB onwards
286fn _comparator_handler_from_second_msb<E: BooleanGates>(
287 evaluator: &mut E,
288 a: &[E::Ciphertext],
289 b: &[E::Ciphertext],
290 mut comp: E::Ciphertext,
291 mut casc: E::Ciphertext,
292 key: &E::Key,
293) -> E::Ciphertext {
294 let n = a.len();
295
296 // handle MSB - 1
297 let mut tmp = evaluator.not(&b[n - 2]);
298 evaluator.and_inplace(&mut tmp, &a[n - 2], key);
299 evaluator.and_inplace(&mut tmp, &casc, key);
300 evaluator.or_inplace(&mut comp, &tmp, key);
301
302 for i in 2..n {
303 // calculate cascading bit
304 let tmp_casc = evaluator.xnor(&a[n - i], &b[n - i], key);
305 evaluator.and_inplace(&mut casc, &tmp_casc, key);
306
307 // calculate computate bit
308 let mut tmp = evaluator.not(&b[n - 1 - i]);
309 evaluator.and_inplace(&mut tmp, &a[n - 1 - i], key);
310 evaluator.and_inplace(&mut tmp, &casc, key);
311 evaluator.or_inplace(&mut comp, &tmp, key);
312 }
313
314 return comp;
315}
316
317/// Signed integer comparison is same as unsigned integer with MSB flipped.
318pub(super) fn arbitrary_signed_bit_comparator<E: BooleanGates>(
319 evaluator: &mut E,
320 a: &[E::Ciphertext],
321 b: &[E::Ciphertext],
322 key: &E::Key,
323) -> E::Ciphertext {
324 assert!(a.len() == b.len());
325 let n = a.len();
326
327 // handle MSB
328 let mut comp = evaluator.not(&a[n - 1]);
329 evaluator.and_inplace(&mut comp, &b[n - 1], key); // comp
330 let casc = evaluator.xnor(&a[n - 1], &b[n - 1], key); // casc
331
332 return _comparator_handler_from_second_msb(evaluator, a, b, comp, casc, key);
333}
334
335pub(super) fn arbitrary_bit_comparator<E: BooleanGates>(
336 evaluator: &mut E,
337 a: &[E::Ciphertext],
338 b: &[E::Ciphertext],
339 key: &E::Key,
340) -> E::Ciphertext {
341 assert!(a.len() == b.len());
342 let n = a.len();
343
344 // handle MSB
345 let mut comp = evaluator.not(&b[n - 1]);
346 evaluator.and_inplace(&mut comp, &a[n - 1], key);
347 let casc = evaluator.xnor(&a[n - 1], &b[n - 1], key);
348
349 return _comparator_handler_from_second_msb(evaluator, a, b, comp, casc, key);
350}