Repository: cloudflare/daphne
Pull request #656 (2024-08-23 18:23:36)

pine: Use a generic `Xof`
Partially addresses #618.

The spec uses `XofTurboShake128`, but we will also need to support `XofHmacSha256Aes128`.


Patch:
diff --git a/crates/dapf/src/acceptance/mod.rs b/crates/dapf/src/acceptance/mod.rs
index a4097af..748f6ce 100644
--- a/crates/dapf/src/acceptance/mod.rs
+++ b/crates/dapf/src/acceptance/mod.rs
@@ -279,7 +279,7 @@ impl Test {
             test_success: prometheus::register_int_gauge_vec!(
                 "daphne_server_acceptance_test",
                 "Counts the number of times tests ran",
-                &["using_mtls", "using_bearer_token", "success"],
+                &["using_mtls", "using_bearer_token"],
             )?,
             test_durations: prometheus::register_histogram_vec!(
                 "daphne_server_acceptance_test_durations",
@@ -710,16 +710,16 @@ impl Test {
 
     pub async fn test_helper(&self, opt: &TestOptions) -> Result<TestDurations> {
         let res = self.test_helper_impl(opt).await;
-        let success = res.is_ok();
         let c = |b: bool| ["F", "T"][b as usize];
-        self.metrics
+        let metric = self
+            .metrics
             .test_success
-            .with_label_values(&[
-                c(self.using_mtls),
-                c(self.bearer_token.is_some()),
-                c(success),
-            ])
-            .inc();
+            .with_label_values(&[c(self.using_mtls), c(self.bearer_token.is_some())]);
+        if res.is_ok() {
+            metric.inc()
+        } else {
+            metric.set(0)
+        }
         if let Ok(TestDurations {
             hpke_config_fetch,
             aggregate_init_req,
diff --git a/crates/daphne/benches/pine.rs b/crates/daphne/benches/pine.rs
index ee66bd6..1a4e2f3 100644
--- a/crates/daphne/benches/pine.rs
+++ b/crates/daphne/benches/pine.rs
@@ -6,10 +6,15 @@ use daphne::pine::Pine;
 use prio::{
     field::random_vector,
     flp::Type,
-    vdaf::{xof::Seed, Aggregator, Client},
+    vdaf::{
+        xof::{Seed, XofTurboShake128},
+        Aggregator, Client,
+    },
 };
 
 fn pine(c: &mut Criterion) {
+    // NOTE We ignore this clippy warning because we may want to benchmark more parameters later.
+    #[allow(clippy::single_element_loop)]
     for (dimension, chunk_len, chunk_len_sq_norm_equal) in [(200_000, 150 * 2, 447 * 18)] {
         let pine =
             Pine::new_64(1 << 15, dimension, 15, chunk_len, chunk_len_sq_norm_equal).unwrap();
@@ -24,7 +29,10 @@ fn pine(c: &mut Criterion) {
             |b, &_d| {
                 b.iter(|| {
                     pine.flp
-                        .encode_with_wr_joint_rand(measurement.iter().copied(), &wr_joint_rand_seed)
+                        .encode_with_wr_joint_rand::<XofTurboShake128, 16>(
+                            measurement.iter().copied(),
+                            &wr_joint_rand_seed,
+                        )
                         .unwrap()
                 });
             },
@@ -33,7 +41,10 @@ fn pine(c: &mut Criterion) {
         {
             let (mut input, wr_test_results) = pine
                 .flp
-                .encode_with_wr_joint_rand(measurement.iter().copied(), &wr_joint_rand_seed)
+                .encode_with_wr_joint_rand::<XofTurboShake128, 16>(
+                    measurement.iter().copied(),
+                    &wr_joint_rand_seed,
+                )
                 .unwrap();
             input.extend_from_slice(&wr_test_results);
             let prove_rand = random_vector(pine.flp_sq_norm_equal.prove_rand_len()).unwrap();
diff --git a/crates/daphne/src/pine/flp.rs b/crates/daphne/src/pine/flp.rs
index 5dd4a1e..a0a34b7 100644
--- a/crates/daphne/src/pine/flp.rs
+++ b/crates/daphne/src/pine/flp.rs
@@ -38,7 +38,7 @@ use prio::{
         FlpError, Gadget, Type,
     },
     vdaf::{
-        xof::{Seed, Xof, XofTurboShake128},
+        xof::{Seed, Xof},
         VdafError,
     },
 };
@@ -67,14 +67,15 @@ impl<F: FftFriendlyFieldElement> PineType<F> {
     ///
     /// Return the encoded measurement `meas` and wraparound test results `wr_test_results`. The
     /// input to the FLP circuit is `meas || wr_test_results`.
-    pub fn encode_with_wr_joint_rand(
+    pub fn encode_with_wr_joint_rand<X: Xof<SEED_SIZE>, const SEED_SIZE: usize>(
         &self,
         gradient: impl Iterator<Item = f64>,
-        wr_joint_rand_seed: &Seed<16>,
+        wr_joint_rand_seed: &Seed<SEED_SIZE>,
     ) -> Result<(Vec<F>, [F; NUM_WR_TESTS]), VdafError> {
         let mut meas = Vec::with_capacity(self.input_len());
         self.append_encoded_gradient(&mut meas, gradient)?;
-        let wr_test_results = self.append_wr_test_results(&mut meas, wr_joint_rand_seed)?;
+        let wr_test_results =
+            self.append_wr_test_results::<X, SEED_SIZE>(&mut meas, wr_joint_rand_seed)?;
         Ok((meas, wr_test_results))
     }
 
@@ -129,12 +130,13 @@ impl<F: FftFriendlyFieldElement> PineType<F> {
     /// `wr_test_results` on its own.
     ///
     /// In case of error, the buffer may have been modified and must be cleared before re-use.
-    pub(crate) fn append_wr_test_results(
+    pub(crate) fn append_wr_test_results<X: Xof<SEED_SIZE>, const SEED_SIZE: usize>(
         &self,
         meas: &mut Vec<F>,
-        wr_joint_rand_seed: &Seed<16>,
+        wr_joint_rand_seed: &Seed<SEED_SIZE>,
     ) -> Result<[F; NUM_WR_TESTS], VdafError> {
-        let wr_test_results = self.run_wr_tests(&meas[..self.cfg.dimension], wr_joint_rand_seed);
+        let wr_test_results =
+            self.run_wr_tests::<X, SEED_SIZE>(&meas[..self.cfg.dimension], wr_joint_rand_seed);
 
         let mut wr_success_count = 0;
         for wr_test_result in wr_test_results {
@@ -169,17 +171,13 @@ impl<F: FftFriendlyFieldElement> PineType<F> {
 
     /// Run the wraparound tests. For each test, compute the dot product of the gradient and a
     /// random `{-1, 0, 1}`-vector derived from the provided seed.
-    pub(crate) fn run_wr_tests(
+    pub(crate) fn run_wr_tests<X: Xof<SEED_SIZE>, const SEED_SIZE: usize>(
         &self,
         gradient: &[F],
-        wr_joint_rand_seed: &Seed<16>,
+        wr_joint_rand_seed: &Seed<SEED_SIZE>,
     ) -> [F; NUM_WR_TESTS] {
         debug_assert_eq!(gradient.len(), self.cfg.dimension);
-        let mut xof = XofTurboShake128::seed_stream(
-            wr_joint_rand_seed,
-            &self.cfg.dst(USAGE_WR_JOINT_RAND),
-            &[],
-        );
+        let mut xof = X::seed_stream(wr_joint_rand_seed, &self.cfg.dst(USAGE_WR_JOINT_RAND), &[]);
         let rand_len_per_test = chunk_count(4, self.cfg.dimension);
         let mut rand = vec![0_u8; rand_len_per_test * NUM_WR_TESTS];
         xof.fill(&mut rand[..]);
@@ -599,6 +597,7 @@ mod tests {
     use prio::{
         field::{Field128, FieldElement, FieldElementWithInteger},
         flp::{test_utils::FlpTest, Type},
+        vdaf::xof::XofTurboShake128,
     };
     use rand::prelude::*;
     use std::iter;
@@ -659,7 +658,7 @@ mod tests {
 
         let (input, wr_test_results) = pine
             .flp
-            .encode_with_wr_joint_rand(gradient, &Seed::generate().unwrap())
+            .encode_with_wr_joint_rand::<XofTurboShake128, 16>(gradient, &Seed::generate().unwrap())
             .unwrap();
         let (_encoded_gradient, rest) = input.split_at(pine.flp.cfg.dimension);
         let (bit_checked, rest) = rest.split_at(pine.flp.cfg.bit_checked_len);
@@ -685,11 +684,14 @@ mod tests {
         }
     }
 
-    impl<F: FftFriendlyFieldElement> Pine<F> {
+    impl<F: FftFriendlyFieldElement> Pine<F, XofTurboShake128, 16> {
         fn run_valid_test_case(&self, gradient: impl Iterator<Item = f64>) {
             let (mut input, wr_test_results) = self
                 .flp
-                .encode_with_wr_joint_rand(gradient, &Seed::generate().unwrap())
+                .encode_with_wr_joint_rand::<XofTurboShake128, 16>(
+                    gradient,
+                    &Seed::generate().unwrap(),
+                )
                 .unwrap();
             input.extend_from_slice(&wr_test_results);
 
@@ -738,7 +740,10 @@ mod tests {
         let pine = Pine::new_128(norm_bound, DIM, 15, 4, 5).unwrap();
         let (mut input, wr_test_results) = pine
             .flp
-            .encode_with_wr_joint_rand([0.0; DIM].into_iter(), &Seed::generate().unwrap())
+            .encode_with_wr_joint_rand::<XofTurboShake128, 16>(
+                [0.0; DIM].into_iter(),
+                &Seed::generate().unwrap(),
+            )
             .unwrap();
         input.extend_from_slice(&wr_test_results);
 
@@ -758,7 +763,10 @@ mod tests {
         let pine = Pine::new_128(norm_bound, DIM, 15, 4, 5).unwrap();
         let (mut input, wr_test_results) = pine
             .flp
-            .encode_with_wr_joint_rand([0.0; DIM].into_iter(), &Seed::generate().unwrap())
+            .encode_with_wr_joint_rand::<XofTurboShake128, 16>(
+                [0.0; DIM].into_iter(),
+                &Seed::generate().unwrap(),
+            )
             .unwrap();
         input.extend_from_slice(&wr_test_results);
 
@@ -779,7 +787,10 @@ mod tests {
         let pine = Pine::new_128(norm_bound, DIM, 15, 4, 5).unwrap();
         let (mut input, wr_test_results) = pine
             .flp
-            .encode_with_wr_joint_rand([0.0; DIM].into_iter(), &Seed::generate().unwrap())
+            .encode_with_wr_joint_rand::<XofTurboShake128, 16>(
+                [0.0; DIM].into_iter(),
+                &Seed::generate().unwrap(),
+            )
             .unwrap();
         input.extend_from_slice(&wr_test_results);
 
@@ -801,7 +812,10 @@ mod tests {
         let pine = Pine::new_128(norm_bound, DIM, 15, 4, 5).unwrap();
         let (mut input, wr_test_results) = pine
             .flp
-            .encode_with_wr_joint_rand([0.0; DIM].into_iter(), &Seed::generate().unwrap())
+            .encode_with_wr_joint_rand::<XofTurboShake128, 16>(
+                [0.0; DIM].into_iter(),
+                &Seed::generate().unwrap(),
+            )
             .unwrap();
         input.extend_from_slice(&wr_test_results);
 
@@ -824,7 +838,10 @@ mod tests {
         let pine = Pine::new_128(norm_bound, DIM, 15, 4, 5).unwrap();
         let (mut input, wr_test_results) = pine
             .flp
-            .encode_with_wr_joint_rand([0.1; DIM].into_iter(), &Seed::generate().unwrap())
+            .encode_with_wr_joint_rand::<XofTurboShake128, 16>(
+                [0.1; DIM].into_iter(),
+                &Seed::generate().unwrap(),
+            )
             .unwrap();
         let i = input.len();
         input.extend_from_slice(&wr_test_results);
diff --git a/crates/daphne/src/pine/mod.rs b/crates/daphne/src/pine/mod.rs
index c1de539..0779f21 100644
--- a/crates/daphne/src/pine/mod.rs
+++ b/crates/daphne/src/pine/mod.rs
@@ -13,9 +13,11 @@
 //!
 //! PINE is based on the scheme from [Rothblum et al. 2023](https://arxiv.org/abs/2311.10237).
 
+use std::marker::PhantomData;
+
 use prio::{
     field::{FftFriendlyFieldElement, Field128, Field64},
-    vdaf::VdafError,
+    vdaf::{xof::XofTurboShake128, VdafError},
 };
 
 use flp::{PineType, PineTypeSquaredNormEqual};
@@ -39,12 +41,13 @@ const USAGE_WR_JOINT_RAND_PART: u16 = 10;
 
 /// The PINE VDAF.
 #[derive(Clone, Debug, Eq, PartialEq)]
-pub struct Pine<F> {
+pub struct Pine<F, X, const SEED_SIZE: usize> {
     pub flp: PineType<F>,
     pub flp_sq_norm_equal: PineTypeSquaredNormEqual<F>,
+    phantom_data: PhantomData<X>,
 }
 
-pub type Pine128 = Pine<Field128>;
+pub type Pine128 = Pine<Field128, XofTurboShake128, 16>;
 
 impl Pine128 {
     /// Construct an instance of [`Pine128`] with the provided parameters.
@@ -67,7 +70,7 @@ impl Pine128 {
     }
 }
 
-pub type Pine64 = Pine<Field64>;
+pub type Pine64 = Pine<Field64, XofTurboShake128, 16>;
 
 impl Pine64 {
     /// Construct an instance of [`Pine64`] with the provided parameters.
@@ -128,7 +131,7 @@ impl<F> PineConfig<F> {
     }
 }
 
-impl<F: FftFriendlyFieldElement> Pine<F> {
+impl<F: FftFriendlyFieldElement, X, const SEED_SIZE: usize> Pine<F, X, SEED_SIZE> {
     /// Construct an instance of the Pine VDAF.
     ///
     /// # Parameters
@@ -276,6 +279,7 @@ impl<F: FftFriendlyFieldElement> Pine<F> {
         Ok(Self {
             flp: PineType { cfg: cfg.clone() },
             flp_sq_norm_equal: PineTypeSquaredNormEqual { cfg },
+            phantom_data: PhantomData,
         })
     }
 }
diff --git a/crates/daphne/src/pine/msg.rs b/crates/daphne/src/pine/msg.rs
index 60b16c0..3d87ee7 100644
--- a/crates/daphne/src/pine/msg.rs
+++ b/crates/daphne/src/pine/msg.rs
@@ -16,12 +16,12 @@ use super::{vdaf::PinePrepState, Pine};
 
 /// The public share sent by the Client to each Aggregator.
 #[derive(Clone, Debug)]
-pub struct PublicShare {
-    pub(crate) wr_joint_rand_parts: [Seed<16>; 2],
-    pub(crate) vf_joint_rand_parts: [Seed<16>; 2],
+pub struct PublicShare<const SEED_SIZE: usize> {
+    pub(crate) wr_joint_rand_parts: [Seed<SEED_SIZE>; 2],
+    pub(crate) vf_joint_rand_parts: [Seed<SEED_SIZE>; 2],
 }
 
-impl std::fmt::Display for PublicShare {
+impl<const SEED_SIZE: usize> std::fmt::Display for PublicShare<SEED_SIZE> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         let msg = self
             .get_encoded()
@@ -31,7 +31,7 @@ impl std::fmt::Display for PublicShare {
     }
 }
 
-impl Encode for PublicShare {
+impl<const SEED_SIZE: usize> Encode for PublicShare<SEED_SIZE> {
     fn encode(&self, bytes: &mut Vec<u8>) -> Result<(), CodecError> {
         self.wr_joint_rand_parts[0].encode(bytes)?;
         self.wr_joint_rand_parts[1].encode(bytes)?;
@@ -41,13 +41,15 @@ impl Encode for PublicShare {
     }
 
     fn encoded_len(&self) -> Option<usize> {
-        Some(16 * 4)
+        Some(SEED_SIZE * 4)
     }
 }
 
-impl<F> ParameterizedDecode<Pine<F>> for PublicShare {
+impl<F, X, const SEED_SIZE: usize> ParameterizedDecode<Pine<F, X, SEED_SIZE>>
+    for PublicShare<SEED_SIZE>
+{
     fn decode_with_param(
-        _pine: &Pine<F>,
+        _pine: &Pine<F, X, SEED_SIZE>,
         bytes: &mut std::io::Cursor<&[u8]>,
     ) -> Result<Self, CodecError> {
         Ok(Self {
@@ -59,26 +61,28 @@ impl<F> ParameterizedDecode<Pine<F>> for PublicShare {
 
 /// An input share sent by the Client to the Leader.
 #[derive(Clone, Debug)]
-pub struct InputShare<F>(pub(crate) InputShareFor<F>);
+pub struct InputShare<F, const SEED_SIZE: usize>(pub(crate) InputShareFor<F, SEED_SIZE>);
 
 #[derive(Clone, Debug)]
-pub(crate) enum InputShareFor<F> {
+pub(crate) enum InputShareFor<F, const SEED_SIZE: usize> {
     Leader {
         meas_share: Vec<F>,
         proofs_share: Vec<F>,
-        wr_blind: Seed<16>,
-        vf_blind: Seed<16>,
+        wr_blind: Seed<SEED_SIZE>,
+        vf_blind: Seed<SEED_SIZE>,
     },
 
     Helper {
-        meas_share: Seed<16>,
-        proofs_share: Seed<16>,
-        wr_blind: Seed<16>,
-        vf_blind: Seed<16>,
+        meas_share: Seed<SEED_SIZE>,
+        proofs_share: Seed<SEED_SIZE>,
+        wr_blind: Seed<SEED_SIZE>,
+        vf_blind: Seed<SEED_SIZE>,
     },
 }
 
-impl<F: FftFriendlyFieldElement> std::fmt::Display for InputShare<F> {
+impl<F: FftFriendlyFieldElement, const SEED_SIZE: usize> std::fmt::Display
+    for InputShare<F, SEED_SIZE>
+{
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         let msg = self
             .get_encoded()
@@ -91,7 +95,7 @@ impl<F: FftFriendlyFieldElement> std::fmt::Display for InputShare<F> {
     }
 }
 
-impl<F: FftFriendlyFieldElement> Encode for InputShare<F> {
+impl<F: FftFriendlyFieldElement, const SEED_SIZE: usize> Encode for InputShare<F, SEED_SIZE> {
     fn encode(&self, bytes: &mut Vec<u8>) -> Result<(), CodecError> {
         match &self.0 {
             InputShareFor::Leader {
@@ -131,16 +135,20 @@ impl<F: FftFriendlyFieldElement> Encode for InputShare<F> {
                 proofs_share,
                 ..
             } => Some(
-                meas_share.len() * F::ENCODED_SIZE + proofs_share.len() * F::ENCODED_SIZE + 16 * 2,
+                meas_share.len() * F::ENCODED_SIZE
+                    + proofs_share.len() * F::ENCODED_SIZE
+                    + SEED_SIZE * 2,
             ),
-            InputShareFor::Helper { .. } => Some(16 * 4),
+            InputShareFor::Helper { .. } => Some(SEED_SIZE * 4),
         }
     }
 }
 
-impl<F: FftFriendlyFieldElement> ParameterizedDecode<(&Pine<F>, usize)> for InputShare<F> {
+impl<F: FftFriendlyFieldElement, X, const SEED_SIZE: usize>
+    ParameterizedDecode<(&Pine<F, X, SEED_SIZE>, usize)> for InputShare<F, SEED_SIZE>
+{
     fn decode_with_param(
-        (pine, agg_id): &(&Pine<F>, usize),
+        (pine, agg_id): &(&Pine<F, X, SEED_SIZE>, usize),
         bytes: &mut std::io::Cursor<&[u8]>,
     ) -> Result<Self, CodecError> {
         match agg_id {
@@ -174,13 +182,15 @@ impl<F: FftFriendlyFieldElement> ParameterizedDecode<(&Pine<F>, usize)> for Inpu
 
 /// A prep share broadcast by one of the Aggregators.
 #[derive(Clone, Debug)]
-pub struct PrepShare<F> {
+pub struct PrepShare<F, const SEED_SIZE: usize> {
     pub(crate) verifiers_share: Vec<F>,
-    pub(crate) wr_joint_rand_part: Seed<16>,
-    pub(crate) vf_joint_rand_part: Seed<16>,
+    pub(crate) wr_joint_rand_part: Seed<SEED_SIZE>,
+    pub(crate) vf_joint_rand_part: Seed<SEED_SIZE>,
 }
 
-impl<F: FftFriendlyFieldElement> std::fmt::Display for PrepShare<F> {
+impl<F: FftFriendlyFieldElement, const SEED_SIZE: usize> std::fmt::Display
+    for PrepShare<F, SEED_SIZE>
+{
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         let msg = self
             .get_encoded()
@@ -190,7 +200,7 @@ impl<F: FftFriendlyFieldElement> std::fmt::Display for PrepShare<F> {
     }
 }
 
-impl<F: FftFriendlyFieldElement> Encode for PrepShare<F> {
+impl<F: FftFriendlyFieldElement, const SEED_SIZE: usize> Encode for PrepShare<F, SEED_SIZE> {
     fn encode(&self, bytes: &mut Vec<u8>) -> Result<(), CodecError> {
         for v in &self.verifiers_share {
             v.encode(bytes)?;
@@ -201,13 +211,15 @@ impl<F: FftFriendlyFieldElement> Encode for PrepShare<F> {
     }
 
     fn encoded_len(&self) -> Option<usize> {
-        Some(self.verifiers_share.len() * F::ENCODED_SIZE + 16 * 2)
+        Some(self.verifiers_share.len() * F::ENCODED_SIZE + SEED_SIZE * 2)
     }
 }
 
-impl<F: FftFriendlyFieldElement> ParameterizedDecode<PinePrepState<F>> for PrepShare<F> {
+impl<F: FftFriendlyFieldElement, const SEED_SIZE: usize>
+    ParameterizedDecode<PinePrepState<F, SEED_SIZE>> for PrepShare<F, SEED_SIZE>
+{
     fn decode_with_param(
-        state: &PinePrepState<F>,
+        state: &PinePrepState<F, SEED_SIZE>,
         bytes: &mut std::io::Cursor<&[u8]>,
     ) -> Result<Self, CodecError> {
         Ok(Self {
@@ -222,12 +234,12 @@ impl<F: FftFriendlyFieldElement> ParameterizedDecode<PinePrepState<F>> for PrepS
 
 /// The prep message combined from the prep shares.
 #[derive(Clone, Debug, Eq, PartialEq)]
-pub struct Prep {
-    pub(crate) wr_joint_rand_seed: Seed<16>,
-    pub(crate) vf_joint_rand_seed: Seed<16>,
+pub struct Prep<const SEED_SIZE: usize> {
+    pub(crate) wr_joint_rand_seed: Seed<SEED_SIZE>,
+    pub(crate) vf_joint_rand_seed: Seed<SEED_SIZE>,
 }
 
-impl std::fmt::Display for Prep {
+impl<const SEED_SIZE: usize> std::fmt::Display for Prep<SEED_SIZE> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         let msg = self
             .get_encoded()
@@ -237,7 +249,7 @@ impl std::fmt::Display for Prep {
     }
 }
 
-impl Encode for Prep {
+impl<const SEED_SIZE: usize> Encode for Prep<SEED_SIZE> {
     fn encode(&self, bytes: &mut Vec<u8>) -> Result<(), CodecError> {
         self.wr_joint_rand_seed.encode(bytes)?;
         self.vf_joint_rand_seed.encode(bytes)?;
@@ -245,13 +257,15 @@ impl Encode for Prep {
     }
 
     fn encoded_len(&self) -> Option<usize> {
-        Some(16 * 2)
+        Some(SEED_SIZE * 2)
     }
 }
 
-impl<F> ParameterizedDecode<PinePrepState<F>> for Prep {
+impl<F, const SEED_SIZE: usize> ParameterizedDecode<PinePrepState<F, SEED_SIZE>>
+    for Prep<SEED_SIZE>
+{
     fn decode_with_param(
-        _state: &PinePrepState<F>,
+        _state: &PinePrepState<F, SEED_SIZE>,
         bytes: &mut std::io::Cursor<&[u8]>,
     ) -> Result<Self, CodecError> {
         Ok(Self {
diff --git a/crates/daphne/src/pine/vdaf.rs b/crates/daphne/src/pine/vdaf.rs
index f46e803..3f81be8 100644
--- a/crates/daphne/src/pine/vdaf.rs
+++ b/crates/daphne/src/pine/vdaf.rs
@@ -5,19 +5,18 @@
 
 use std::{borrow::Cow, iter};
 
+use crate::pine::USAGE_QUERY_RAND;
 use prio::{
     codec::{CodecError, Decode, Encode, ParameterizedDecode},
     field::FftFriendlyFieldElement,
     flp::Type,
     vdaf::{
-        xof::{IntoFieldVec, Seed, Xof, XofTurboShake128},
+        xof::{IntoFieldVec, Seed, Xof},
         Aggregatable, Aggregator, Client, Collector, PrepareTransition, Vdaf, VdafError,
     },
 };
 use rand::prelude::*;
 
-use crate::pine::USAGE_QUERY_RAND;
-
 use super::{
     field_to_f64, msg, Pine, USAGE_MEAS_SHARE, USAGE_PROOF_SHARE, USAGE_PROVE_RAND,
     USAGE_VF_JOINT_RAND, USAGE_VF_JOINT_RAND_PART, USAGE_VF_JOINT_RAND_SEED,
@@ -43,9 +42,11 @@ impl<F: FftFriendlyFieldElement> Encode for PineVec<F> {
     }
 }
 
-impl<F: FftFriendlyFieldElement> ParameterizedDecode<(&Pine<F>, &())> for PineVec<F> {
+impl<F: FftFriendlyFieldElement, X, const SEED_SIZE: usize>
+    ParameterizedDecode<(&Pine<F, X, SEED_SIZE>, &())> for PineVec<F>
+{
     fn decode_with_param(
-        (pine, _agg_param): &(&Pine<F>, &()),
+        (pine, _agg_param): &(&Pine<F, X, SEED_SIZE>, &()),
         bytes: &mut std::io::Cursor<&[u8]>,
     ) -> Result<Self, CodecError> {
         iter::repeat_with(|| F::decode(bytes))
@@ -75,12 +76,14 @@ impl<F: FftFriendlyFieldElement> Aggregatable for PineVec<F> {
 
 // Vdaf
 
-impl<F: FftFriendlyFieldElement> Vdaf for Pine<F> {
+impl<F: FftFriendlyFieldElement, X: Xof<SEED_SIZE>, const SEED_SIZE: usize> Vdaf
+    for Pine<F, X, SEED_SIZE>
+{
     type Measurement = Vec<f64>;
     type AggregateResult = Vec<f64>;
     type AggregationParam = ();
-    type PublicShare = msg::PublicShare;
-    type InputShare = msg::InputShare<F>;
+    type PublicShare = msg::PublicShare<SEED_SIZE>;
+    type InputShare = msg::InputShare<F, SEED_SIZE>;
     type OutputShare = PineVec<F>;
     type AggregateShare = PineVec<F>;
 
@@ -93,16 +96,16 @@ impl<F: FftFriendlyFieldElement> Vdaf for Pine<F> {
     }
 }
 
-impl<F: FftFriendlyFieldElement> Pine<F> {
+impl<F: FftFriendlyFieldElement, X: Xof<SEED_SIZE>, const SEED_SIZE: usize> Pine<F, X, SEED_SIZE> {
     fn joint_rand_part(
         &self,
         agg_id: u8,
-        blind: &[u8; 16],
+        blind: &[u8; SEED_SIZE],
         nonce: &[u8; 16],
         meas_share: &[F],
         usage: u16,
-    ) -> Seed<16> {
-        let mut xof = XofTurboShake128::init(blind, &self.flp.cfg.dst(usage));
+    ) -> Seed<SEED_SIZE> {
+        let mut xof = X::init(blind, &self.flp.cfg.dst(usage));
         xof.update(&[agg_id]);
         xof.update(nonce);
         let mut buf = Vec::with_capacity(F::ENCODED_SIZE);
@@ -114,22 +117,26 @@ impl<F: FftFriendlyFieldElement> Pine<F> {
         xof.into_seed()
     }
 
-    fn joint_rand_seed(&self, joint_rand_parts: &[Seed<16>; 2], usage: u16) -> Seed<16> {
-        let mut xof = XofTurboShake128::init(&[0; 16], &self.flp.cfg.dst(usage));
+    fn joint_rand_seed(
+        &self,
+        joint_rand_parts: &[Seed<SEED_SIZE>; 2],
+        usage: u16,
+    ) -> Seed<SEED_SIZE> {
+        let mut xof = X::init(&[0; SEED_SIZE], &self.flp.cfg.dst(usage));
         xof.update(joint_rand_parts[0].as_ref());
         xof.update(joint_rand_parts[1].as_ref());
         xof.into_seed()
     }
 
-    fn helper_meas_share(&self, seed: &[u8; 16]) -> Vec<F> {
-        let mut xof = XofTurboShake128::init(seed, &self.flp.cfg.dst(USAGE_MEAS_SHARE));
+    fn helper_meas_share(&self, seed: &[u8; SEED_SIZE]) -> Vec<F> {
+        let mut xof = X::init(seed, &self.flp.cfg.dst(USAGE_MEAS_SHARE));
         xof.update(&[1]); // agg_id
         xof.into_seed_stream()
             .into_field_vec(self.flp.cfg.encoded_input_len)
     }
 
-    fn helper_proof_share(&self, seed: &[u8; 16]) -> Vec<F> {
-        let mut xof = XofTurboShake128::init(seed, &self.flp.cfg.dst(USAGE_PROOF_SHARE));
+    fn helper_proof_share(&self, seed: &[u8; SEED_SIZE]) -> Vec<F> {
+        let mut xof = X::init(seed, &self.flp.cfg.dst(USAGE_PROOF_SHARE));
         xof.update(&[self.flp.cfg.num_proofs, 1]); // num_proofs, agg_id
         xof.into_seed_stream().into_field_vec(
             self.flp_sq_norm_equal.proof_len()
@@ -140,24 +147,42 @@ impl<F: FftFriendlyFieldElement> Pine<F> {
 
 // Client
 
-impl<F: FftFriendlyFieldElement> Client<16> for Pine<F> {
+impl<F: FftFriendlyFieldElement, X: Xof<SEED_SIZE>, const SEED_SIZE: usize> Client<16>
+    for Pine<F, X, SEED_SIZE>
+{
     fn shard(
         &self,
         gradient: &Vec<f64>,
         nonce: &[u8; 16],
-    ) -> Result<(msg::PublicShare, Vec<msg::InputShare<F>>), VdafError> {
+    ) -> Result<
+        (
+            msg::PublicShare<SEED_SIZE>,
+            Vec<msg::InputShare<F, SEED_SIZE>>,
+        ),
+        VdafError,
+    > {
         let mut rng = thread_rng();
-        self.shard_with_rand(gradient, nonce, rng.gen())
+        let mut rand = [[0; SEED_SIZE]; NUM_SHARD_RAND_SEEDS];
+        for seed in &mut rand {
+            rng.fill(&mut seed[..]);
+        }
+        self.shard_with_rand(gradient, nonce, rand)
     }
 }
 
-impl<F: FftFriendlyFieldElement> Pine<F> {
+impl<F: FftFriendlyFieldElement, X: Xof<SEED_SIZE>, const SEED_SIZE: usize> Pine<F, X, SEED_SIZE> {
     pub(crate) fn shard_with_rand(
         &self,
         gradient: &[f64],
         nonce: &[u8; 16],
-        rand: [[u8; 16]; NUM_SHARD_RAND_SEEDS],
-    ) -> Result<(msg::PublicShare, Vec<msg::InputShare<F>>), VdafError> {
+        rand: [[u8; SEED_SIZE]; NUM_SHARD_RAND_SEEDS],
+    ) -> Result<
+        (
+            msg::PublicShare<SEED_SIZE>,
+            Vec<msg::InputShare<F, SEED_SIZE>>,
+        ),
+        VdafError,
+    > {
         let mut meas = Vec::with_capacity(self.flp.input_len());
         self.flp
             .append_encoded_gradient(&mut meas, gradient.iter().copied())?;
@@ -201,7 +226,7 @@ impl<F: FftFriendlyFieldElement> Pine<F> {
             let i = meas.len();
             let wr_test_results = self
                 .flp
-                .append_wr_test_results(&mut meas, &wr_joint_rand_seed)?;
+                .append_wr_test_results::<X, SEED_SIZE>(&mut meas, &wr_joint_rand_seed)?;
 
             for (m, s1) in meas[i..].iter().zip(meas_share_1[i..].iter()) {
                 meas_share_0.push(*m - *s1);
@@ -235,7 +260,7 @@ impl<F: FftFriendlyFieldElement> Pine<F> {
                 let vf_joint_rand_seed =
                     self.joint_rand_seed(&vf_joint_rand_parts, USAGE_VF_JOINT_RAND_SEED);
 
-                XofTurboShake128::seed_stream(
+                X::seed_stream(
                     &vf_joint_rand_seed,
                     &self.flp.cfg.dst(USAGE_VF_JOINT_RAND),
                     &[self.flp.cfg.num_proofs],
@@ -244,8 +269,7 @@ impl<F: FftFriendlyFieldElement> Pine<F> {
             };
 
             let prove_rands = {
-                let mut xof =
-                    XofTurboShake128::init(&prove_seed, &self.flp.cfg.dst(USAGE_PROVE_RAND));
+                let mut xof = X::init(&prove_seed, &self.flp.cfg.dst(USAGE_PROVE_RAND));
                 xof.update(&[self.flp.cfg.num_proofs]);
                 xof.into_seed_stream().into_field_vec(
                     self.flp_sq_norm_equal.prove_rand_len()
@@ -305,35 +329,37 @@ impl<F: FftFriendlyFieldElement> Pine<F> {
 // Aggregator
 
 #[derive(Clone, Debug)]
-pub struct PinePrepState<F> {
+pub struct PinePrepState<F, const SEED_SIZE: usize> {
     gradient_share: Vec<F>,
-    corrected_wr_joint_rand_seed: Seed<16>,
-    corrected_vf_joint_rand_seed: Seed<16>,
+    corrected_wr_joint_rand_seed: Seed<SEED_SIZE>,
+    corrected_vf_joint_rand_seed: Seed<SEED_SIZE>,
     pub(crate) verifiers_len: usize,
 }
 
-impl<F: FftFriendlyFieldElement> PartialEq for PinePrepState<F> {
+impl<F: FftFriendlyFieldElement, const SEED_SIZE: usize> PartialEq for PinePrepState<F, SEED_SIZE> {
     fn eq(&self, _other: &Self) -> bool {
         todo!("constant time compare")
     }
 }
 
-impl<F: FftFriendlyFieldElement> Eq for PinePrepState<F> {}
+impl<F: FftFriendlyFieldElement, const SEED_SIZE: usize> Eq for PinePrepState<F, SEED_SIZE> {}
 
-impl<F: FftFriendlyFieldElement> Aggregator<16, 16> for Pine<F> {
-    type PrepareState = PinePrepState<F>;
-    type PrepareShare = msg::PrepShare<F>;
-    type PrepareMessage = msg::Prep;
+impl<F: FftFriendlyFieldElement, X: Xof<SEED_SIZE>, const SEED_SIZE: usize>
+    Aggregator<SEED_SIZE, 16> for Pine<F, X, SEED_SIZE>
+{
+    type PrepareState = PinePrepState<F, SEED_SIZE>;
+    type PrepareShare = msg::PrepShare<F, SEED_SIZE>;
+    type PrepareMessage = msg::Prep<SEED_SIZE>;
 
     fn prepare_init(
         &self,
-        verify_key: &[u8; 16],
+        verify_key: &[u8; SEED_SIZE],
         agg_id: usize,
         _agg_param: &(),
         nonce: &[u8; 16],
-        public_share: &msg::PublicShare,
-        input_share: &msg::InputShare<F>,
-    ) -> Result<(PinePrepState<F>, msg::PrepShare<F>), VdafError> {
+        public_share: &msg::PublicShare<SEED_SIZE>,
+        input_share: &msg::InputShare<F, SEED_SIZE>,
+    ) -> Result<(PinePrepState<F, SEED_SIZE>, msg::PrepShare<F, SEED_SIZE>), VdafError> {
         let msg::PublicShare {
             wr_joint_rand_parts,
             vf_joint_rand_parts,
@@ -397,7 +423,7 @@ impl<F: FftFriendlyFieldElement> Aggregator<16, 16> for Pine<F> {
 
         // Run the wraparound tests.
         {
-            let wr_test_results = self.flp.run_wr_tests(
+            let wr_test_results = self.flp.run_wr_tests::<X, SEED_SIZE>(
                 &meas_share[..self.flp.cfg.dimension],
                 &corrected_wr_joint_rand_seed,
             );
@@ -422,7 +448,7 @@ impl<F: FftFriendlyFieldElement> Aggregator<16, 16> for Pine<F> {
         // Query the proofs.
         let verifiers_share = {
             let corrected_vf_joint_rands = {
-                XofTurboShake128::seed_stream(
+                X::seed_stream(
                     &corrected_vf_joint_rand_seed,
                     &self.flp.cfg.dst(USAGE_VF_JOINT_RAND),
                     &[self.flp.cfg.num_proofs],
@@ -431,8 +457,7 @@ impl<F: FftFriendlyFieldElement> Aggregator<16, 16> for Pine<F> {
             };
 
             let query_rands = {
-                let mut xof =
-                    XofTurboShake128::init(verify_key, &self.flp.cfg.dst(USAGE_QUERY_RAND));
+                let mut xof = X::init(verify_key, &self.flp.cfg.dst(USAGE_QUERY_RAND));
                 xof.update(&[self.flp.cfg.num_proofs]);
                 xof.update(nonce);
                 xof.into_seed_stream().into_field_vec(
@@ -487,11 +512,11 @@ impl<F: FftFriendlyFieldElement> Aggregator<16, 16> for Pine<F> {
         ))
     }
 
-    fn prepare_shares_to_prepare_message<M: IntoIterator<Item = msg::PrepShare<F>>>(
+    fn prepare_shares_to_prepare_message<M: IntoIterator<Item = msg::PrepShare<F, SEED_SIZE>>>(
         &self,
         _agg_param: &(),
         inputs: M,
-    ) -> Result<msg::Prep, VdafError> {
+    ) -> Result<msg::Prep<SEED_SIZE>, VdafError> {
         let [prep_share_0, prep_share_1] = inputs
             .into_iter()
             .collect::<Vec<_>>()
@@ -549,9 +574,9 @@ impl<F: FftFriendlyFieldElement> Aggregator<16, 16> for Pine<F> {
 
     fn prepare_next(
         &self,
-        state: PinePrepState<F>,
-        input: msg::Prep,
-    ) -> Result<PrepareTransition<Self, 16, 16>, VdafError> {
+        state: PinePrepState<F, SEED_SIZE>,
+        input: msg::Prep<SEED_SIZE>,
+    ) -> Result<PrepareTransition<Self, SEED_SIZE, 16>, VdafError> {
         if state.corrected_wr_joint_rand_seed != input.wr_joint_rand_seed {
             return Err(VdafError::Uncategorized(
                 "wraparound joint randomness check failed".into(),
@@ -583,7 +608,9 @@ impl<F: FftFriendlyFieldElement> Aggregator<16, 16> for Pine<F> {
 
 // Collector
 
-impl<F: FftFriendlyFieldElement> Collector for Pine<F> {
+impl<F: FftFriendlyFieldElement, X: Xof<SEED_SIZE>, const SEED_SIZE: usize> Collector
+    for Pine<F, X, SEED_SIZE>
+{
     fn unshard<M: IntoIterator<Item = PineVec<F>>>(
         &self,
         _agg_param: &(),
diff --git a/crates/daphne/src/vdaf/mod.rs b/crates/daphne/src/vdaf/mod.rs
index 71b7043..fb0caef 100644
--- a/crates/daphne/src/vdaf/mod.rs
+++ b/crates/daphne/src/vdaf/mod.rs
@@ -207,7 +207,7 @@ pub enum VdafPrepState {
         out_share: Vec<Field64>,
     },
     #[cfg(feature = "experimental")]
-    Pine128(PinePrepState<Field128>),
+    Pine128(PinePrepState<Field128, 16>),
 }
 
 #[cfg(any(test, feature = "test-utils"))]
@@ -277,7 +277,7 @@ pub enum VdafPrepMessage {
     #[cfg(feature = "experimental")]
     MasticShare(Field64),
     #[cfg(feature = "experimental")]
-    Pine128Share(crate::pine::msg::PrepShare<Field128>),
+    Pine128Share(crate::pine::msg::PrepShare<Field128, 16>),
 }
 
 #[cfg(any(test, feature = "test-utils"))]
diff --git a/crates/daphne/src/vdaf/pine.rs b/crates/daphne/src/vdaf/pine.rs
index b769119..024f56b 100644
--- a/crates/daphne/src/vdaf/pine.rs
+++ b/crates/daphne/src/vdaf/pine.rs
@@ -11,7 +11,11 @@ use crate::{
 use super::{
     shard_then_encode, VdafAggregateShare, VdafError, VdafPrepMessage, VdafPrepState, VdafVerifyKey,
 };
-use prio::{codec::ParameterizedDecode, field::FftFriendlyFieldElement, vdaf::Aggregator};
+use prio::{
+    codec::ParameterizedDecode,
+    field::FftFriendlyFieldElement,
+    vdaf::{xof::Xof, Aggregator},
+};
 use serde::{Deserialize, Serialize};
 
 /// [Pine](crate::pine::Pine) parameters.
@@ -258,14 +262,14 @@ impl PineConfig {
     }
 }
 
-fn prep_init<F: FftFriendlyFieldElement>(
-    vdaf: Pine<F>,
-    verify_key: &[u8; 16],
+fn prep_init<F: FftFriendlyFieldElement, X: Xof<SEED_SIZE>, const SEED_SIZE: usize>(
+    vdaf: Pine<F, X, SEED_SIZE>,
+    verify_key: &[u8; SEED_SIZE],
     agg_id: usize,
     nonce: &[u8; 16],
     public_share_data: &[u8],
     input_share_data: &[u8],
-) -> Result<(PinePrepState<F>, msg::PrepShare<F>), VdafError> {
+) -> Result<(PinePrepState<F, SEED_SIZE>, msg::PrepShare<F, SEED_SIZE>), VdafError> {
     // Parse the public share.
     let public_share = msg::PublicShare::get_decoded_with_param(&vdaf, public_share_data)?;
 


Test patch:
diff --git a/crates/daphne/src/pine/test_vec/mod.rs b/crates/daphne/src/pine/test_vec/mod.rs
index d5089e6..c7a41e0 100644
--- a/crates/daphne/src/pine/test_vec/mod.rs
+++ b/crates/daphne/src/pine/test_vec/mod.rs
@@ -5,7 +5,7 @@
 use prio::{
     codec::Encode,
     field::FftFriendlyFieldElement,
-    vdaf::{Aggregator, Collector, PrepareTransition},
+    vdaf::{xof::XofTurboShake128, Aggregator, Collector, PrepareTransition},
 };
 use serde::Deserialize;
 
@@ -54,7 +54,7 @@ struct TestVec {
     agg_result: Vec<f64>,
 }
 
-impl<F: FftFriendlyFieldElement> Pine<F> {
+impl<F: FftFriendlyFieldElement> Pine<F, XofTurboShake128, 16> {
     fn run_test_vec(&self, test_vec: &TestVec) {
         // Check that the test vector parameters have the values we expect.
         //