Skip to main content

binius_math/field_buffer/
structured.rs

1// Copyright 2026 The Binius Developers
2
3//! A field buffer whose nonzero values may be confined to one aligned block.
4
5use std::ops::Deref;
6
7use binius_compute::{Allocator, VecLike};
8use binius_field::PackedField;
9
10use super::{FieldBuffer, FieldSliceMut};
11
12/// A field buffer that may hold explicit values for only one aligned block, and zero elsewhere.
13///
14/// Passing the structure along lets a consumer that understands it skip the zeros. A consumer
15/// that does not calls [`Self::materialize`].
16#[derive(Debug, Clone)]
17pub enum StructuredBuffer<P: PackedField, Data: Deref<Target = [P]>> {
18	/// Every value held explicitly.
19	Buffer(FieldBuffer<P, Data>),
20	/// `inner` placed in block `index` of `2^log_n_blocks` equal blocks, and zero everywhere else.
21	ZeroPadded {
22		inner: Box<Self>,
23		log_n_blocks: usize,
24		/// Index in the range `0..1 << log_n_blocks`.
25		index: usize,
26	},
27}
28
29impl<P: PackedField, Data: Deref<Target = [P]>> StructuredBuffer<P, Data> {
30	/// Returns the base-2 logarithm of the number of field elements.
31	pub fn log_len(&self) -> usize {
32		match self {
33			Self::Buffer(buffer) => buffer.log_len(),
34			Self::ZeroPadded {
35				inner,
36				log_n_blocks,
37				..
38			} => inner.log_len() + log_n_blocks,
39		}
40	}
41
42	/// Writes every value into `dst`, which must hold zeros outside the explicit block.
43	///
44	/// # Panics
45	///
46	/// Panics if `dst` is not the same size as `self`, or any `ZeroPadded` index is out of range.
47	fn write_into(self, mut dst: FieldSliceMut<'_, P>) {
48		match self {
49			Self::Buffer(buffer) => {
50				assert_eq!(buffer.log_len(), dst.log_len()); // precondition
51				if buffer.log_len() < P::LOG_WIDTH {
52					dst.as_mut()[0] = P::from_scalars(buffer.iter_scalars());
53				} else {
54					dst.as_mut().copy_from_slice(buffer.as_ref());
55				}
56			}
57			Self::ZeroPadded {
58				inner,
59				log_n_blocks,
60				index,
61			} => {
62				let mut block = dst.chunk_mut(dst.log_len() - log_n_blocks, index);
63				inner.write_into(block.chunk());
64			}
65		}
66	}
67}
68
69impl<P: PackedField, Data: VecLike<P>> StructuredBuffer<P, Data> {
70	/// Writes out every value, zeros included, as one buffer.
71	///
72	/// A buffer already holding every value is returned with no copy.
73	pub fn materialize<A>(self, alloc: &A) -> FieldBuffer<P, Data>
74	where
75		A: Allocator<Vec<P> = Data>,
76	{
77		match self {
78			Self::Buffer(buffer) => buffer,
79			padded => {
80				let mut buffer = FieldBuffer::zeros_in(alloc, padded.log_len());
81				padded.write_into(buffer.as_mut_view());
82				buffer
83			}
84		}
85	}
86}
87
88impl<P: PackedField, Data: Deref<Target = [P]>> From<FieldBuffer<P, Data>>
89	for StructuredBuffer<P, Data>
90{
91	fn from(buffer: FieldBuffer<P, Data>) -> Self {
92		Self::Buffer(buffer)
93	}
94}
95
96#[cfg(test)]
97mod tests {
98	use binius_compute::GlobalAllocator;
99	use binius_field::{Field, PackedField, PackedGhash1x128b, PackedGhash4x128b};
100	use rand::{SeedableRng, rngs::StdRng};
101
102	use super::StructuredBuffer;
103	use crate::{FieldBuffer, test_utils::random_field_buffer};
104
105	/// Materializing a nested zero-padding matches placing the values by hand.
106	fn check<P: PackedField>(log_inner: usize) {
107		let mut rng = StdRng::seed_from_u64(0);
108		let inner = random_field_buffer::<P>(&mut rng, log_inner);
109
110		// The inner buffer at block 1 of 2, and that at block 2 of 4.
111		let structured = StructuredBuffer::ZeroPadded {
112			inner: Box::new(StructuredBuffer::ZeroPadded {
113				inner: Box::new(inner.clone().into()),
114				log_n_blocks: 1,
115				index: 1,
116			}),
117			log_n_blocks: 2,
118			index: 2,
119		};
120		assert_eq!(structured.log_len(), log_inner + 3);
121
122		let offset = (2 * 2 + 1) << log_inner;
123		let expected = FieldBuffer::<P>::from_values(
124			&(0..1 << (log_inner + 3))
125				.map(|i| {
126					if (offset..offset + (1 << log_inner)).contains(&i) {
127						inner.get(i - offset)
128					} else {
129						P::Scalar::ZERO
130					}
131				})
132				.collect::<Vec<_>>(),
133		);
134		assert_eq!(structured.materialize(&GlobalAllocator), expected);
135	}
136
137	#[test]
138	fn materialize_matches_naive_placement() {
139		for log_inner in 0..4 {
140			check::<PackedGhash1x128b>(log_inner);
141			check::<PackedGhash4x128b>(log_inner);
142		}
143	}
144}