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

@@ -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<F: Field> {
@@ -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");

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

@@ -36,10 +36,12 @@ pub trait StepNarrow<S: Step + ?Sized> {
 pub trait Step: AsRef<str> {}
 
 // 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.

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<F, const N: usize>(
     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();
     }
 

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},

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::<Vec<_>>()
             .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::<Vec<_>>()
             .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]
     }