Change visibility and minor refactors

benches/expand_from_coeff.rs

@@ -1,5 +1,8 @@
 use divan::{black_box, AllocProfiler, Bencher};
-use whir::algebra::{fields::Field64, ntt, random_vector};
+use whir::{
+    algebra::{fields::Field64, ntt, random_vector},
+    buffer::{Buffer, BufferMath, BufferOps},
+};
 
 #[global_allocator]
 static ALLOC: AllocProfiler = AllocProfiler::system();
@@ -28,15 +31,18 @@ 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<Buffer<Field64>> = (0..num_messages)
+                .map(|_| Buffer::from(random_vector(&mut rng, message_length)))
                 .collect();
-            (coeffs, expansion, coset_sz)
+            let masks = Buffer::zeros(0);
+            (coeffs, masks, expansion, coset_sz)
         })
-        .bench_values(|(coeffs, expansion, _coset_sz)| {
-            let coeffs_refs = coeffs.iter().map(|v| v.as_slice()).collect::<Vec<_>>();
+        .bench_values(|(coeffs, masks, expansion, _coset_sz)| {
+            let coeffs_refs = coeffs.iter().collect::<Vec<_>>();
+            let messages = ntt::Messages::new(&coeffs_refs, coeffs[0].len(), 1);
             black_box(ntt::interleaved_rs_encode(
-                &coeffs_refs,
+                messages,
+                &masks,
                 coeffs[0].len() * expansion,
             ))
         });

src/algebra/linear_form/univariate_evaluation.rs

@@ -1,9 +1,7 @@
 use ark_ff::Field;
 
 use super::LinearForm;
-use crate::algebra::{
-    embedding::Embedding, geometric_accumulate, linear_form::Evaluate, mixed_univariate_evaluate,
-};
+use crate::algebra::{embedding::Embedding, linear_form::Evaluate, mixed_univariate_evaluate};
 
 /// Linear form to represent univariate polynomial evaluation.
 ///
@@ -21,21 +19,6 @@ impl<F: Field> UnivariateEvaluation<F> {
     pub const fn new(point: F, size: usize) -> Self {
         Self { size, point }
     }
-
-    /// Batched version of [`LinearForm::accumulate`] for many [`UnivariateEvaluation`]s.
-    pub fn accumulate_many(evaluators: &[Self], accumulator: &mut [F], scalars: &[F]) {
-        assert_eq!(evaluators.len(), scalars.len());
-        let Some(size) = evaluators.first().map(|e| e.size) else {
-            return;
-        };
-        assert_eq!(accumulator.len(), size);
-        for evaluator in evaluators {
-            assert_eq!(evaluator.size, size);
-        }
-        let points = evaluators.iter().map(|e| e.point).collect::<Vec<F>>();
-        let scalars = scalars.to_vec();
-        geometric_accumulate(accumulator, scalars, &points);
-    }
 }
 
 impl<F: Field> LinearForm<F> for UnivariateEvaluation<F> {
@@ -55,7 +38,6 @@ impl<F: Field> LinearForm<F> for UnivariateEvaluation<F> {
         result
     }
 
-    /// See also [`Self::accumulate_many`] for a more efficient batched version.
     fn accumulate(&self, accumulator: &mut [F], scalar: F) {
         assert_eq!(accumulator.len(), self.size);
         let mut power = scalar;

src/algebra/ntt/cooley_tukey.rs

@@ -14,11 +14,15 @@ 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::Buffer};
+use crate::{
+    algebra::ntt::utils::divisors,
+    buffer::{Buffer, BufferOps},
+    utils::{chunks_exact_or_empty, zip_strict},
+};
 
 // Supported primes
 const PRIMES: [usize; 2] = [2, 3];
@@ -398,19 +402,41 @@ impl<F: Field> ReedSolomon<F> for NttEngine<F> {
         result
     }
 
-    #[cfg_attr(feature = "tracing", instrument(skip(self, polys), fields(
-        num_polys = polys.len(),
-        poly_length = polys.first().map(|p| p.len()),
+    #[cfg_attr(feature = "tracing", instrument(skip(self, messages, masks), fields(
+        num_polys = messages.vectors.len() * messages.interleaving_depth,
+        message_length = messages.message_length,
         codeword_length = codeword_length,
     )))]
