From 10040b83d0bc5620a90fc99b98cba7adb54b103a Mon Sep 17 00:00:00 2001
From: Andy Leiserson <aleiserson@mozilla.com>
Date: Thu, 4 Apr 2024 13:46:57 -0700
Subject: [PATCH] Add BooleanProtocols rollup trait

---
 ipa-core/src/protocol/basics/mod.rs           | 21 +++++++++++++++++--
 .../boolean_ops/addition_sequential.rs        | 17 +++++++++------
 .../comparison_and_subtraction_sequential.rs  | 16 +++++++-------
 .../boolean_ops/share_conversion_aby.rs       | 18 +++++++---------
 ipa-core/src/protocol/ipa_prf/prf_eval.rs     |  7 +++----
 5 files changed, 48 insertions(+), 31 deletions(-)

diff --git a/ipa-core/src/protocol/basics/mod.rs b/ipa-core/src/protocol/basics/mod.rs
index f7536d606..5c4cb7929 100644
--- a/ipa-core/src/protocol/basics/mod.rs
+++ b/ipa-core/src/protocol/basics/mod.rs
@@ -8,6 +8,8 @@ mod reveal;
 mod share_known_value;
 pub mod sum_of_product;
 
+use std::ops::Not;
+
 #[cfg(feature = "descriptive-gate")]
 pub use check_zero::check_zero;
 pub use if_else::{if_else, select};
