Slightly less compact gate

ipa-core/Cargo.toml

@@ -11,7 +11,7 @@ default = [
     # by default remove all TRACE, DEBUG spans from release builds
     "tracing/max_level_trace",
     "tracing/release_max_level_info",
-    "descriptive-gate",
+    "compact-gate",
     "aggregate-circuit",
     "stall-detection",
     "aggregate-circuit",
@@ -140,6 +140,7 @@ sha2 = "0.10"
 shuttle-crate = { package = "shuttle", version = "0.6.1", optional = true }
 thiserror = "1.0"
 time = { version = "0.3", optional = true }
+tinyvec = { version = "1.6" }
 tokio = { version = "1.35", features = ["fs", "rt", "rt-multi-thread", "macros"] }
 # TODO: axum-server holds onto 0.24 and we can't upgrade until they do. Or we move away from axum-server
 tokio-rustls = { version = "0.24", optional = true }
@@ -214,7 +215,7 @@ required-features = ["enable-benches", "descriptive-gate"]
 name = "oneshot_ipa"
 path = "benches/oneshot/ipa.rs"
 harness = false
-required-features = ["enable-benches", "descriptive-gate"]
+required-features = ["enable-benches"]
 
 [[test]]
 name = "helper_networks"

ipa-core/src/helpers/error.rs

@@ -75,7 +75,7 @@ impl Error {
     ) -> Error {
         Self::SerializationError {
             record_id,
-            step: String::from(gate.as_ref()),
+            step: gate.to_string(),
             inner: inner.into(),
         }
     }

ipa-core/src/helpers/gateway/receive.rs

@@ -44,7 +44,7 @@ impl<M: Message> ReceivingEnd<M> {
     /// ## Panics
     /// This will panic if message size does not fit into 8 bytes and it somehow got serialized
     /// and sent to this helper.
-    #[tracing::instrument(level = "trace", "receive", skip_all, fields(i = %record_id, from = ?self.channel_id.role, gate = ?self.channel_id.gate.as_ref()))]
+    //#[tracing::instrument(level = "trace", "receive", skip_all, fields(i = %record_id, from = ?self.channel_id.role, gate = ?self.channel_id.gate.to_string()))]
     pub async fn receive(&self, record_id: RecordId) -> Result<M, Error> {
         self.unordered_rx
             .recv::<M, _>(record_id)

ipa-core/src/helpers/gateway/send.rs

@@ -113,19 +113,8 @@ impl<M: Message> SendingEnd<M> {
     /// call.
     ///
     /// [`set_total_records`]: crate::protocol::context::Context::set_total_records
-    #[tracing::instrument(level = "trace", "send", skip_all, fields(i = %record_id, total = %self.inner.total_records, to = ?self.channel_id.role, gate = ?self.channel_id.gate.as_ref()))]
     pub async fn send<B: Borrow<M>>(&self, record_id: RecordId, msg: B) -> Result<(), Error> {
-        let r = self.inner.send(record_id, msg).await;
-        metrics::increment_counter!(RECORDS_SENT,
-            STEP => self.channel_id.gate.as_ref().to_string(),
-            ROLE => self.sender_role.as_static_str()
-        );
-        metrics::counter!(BYTES_SENT, M::Size::U64,
-            STEP => self.channel_id.gate.as_ref().to_string(),
-            ROLE => self.sender_role.as_static_str()
-        );
-
-        r
+        self.inner.send(record_id, msg).await
     }
 }
 

ipa-core/src/helpers/mod.rs

@@ -421,7 +421,7 @@ impl ChannelId {
 
 impl Debug for ChannelId {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        write!(f, "channel[{:?},{:?}]", self.role, self.gate.as_ref())
+        write!(f, "channel[{:?},{:?}]", self.role, self.gate.to_string())
     }
 }
 

ipa-core/src/helpers/prss_protocol.rs

@@ -1,4 +1,7 @@
+use std::fmt::Display;
+
 use futures_util::future::try_join4;
+use generic_array::GenericArray;
 use rand_core::{CryptoRng, RngCore};
 use x25519_dalek::PublicKey;
 
@@ -13,13 +16,21 @@ use crate::{
 
 pub struct PrssExchangeStep;
 
-impl AsRef<str> for PrssExchangeStep {
-    fn as_ref(&self) -> &str {
-        "prss_exchange"
+impl Display for PrssExchangeStep {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str("prss_exchange")
     }
 }
 
-impl Step for PrssExchangeStep {}
+impl Step for PrssExchangeStep {
+    #[cfg(feature = "compact-gate")]
+    type Length = generic_array::typenum::U1;
+
+    #[cfg(feature = "compact-gate")]
+    fn as_bytes(&self) -> GenericArray<u8, Self::Length> {
+        [0u8].into()
+    }
+}
 
 /// establish the prss endpoint by exchanging public keys with the other helpers
 /// # Errors

ipa-core/src/helpers/transport/query/mod.rs

@@ -5,6 +5,7 @@ use std::{
     num::NonZeroU32,
 };
 
+use generic_array::GenericArray;
 use serde::{Deserialize, Deserializer, Serialize};
 
 use crate::{
@@ -213,17 +214,26 @@ impl QueryType {
 }
 
 /// TODO: should this `AsRef` impl (used for `Substep`) take into account config of IPA?
-impl AsRef<str> for QueryType {
-    fn as_ref(&self) -> &str {
-        match self {
+impl Display for QueryType {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        let id = match self {
             #[cfg(any(test, feature = "cli", feature = "test-fixture"))]
             QueryType::TestMultiply => Self::TEST_MULTIPLY_STR,
             QueryType::OprfIpa(_) => Self::OPRF_IPA_STR,
-        }
+        };
+        f.write_str(id)
     }
 }
 
-impl Step for QueryType {}
+impl Step for QueryType {
+    #[cfg(feature = "compact-gate")]
+    type Length = generic_array::typenum::U1;
+
+    #[cfg(feature = "compact-gate")]
+    fn as_bytes(&self) -> GenericArray<u8, Self::Length> {
+        [1u8].into()
+    }
+}
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
 #[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]

ipa-core/src/net/http_serde.rs

@@ -142,7 +142,7 @@ pub mod query {
             write!(
                 f,
                 "query_type={qt}&field_type={f:?}&size={size}",
-                qt = self.query_type.as_ref(),
+                qt = self.query_type.to_string(),
                 f = self.field_type,
                 size = self.size
             )?;
@@ -422,7 +422,7 @@ pub mod query {
                         "{}/{}/step/{}",
                         BASE_AXUM_PATH,
                         self.query_id.as_ref(),
-                        self.gate.as_ref()
+                        self.gate.to_string()
                     ))
                     .build()?;
                 Ok(hyper::Request::post(uri).body(self.body)?)

ipa-core/src/protocol/context/prss.rs

@@ -39,11 +39,11 @@ impl SharedRandomness for InstrumentedIndexedSharedRandomness<'_> {
         &self,
         index: I,
     ) -> (GenericArray<u128, N>, GenericArray<u128, N>) {
-        let step = self.step.as_ref().to_string();
+        //let step = self.step.to_string(); // TODO: not OK here!
         // TODO: what we really want here is a gauge indicating the maximum index used to generate
         // PRSS. Gauge infrastructure is not supported yet, `Metrics` struct needs to be able to
         // handle gauges
-        metrics::increment_counter!(INDEXED_PRSS_GENERATED, STEP => step, ROLE => self.role.as_static_str());
+        //metrics::increment_counter!(INDEXED_PRSS_GENERATED, STEP => step, ROLE => self.role.as_static_str());
         self.inner.generate_arrays(index)
     }
 }
@@ -73,8 +73,8 @@ impl RngCore for InstrumentedSequentialSharedRandomness<'_> {
     }
 
     fn next_u64(&mut self) -> u64 {
-        let step = self.step.as_ref().to_string();
-        metrics::increment_counter!(SEQUENTIAL_PRSS_GENERATED, STEP => step, ROLE => self.role.as_static_str());
+        //let step = self.step.to_string(); // TODO: not OK here!
+        //metrics::increment_counter!(SEQUENTIAL_PRSS_GENERATED, STEP => step, ROLE => self.role.as_static_str());
         self.inner.next_u64()
     }
 

ipa-core/src/protocol/context/validator.rs

@@ -71,15 +71,17 @@ impl<F: ExtendableField> Debug for SemiHonest<'_, F> {
 
 /// Steps used by the validation component of malicious protocol execution.
 /// In addition to these, an implicit step is used to initialize the value of `r`.
-#[cfg_attr(feature = "descriptive-gate", derive(ipa_macros::Step))]
+//#[cfg_attr(feature = "descriptive-gate", derive(ipa_macros::Step))]
+#[derive(ipa_macros::Step)]
 pub(crate) enum Step {
     /// For the execution of the malicious protocol.
     MaliciousProtocol,
     /// The final validation steps.
     Validate,
 }
 
-#[cfg_attr(feature = "descriptive-gate", derive(ipa_macros::Step))]
+//#[cfg_attr(feature = "descriptive-gate", derive(ipa_macros::Step))]
+#[derive(ipa_macros::Step)]
 pub(crate) enum ValidateStep {
     /// Propagate the accumulated values of `u` and `w`.
     PropagateUAndW,
@@ -220,7 +222,7 @@ impl<'a, F: ExtendableField> Validator<MaliciousContext<'a>, F> for Malicious<'a
     ///
     /// ## Panics
     /// Will panic if the mutex is poisoned
-    #[tracing::instrument(name = "validate", skip_all, fields(gate = %self.validate_ctx.gate().as_ref()))]
+    #[tracing::instrument(name = "validate", skip_all, fields(gate = %self.validate_ctx.gate().to_string()))]
     async fn validate<D: DowngradeMalicious>(self, values: D) -> Result<D::Target, Error> {
         // send our `u_i+1` value to the helper on the right
         let (u_share, w_share) = self.propagate_u_and_w().await?;

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

@@ -38,6 +38,7 @@ mod shuffle;
 
 #[derive(Step)]
 pub(crate) enum Step {
+    PRFKeyGen,
     ConvertFp25519,
     EvalPrf,
     ConvertInputRowsToPrf,
@@ -226,7 +227,7 @@ where
     let convert_ctx = ctx.narrow(&Step::ConvertFp25519);
     let eval_ctx = ctx.narrow(&Step::EvalPrf);
 
-    let prf_key = gen_prf_key(&convert_ctx);
+    let prf_key = gen_prf_key(&ctx.narrow(&Step::PRFKeyGen));
 
     ctx.try_join(input_rows.into_iter().enumerate().map(|(idx, record)| {
         let convert_ctx = convert_ctx.clone();

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

@@ -17,7 +17,6 @@ use crate::{
 
 #[derive(Step)]
 pub(crate) enum Step {
-    PRFKeyGen,
     GenRandomMask,
     MultMaskWithPRFInput,
     RevealR,
@@ -57,7 +56,7 @@ pub fn gen_prf_key<C>(ctx: &C) -> AdditiveShare<Fp25519>
 where
     C: Context,
 {
-    ctx.narrow(&Step::PRFKeyGen).prss().generate(RecordId(0))
+    ctx.prss().generate(RecordId(0))
 }
 
 /// evaluates the Dodis-Yampolski PRF g^(1/(k+x))

ipa-core/src/protocol/modulus_conversion/convert_shares.rs

@@ -298,7 +298,7 @@ where
 /// Propagates errors from convert shares
 /// # Panics
 /// If the total record count on the context is unspecified.
-#[tracing::instrument(name = "modulus_conversion", skip_all, fields(bits = ?bit_range, gate = %ctx.gate().as_ref()))]
+#[tracing::instrument(name = "modulus_conversion", skip_all, fields(bits = ?bit_range, gate = %ctx.gate().to_string()))]
 pub fn convert_bits<'a, F, V, C, S, VS>(
     ctx: C,
     binary_shares: VS,
@@ -319,7 +319,7 @@ where
 /// A version of `convert_bits` that allows for the retention of unconverted fields in the input.
 /// Note that unconverted fields are not upgraded, so they might need to be upgraded either before or
 /// after invoking this function.
-#[tracing::instrument(name = "modulus_conversion", skip_all, fields(bits = ?bit_range, gate = %ctx.gate().as_ref()))]
+#[tracing::instrument(name = "modulus_conversion", skip_all, fields(bits = ?bit_range, gate = %ctx.gate().to_string()))]
 pub fn convert_selected_bits<'a, F, V, C, S, VS, R>(
     ctx: C,
     binary_shares: VS,

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

@@ -15,7 +15,7 @@ use x25519_dalek::PublicKey;
 
 use super::step::Gate;
 use crate::{
-    protocol::RecordId,
+    protocol::{step::GateId, RecordId},
     rand::{CryptoRng, RngCore},
     sync::{Arc, Mutex},
 };
@@ -26,13 +26,13 @@ use crate::{
 /// a given step.
 #[cfg(debug_assertions)]
 struct UsedSet {
-    key: Gate,
+    key: GateId,
     used: Arc<Mutex<HashSet<usize>>>,
 }
 
 #[cfg(debug_assertions)]
 impl UsedSet {
-    fn new(key: Gate) -> Self {
+    fn new(key: GateId) -> Self {
         Self {
             key,
             used: Arc::new(Mutex::new(HashSet::new())),
@@ -51,7 +51,7 @@ impl UsedSet {
         } else {
             assert!(
                 self.used.lock().unwrap().insert(raw_index as usize),
-                "Generated randomness for index '{index}' twice using the same key '{}'",
+                "Generated randomness for index '{index}' twice using the same key '{:?}'",
                 self.key,
             );
         }
@@ -61,7 +61,7 @@ impl UsedSet {
 #[cfg(debug_assertions)]
 impl Debug for UsedSet {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        write!(f, "IndicesSet(key={})", self.key)
+        write!(f, "IndicesSet(key={:?})", self.key)
     }
 }
 
@@ -270,45 +270,47 @@ enum EndpointItem {
 struct EndpointInner {
     left: GeneratorFactory,
     right: GeneratorFactory,
-    items: HashMap<Gate, EndpointItem>,
+    items: HashMap<GateId, EndpointItem>,
 }
 
 impl EndpointInner {
     pub fn indexed(&mut self, key: &Gate) -> Arc<IndexedSharedRandomness> {
         // The second arm of this statement would be fine, except that `HashMap::entry()`
         // only takes an owned value as an argument.
         // This makes the lookup perform an allocation, which is very much suboptimal.
-        let item = if let Some(item) = self.items.get(key) {
+        let key = key.id().clone();
+        let item = if let Some(item) = self.items.get(&key) {
             item
         } else {
             self.items.entry(key.clone()).or_insert_with_key(|k| {
                 EndpointItem::Indexed(Arc::new(IndexedSharedRandomness {
-                    left: self.left.generator(k.as_ref().as_bytes()),
-                    right: self.right.generator(k.as_ref().as_bytes()),
+                    left: self.left.generator(&k),
+                    right: self.right.generator(&k),
                     #[cfg(debug_assertions)]
-                    used: UsedSet::new(key.clone()),
+                    used: UsedSet::new(k.clone()),
                 }))
             })
         };
         if let EndpointItem::Indexed(idxd) = item {
             Arc::clone(idxd)
         } else {
-            panic!("Attempt to get an indexed PRSS for {key} after retrieving a sequential PRSS");
+            panic!("Attempt to get an indexed PRSS for {key:?} after retrieving a sequential PRSS");
         }
     }
 
     pub fn sequential(
         &mut self,
         key: &Gate,
     ) -> (SequentialSharedRandomness, SequentialSharedRandomness) {
+        let key = key.id().clone();
         let prev = self.items.insert(key.clone(), EndpointItem::Sequential);
         assert!(
             prev.is_none(),
-            "Attempt access a sequential PRSS for {key} after another access"
+            "Attempt access a sequential PRSS for {key:?} after another access"
         );
         (
-            SequentialSharedRandomness::new(self.left.generator(key.as_ref().as_bytes())),
-            SequentialSharedRandomness::new(self.right.generator(key.as_ref().as_bytes())),
+            SequentialSharedRandomness::new(self.left.generator(key.as_slice())),
+            SequentialSharedRandomness::new(self.right.generator(key.as_slice())),
         )
     }
 }