-    fn interleaved_encode(&self, polys: &[&[F]], codeword_length: usize) -> Buffer<F> {
+    fn interleaved_encode(
+        &self,
+        messages: Messages<'_, F>,
+        masks: &Buffer<F>,
+        codeword_length: usize,
+    ) -> Buffer<F> {
+        let vectors = messages
+            .vectors
+            .iter()
+            .map(|vector| vector.to_slice())
+            .collect::<Vec<_>>();
+        let messages = vectors
+            .iter()
+            .flat_map(|vector| {
+                chunks_exact_or_empty(vector, messages.message_length, messages.interleaving_depth)
+            })
+            .collect::<Vec<_>>();
         assert!(self.order.is_multiple_of(codeword_length));
-        if polys.is_empty() {
+        if messages.is_empty() {
+            assert!(masks.is_empty());
             return Buffer::from(Vec::new());
         }
-        let num_polys = polys.len();
-        let poly_length = polys[0].len();
-        assert!(polys.iter().all(|p| p.len() == poly_length));
+        let num_polys = messages.len();
+        let message_length = messages[0].len();
+        assert!(messages
+            .iter()
+            .all(|message| message.len() == message_length));
+        assert!(masks.len().is_multiple_of(num_polys));
+        let mask_length = masks.len() / num_polys;
+        let poly_length = message_length + mask_length;
         assert!(poly_length <= codeword_length);
 
         // Coset-NTT: instead of doing one codeword-length NTT on mostly zeros,
@@ -434,10 +460,14 @@ impl<F: Field> ReedSolomon<F> for NttEngine<F> {
         // Lay out twisted coefficients in contiguous coset blocks of length
         // `coset_size`, zero-padding each block as needed.
         let mut result = Vec::with_capacity(num_polys * codeword_length);
-        for poly in polys {
+        for (message, mask) in zip_strict(
+            messages,
+            chunks_exact_or_empty(masks.to_slice(), mask_length, num_polys),
+        ) {
             // FFT[a 0 0 0] = [a a a a], so just replicate input in coset dimension.
             for _ in 0..num_cosets {
-                result.extend_from_slice(poly);
+                result.extend_from_slice(message);
+                result.extend_from_slice(mask);
                 result.resize(result.len() + coset_padding, F::ZERO);
             }
         }

src/algebra/ntt/mod.rs

@@ -22,7 +22,7 @@ pub use self::{
 };
 use crate::{
     algebra::fields,
-    buffer::{Buffer, DefaultRs},
+    buffer::{Buffer, BufferOps, DefaultRs},
     type_map::{self, TypeMap},
 };
 
@@ -62,12 +62,42 @@ impl type_map::Family for NttFamily {
     type Dyn<F: 'static> = dyn ReedSolomon<F>;
 }
 
+/// Buffer-native description of interleaved Reed-Solomon messages.
+///
+/// Each vector contains `interleaving_depth` consecutive messages of
+/// `message_length` elements. The encoder decides how to gather those chunks,
+/// append masks, and pad them for the NTT.
+// TODO: Generalize `Messages` plus the separate mask buffer into an
+// IRS-agnostic segmented-polynomial batch so the encoder sees only resident
+// prefix/suffix segments, not message/mask protocol semantics.
+#[derive(Clone)]
+pub struct Messages<'a, F> {
+    pub vectors: &'a [&'a Buffer<F>],
+    pub message_length: usize,
+    pub interleaving_depth: usize,
+}
+
+impl<'a, F: Copy> Messages<'a, F> {
+    pub fn new(
+        vectors: &'a [&'a Buffer<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,
+        }
+    }
+}
+
 /// Reed-Solomon encoder for a given field `F`.
 ///
-/// Pure-NTT abstraction: encodes polynomials, knows nothing about how callers
-/// structure those polynomials (whir's IRS, for example, concatenates a
-/// message and a mask into a single polynomial before calling this trait —
-/// that split lives entirely on the caller side).
+/// The input remains buffer-native so each backend can gather messages, append
+/// masks, and pad without forcing protocol code through host slices.
 pub trait ReedSolomon<F>: Debug + Send + Sync {
     /// Smallest supported codeword length `>= size`, or `None` if `size`
     /// exceeds the engine's maximum order. The returned length is always
@@ -96,16 +126,19 @@ pub trait ReedSolomon<F>: Debug + Send + Sync {
         indices: &[usize],
     ) -> Vec<F>;
 
-    /// Batch-encode polynomials in parallel.
+    /// Batch-encode masked polynomials in parallel.
     ///
-    /// All `polys[i]` must have the same length. Output is a flat buffer of
-    /// `polys.len() * codeword_length` elements in row-major
-    /// `(eval_index, poly)` layout: `result[i * polys.len() + j]` is poly
-    /// `j`'s value at the `i`-th evaluation point.
+    /// Each logical polynomial is a message chunk followed by its mask row.
+    /// Output is a flat buffer in row-major `(eval_index, polynomial)` layout.
     ///
     /// `codeword_length` must be NTT-smooth for this engine and at least the
     /// polynomial length.
-    fn interleaved_encode(&self, polys: &[&[F]], codeword_length: usize) -> Buffer<F>;
+    fn interleaved_encode(
+        &self,
+        messages: Messages<'_, F>,
+        masks: &Buffer<F>,
+        codeword_length: usize,
+    ) -> Buffer<F>;
 }
 
 assert_obj_safe!(ReedSolomon<crate::algebra::fields::Field256>);
@@ -126,10 +159,14 @@ pub fn evaluation_points<F: 'static>(
         .evaluation_points(poly_length, codeword_length, indices)
 }
 
-pub fn interleaved_rs_encode<F: 'static>(polys: &[&[F]], codeword_length: usize) -> Buffer<F> {
+pub fn interleaved_rs_encode<F: 'static>(
+    messages: Messages<'_, F>,
+    masks: &Buffer<F>,
+    codeword_length: usize,
+) -> Buffer<F> {
     NTT.get::<F>()
         .expect("Unsupported NTT field.")
-        .interleaved_encode(polys, codeword_length)
+        .interleaved_encode(messages, masks, codeword_length)
 }
 
 pub fn generator<F: 'static>(codeword_length: usize) -> F {