@@ -18,8 +20,8 @@ pub use share_known_value::ShareKnownValue;
 pub use sum_of_product::SumOfProducts;
 
 use crate::{
-    ff::Field,
-    protocol::context::Context,
+    ff::{boolean::Boolean, Field},
+    protocol::{context::Context, ipa_prf::PRF_CHUNK},
     secret_sharing::{
         replicated::semi_honest::AdditiveShare, SecretSharing, SharedValue, Vectorizable,
     },
@@ -42,8 +44,23 @@ pub trait BasicProtocols<C: Context, V: SharedValue + Vectorizable<N>, const N:
 {
 }
 
+pub trait BooleanProtocols<C: Context, V: SharedValue + Vectorizable<N>, const N: usize = 1>:
+    SecretSharing<V>
+    + Reveal<C, N, Output = <V as Vectorizable<N>>::Array>
+    + SecureMul<C>
+    + Not<Output = Self>
+{
+}
+
 impl<C: Context, F: Field> BasicProtocols<C, F> for AdditiveShare<F> {}
 
+impl<C: Context> BooleanProtocols<C, Boolean, 1> for AdditiveShare<Boolean> {}
+
+impl<C: Context> BooleanProtocols<C, Boolean, PRF_CHUNK> for AdditiveShare<Boolean, PRF_CHUNK> {}
+
+// Used by semi_honest_compare_gt_vec test.
+impl<C: Context> BooleanProtocols<C, Boolean, 256> for AdditiveShare<Boolean, 256> {}
+
 #[cfg(feature = "descriptive-gate")]
 impl<'a, F: ExtendableField> BasicProtocols<UpgradedMaliciousContext<'a, F>, F>
     for MaliciousReplicated<F>
diff --git a/ipa-core/src/protocol/ipa_prf/boolean_ops/addition_sequential.rs b/ipa-core/src/protocol/ipa_prf/boolean_ops/addition_sequential.rs
index b4dc09bce..1cd04e2b9 100644
--- a/ipa-core/src/protocol/ipa_prf/boolean_ops/addition_sequential.rs
+++ b/ipa-core/src/protocol/ipa_prf/boolean_ops/addition_sequential.rs
@@ -1,4 +1,4 @@
-use std::{borrow::Borrow, iter::repeat, ops::Not};
+use std::{borrow::Borrow, iter::repeat};
 
 #[cfg(all(test, unit_test))]
 use ipa_macros::Step;
@@ -6,7 +6,12 @@ use ipa_macros::Step;
 use crate::{
     error::Error,
     ff::{ArrayAccessRef, ArrayBuild, ArrayBuilder, Field},
-    protocol::{basics::SecureMul, context::Context, step::BitOpStep, RecordId},
+    protocol::{
+        basics::{BooleanProtocols, SecureMul},
+        context::Context,
+        step::BitOpStep,
+        RecordId,
+    },
     secret_sharing::{replicated::semi_honest::AdditiveShare, FieldSimd},
 };
 #[cfg(all(test, unit_test))]
@@ -41,7 +46,7 @@ where
     F: Field + FieldSimd<N>,
     XS: ArrayAccessRef<Element = AdditiveShare<F, N>> + ArrayBuild<Input = AdditiveShare<F, N>>,
     YS: ArrayAccessRef<Element = AdditiveShare<F, N>>,
-    AdditiveShare<F, N>: SecureMul<C> + Not<Output = AdditiveShare<F, N>>,
+    AdditiveShare<F, N>: BooleanProtocols<C, F, N>,
 {
     let mut carry = AdditiveShare::<F, N>::ZERO;
     let sum = addition_circuit(ctx, record_id, x, y, &mut carry).await?;
@@ -65,7 +70,7 @@ where
     C: Context,
     S: SharedValue + CustomArray<Element = F>,
     AdditiveShare<S>: From<AdditiveShare<F, N>> + Into<AdditiveShare<F, N>>,
-    AdditiveShare<F>: Not<Output = AdditiveShare<F>>,
+    AdditiveShare<F>: BooleanProtocols<C, F>,
 {
     use crate::{ff::Expand, protocol::basics::if_else};
     let mut carry = AdditiveShare::<F>::ZERO;
@@ -114,7 +119,7 @@ where
     F: Field + FieldSimd<N>,
     XS: ArrayAccessRef<Element = AdditiveShare<F, N>> + ArrayBuild<Input = AdditiveShare<F, N>>,
     YS: ArrayAccessRef<Element = AdditiveShare<F, N>>,
-    AdditiveShare<F, N>: SecureMul<C> + Not<Output = AdditiveShare<F, N>>,
+    AdditiveShare<F, N>: BooleanProtocols<C, F, N>,
 {
     let x = x.iter();
     let y = y.iter();
@@ -168,7 +173,7 @@ async fn bit_adder<C, F, const N: usize>(
 where
     C: Context,
     F: Field + FieldSimd<N>,
-    AdditiveShare<F, N>: SecureMul<C> + Not<Output = AdditiveShare<F, N>>,
+    AdditiveShare<F, N>: BooleanProtocols<C, F, N>,
 {
     let output = x + y + &*carry;
 
diff --git a/ipa-core/src/protocol/ipa_prf/boolean_ops/comparison_and_subtraction_sequential.rs b/ipa-core/src/protocol/ipa_prf/boolean_ops/comparison_and_subtraction_sequential.rs
index cee540964..cba927c04 100644
--- a/ipa-core/src/protocol/ipa_prf/boolean_ops/comparison_and_subtraction_sequential.rs
+++ b/ipa-core/src/protocol/ipa_prf/boolean_ops/comparison_and_subtraction_sequential.rs
@@ -4,7 +4,7 @@
 //! the bit-width of the first (x) operand, then the excess bits of y must be zero. This condition
 //! is abbreviated below as `length(x) >= log2(y)`.
 
-use std::{borrow::Borrow, iter::repeat, ops::Not};
+use std::{borrow::Borrow, iter::repeat};
 
 #[cfg(all(test, unit_test))]
 use ipa_macros::Step;
@@ -15,7 +15,7 @@ use crate::{
     error::Error,
     ff::{ArrayAccessRef, ArrayBuild, ArrayBuilder, Field},
     protocol::{
-        basics::{SecureMul, ShareKnownValue},
+        basics::{BooleanProtocols, SecureMul, ShareKnownValue},
         context::Context,
         step::BitOpStep,
         RecordId,
@@ -52,7 +52,7 @@ where
     F: Field,
     XS: ArrayAccessRef<Element = AdditiveShare<F>> + ArrayBuild<Input = AdditiveShare<F>>,
     YS: ArrayAccessRef<Element = AdditiveShare<F>>,
-    AdditiveShare<F>: SecureMul<C> + Not<Output = AdditiveShare<F>>,
+    AdditiveShare<F>: BooleanProtocols<C, F>,
 {
     // we need to initialize carry to 1 for x>=y,
     let mut carry = AdditiveShare::<F>::share_known_value(&ctx, F::ONE);
@@ -77,7 +77,7 @@ where
     F: Field + FieldSimd<N>,
     XS: ArrayAccessRef<Element = AdditiveShare<F, N>> + ArrayBuild<Input = AdditiveShare<F, N>>,
     YS: ArrayAccessRef<Element = AdditiveShare<F, N>>,
-    AdditiveShare<F, N>: SecureMul<C> + Not<Output = AdditiveShare<F, N>>,
+    AdditiveShare<F, N>: BooleanProtocols<C, F, N>,
 {
     // we need to initialize carry to 0 for x>y
     let mut carry = AdditiveShare::<F, N>::ZERO;
@@ -102,7 +102,7 @@ where
     F: Field,
     XS: ArrayAccessRef<Element = AdditiveShare<F>> + ArrayBuild<Input = AdditiveShare<F>>,
     YS: ArrayAccessRef<Element = AdditiveShare<F>>,
-    AdditiveShare<F>: SecureMul<C> + Not<Output = AdditiveShare<F>>,
+    AdditiveShare<F>: BooleanProtocols<C, F>,
 {
     // we need to initialize carry to 1 for a subtraction
     let mut carry = AdditiveShare::<F>::share_known_value(&ctx, F::ONE);
@@ -127,7 +127,7 @@ where
     S: SharedValue + CustomArray<Element = F>,
     AdditiveShare<S>:
         ArrayAccessRef<Element = AdditiveShare<F>> + ArrayBuild<Input = AdditiveShare<F>>,
-    AdditiveShare<F>: SecureMul<C> + Not<Output = AdditiveShare<F>>,
+    AdditiveShare<F>: BooleanProtocols<C, F>,
     AdditiveShare<S>: From<AdditiveShare<F, N>> + Into<AdditiveShare<F, N>>,
 {
     let mut carry = AdditiveShare::<F>::share_known_value(&ctx, F::ONE);
@@ -170,7 +170,7 @@ where
     F: Field + FieldSimd<N>,
     XS: ArrayAccessRef<Element = AdditiveShare<F, N>> + ArrayBuild<Input = AdditiveShare<F, N>>,
     YS: ArrayAccessRef<Element = AdditiveShare<F, N>>,
-    AdditiveShare<F, N>: SecureMul<C> + Not<Output = AdditiveShare<F, N>>,
+    AdditiveShare<F, N>: BooleanProtocols<C, F, N>,
 {
     let x = x.iter();
     let y = y.iter();
@@ -220,7 +220,7 @@ async fn bit_subtractor<C, F, const N: usize>(
 where
     C: Context,
     F: Field + FieldSimd<N>,
-    AdditiveShare<F, N>: SecureMul<C> + Not<Output = AdditiveShare<F, N>>,
+    AdditiveShare<F, N>: BooleanProtocols<C, F, N>,
 {
     let output = x + !(y + &*carry);
 
diff --git a/ipa-core/src/protocol/ipa_prf/boolean_ops/share_conversion_aby.rs b/ipa-core/src/protocol/ipa_prf/boolean_ops/share_conversion_aby.rs
index b75660d4d..c92b430b2 100644
--- a/ipa-core/src/protocol/ipa_prf/boolean_ops/share_conversion_aby.rs
+++ b/ipa-core/src/protocol/ipa_prf/boolean_ops/share_conversion_aby.rs
@@ -1,8 +1,4 @@
-use std::{
-    borrow::Borrow,
-    convert::Infallible,
-    ops::{Neg, Not},
-};
+use std::{borrow::Borrow, convert::Infallible, ops::Neg};
 
 use ipa_macros::Step;
 
@@ -14,7 +10,7 @@ use crate::{
     },
     helpers::Role,
     protocol::{
-        basics::{partial_reveal, Reveal, SecureMul},
+        basics::{partial_reveal, BooleanProtocols},
         context::Context,
         ipa_prf::boolean_ops::addition_sequential::integer_add,
         prss::{FromPrss, SharedRandomness},
@@ -22,7 +18,7 @@ use crate::{
     },
     secret_sharing::{
         replicated::{semi_honest::AdditiveShare, ReplicatedSecretSharing},
-        FieldSimd, FieldVectorizable, SharedValue, SharedValueArray, TransposeFrom, Vectorizable,
+        FieldSimd, SharedValue, SharedValueArray, TransposeFrom, Vectorizable,
     },
 };
 
@@ -119,9 +115,7 @@ where
     YS: ArrayAccessRef<Element = AdditiveShare<Boolean, N>>
         + ArrayBuild<Input = AdditiveShare<Boolean, N>>
         + FromPrss<usize>,
-    AdditiveShare<Boolean, N>: SecureMul<C>
-        + Reveal<C, N, Output = <Boolean as FieldVectorizable<N>>::ArrayAlias>
-        + Not<Output = AdditiveShare<Boolean, N>>,
+    AdditiveShare<Boolean, N>: BooleanProtocols<C, Boolean, N>,
     Vec<AdditiveShare<BA256>>: for<'a> TransposeFrom<&'a YS>,
     Vec<BA256>:
         for<'a> TransposeFrom<&'a [<Boolean as Vectorizable<N>>::Array; 256], Error = Infallible>,
@@ -354,6 +348,7 @@ mod tests {
     use crate::{
         ff::{boolean_array::BA64, Serializable},
         helpers::stream::{ProcessChunks, TryFlattenItersExt},
+        protocol::context::SemiHonestContext,
         rand::thread_rng,
         seq_join::{seq_join, SeqJoin},
         test_executor::run,
@@ -379,7 +374,8 @@ mod tests {
             + 'static,
         for<'a> <XS as ArrayAccessRef>::Ref<'a>: Send,
         for<'a> <YS as ArrayAccessRef>::Ref<'a>: Send,
-        AdditiveShare<Boolean, CHUNK>: Not<Output = AdditiveShare<Boolean, CHUNK>>,
+        AdditiveShare<Boolean, CHUNK>:
+            for<'a> BooleanProtocols<SemiHonestContext<'a>, Boolean, CHUNK>,
         Vec<AdditiveShare<BA256>>: for<'a> TransposeFrom<&'a YS>,
         Vec<BA256>: for<'a> TransposeFrom<
             &'a [<Boolean as Vectorizable<CHUNK>>::Array; 256],
diff --git a/ipa-core/src/protocol/ipa_prf/prf_eval.rs b/ipa-core/src/protocol/ipa_prf/prf_eval.rs
index a10a9a0e2..854e0baaa 100644
--- a/ipa-core/src/protocol/ipa_prf/prf_eval.rs
+++ b/ipa-core/src/protocol/ipa_prf/prf_eval.rs
@@ -6,14 +6,13 @@ use crate::{
     error::Error,
     ff::{boolean::Boolean, curve_points::RP25519, ec_prime_field::Fp25519, Expand},
     protocol::{
-        basics::{Reveal, SecureMul},
+        basics::{BooleanProtocols, Reveal, SecureMul},
         context::Context,
         prss::{FromPrss, SharedRandomness},
         RecordId,
     },
     secret_sharing::{
-        replicated::semi_honest::AdditiveShare, FieldSimd, FieldVectorizable, Sendable, StdArray,
-        Vectorizable,
+        replicated::semi_honest::AdditiveShare, FieldSimd, Sendable, StdArray, Vectorizable,
     },
 };
 
@@ -92,7 +91,7 @@ where
     Fp25519: Vectorizable<N>,
     RP25519: Vectorizable<N, Array = StdArray<RP25519, N>>,
     Boolean: FieldSimd<N>,
-    AdditiveShare<Boolean, N>: Reveal<C, N, Output = <Boolean as FieldVectorizable<N>>::ArrayAlias>,
+    AdditiveShare<Boolean, N>: BooleanProtocols<C, Boolean, N>,
     AdditiveShare<Fp25519, N>: SecureMul<C> + FromPrss,
     StdArray<RP25519, N>: Sendable,
 {