From beca008ac325429c19e1b8b2e978551a9a8e8cba Mon Sep 17 00:00:00 2001 From: Andy Leiserson Date: Wed, 6 Mar 2024 15:30:28 -0800 Subject: [PATCH] Define real `Step` in arith circuit benchmark and TestWorld This allows limiting strings-as-steps to `cfg(test)` only. --- ipa-core/src/protocol/context/mod.rs | 6 +++--- ipa-core/src/protocol/step/mod.rs | 6 ++++-- ipa-core/src/test_fixture/circuit.rs | 9 ++++++++- ipa-core/src/test_fixture/mod.rs | 2 +- ipa-core/src/test_fixture/world.rs | 19 ++++++++++++------- 5 files changed, 28 insertions(+), 14 deletions(-) diff --git a/ipa-core/src/protocol/context/mod.rs b/ipa-core/src/protocol/context/mod.rs index ca1965135..4a54302ce 100644 --- a/ipa-core/src/protocol/context/mod.rs +++ b/ipa-core/src/protocol/context/mod.rs @@ -300,7 +300,7 @@ mod tests { telemetry::metrics::{ BYTES_SENT, INDEXED_PRSS_GENERATED, RECORDS_SENT, SEQUENTIAL_PRSS_GENERATED, }, - test_fixture::{Reconstruct, Runner, TestWorld, TestWorldConfig}, + test_fixture::{Reconstruct, Runner, TestExecutionStep, TestWorld, TestWorldConfig}, }; trait ReplicatedLeftValue { @@ -391,7 +391,7 @@ mod tests { let input_size = input.len(); let snapshot = world.metrics_snapshot(); let metrics_step = Gate::default() - .narrow(&TestWorld::execution_step(0)) + .narrow(&TestExecutionStep::Iter(0)) .narrow("metrics"); // for semi-honest protocols, amplification factor per helper is 1. @@ -448,7 +448,7 @@ mod tests { .await; let metrics_step = Gate::default() - .narrow(&TestWorld::execution_step(0)) + .narrow(&TestExecutionStep::Iter(0)) // TODO: leaky abstraction, test world should tell us the exact step .narrow(&MaliciousProtocol) .narrow("metrics"); diff --git a/ipa-core/src/protocol/step/mod.rs b/ipa-core/src/protocol/step/mod.rs index 81ede8472..73815967b 100644 --- a/ipa-core/src/protocol/step/mod.rs +++ b/ipa-core/src/protocol/step/mod.rs @@ -36,10 +36,12 @@ pub trait StepNarrow { pub trait Step: AsRef {} // In test code, allow a string (or string reference) to be used as a `Step`. -#[cfg(feature = "descriptive-gate")] +// Note: Since the creation of the `derive(Step)` macro, hardly any code is +// required to define a step. Doing so is highly encouraged, even in tests. +#[cfg(test)] impl Step for String {} -#[cfg(feature = "descriptive-gate")] +#[cfg(test)] impl Step for str {} /// A step generator for bitwise secure operations. diff --git a/ipa-core/src/test_fixture/circuit.rs b/ipa-core/src/test_fixture/circuit.rs index 278d57195..d5c983140 100644 --- a/ipa-core/src/test_fixture/circuit.rs +++ b/ipa-core/src/test_fixture/circuit.rs @@ -1,6 +1,7 @@ use std::{array, num::NonZeroUsize}; use futures::{future::join3, stream, StreamExt}; +use ipa_macros::Step; use rand::distributions::{Distribution, Standard}; use crate::{ @@ -121,6 +122,12 @@ pub async fn arithmetic( assert_eq!(sum, u128::from(width)); } +#[derive(Step)] +enum Step { + #[dynamic(1024)] + Stripe(usize), +} + async fn circuit<'a, F, const N: usize>( ctx: SemiHonestContext<'a>, record_id: RecordId, @@ -134,7 +141,7 @@ where { assert_eq!(b.len(), usize::from(depth)); for (stripe_ix, stripe) in b.iter().enumerate() { - let stripe_ctx = ctx.narrow(&format!("s{stripe_ix}")); + let stripe_ctx = ctx.narrow(&Step::Stripe(stripe_ix)); a = a.multiply(stripe, stripe_ctx, record_id).await.unwrap(); } diff --git a/ipa-core/src/test_fixture/mod.rs b/ipa-core/src/test_fixture/mod.rs index 0f96aee7b..e54e5eba4 100644 --- a/ipa-core/src/test_fixture/mod.rs +++ b/ipa-core/src/test_fixture/mod.rs @@ -25,7 +25,7 @@ use rand::{distributions::Standard, prelude::Distribution, rngs::mock::StepRng}; use rand_core::{CryptoRng, RngCore}; pub use sharing::{get_bits, into_bits, Reconstruct, ReconstructArr}; #[cfg(feature = "in-memory-infra")] -pub use world::{Runner, TestWorld, TestWorldConfig}; +pub use world::{Runner, TestExecutionStep, TestWorld, TestWorldConfig}; use crate::{ ff::{Field, U128Conversions}, diff --git a/ipa-core/src/test_fixture/world.rs b/ipa-core/src/test_fixture/world.rs index 14734c352..1b5a691d2 100644 --- a/ipa-core/src/test_fixture/world.rs +++ b/ipa-core/src/test_fixture/world.rs @@ -2,6 +2,7 @@ use std::{fmt::Debug, io::stdout, iter::zip}; use async_trait::async_trait; use futures::{future::join_all, Future}; +use ipa_macros::Step; use rand::{distributions::Standard, prelude::Distribution, rngs::StdRng}; use rand_core::{RngCore, SeedableRng}; use tracing::{Instrument, Level, Span}; @@ -31,6 +32,14 @@ use crate::{ }, }; +// This is used by the metrics tests in `protocol::context`. It otherwise would/should not be pub. +#[derive(Step)] +pub enum TestExecutionStep { + /// Provides a unique per-iteration context in tests. + #[dynamic(1024)] + Iter(usize), +} + /// Test environment for protocols to run tests that require communication between helpers. /// For now the messages sent through it never leave the test infra memory perimeter, so /// there is no need to associate each of them with `QueryId`, but this API makes it possible @@ -139,7 +148,7 @@ impl TestWorld { zip(&self.participants, &self.gateways) .map(|(participant, gateway)| { SemiHonestContext::new(participant, gateway) - .narrow(&Self::execution_step(execution)) + .narrow(&TestExecutionStep::Iter(execution)) }) .collect::>() .try_into() @@ -155,7 +164,8 @@ impl TestWorld { let execution = self.executions.fetch_add(1, Ordering::Relaxed); zip(&self.participants, &self.gateways) .map(|(participant, gateway)| { - MaliciousContext::new(participant, gateway).narrow(&Self::execution_step(execution)) + MaliciousContext::new(participant, gateway) + .narrow(&TestExecutionStep::Iter(execution)) }) .collect::>() .try_into() @@ -167,11 +177,6 @@ impl TestWorld { self.metrics_handle.snapshot() } - #[must_use] - pub fn execution_step(execution: usize) -> String { - format!("run-{execution}") - } - pub fn gateway(&self, role: Role) -> &Gateway { &self.gateways[role] }