@@ -149,8 +186,8 @@ mod tests {
 
     use super::*;
     use crate::{
-        algebra::{random_vector, univariate_evaluate},
-        buffer::BufferOps,
+        algebra::univariate_evaluate,
+        buffer::{BufferMath, BufferOps},
         utils::zip_strict,
     };
 
@@ -191,19 +228,12 @@ mod tests {
         )| {
             let mut rng = StdRng::seed_from_u64(seed);
             let messages = (0..num_messages)
-                .map(|_| random_vector(&mut rng, message_length))
+                .map(|_| Buffer::random(&mut rng, message_length))
                 .collect::<Vec<_>>();
-            let masks: Vec<Vec<F>> = (0..num_messages)
-                .map(|_| random_vector(&mut rng, mask_length))
-                .collect();
-            // Build each polynomial as `message || mask`. The engine takes
-            // unified polynomial slices; the message/mask split is purely a
-            // caller-side concept.
-            let polys: Vec<Vec<F>> = (0..num_messages)
-                .map(|i| messages[i].iter().chain(masks[i].iter()).copied().collect())
-                .collect();
-            let poly_refs: Vec<&[F]> = polys.iter().map(Vec::as_slice).collect();
-            let codeword = ntt.interleaved_encode(&poly_refs, codeword_length);
+            let masks = Buffer::random(&mut rng, mask_length * num_messages);
+            let message_refs = messages.iter().collect::<Vec<_>>();
+            let input = Messages::new(&message_refs, message_length, 1);
+            let codeword = ntt.interleaved_encode(input, &masks, codeword_length);
 
             // Output must be the right size.
             assert_eq!(codeword.len(), codeword_length * num_messages);
@@ -213,9 +243,10 @@ mod tests {
             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];
-                for ((message, mask), value) in zip_strict(zip_strict(&messages, &masks), evaluations) {
+                for (poly_index, (message, value)) in zip_strict(&messages, evaluations).enumerate() {
+                    let mask = &masks.to_slice()[poly_index * mask_length..(poly_index + 1) * mask_length];
                     assert_eq!(*value,
-                        univariate_evaluate(message, evaluation_point)
+                        univariate_evaluate(message.to_slice(), evaluation_point)
                         + evaluation_point.pow([message_length as u64])
                         * univariate_evaluate(mask, evaluation_point));
                 }