Fix the saturated addition panic inside IPA

.github/workflows/check.yml

@@ -145,8 +145,11 @@ jobs:
  - name: Run arithmetic bench
  run: cargo bench --bench oneshot_arithmetic --no-default-features --features "enable-benches compact-gate"
 
- - name: Run compact gate tests
+ - name: Run compact gate tests for HTTP stack
  run: cargo test --no-default-features --features "cli web-app real-world-infra test-fixture compact-gate"
+
+ - name: Run in-memory compact gate tests
+ run: cargo test --features "compact-gate"
  slow:
  name: Slow tests
  env:

ipa-core/build.rs

@@ -27,7 +27,6 @@ track_steps!(
  dp::step,
  step,
  },
- test_fixture::step
 );
 
 fn main() {

ipa-core/src/lib.rs

@@ -118,7 +118,7 @@ pub(crate) mod test_executor {
  }
 }
 
-#[cfg(all(test, unit_test, not(feature = "shuttle")))]
+#[cfg(all(test, not(feature = "shuttle")))]
 pub(crate) mod test_executor {
  use std::future::Future;
 
@@ -137,6 +137,7 @@ pub(crate) mod test_executor {
  .block_on(f())
  }
 
+ #[allow(dead_code)]
  pub fn run<F, Fut, T>(f: F) -> T
  where
  F: Fn() -> Fut + Send + Sync + 'static,

ipa-core/src/net/client/mod.rs

@@ -468,10 +468,10 @@ pub(crate) mod tests {
  RequestHandler, RoleAssignment, Transport, MESSAGE_PAYLOAD_SIZE_BYTES,
  },
  net::test::TestServer,
+ protocol::step::TestExecutionStep,
  query::ProtocolResult,
  secret_sharing::replicated::semi_honest::AdditiveShare as Replicated,
  sync::Arc,
- test_fixture::step::TestExecutionStep,
  };
 
  #[tokio::test]

ipa-core/src/protocol/basics/mod.rs

@@ -149,6 +149,11 @@ impl<'a, B: ShardBinding> BooleanProtocols<UpgradedSemiHonestContext<'a, B, Bool
 {
 }
 
