binius_math/field_buffer/
structured.rs1use std::ops::Deref;
6
7use binius_compute::{Allocator, VecLike};
8use binius_field::PackedField;
9
10use super::{FieldBuffer, FieldSliceMut};
11
12#[derive(Debug, Clone)]
17pub enum StructuredBuffer<P: PackedField, Data: Deref<Target = [P]>> {
18 Buffer(FieldBuffer<P, Data>),
20 ZeroPadded {
22 inner: Box<Self>,
23 log_n_blocks: usize,
24 index: usize,
26 },
27}
28
29impl<P: PackedField, Data: Deref<Target = [P]>> StructuredBuffer<P, Data> {
30 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 fn write_into(self, mut dst: FieldSliceMut<'_, P>) {
48 match self {
49 Self::Buffer(buffer) => {
50 assert_eq!(buffer.log_len(), dst.log_len()); 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 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 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 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}