Resolve conflicts with buffer abstraction

benches/expand_from_coeff.rs

@@ -1,5 +1,11 @@
 use divan::{black_box, AllocProfiler, Bencher};
-use whir::algebra::{fields::Field64, ntt, random_vector};
+use whir::{
+    algebra::{
+        fields::Field64,
+        ntt::{Messages, NttEngine, ReedSolomon},
+    },
+    buffer::{ActiveBuffer, Buffer, BufferOps},
+};
 
 #[global_allocator]
 static ALLOC: AllocProfiler = AllocProfiler::system();
@@ -28,18 +34,17 @@ fn interleaved_rs_encode(bencher: Bencher, case: &(usize, usize, usize)) {
             let message_length = 1 << (exp - coset_sz);
             let num_messages = 1 << coset_sz;
             let mut rng = ark_std::rand::thread_rng();
-            let coeffs: Vec<Vec<Field64>> = (0..num_messages)
-                .map(|_| random_vector(&mut rng, message_length))
+            let coeffs: Vec<ActiveBuffer<Field64>> = (0..num_messages)
+                .map(|_| ActiveBuffer::random(&mut rng, message_length))
                 .collect();
-            (coeffs, expansion, coset_sz)
+            let engine = NttEngine::<Field64>::new_from_fftfield();
+            (engine, coeffs, expansion)
         })
-        .bench_values(|(coeffs, expansion, _coset_sz)| {
-            let coeffs_refs = coeffs.iter().map(|v| v.as_slice()).collect::<Vec<_>>();
-            black_box(ntt::interleaved_rs_encode(
-                &coeffs_refs,
-                &[],
-                coeffs[0].len() * expansion,
-            ))
+        .bench_values(|(engine, coeffs, expansion)| {
+            let coeffs_refs = coeffs.iter().collect::<Vec<_>>();
+            let messages = Messages::new(&coeffs_refs, coeffs[0].len(), 1);
+            let masks = ActiveBuffer::from_slice(&[]);
+            black_box(engine.interleaved_encode(messages, &masks, coeffs[0].len() * expansion))
         });
 }
 

src/algebra/ntt/cooley_tukey.rs

@@ -14,12 +14,13 @@ use {crate::utils::workload_size, rayon::prelude::*, std::cmp::max};
 use super::{
     transpose,
     utils::{lcm, sqrt_factor},
-    ReedSolomon,
+    Messages, ReedSolomon,
 };
 #[cfg(not(feature = "rs_in_order"))]
 use crate::algebra::ntt::transpose::transpose_permute;
 use crate::{
     algebra::ntt::utils::divisors,
+    buffer::{ActiveBuffer, BufferOps},
     utils::{chunks_exact_or_empty, zip_strict},
 };
 
@@ -402,17 +403,33 @@ impl<F: Field> ReedSolomon<F> for NttEngine<F> {
     }
 
     #[cfg_attr(feature = "tracing", instrument(skip(self, messages, masks), fields(
-        num_messages = messages.len(),
-        message_len = messages.first().map(|c| c.len()),
+        num_messages = messages.vectors.len() * messages.interleaving_depth,
+        message_len = messages.message_length,
         codeword_length = codeword_length,
-        mask_len = masks.len().checked_div(messages.len())
+        mask_len = masks.len().checked_div(messages.vectors.len() * messages.interleaving_depth)
 
     )))]