+impl<'a, B: ShardBinding> BooleanProtocols<UpgradedSemiHonestContext<'a, B, Boolean>, 3>
+ for AdditiveShare<Boolean, 3>
+{
+}
+
 impl<'a, B: ShardBinding> BooleanProtocols<DZKPUpgradedSemiHonestContext<'a, B>, 32>
  for AdditiveShare<Boolean, 32>
 {

ipa-core/src/protocol/boolean/or.rs

@@ -1,9 +1,11 @@
 use std::iter::zip;
 
+use ipa_step::StepNarrow;
+
 use crate::{
  error::Error,
  ff::{boolean::Boolean, Field},
- protocol::{basics::SecureMul, boolean::step::SixteenBitStep, context::Context, RecordId},
+ protocol::{basics::SecureMul, boolean::NBitStep, context::Context, Gate, RecordId},
  secret_sharing::{
  replicated::semi_honest::AdditiveShare, BitDecomposed, FieldSimd,
  Linear as LinearSecretSharing,
@@ -34,25 +36,27 @@ pub async fn or<F: Field, C: Context, S: LinearSecretSharing<F> + SecureMul<C>>(
 //
 // Supplying an iterator saves constructing a complete copy of the argument
 // in memory when it is a uniform constant.
-pub async fn bool_or<'a, C, BI, const N: usize>(
+pub async fn bool_or<'a, C, S, BI, const N: usize>(
  ctx: C,
  record_id: RecordId,
  a: &BitDecomposed<AdditiveShare<Boolean, N>>,
  b: BI,
 ) -> Result<BitDecomposed<AdditiveShare<Boolean, N>>, Error>
 where
  C: Context,
+ S: NBitStep,
  BI: IntoIterator,
  <BI as IntoIterator>::IntoIter: ExactSizeIterator<Item = &'a AdditiveShare<Boolean, N>> + Send,
  Boolean: FieldSimd<N>,
  AdditiveShare<Boolean, N>: SecureMul<C>,
+ Gate: StepNarrow<S>,
 {
  let b = b.into_iter();
  assert_eq!(a.len(), b.len());
 
  BitDecomposed::try_from(
  ctx.parallel_join(zip(a.iter(), b).enumerate().map(|(i, (a, b))| {
- let ctx = ctx.narrow(&SixteenBitStep::from(i));
+ let ctx = ctx.narrow(&S::from(i));
  async move {
  let ab = a.multiply(b, ctx, record_id).await?;
  Ok::<_, Error>(-ab + a + b)

ipa-core/src/protocol/ipa_prf/boolean_ops/addition_sequential.rs

@@ -73,7 +73,7 @@ where
  .await?;
 
  // if carry==1 then {all ones} else {result}
- bool_or(
+ bool_or::<_, S, _, N>(
  ctx.narrow::<Step>(&Step::Select),
  record_id,
  &result,

ipa-core/src/protocol/ipa_prf/boolean_ops/step.rs

@@ -1,15 +1,15 @@
 use ipa_step_derive::CompactStep;
 
 /// FIXME: This step is not generic enough to be used in the `saturated_addition` protocol.
-/// It constrains the input to be at most 2 bytes and it will panic in runtime if it is greater
+/// It constrains the input to be at most 4 bytes and it will panic in runtime if it is greater
 /// than that. The issue is that compact gate requires concrete type to be put as child.
 /// If we ever see it being an issue, we should make a few implementations of this similar to what
 /// we've done for bit steps
 #[derive(CompactStep)]
 pub(crate) enum SaturatedAdditionStep {
- #[step(child = crate::protocol::boolean::step::SixteenBitStep)]
+ #[step(child = crate::protocol::boolean::step::ThirtyTwoBitStep)]
  Add,
- #[step(child = crate::protocol::boolean::step::SixteenBitStep)]
+ #[step(child = crate::protocol::boolean::step::ThirtyTwoBitStep)]
  Select,
 }
 

ipa-core/src/protocol/ipa_prf/mod.rs

@@ -721,3 +721,109 @@ pub mod tests {
  });
  }
 }
+
+#[cfg(all(test, all(feature = "compact-gate", feature = "in-memory-infra")))]
+mod compact_gate_tests {
+
+ use ipa_step::StepNarrow;
+
+ use crate::{
+ ff::{
+ boolean_array::{BA20, BA5, BA8},
+ U128Conversions,
+ },
+ helpers::query::DpMechanism,
+ protocol::{
+ ipa_prf::{oprf_ipa, oprf_padding::PaddingParameters},
+ step::{ProtocolGate, ProtocolStep},
+ },
+ test_executor::run,
+ test_fixture::{ipa::TestRawDataRecord, Reconstruct, Runner, TestWorld, TestWorldConfig},
+ };
+
+ #[test]
+ fn saturated_agg() {
+ const EXPECTED: &[u128] = &[0, 255, 255, 0, 0, 0, 0, 0];
+
+ run(|| async {
+ let world = TestWorld::new_with(TestWorldConfig {
+ initial_gate: Some(ProtocolGate::default().narrow(&ProtocolStep::IpaPrf)),
+ ..Default::default()
+ });
+
+ let records: Vec<TestRawDataRecord> = vec![
+ TestRawDataRecord {
+ timestamp: 0,
+ user_id: 12345,
+ is_trigger_report: false,
+ breakdown_key: 1,
+ trigger_value: 0,
+ },
+ TestRawDataRecord {
+ timestamp: 5,
+ user_id: 12345,
+ is_trigger_report: false,
+ breakdown_key: 2,
+ trigger_value: 0,
+ },
+ TestRawDataRecord {
+ timestamp: 10,
+ user_id: 12345,
+ is_trigger_report: true,
+ breakdown_key: 0,
+ trigger_value: 255,
+ },
+ TestRawDataRecord {
+ timestamp: 20,
+ user_id: 12345,
+ is_trigger_report: true,
+ breakdown_key: 0,
+ trigger_value: 255,
+ },
+ TestRawDataRecord {
+ timestamp: 30,
+ user_id: 12345,
+ is_trigger_report: true,
+ breakdown_key: 0,
+ trigger_value: 255,
+ },
+ TestRawDataRecord {
+ timestamp: 0,
+ user_id: 68362,
+ is_trigger_report: false,
+ breakdown_key: 1,
+ trigger_value: 0,
+ },
+ TestRawDataRecord {
+ timestamp: 20,
+ user_id: 68362,
+ is_trigger_report: true,
+ breakdown_key: 1,
+ trigger_value: 255,
+ },
+ ];
+ let dp_params = DpMechanism::NoDp;
+ let padding_params = PaddingParameters::relaxed();
+
+ let mut result: Vec<_> = world
+ .semi_honest(records.into_iter(), |ctx, input_rows| async move {
+ oprf_ipa::<_, BA5, BA8, BA8, BA20, 5, 32>(
+ ctx,
+ input_rows,
+ None,
+ dp_params,
+ padding_params,
+ )
+ .await
+ .unwrap()
+ })
+ .await
+ .reconstruct();
+ result.truncate(EXPECTED.len());
+ assert_eq!(
+ result.iter().map(|&v| v.as_u128()).collect::<Vec<_>>(),
+ EXPECTED,
+ );
+ });
+ }
+}

ipa-core/src/protocol/ipa_prf/prf_sharding/mod.rs

@@ -320,14 +320,14 @@ pub struct AttributionOutputs<BK, TV> {
 pub type SecretSharedAttributionOutputs<BK, TV> =
  AttributionOutputs<Replicated<BK>, Replicated<TV>>;
 
-#[cfg(all(test, any(unit_test, feature = "shuttle")))]
+#[cfg(test)]
 #[derive(Debug, Clone, Ord, PartialEq, PartialOrd, Eq)]
 pub struct AttributionOutputsTestInput<BK: BooleanArray, TV: BooleanArray> {
  pub bk: BK,
  pub tv: TV,
 }
 
-#[cfg(all(test, any(unit_test, feature = "shuttle")))]
+#[cfg(test)]
 impl<BK, TV> crate::secret_sharing::IntoShares<(Replicated<BK>, Replicated<TV>)>
  for AttributionOutputsTestInput<BK, TV>
 where

ipa-core/src/protocol/step.rs

@@ -9,9 +9,11 @@ pub enum ProtocolStep {
  IpaPrf,
  Multiply,
  PrimeFieldAddition,
- #[cfg(any(test, feature = "test-fixture"))]
- #[step(count = 10, child = crate::test_fixture::step::TestExecutionStep)]
- Test(usize),
+ /// Steps used in unit tests are grouped under this one. Ideally it should be
+ /// gated behind test configuration, but it does not work with build.rs that
+ /// does not enable any features when creating protocol gate file
+ #[step(child = TestExecutionStep)]
+ Test,
 
  /// This step includes all the steps that are currently not linked into a top-level protocol.
  ///
@@ -46,3 +48,10 @@ pub enum DeadCodeStep {
  #[step(child = crate::protocol::ipa_prf::boolean_ops::step::MultiplicationStep)]
  Multiplication,
 }
+
+/// Provides a unique per-iteration context in tests.
+#[derive(CompactStep)]
+pub enum TestExecutionStep {
+ #[step(count = 999)]
+ Iter(usize),
+}

ipa-core/src/secret_sharing/into_shares.rs

@@ -41,7 +41,7 @@ where
  }
 }
 
-#[cfg(all(test, unit_test))]
+#[cfg(test)]
 impl<U, T> IntoShares<Result<T, crate::error::Error>> for Result<U, crate::error::Error>
 where
  U: IntoShares<T>,

ipa-core/src/test_fixture/circuit.rs

@@ -10,13 +10,13 @@ use crate::{
  protocol::{
  basics::SecureMul,
  context::{Context, SemiHonestContext},
- step::ProtocolStep,
+ step::{ProtocolStep, TestExecutionStep as Step},
  Gate, RecordId,
  },
  rand::thread_rng,
  secret_sharing::{replicated::semi_honest::AdditiveShare as Replicated, FieldSimd, IntoShares},
  seq_join::seq_join,
- test_fixture::{step::TestExecutionStep as Step, ReconstructArr, TestWorld, TestWorldConfig},
+ test_fixture::{ReconstructArr, TestWorld, TestWorldConfig},
  utils::array::zip3,
 };
 
@@ -82,7 +82,7 @@ pub async fn arithmetic<F, const N: usize>(
  active,
  ..Default::default()
  },
- initial_gate: Some(Gate::default().narrow(&ProtocolStep::Test(0))),
+ initial_gate: Some(Gate::default().narrow(&ProtocolStep::Test)),
  ..Default::default()
  };
  let world = TestWorld::new_with(config);

ipa-core/src/test_fixture/mod.rs

@@ -15,7 +15,6 @@ pub mod hybrid;
 pub mod ipa;
 pub mod logging;
 pub mod metrics;
-pub(crate) mod step;
 #[cfg(feature = "in-memory-infra")]
 mod test_gate;
 

ipa-core/src/test_fixture/test_gate.rs

@@ -2,7 +2,7 @@ use std::sync::atomic::{AtomicUsize, Ordering};
 
 use ipa_step::StepNarrow;
 
-use crate::{protocol::Gate, test_fixture::step::TestExecutionStep};
+use crate::protocol::{step::TestExecutionStep, Gate};
 
 /// This manages the gate information for test runs. Most unit tests want to have multiple runs
 /// using the same instance of [`TestWorld`], but they don't care about the name of that particular

ipa-step/src/gate.rs

@@ -29,7 +29,7 @@
  let short_name = t.rsplit_once("::").map_or_else(|| t.as_ref(), |(_a, b)| b);
  let msg = format!("unexpected narrow for {gate_name}({{s}}) => {short_name}({{ss}})");
  syntax.extend(quote! {
- #[allow(clippy::too_many_lines)]
+ #[allow(clippy::too_many_lines, clippy::unreadable_literal)]
  impl ::ipa_step::StepNarrow<#ty> for #ident {
  fn narrow(&self, step: &#ty) -> Self {
  match self.0 {
@@ -87,6 +87,7 @@
  let gate_lookup_type = step_hasher.lookup_type();
  let mut syntax = quote! {
 
+ #[allow(clippy::unreadable_literal)]
  static STR_LOOKUP: [&str; #step_count] = [#(#gate_names),*];
  static GATE_LOOKUP: #gate_lookup_type = #step_hasher
 

ipa-step/src/hash.rs

@@ -62,6 +62,7 @@ impl ToTokens for HashingSteps {
  };
 
  struct #lookup_type {
+ #[allow(clippy::unreadable_literal)]
  inner: [(u64, u32); #sz]
  }
 

scripts/coverage-ci

@@ -19,4 +19,7 @@ done
 cargo test --bench oneshot_ipa --no-default-features --features "enable-benches compact-gate" -- -n 62 -c 16
 cargo test --bench criterion_arithmetic --no-default-features --features "enable-benches compact-gate"
 
+# compact gate + in-memory-infra
+cargo test --features "compact-gate"
+
 cargo llvm-cov report "$@"