-    fn interleaved_encode(&self, messages: &[&[F]], masks: &[F], codeword_length: usize) -> Vec<F> {
+    fn interleaved_encode(
+        &self,
+        messages: Messages<'_, F>,
+        masks: &ActiveBuffer<F>,
+        codeword_length: usize,
+    ) -> ActiveBuffer<F> {
+        let vectors = messages
+            .vectors
+            .iter()
+            .map(|v| v.to_slice())
+            .collect::<Vec<_>>();
+        let messages = vectors
+            .iter()
+            .flat_map(|v| {
+                chunks_exact_or_empty(v, messages.message_length, messages.interleaving_depth)
+            })
+            .collect::<Vec<_>>();
         assert!(self.order.is_multiple_of(codeword_length));
         if messages.is_empty() {
             assert!(masks.is_empty());
-            return Vec::new();
+            return ActiveBuffer::from_vec(Vec::new());
         }
         let num_messages = messages.len();
         let message_len = messages[0].len();
@@ -445,7 +462,7 @@ impl<F: Field> ReedSolomon<F> for NttEngine<F> {
         let mut result = Vec::with_capacity(num_messages * codeword_length);
         for (message, mask) in zip_strict(
             messages,
-            chunks_exact_or_empty(masks, mask_length, num_messages),
+            chunks_exact_or_empty(masks.to_slice(), mask_length, num_messages),
         ) {
             // FFT[a 0 0 0] = [a a a a], so just replicate input in coset dimension.
             for _ in 0..num_cosets {
@@ -473,7 +490,7 @@ impl<F: Field> ReedSolomon<F> for NttEngine<F> {
 
         // Transpose to row-major order with vectors stacked horizontally.
         transpose(&mut result, num_messages, codeword_length);
-        result
+        ActiveBuffer::from_vec(result)
     }
 }
 

src/algebra/ntt/mod.rs

@@ -22,34 +22,35 @@ pub use self::{
 };
 use crate::{
     algebra::fields,
+    buffer::{ActiveBuffer, BufferOps, DefaultRs},
     type_map::{self, TypeMap},
 };
 
 pub static NTT: LazyLock<TypeMap<NttFamily>> = LazyLock::new(|| {
     let map = TypeMap::new();
     map.insert(
-        Arc::new(NttEngine::<fields::Field64>::new_from_fftfield()) as Arc<dyn ReedSolomon<_>>
+        Arc::new(DefaultRs::<fields::Field64>::new_from_fftfield()) as Arc<dyn ReedSolomon<_>>
     );
     map.insert(
-        Arc::new(NttEngine::<fields::Field128>::new_from_fftfield()) as Arc<dyn ReedSolomon<_>>
+        Arc::new(DefaultRs::<fields::Field128>::new_from_fftfield()) as Arc<dyn ReedSolomon<_>>
     );
     map.insert(
-        Arc::new(NttEngine::<fields::Field192>::new_from_fftfield()) as Arc<dyn ReedSolomon<_>>
+        Arc::new(DefaultRs::<fields::Field192>::new_from_fftfield()) as Arc<dyn ReedSolomon<_>>
     );
     map.insert(
-        Arc::new(NttEngine::<fields::Field256>::new_from_fftfield()) as Arc<dyn ReedSolomon<_>>
+        Arc::new(DefaultRs::<fields::Field256>::new_from_fftfield()) as Arc<dyn ReedSolomon<_>>
     );
     map.insert(
-        Arc::new(NttEngine::<fields::Field64_2>::new_from_fftfield()) as Arc<dyn ReedSolomon<_>>,
+        Arc::new(DefaultRs::<fields::Field64_2>::new_from_fftfield()) as Arc<dyn ReedSolomon<_>>,
     );
     map.insert(
-        Arc::new(NttEngine::<fields::Field64_3>::new_from_fftfield()) as Arc<dyn ReedSolomon<_>>,
+        Arc::new(DefaultRs::<fields::Field64_3>::new_from_fftfield()) as Arc<dyn ReedSolomon<_>>,
     );
     map.insert(Arc::new(
-        NttEngine::<<fields::Field64_2 as Field>::BasePrimeField>::new_from_fftfield(),
+        DefaultRs::<<fields::Field64_2 as Field>::BasePrimeField>::new_from_fftfield(),
     ) as Arc<dyn ReedSolomon<_>>);
     map.insert(Arc::new(
-        NttEngine::<<fields::Field64_3 as Field>::BasePrimeField>::new_from_fftfield(),
+        DefaultRs::<<fields::Field64_3 as Field>::BasePrimeField>::new_from_fftfield(),
     ) as Arc<dyn ReedSolomon<_>>);
     map
 });
@@ -61,6 +62,37 @@ impl type_map::Family for NttFamily {
     type Dyn<F: 'static> = dyn ReedSolomon<F>;
 }
 
+/// Interleaved Reed-Solomon input buffers.
+///
+/// Chunking buffer-backed vectors into per-message views can require backend work. This descriptor
+/// keeps the encoding input compact and lets each backend derive the chunked messages internally.
+///
+/// Each vector stores `interleaving_depth` consecutive messages of `message_length` field
+/// elements.
+#[derive(Clone, Copy)]
+pub struct Messages<'a, F> {
+    pub(crate) vectors: &'a [&'a ActiveBuffer<F>],
+    pub(crate) message_length: usize,
+    pub(crate) interleaving_depth: usize,
+}
+
+impl<'a, F: Field> Messages<'a, F> {
+    pub fn new(
+        vectors: &'a [&'a ActiveBuffer<F>],
+        message_length: usize,
+        interleaving_depth: usize,
+    ) -> Self {
+        assert!(vectors
+            .iter()
+            .all(|vector| vector.len() == message_length * interleaving_depth));
+        Self {
+            vectors,
+            message_length,
+            interleaving_depth,
+        }
+    }
+}
+
 /// Trait for a Reed-Solomon encoder implementation for a given field `F`.
 pub trait ReedSolomon<F>: Debug + Send + Sync {
     /// Returns the next supported order equal or larger than `size`.
@@ -88,15 +120,20 @@ pub trait ReedSolomon<F>: Debug + Send + Sync {
 
     /// Compute a masked interleaved Reed-Solomon encoding.
     ///
-    /// `messages` are `num_messages` slices of `message_length` elements.
-    /// `masks` is a `num_messages` × `mask_length` matrix of blinding coefficients.
+    /// `messages` contains `num_vectors` buffers, each holding `interleaving_depth` messages of
+    /// `message_length` elements laid out contiguously. `masks` is a
+    /// `(num_vectors * interleaving_depth) × mask_length` matrix of blinding coefficients.
     /// `codeword_length` must be an NTT-smooth number >= `message_length + mask_length`.
-    /// returns an `codeword_length × num_messages` matrix.
+    /// Returns a `codeword_length × (num_vectors * interleaving_depth)` matrix.
     ///
     /// Each output value is the univariate polynomial evaluation in the evaluation point
     /// corresponding with the index of a coefficient list formed by concatenating message and mask.
-    ///
-    fn interleaved_encode(&self, messages: &[&[F]], masks: &[F], codeword_length: usize) -> Vec<F>;
+    fn interleaved_encode(
+        &self,
+        messages: Messages<'_, F>,
+        masks: &ActiveBuffer<F>,
+        codeword_length: usize,
+    ) -> ActiveBuffer<F>;
 }
 
 assert_obj_safe!(ReedSolomon<crate::algebra::fields::Field256>);
@@ -118,10 +155,10 @@ pub fn evaluation_points<F: 'static>(
 }
 
 pub fn interleaved_rs_encode<F: 'static>(
-    messages: &[&[F]],
-    masks: &[F],
+    messages: Messages<'_, F>,
+    masks: &ActiveBuffer<F>,
     codeword_length: usize,
-) -> Vec<F> {
+) -> ActiveBuffer<F> {
     NTT.get::<F>()
         .expect("Unsupported NTT field.")
         .interleaved_encode(messages, masks, codeword_length)
@@ -145,6 +182,7 @@ mod tests {
     use super::*;
     use crate::{
         algebra::{random_vector, univariate_evaluate},
+        buffer::BufferOps,
         utils::{chunks_exact_or_empty, zip_strict},
     };
 
@@ -188,10 +226,16 @@ mod tests {
                 .map(|_| random_vector(&mut rng, message_length))
                 .collect::<Vec<_>>();
             let masks = random_vector(&mut rng, mask_length * num_messages);
-            let message_refs = messages.iter().map(|v| v.as_slice()).collect::<Vec<_>>();
+            let vectors = messages
+                .iter()
+                .map(|message| ActiveBuffer::from_slice(message))
+                .collect::<Vec<_>>();
+            let vector_refs = vectors.iter().collect::<Vec<_>>();
+            let mask_buffer = ActiveBuffer::from_slice(&masks);
+            let rs_messages = Messages::new(&vector_refs, message_length, 1);
             let codeword = ntt.interleaved_encode(
-                &message_refs,
-                &masks,
+                rs_messages,
+                &mask_buffer,
                 codeword_length,
             );
 
@@ -200,6 +244,7 @@ mod tests {
 
             // Output values are polynomial evaluations in the evaluation points.
             let mut evaluation_points = ntt.evaluation_points(message_length + mask_length, codeword_length, &sampled_indices);
+            let codeword = codeword.to_slice();
             for (&index, &evaluation_point) in zip_strict(&sampled_indices, &evaluation_points) {
                 let evaluations = &codeword[index * num_messages.. (index + 1) * num_messages];
                 let masks = chunks_exact_or_empty(&masks, mask_length, num_messages);

src/bin/benchmark.rs

@@ -15,6 +15,7 @@ use whir::{
         linear_form::{Evaluate, LinearForm, MultilinearExtension},
     },
     bits::Bits,
+    buffer::{ActiveBuffer, BufferOps},
     cmdline_utils::{AvailableFields, AvailableHash},
     hash::HASH_COUNTER,
     parameters::ProtocolParameters,
@@ -162,12 +163,13 @@ where
 
         HASH_COUNTER.reset();
 
-        let witness = params.commit(&mut prover_state, &[&vector]);
+        let vector_buffer = ActiveBuffer::from_slice(&vector);
+        let witness = params.commit(&mut prover_state, &[&vector_buffer]);
 
         let _ = params.prove(
             &mut prover_state,
-            vec![Cow::Borrowed(vector.as_slice())],
-            vec![Cow::Owned(witness)],
+            &[&vector_buffer],
+            vec![&witness],
             vec![],
             Cow::Owned(vec![]),
         );
@@ -239,7 +241,8 @@ where
         HASH_COUNTER.reset();
         let whir_prover_time = Instant::now();
 
-        let witness = params.commit(&mut prover_state, &[&vector]);
+        let vector_buffer = ActiveBuffer::from_slice(&vector);
+        let witness = params.commit(&mut prover_state, &[&vector_buffer]);
 
         let prove_linear_forms: Vec<Box<dyn LinearForm<M::Target>>> = points
             .iter()
@@ -250,8 +253,8 @@ where
 
         let _ = params.prove(
             &mut prover_state,
-            vec![Cow::Borrowed(vector.as_slice())],
-            vec![Cow::Owned(witness)],
+            &[&vector_buffer],
+            vec![&witness],
             prove_linear_forms,
             Cow::Borrowed(evaluations.as_slice()),
         );

src/bin/main.rs

@@ -11,6 +11,7 @@ use whir::{
         linear_form::{Covector, Evaluate, LinearForm, MultilinearExtension},
     },
     bits::Bits,
+    buffer::{ActiveBuffer, BufferOps},
     cmdline_utils::{AvailableFields, AvailableHash},
     hash::HASH_COUNTER,
     parameters::ProtocolParameters,
@@ -149,9 +150,10 @@ where
     }
 
     let vector = (0..num_coeffs).map(M::Source::from).collect::<Vec<_>>();
+    let vector_buffer = ActiveBuffer::from_slice(&vector);
 
     let whir_commit_time = Instant::now();
-    let witness = params.commit(&mut prover_state, &[&vector]);
+    let witness = params.commit(&mut prover_state, &[&vector_buffer]);
     let whir_commit_time = whir_commit_time.elapsed();
 
     // Allocate constraints
@@ -184,8 +186,8 @@ where
     let whir_prove_time = Instant::now();
     let _ = params.prove(
         &mut prover_state,
-        vec![Cow::Borrowed(vector.as_slice())],
-        vec![Cow::Owned(witness)],
+        &[&vector_buffer],
+        vec![&witness],
         prove_linear_forms,
         Cow::Borrowed(evaluations.as_slice()),
     );
@@ -315,13 +317,14 @@ where
     }
 
     let whir_commit_time = Instant::now();
-    let witness = params.commit(&mut prover_state, &[vector.as_slice()]);
+    let vector_buffer = ActiveBuffer::from_slice(&vector);
+    let witness = params.commit(&mut prover_state, &[&vector_buffer]);
     let whir_commit_time = whir_commit_time.elapsed();
 
     let whir_prove_time = Instant::now();
     let _ = params.prove(
         &mut prover_state,
-        vec![Cow::Borrowed(&vector)],
+        &[&vector_buffer],
         witness,
         prove_linear_forms,
         Cow::Borrowed(&evaluations),