Skip to main content

binius_math/field_buffer/
mod.rs

1// Copyright 2025 Irreducible Inc.
2// Copyright 2026 The Binius Developers
3
4//! A power-of-two-sized buffer of packed field elements.
5//!
6//! The buffer lives in this file, with its methods and the guard its halving methods hand out.
7//! `SplitMut` is built in exactly one place, so it needs no visibility beyond this module.
8//!
9//! The rest sits in modules of their own:
10//!
11//! ```text
12//! chunks      the buffer's chunk methods, and the iterators and guard they hand out
13//! view        the borrowed aliases, which name the buffer rather than being produced by it
14//! ```
15//!
16//! # Why the backing store is a type parameter
17//!
18//! A store's length fixes the element count only once the buffer fills a whole packed word.
19//! Below one word, the same single word backs any length:
20//!
21//! ```text
22//! WIDTH = 4 lanes per word
23//!
24//! 1 element   ->  [ x . . . ]   one word
25//! 2 elements  ->  [ x x . . ]   one word
26//! 4 elements  ->  [ x x x x ]   one word
27//! ```
28//!
29//! So the logical length is irreducible extra state, carried alongside the store.
30//! That rules out a transparent wrapper around a packed slice.
31//! With it goes the deref-coercion design a vector and its two slice types enjoy.
32
33use std::{
34	mem::MaybeUninit,
35	ops::{Deref, DerefMut},
36	slice,
37};
38
39use binius_compute::{Allocator, BufferData, VecLike};
40use binius_field::{
41	Field, PackedField,
42	packed::{get_packed_slice_unchecked, set_packed_slice_unchecked},
43};
44use binius_utils::{
45	checked_arithmetics::strict_log_2,
46	rayon::{
47		prelude::*,
48		task_size::{IndexedParallelIteratorExt, task_chunk_len},
49	},
50};
51use bytemuck::zeroed_vec;
52
53mod chunks;
54mod view;
55
56pub use chunks::{ChunkMut, Chunks, ChunksMut};
57pub use view::{FieldSlice, FieldSliceData, FieldSliceMut, FieldVec};
58
59/// A power-of-two-sized buffer containing field elements, stored in packed fields.
60///
61/// The backing store length is fully determined by `log_len`:
62///
63/// ```text
64/// words.len() == 1 << log_len.saturating_sub(P::LOG_WIDTH)
65/// ```
66///
67/// A buffer shorter than one packed word still occupies a whole word.
68/// The lanes at and past the logical length are then not elements of the buffer:
69///
70/// ```text
71/// WIDTH = 4, log_len = 1
72///
73/// word:  [ s_0, s_1, dead, dead ]
74///          ^^^^^^^^  live prefix
75/// ```
76///
77/// Nothing reads a dead lane as an element.
78/// Equality and the halving routines touch the live prefix only.
79///
80/// Truncation additionally zeros the dead lanes it creates.
81/// Other constructors may leave whatever the caller's store held there.
82// The buffer's methods are grouped by topic, one module per group, so its inherent impls span
83// more than one file.
84#[allow(clippy::multiple_inherent_impl)]
85#[derive(Debug, Clone, Eq)]
86pub struct FieldBuffer<P: PackedField, Data: Deref<Target = [P]> = Vec<P>> {
87	/// log2 the number over elements in the buffer.
88	log_len: usize,
89	/// The packed words.
90	words: Data,
91}
92
93impl<P: PackedField, Data: Deref<Target = [P]> + Copy> Copy for FieldBuffer<P, Data> {}
94
95impl<P: PackedField, Data: Deref<Target = [P]>> PartialEq for FieldBuffer<P, Data> {
96	fn eq(&self, other: &Self) -> bool {
97		// Equality compares only the live scalars, never the raw packed backing store.
98		//
99		// Invariant: buffers of different lengths are never equal.
100		//
101		// The length-below-width branch compares only the live prefix of a single packed word,
102		// never the dead lanes past it: shorter and longer buffers sharing that prefix would
103		// otherwise compare equal, and the dead lanes may hold unrelated data.
104		if self.log_len != other.log_len {
105			return false;
106		}
107		if self.log_len < P::LOG_WIDTH {
108			let iter_1 = self
109				.words
110				.first()
111				.expect("len >= 1")
112				.iter()
113				.take(1 << self.log_len);
114			let iter_2 = other
115				.words
116				.first()
117				.expect("len >= 1")
118				.iter()
119				.take(1 << self.log_len);
120			iter_1.eq(iter_2)
121		} else {
122			let prefix = 1 << (self.log_len - P::LOG_WIDTH);
123			self.words[..prefix] == other.words[..prefix]
124		}
125	}
126}
127
128impl<P: PackedField> FieldBuffer<P> {
129	/// Create a new FieldBuffer from a vector of values.
130	///
131	/// # Preconditions
132	///
133	/// * `values.len()` must be a power of two.
134	#[track_caller]
135	pub fn from_values(values: &[P::Scalar]) -> Self {
136		let log_len =
137			strict_log_2(values.len()).expect("precondition: values.len() must be a power of two");
138
139		let packed_len = 1 << log_len.saturating_sub(P::LOG_WIDTH);
140		let mut words = Vec::with_capacity(packed_len);
141		words.extend(
142			values
143				.chunks(P::WIDTH)
144				.map(|chunk| P::from_scalars(chunk.iter().copied())),
145		);
146
147		Self { log_len, words }
148	}
149
150	/// Builds a buffer of `2^log_len` zeros.
151	pub fn zeros(log_len: usize) -> Self {
152		let packed_len = 1 << log_len.saturating_sub(P::LOG_WIDTH);
153		let words = zeroed_vec(packed_len);
154		Self { log_len, words }
155	}
156
157	/// Builds a one-element buffer holding `value`, with room reserved to grow.
158	///
159	/// The store is sized for `2^log_capacity` elements up front.
160	/// Growing the buffer to that many therefore never reallocates.
161	pub fn scalar_with_capacity(value: P::Scalar, log_capacity: usize) -> Self {
162		let mut words = Vec::with_capacity(1 << log_capacity.saturating_sub(P::LOG_WIDTH));
163		words.push(P::from_scalars([value]));
164		Self { log_len: 0, words }
165	}
166}
167
168impl<P: PackedField, Data: VecLike<P>> FieldBuffer<P, Data> {
169	/// Builds a zeroed buffer of `2^log_len` elements, backed by memory drawn from `alloc`.
170	///
171	/// The allocator-aware counterpart to the plain zeroed constructor.
172	/// Under a pool the result is a recyclable buffer rather than a fresh allocation.
173	pub fn zeros_in<A>(alloc: &A, log_len: usize) -> Self
174	where
175		A: Allocator<Vec<P> = Data>,
176	{
177		let packed_len = 1 << log_len.saturating_sub(P::LOG_WIDTH);
178		// An allocator hands out an empty buffer, so the words are both created and zeroed here.
179		let mut words = alloc.alloc::<P>(packed_len);
180		words.resize(packed_len, P::default());
181		FieldBuffer::new(log_len, words)
182	}
183
184	/// Copies a borrowed buffer into memory drawn from `alloc`.
185	///
186	/// Whole packed words are copied, dead lanes and all, so the copy is bit-identical.
187	/// A long source spreads across workers.
188	pub fn from_view_in<A>(alloc: &A, src: FieldSlice<'_, P>) -> Self
189	where
190		A: Allocator<Vec<P> = Data>,
191	{
192		Self::from_view_with_capacity_in(alloc, src, src.log_len())
193	}
194
195	/// Copies a borrowed buffer into memory drawn from `alloc`, with room reserved to grow.
196	///
197	/// The buffer spans `src.log_len()` elements and its store is sized for `2^log_capacity`.
198	/// Growing it to that many elements therefore never reallocates.
199	/// [`Self::repeat_extend`] is the growth this reserves for.
200	///
201	/// Whole packed words are copied, dead lanes and all, so the copy is bit-identical.
202	/// A long source spreads across workers.
203	///
204	/// # Panics
205	///
206	/// Panics if `log_capacity` is shorter than what `src` spans.
207	#[track_caller]
208	pub fn from_view_with_capacity_in<A>(
209		alloc: &A,
210		src: FieldSlice<'_, P>,
211		log_capacity: usize,
212	) -> Self
213	where
214		A: Allocator<Vec<P> = Data>,
215	{
216		assert!(
217			log_capacity >= src.log_len(),
218			"precondition: log_capacity must be at least src.log_len()"
219		);
220
221		let mut words = alloc.alloc::<P>(1 << log_capacity.saturating_sub(P::LOG_WIDTH));
222
223		// The allocator rounds its blocks up, so the fill is bounded to the words that are live.
224		let source = src.as_ref();
225		let head = &mut words.spare_capacity_mut()[..source.len()];
226
227		// The floor holds a worker's share at one task's byte budget, so a short source stays put.
228		(head.par_iter_mut(), source.par_iter())
229			.into_par_iter()
230			.with_min_task_bytes::<P>()
231			.for_each(|(dst, src)| {
232				dst.write(*src);
233			});
234
235		// SAFETY: the loop above wrote every word of the source.
236		unsafe { words.set_len(source.len()) };
237
238		FieldBuffer::new(src.log_len(), words)
239	}
240
241	/// Grows the buffer to span `2^log_len` elements, repeating what it already holds.
242	///
243	/// ```text
244	///     [x]  ->  [x | x | x | x]
245	/// ```
246	///
247	/// This is the buffer's [`Vec::extend_from_within`], split across workers.
248	/// Every copy is drawn from the live prefix, so one pass fills the whole buffer.
249	///
250	/// Returns the buffer untouched when it already spans that many elements.
251	///
252	/// # Panics
253	///
254	/// Panics if `log_len` is shorter than what the buffer already spans.
255	/// Panics if the store has no room reserved for `2^log_len` elements.
256	#[track_caller]
257	pub fn repeat_extend(&mut self, log_len: usize) {
258		assert!(
259			log_len >= self.log_len,
260			"precondition: log_len must be at least the buffer's own log_len"
261		);
262		if log_len == self.log_len {
263			return;
264		}
265
266		let total = 1 << log_len.saturating_sub(P::LOG_WIDTH);
267		assert!(
268			total <= self.words.capacity(),
269			"precondition: the store must have room reserved for 2^log_len elements"
270		);
271
272		if self.log_len < P::LOG_WIDTH {
273			// The live elements share one word, so repeating them cycles lanes, not words.
274			let word = P::from_scalars(self.iter_scalars().cycle());
275			self.words.clear();
276			self.words.resize(total, word);
277		} else {
278			let prefix = self.words.len();
279
280			// A run is the words one worker copies at a time.
281			// It is a power of two at most the prefix, so it divides the prefix evenly.
282			let run = prefix.min(task_chunk_len::<P>().next_power_of_two());
283
284			// One borrow has to span both the prefix the copies read and the room they fill.
285			// Emptying the buffer first is what puts the whole store behind that single borrow.
286			//
287			// SAFETY: a packed field has no destructor, so forgetting the live words is a no-op.
288			unsafe { self.words.set_len(0) };
289			let store = &mut self.words.spare_capacity_mut()[..total];
290			let (head, tail) = store.split_at_mut(prefix);
291
292			// SAFETY: the split point is the length the buffer just had.
293			// So every word of `head` was written before this call.
294			let head = unsafe { &*(head as *const [MaybeUninit<P>] as *const [P]) };
295
296			repeat_words(head, tail, run);
297
298			// SAFETY: the call above wrote every word between the prefix and `total`.
299			unsafe { self.words.set_len(total) };
300		}
301
302		self.log_len = log_len;
303	}
304
305	/// Builds a buffer from scalar values, directly into memory drawn from `alloc`.
306	///
307	/// The allocator-aware counterpart to building from a scalar slice.
308	/// Packing happens straight into the allocator's buffer, so no intermediate vector is copied.
309	/// Under a pool the result is a recyclable buffer.
310	///
311	/// # Preconditions
312	///
313	/// * `values.len()` must be a power of two.
314	#[track_caller]
315	pub fn from_values_in<A>(alloc: &A, values: &[P::Scalar]) -> Self
316	where
317		A: Allocator<Vec<P> = Data>,
318	{
319		let log_len =
320			strict_log_2(values.len()).expect("precondition: values.len() must be a power of two");
321
322		let packed_len = 1 << log_len.saturating_sub(P::LOG_WIDTH);
323		let mut words = alloc.alloc::<P>(packed_len);
324		words.extend(
325			values
326				.chunks(P::WIDTH)
327				.map(|chunk| P::from_scalars(chunk.iter().copied())),
328		);
329
330		FieldBuffer::new(log_len, words)
331	}
332
333	/// Grows the buffer to span `log_len` variables, padding the new positions with zeros.
334	///
335	/// Returns the buffer untouched when it already spans that many.
336	/// Padding the shorter of two buffers to match the longer hits that case when both are equal.
337	///
338	/// New memory comes from `alloc`.
339	/// A pooled buffer therefore stays pooled, instead of escaping the pool on a reallocation.
340	///
341	/// # Panics
342	///
343	/// Panics if `log_len` is shorter than what the buffer already spans.
344	/// Shrinking is a separate operation.
345	/// Silently returning a narrower buffer than asked for would hide the mistake.
346	#[track_caller]
347	pub fn zero_extend_in<A>(self, alloc: &A, log_len: usize) -> Self
348	where
349		A: Allocator<Vec<P> = Data>,
350	{
351		assert!(
352			log_len >= self.log_len,
353			"precondition: log_len must be at least the buffer's own log_len"
354		);
355		// Nothing to pad, so the caller's buffer is handed straight back with no copy.
356		if log_len == self.log_len {
357			return self;
358		}
359
360		// The target starts fully zeroed, so only the occupied words need writing.
361		let mut extended = Self::zeros_in(alloc, log_len);
362
363		if self.log_len < P::LOG_WIDTH {
364			// The source's single word also spans lanes that the padding covers.
365			// Nothing sitting there is an element of it, so only the live scalars carry over.
366			//
367			//     source   [ s_0, s_1, x, y ]
368			//     result   [ s_0, s_1, 0, 0 ] [ 0, 0, 0, 0 ]
369			extended.as_mut()[0] = P::from_scalars(self.iter_scalars());
370		} else {
371			// Every lane of every source word is live, so whole words move across untouched.
372			extended.as_mut()[..self.as_ref().len()].copy_from_slice(self.as_ref());
373		}
374
375		extended
376	}
377}
378
379#[allow(clippy::len_without_is_empty)]
380impl<P: PackedField, Data: Deref<Target = [P]>> FieldBuffer<P, Data> {
381	/// Create a new FieldBuffer from a slice of packed words.
382	///
383	/// # Preconditions
384	///
385	/// * `words.len()` must equal the expected packed length for `log_len`.
386	#[track_caller]
387	pub fn new(log_len: usize, words: Data) -> Self {
388		let expected_packed_len = 1 << log_len.saturating_sub(P::LOG_WIDTH);
389		assert!(
390			words.len() == expected_packed_len,
391			"precondition: words.len() must equal expected packed length"
392		);
393
394		Self { log_len, words }
395	}
396
397	/// Consumes the buffer and returns its backing data store.
398	pub fn into_inner(self) -> Data {
399		self.words
400	}
401
402	/// Returns log2 the number of field elements.
403	pub const fn log_len(&self) -> usize {
404		self.log_len
405	}
406
407	/// Returns the number of field elements.
408	pub const fn len(&self) -> usize {
409		1 << self.log_len
410	}
411
412	/// Borrows the whole buffer as a shared view.
413	pub fn as_view(&self) -> FieldSlice<'_, P> {
414		FieldSlice::from_slice(self.log_len, self.as_ref())
415	}
416
417	/// Get a field element at the given index.
418	///
419	/// # Preconditions
420	///
421	/// * the index is in the range `0..self.len()`
422	#[track_caller]
423	pub fn get(&self, index: usize) -> P::Scalar {
424		assert!(
425			index < self.len(),
426			"precondition: index {index} must be less than len {}",
427			self.len()
428		);
429
430		// Safety: bound check on index performed above. The buffer length is at least
431		// `self.len() >> P::LOG_WIDTH` by struct invariant.
432		unsafe { get_packed_slice_unchecked(&self.words, index) }
433	}
434
435	/// Returns an iterator over the scalar elements in the buffer.
436	pub fn iter_scalars(&self) -> impl Iterator<Item = P::Scalar> + Send + Clone + '_ {
437		P::iter_slice(self.as_ref()).take(self.len())
438	}
439
440	/// Returns an iterator over the packed words the elements occupy.
441	///
442	/// The run covers the live words only, never a store's spare room past them.
443	/// A buffer shorter than one packed word yields that one word whole, dead lanes and all.
444	#[inline]
445	pub fn iter_packed(&self) -> slice::Iter<'_, P> {
446		self.as_ref().iter()
447	}
448
449	/// Splits the buffer in half and returns a pair of borrowed slices.
450	///
451	/// # Preconditions
452	///
453	/// * `self.log_len()` must be greater than 0.
454	#[track_caller]
455	pub fn split_half(&self) -> (FieldSlice<'_, P>, FieldSlice<'_, P>) {
456		assert!(self.log_len > 0, "precondition: cannot split a buffer of length 1");
457
458		let new_log_len = self.log_len - 1;
459		if new_log_len < P::LOG_WIDTH {
460			// The result will be two Single variants
461			// We have exactly one packed element that needs to be split
462			let packed = self.words[0];
463			let zeros = P::default();
464
465			let (first_half, second_half) = packed.interleave(zeros, new_log_len);
466
467			let first = FieldBuffer {
468				log_len: new_log_len,
469				words: FieldSliceData::Single(first_half),
470			};
471			let second = FieldBuffer {
472				log_len: new_log_len,
473				words: FieldSliceData::Single(second_half),
474			};
475
476			(first, second)
477		} else {
478			// Split the packed word slice in half
479			let half_len = 1 << (new_log_len - P::LOG_WIDTH);
480			let (first_half, second_half) = self.words.split_at(half_len);
481			let second_half = &second_half[..half_len];
482
483			let first = FieldBuffer {
484				log_len: new_log_len,
485				words: FieldSliceData::Slice(first_half),
486			};
487			let second = FieldBuffer {
488				log_len: new_log_len,
489				words: FieldSliceData::Slice(second_half),
490			};
491
492			(first, second)
493		}
494	}
495}
496
497impl<P: PackedField, Data: DerefMut<Target = [P]>> FieldBuffer<P, Data> {
498	/// Borrows the whole buffer as a mutable view.
499	pub fn as_mut_view(&mut self) -> FieldSliceMut<'_, P> {
500		FieldSliceMut::from_slice(self.log_len, self.as_mut())
501	}
502
503	/// Set a field element at the given index.
504	///
505	/// # Preconditions
506	///
507	/// * the index is in the range `0..self.len()`
508	#[track_caller]
509	pub fn set(&mut self, index: usize, value: P::Scalar) {
510		assert!(
511			index < self.len(),
512			"precondition: index {index} must be less than len {}",
513			self.len()
514		);
515
516		// Safety: bound check on index performed above. The buffer length is at least
517		// `self.len() >> P::LOG_WIDTH` by struct invariant.
518		unsafe { set_packed_slice_unchecked(&mut self.words, index, value) };
519	}
520
521	/// Returns a mutable iterator over the packed words the elements occupy.
522	///
523	/// The mutable counterpart of the shared word iterator, covering the same run of words.
524	/// A final word below the packing width is lent out whole, dead lanes included.
525	#[inline]
526	pub fn iter_packed_mut(&mut self) -> slice::IterMut<'_, P> {
527		self.as_mut().iter_mut()
528	}
529
530	/// Consumes the buffer and halves it, returning a guard that owns the store.
531	///
532	/// The guard lends out the two halves mutably.
533	///
534	/// When both halves live inside one packed word, the guard holds detached copies of them.
535	/// Those are merged back into the original word when it drops.
536	///
537	/// # Preconditions
538	///
539	/// * `self.log_len()` must be greater than 0.
540	#[track_caller]
541	pub fn into_split_half(self) -> SplitMut<P, Data> {
542		assert!(self.log_len > 0, "precondition: cannot split a buffer of length 1");
543
544		// Each half holds one variable fewer than the buffer, and the guard decides from that
545		// alone whether the halves have to be detached from the store.
546		SplitMut::new(self.log_len - 1, self.words)
547	}
548
549	/// Halves the buffer through a mutable borrow, rather than consuming it.
550	///
551	/// Equivalent to borrowing the buffer as a mutable view and halving that.
552	///
553	/// # Preconditions
554	///
555	/// * `self.log_len()` must be greater than 0.
556	#[track_caller]
557	pub fn split_half_mut(&mut self) -> SplitMut<P, &'_ mut [P]> {
558		self.as_mut_view().into_split_half()
559	}
560}
561
562impl<P: PackedField, Data: BufferData<P>> FieldBuffer<P, Data> {
563	/// Truncates the buffer to a shorter length, shrinking the backing store to match.
564	///
565	/// Asking for a length at or above the current one does nothing.
566	///
567	/// A result shorter than one packed word has the dead lanes of its final word zeroed.
568	/// That upholds the invariant that lanes past the logical length are zero.
569	pub fn truncate(&mut self, new_log_len: usize) {
570		if new_log_len >= self.log_len {
571			return;
572		}
573		self.log_len = new_log_len;
574
575		// Zero the lanes past the new logical length in the final word, so a sub-packing-width
576		// result carries no stale scalars. Wider results drop whole words and need no masking.
577		if new_log_len < P::LOG_WIDTH {
578			for i in 1 << new_log_len..P::WIDTH {
579				self.words[0].set(i, <P::Scalar as Field>::ZERO);
580			}
581		}
582
583		self.words
584			.truncate(1 << new_log_len.saturating_sub(P::LOG_WIDTH));
585	}
586}
587
588impl<P: PackedField, Data: Deref<Target = [P]>> AsRef<[P]> for FieldBuffer<P, Data> {
589	#[inline]
590	fn as_ref(&self) -> &[P] {
591		&self.words[..1 << self.log_len.saturating_sub(P::LOG_WIDTH)]
592	}
593}
594
595impl<P: PackedField, Data: DerefMut<Target = [P]>> AsMut<[P]> for FieldBuffer<P, Data> {
596	#[inline]
597	fn as_mut(&mut self) -> &mut [P] {
598		&mut self.words[..1 << self.log_len.saturating_sub(P::LOG_WIDTH)]
599	}
600}
601
602impl<P: PackedField> FromIterator<P::Scalar> for FieldBuffer<P> {
603	/// Builds a buffer over a fresh vector, packing the elements as they arrive.
604	///
605	/// # Panics
606	///
607	/// Panics unless the number of elements is a power of two.
608	/// An empty iterator panics too, matching what building from an empty scalar slice does.
609	/// The count is known only once the iterator runs dry, so a bad one can only panic.
610	#[track_caller]
611	fn from_iter<I: IntoIterator<Item = P::Scalar>>(iter: I) -> Self {
612		let mut iter = iter.into_iter();
613		// The lower bound is the whole count whenever the iterator knows its own length.
614		let mut words = Vec::with_capacity(iter.size_hint().0.div_ceil(P::WIDTH));
615
616		// The length check needs the total, which only the elements themselves can give.
617		let mut len = 0usize;
618		loop {
619			// Fill one word from at most one packing width of elements.
620			// That bound is a constant, so the lanes are written at constant indices rather
621			// than at a running offset into the stream.
622			let mut filled = 0usize;
623			let word = P::from_scalars(iter.by_ref().take(P::WIDTH).inspect(|_| filled += 1));
624
625			// A dry iterator yields an all-zero word belonging to no element.
626			if filled == 0 {
627				break;
628			}
629			words.push(word);
630			len += filled;
631
632			// A short group can only be the last, and its unwritten lanes are already zero,
633			// which is what the invariant on lanes past the logical length asks for.
634			if filled < P::WIDTH {
635				break;
636			}
637		}
638
639		let log_len =
640			strict_log_2(len).expect("precondition: element count must be a power of two");
641
642		Self { log_len, words }
643	}
644}
645
646/// Fills `tail` with copies of `head`, `run` words at a time.
647///
648/// ```text
649///     head            tail
650///     [x]      ->     [x | x | x]
651/// ```
652///
653/// Every destination run draws from one contiguous run of `head`.
654/// So a worker copying one run never reads outside it.
655///
656/// # Preconditions
657///
658/// * `head.len()` is a power of two, and `run` is a power of two at most that
659/// * `tail.len()` is a multiple of `head.len()`
660fn repeat_words<P: PackedField>(head: &[P], tail: &mut [MaybeUninit<P>], run: usize) {
661	debug_assert!(head.len().is_power_of_two());
662	debug_assert!(run.is_power_of_two() && run <= head.len());
663	debug_assert_eq!(tail.len() % head.len(), 0);
664
665	tail.par_chunks_mut(run).enumerate().for_each(|(i, dst)| {
666		// Run `i` of the tail sits one whole head past run `i` of the buffer.
667		// The head length is a power of two, so masking that position gives its source.
668		let source = (i * run) & (head.len() - 1);
669		dst.write_copy_of_slice(&head[source..source + dst.len()]);
670	});
671}
672
673/// Guards the two halves of a buffer that was split along its highest variable.
674///
675/// Holds the parent store, and lends each half out as a mutable view.
676#[derive(Debug)]
677pub struct SplitMut<P: PackedField, Data: DerefMut<Target = [P]>> {
678	/// Element count of each half, as a base-2 logarithm.
679	log_len: usize,
680	/// The detached halves, present only when a half is narrower than a packed word.
681	singles: Option<[P; 2]>,
682	/// The store both halves come from, and the one a detached pair merges back into.
683	data: Data,
684}
685
686impl<P: PackedField, Data: DerefMut<Target = [P]>> SplitMut<P, Data> {
687	/// Guards the two halves of `data`, each holding `2^log_len` elements.
688	fn new(log_len: usize, data: Data) -> Self {
689		// Halves this narrow share word 0, so neither can be lent out as a slice of that store.
690		// Interleaving against zero lifts each half into a word of its own, starting at lane 0.
691		let singles = (log_len < P::LOG_WIDTH).then(|| {
692			let (lo_half, hi_half) = data[0].interleave(P::default(), log_len);
693			[lo_half, hi_half]
694		});
695
696		Self {
697			log_len,
698			singles,
699			data,
700		}
701	}
702
703	/// Lends the two halves out as mutable views, the low half first.
704	///
705	/// A half of whole words is a view straight onto the store, so edits land at once.
706	/// A narrower half is a view onto a detached word, so edits land when this guard drops.
707	pub fn halves(&mut self) -> (FieldSliceMut<'_, P>, FieldSliceMut<'_, P>) {
708		match &mut self.singles {
709			Some([lo_half, hi_half]) => (
710				FieldBuffer {
711					log_len: self.log_len,
712					words: slice::from_mut(lo_half),
713				},
714				FieldBuffer {
715					log_len: self.log_len,
716					words: slice::from_mut(hi_half),
717				},
718			),
719			None => {
720				let half_len = 1 << (self.log_len - P::LOG_WIDTH);
721				let (lo_half, hi_half) = self.data.split_at_mut(half_len);
722				(
723					FieldBuffer {
724						log_len: self.log_len,
725						words: lo_half,
726					},
727					FieldBuffer {
728						log_len: self.log_len,
729						words: hi_half,
730					},
731				)
732			}
733		}
734	}
735}
736
737impl<P: PackedField, Data: DerefMut<Target = [P]>> Drop for SplitMut<P, Data> {
738	fn drop(&mut self) {
739		// Detached halves are the only shape with anything to write back.
740		if let Some([lo_half, hi_half]) = self.singles {
741			(self.data[0], _) = lo_half.interleave(hi_half, self.log_len);
742		}
743	}
744}
745
746#[cfg(test)]
747mod tests {
748	use binius_compute::{BufferPool, GlobalAllocator};
749	use binius_field::packed::get_packed_slice;
750	use binius_utils::rayon::task_size::min_len_for_bytes;
751	use proptest::prelude::*;
752	use rand::{SeedableRng, rngs::StdRng};
753
754	use super::*;
755	use crate::test_utils::{B128, Packed128b, random_field_buffer};
756
757	type P = Packed128b;
758	type F = B128;
759
760	/// Runs every allocator-backed constructor and the zero-padding grow against one allocator.
761	///
762	/// Both a plain heap allocator and a recycling pool must satisfy the same contract.
763	fn check_alloc_constructors<A: Allocator>(alloc: &A) {
764		// Fixture state: a packed word holds 4 lanes, so lengths map to backing words as
765		//
766		//     16 scalars -> log_len 4 -> 4 packed words, every lane live
767		//      2 scalars -> log_len 1 -> 1 packed word,  2 lanes live and 2 dead
768		let scalars: Vec<F> = (0..16).map(F::new).collect();
769		let src = FieldBuffer::<P>::from_values(&scalars);
770
771		// A copy drawn from the allocator carries over both the length and the scalars.
772		let cloned: FieldVec<P, A> = FieldBuffer::from_view_in(alloc, src.as_view());
773		assert_eq!(cloned.log_len(), 4);
774		assert_eq!(cloned.as_view(), src.as_view());
775
776		// Copying a source shorter than one packed word keeps the two live lanes, not four.
777		//
778		// Invariant: a lane past the logical length is not an element, whatever it holds.
779		//
780		// Fixture state: two live scalars, and two dead lanes carrying unrelated values.
781		//
782		//     word   [ 0, 1 | 0xdead, 0xbeef ]
783		//              ^^^^  live prefix
784		//
785		// Packing from scalars would zero those lanes, so the word is supplied whole instead.
786		let small = FieldBuffer::<P>::new(
787			1,
788			vec![P::from_scalars([
789				scalars[0],
790				scalars[1],
791				F::new(0xdead),
792				F::new(0xbeef),
793			])],
794		);
795		let cloned_small: FieldVec<P, A> = FieldBuffer::from_view_in(alloc, small.as_view());
796		assert_eq!(cloned_small.log_len(), 1);
797		assert_eq!(cloned_small.as_view(), small.as_view());
798
799		// 32 elements requested, 32 zeros readable: the buffer arrives at full length, not empty.
800		let mut zeros: FieldVec<P, A> = FieldBuffer::zeros_in(alloc, 5);
801		assert_eq!(zeros.log_len(), 5);
802		assert!(zeros.iter_scalars().all(|scalar| scalar == F::ZERO));
803
804		// Index 31 is the last live element, so writing it stays inside the allocated words.
805		zeros.set(31, F::new(7));
806		assert_eq!(zeros.get(31), F::new(7));
807
808		// Two elements still need the one word that spans them, rounded up from half a word.
809		let zeros_small: FieldVec<P, A> = FieldBuffer::zeros_in(alloc, 1);
810		assert_eq!(zeros_small.as_ref().len(), 1);
811		assert!(zeros_small.iter_scalars().all(|scalar| scalar == F::ZERO));
812
813		// Padding to a wider length keeps the live scalars and zeros every position above them.
814		//
815		//     source (log_len 4)    [0 .. 16]
816		//     result (log_len 5)    [0 .. 16] unchanged, [16 .. 32] zero
817		//
818		// Under a pool the target may sit on memory a prior buffer dirtied.
819		// So the zeros above the copied words have to be written, never assumed.
820		let src_vec: FieldVec<P, A> = FieldBuffer::from_view_in(alloc, src.as_view());
821		let extended = src_vec.zero_extend_in(alloc, 5);
822		assert_eq!(extended.log_len(), 5);
823		for i in 0..16 {
824			assert_eq!(extended.get(i), F::new(i as u128));
825		}
826		assert!((16..32).all(|i| extended.get(i) == F::ZERO));
827
828		// An already-wide-enough buffer comes back unchanged, with no copy taken.
829		let same: FieldVec<P, A> = FieldBuffer::from_view_in(alloc, src.as_view());
830		let same = same.zero_extend_in(alloc, 4);
831		assert_eq!(same.log_len(), 4);
832		assert_eq!(same.as_view(), src.as_view());
833
834		// A source narrower than one packed word pads out of that word's dead lanes.
835		//
836		//     source (log_len 1)    [0, 1 | dead, dead]
837		//     result (log_len 3)    [0, 1 | 0, 0] [0, 0 | 0, 0]
838		//
839		// The two dead lanes become live, so they read back as zeros, not stale scalars.
840		let small_vec: FieldVec<P, A> = FieldBuffer::from_view_in(alloc, small.as_view());
841		let widened = small_vec.zero_extend_in(alloc, 3);
842		assert_eq!(widened.log_len(), 3);
843		assert_eq!(widened.get(0), F::new(0));
844		assert_eq!(widened.get(1), F::new(1));
845		assert!((2..8).all(|i| widened.get(i) == F::ZERO));
846
847		// Reserving room copies the source and leaves the buffer at the source's own length.
848		let reserved: FieldVec<P, A> =
849			FieldBuffer::from_view_with_capacity_in(alloc, src.as_view(), 6);
850		assert_eq!(reserved.log_len(), 4);
851		assert_eq!(reserved.as_view(), src.as_view());
852
853		// Growing into that reserved room repeats the source, four copies over.
854		let mut repeated = reserved;
855		repeated.repeat_extend(6);
856		assert_eq!(repeated.log_len(), 6);
857		assert!((0..64).all(|i| repeated.get(i) == F::new((i % 16) as u128)));
858
859		// A source narrower than one packed word repeats its live lanes, never its dead ones.
860		//
861		//     source (log_len 1)    [0, 1 | dead, dead]
862		//     result (log_len 3)    [0, 1 | 0, 1] [0, 1 | 0, 1]
863		let mut cycled: FieldVec<P, A> =
864			FieldBuffer::from_view_with_capacity_in(alloc, small.as_view(), 3);
865		cycled.repeat_extend(3);
866		assert_eq!(cycled.log_len(), 3);
867		assert!((0..8).all(|i| cycled.get(i) == F::new((i % 2) as u128)));
868	}
869
870	/// Pins `repeat_words` against the definition: word `i` of the tail is word `i mod n` of head.
871	fn check_repeat_words(log_head: usize, log_copies: usize) {
872		let head: Vec<P> = (0..1 << log_head)
873			.map(|i| P::broadcast(F::new(i)))
874			.collect();
875
876		// Every run width the caller can pass, from one word up to the whole head.
877		// A run under the head length is what splits a fill across workers.
878		for log_run in 0..=log_head {
879			let mut tail = vec![MaybeUninit::uninit(); head.len() * ((1 << log_copies) - 1)];
880			repeat_words(&head, &mut tail, 1 << log_run);
881
882			for (i, word) in tail.iter().enumerate() {
883				// SAFETY: the call above wrote every word of the tail.
884				let word = unsafe { word.assume_init() };
885				assert_eq!(word, head[i % head.len()], "log_run={log_run} i={i}");
886			}
887		}
888	}
889
890	#[test]
891	fn repeat_words_tiles_the_head_at_every_run_width() {
892		// Heads of 1 to 16 words, each repeated 1 to 8 times over.
893		for log_head in 0..5 {
894			for log_copies in 0..4 {
895				check_repeat_words(log_head, log_copies);
896			}
897		}
898	}
899
900	/// Pins `repeat_extend` against the definition, in scalars rather than words.
901	fn check_repeat_extend(log_src: usize, log_dst: usize) {
902		let mut rng = StdRng::seed_from_u64(0);
903		let src = random_field_buffer::<P>(&mut rng, log_src);
904
905		let mut buffer: FieldVec<P, GlobalAllocator> =
906			FieldBuffer::from_view_with_capacity_in(&GlobalAllocator, src.as_view(), log_dst);
907		buffer.repeat_extend(log_dst);
908
909		assert_eq!(buffer.log_len(), log_dst);
910		for i in 0..1 << log_dst {
911			let expected = src.get(i % (1 << log_src));
912			assert_eq!(buffer.get(i), expected, "log_src={log_src} log_dst={log_dst} i={i}");
913		}
914	}
915
916	#[test]
917	fn repeat_extend_repeats_the_live_scalars() {
918		// A packed word holds 4 lanes here, so `log_src` 0 and 1 exercise the sub-word path,
919		// and a `log_dst` that stays under the width exercises repeating inside one word.
920		for log_src in 0..5 {
921			for log_dst in log_src..7 {
922				check_repeat_extend(log_src, log_dst);
923			}
924		}
925	}
926
927	#[test]
928	#[should_panic(expected = "precondition: log_len must be at least the buffer's own log_len")]
929	fn repeat_extend_rejects_shrinking() {
930		let mut buffer: FieldVec<P, GlobalAllocator> = FieldBuffer::from_view_with_capacity_in(
931			&GlobalAllocator,
932			FieldBuffer::<P>::zeros(4).as_view(),
933			5,
934		);
935		buffer.repeat_extend(3);
936	}
937
938	#[test]
939	#[should_panic(expected = "precondition: the store must have room reserved")]
940	fn repeat_extend_rejects_a_store_without_room() {
941		let mut buffer: FieldVec<P, GlobalAllocator> = FieldBuffer::from_view_with_capacity_in(
942			&GlobalAllocator,
943			FieldBuffer::<P>::zeros(4).as_view(),
944			4,
945		);
946		buffer.repeat_extend(5);
947	}
948
949	#[test]
950	fn a_copy_reproduces_every_word_on_both_sides_of_the_task_floor() {
951		// The floor is the fewest words a worker takes, set by one task's byte budget.
952		//
953		//     at the floor   [--------]              one task
954		//     above it       [--------][--------]    split across workers
955		let floor = min_len_for_bytes::<P>();
956		for words in [floor, 2 * floor] {
957			let log_len = words.ilog2() as usize + P::LOG_WIDTH;
958			let src = random_field_buffer::<P>(&mut StdRng::seed_from_u64(0), log_len);
959			let copy: FieldVec<P, GlobalAllocator> =
960				FieldBuffer::from_view_in(&GlobalAllocator, src.as_view());
961			assert_eq!(copy.log_len(), log_len);
962			assert_eq!(copy.as_ref(), src.as_ref(), "words={words}");
963		}
964	}
965
966	#[test]
967	#[should_panic(expected = "precondition: log_capacity must be at least src.log_len()")]
968	fn from_view_with_capacity_rejects_a_capacity_below_the_source() {
969		let src = FieldBuffer::<P>::zeros(4);
970		let _: FieldVec<P, GlobalAllocator> =
971			FieldBuffer::from_view_with_capacity_in(&GlobalAllocator, src.as_view(), 3);
972	}
973
974	#[test]
975	fn zeros() {
976		// Make a buffer with `zeros()` and check that all elements are zero.
977		// Test with log_len >= LOG_WIDTH
978		let buffer = FieldBuffer::<P>::zeros(6); // 64 elements
979		assert_eq!(buffer.log_len(), 6);
980		assert_eq!(buffer.len(), 64);
981
982		// Check all elements are zero
983		for i in 0..64 {
984			assert_eq!(buffer.get(i), F::ZERO);
985		}
986
987		// Test with log_len < LOG_WIDTH
988		let buffer = FieldBuffer::<P>::zeros(1); // 2 elements
989		assert_eq!(buffer.log_len(), 1);
990		assert_eq!(buffer.len(), 2);
991
992		// Check all elements are zero
993		for i in 0..2 {
994			assert_eq!(buffer.get(i), F::ZERO);
995		}
996	}
997
998	#[test]
999	fn alloc_constructors_global() {
1000		// Every allocation is an independent heap buffer, freed on drop.
1001		check_alloc_constructors(&GlobalAllocator);
1002	}
1003
1004	#[test]
1005	fn alloc_constructors_pooled() {
1006		// A pool reuses freed blocks, so a buffer may start life on memory a prior buffer dirtied.
1007		let pool = BufferPool::new();
1008		check_alloc_constructors(&&pool);
1009	}
1010
1011	#[test]
1012	fn from_values_below_packing_width() {
1013		// Make a buffer using `from_values()`, where the number of scalars is below the packing
1014		// width
1015		// P::LOG_WIDTH = 2, so P::WIDTH = 4
1016		let values = vec![F::new(1), F::new(2)]; // 2 elements < 4
1017		let buffer = FieldBuffer::<P>::from_values(&values);
1018
1019		assert_eq!(buffer.log_len(), 1); // log2(2) = 1
1020		assert_eq!(buffer.len(), 2);
1021
1022		// Verify the values
1023		assert_eq!(buffer.get(0), F::new(1));
1024		assert_eq!(buffer.get(1), F::new(2));
1025	}
1026
1027	#[test]
1028	fn from_values_above_packing_width() {
1029		// Make a buffer using `from_values()`, where the number of scalars is above the packing
1030		// width
1031		// P::LOG_WIDTH = 2, so P::WIDTH = 4
1032		let values: Vec<F> = (0..16).map(F::new).collect(); // 16 elements > 4
1033		let buffer = FieldBuffer::<P>::from_values(&values);
1034
1035		assert_eq!(buffer.log_len(), 4); // log2(16) = 4
1036		assert_eq!(buffer.len(), 16);
1037
1038		// Verify all values
1039		for i in 0..16 {
1040			assert_eq!(buffer.get(i), F::new(i as u128));
1041		}
1042	}
1043
1044	#[test]
1045	#[should_panic(expected = "power of two")]
1046	fn from_values_non_power_of_two() {
1047		let values: Vec<F> = (0..7).map(F::new).collect(); // 7 is not a power of two
1048		let _ = FieldBuffer::<P>::from_values(&values);
1049	}
1050
1051	#[test]
1052	#[should_panic(expected = "power of two")]
1053	fn from_values_empty() {
1054		let values: Vec<F> = vec![];
1055		let _ = FieldBuffer::<P>::from_values(&values);
1056	}
1057
1058	#[test]
1059	fn new_below_packing_width() {
1060		// Make a buffer using `new()`, where the number of scalars is below the packing
1061		// width
1062		// P::LOG_WIDTH = 2, so P::WIDTH = 4
1063		// For log_len = 1 (2 elements), we need 1 packed value
1064		let mut packed_values = vec![P::default()];
1065		let mut buffer = FieldBuffer::new(1, packed_values.as_mut_slice());
1066
1067		assert_eq!(buffer.log_len(), 1);
1068		assert_eq!(buffer.len(), 2);
1069
1070		// Set and verify values
1071		buffer.set(0, F::new(10));
1072		buffer.set(1, F::new(20));
1073		assert_eq!(buffer.get(0), F::new(10));
1074		assert_eq!(buffer.get(1), F::new(20));
1075	}
1076
1077	#[test]
1078	fn new_above_packing_width() {
1079		// Make a buffer using `new()`, where the number of scalars is above the packing
1080		// width
1081		// P::LOG_WIDTH = 2, so P::WIDTH = 4
1082		// For log_len = 4 (16 elements), we need 4 packed values
1083		let mut packed_values = vec![P::default(); 4];
1084		let mut buffer = FieldBuffer::new(4, packed_values.as_mut_slice());
1085
1086		assert_eq!(buffer.log_len(), 4);
1087		assert_eq!(buffer.len(), 16);
1088
1089		// Set and verify values
1090		for i in 0..16 {
1091			buffer.set(i, F::new(i as u128 * 10));
1092		}
1093		for i in 0..16 {
1094			assert_eq!(buffer.get(i), F::new(i as u128 * 10));
1095		}
1096	}
1097
1098	#[test]
1099	#[should_panic(expected = "precondition")]
1100	fn new_wrong_packed_length() {
1101		let packed_values = vec![P::default(); 3]; // Wrong: should be 4 for log_len=4
1102		let _ = FieldBuffer::new(4, packed_values.as_slice());
1103	}
1104
1105	#[test]
1106	fn get_set() {
1107		let mut buffer = FieldBuffer::<P>::zeros(3); // 8 elements
1108
1109		// Set some values
1110		for i in 0..8 {
1111			buffer.set(i, F::new(i as u128));
1112		}
1113
1114		// Get them back
1115		for i in 0..8 {
1116			assert_eq!(buffer.get(i), F::new(i as u128));
1117		}
1118	}
1119
1120	#[test]
1121	#[should_panic(expected = "precondition")]
1122	fn get_out_of_bounds() {
1123		let buffer = FieldBuffer::<P>::zeros(3); // 8 elements
1124		let _ = buffer.get(8);
1125	}
1126
1127	#[test]
1128	#[should_panic(expected = "precondition")]
1129	fn set_out_of_bounds() {
1130		let mut buffer = FieldBuffer::<P>::zeros(3); // 8 elements
1131		buffer.set(8, F::new(0));
1132	}
1133
1134	#[test]
1135	fn borrowed_views() {
1136		let mut buffer = FieldBuffer::<P>::zeros(3);
1137
1138		// Test the shared view
1139		let slice_ref = buffer.as_view();
1140		assert_eq!(slice_ref.len(), buffer.len());
1141		assert_eq!(slice_ref.log_len(), buffer.log_len());
1142		assert_eq!(slice_ref.as_ref().len(), 1 << slice_ref.log_len().saturating_sub(P::LOG_WIDTH));
1143
1144		// Test the mutable view
1145		let mut slice_mut = buffer.as_mut_view();
1146		slice_mut.set(0, F::new(123));
1147		assert_eq!(slice_mut.as_mut().len(), 1 << slice_mut.log_len().saturating_sub(P::LOG_WIDTH));
1148		assert_eq!(buffer.get(0), F::new(123));
1149	}
1150
1151	#[test]
1152	fn iter_scalars() {
1153		// Test with buffer size below packing width
1154		// P::LOG_WIDTH = 2, so P::WIDTH = 4
1155		let values = vec![F::new(10), F::new(20)]; // 2 elements < 4
1156		let buffer = FieldBuffer::<P>::from_values(&values);
1157
1158		let collected: Vec<F> = buffer.iter_scalars().collect();
1159		assert_eq!(collected, values);
1160
1161		// Verify it matches individual get calls
1162		for (i, &val) in collected.iter().enumerate() {
1163			assert_eq!(val, buffer.get(i));
1164		}
1165
1166		// Test with buffer size equal to packing width
1167		let values = vec![F::new(1), F::new(2), F::new(3), F::new(4)]; // 4 elements = P::WIDTH
1168		let buffer = FieldBuffer::<P>::from_values(&values);
1169
1170		let collected: Vec<F> = buffer.iter_scalars().collect();
1171		assert_eq!(collected, values);
1172
1173		// Test with buffer size above packing width
1174		let values: Vec<F> = (0..16).map(F::new).collect(); // 16 elements > 4
1175		let buffer = FieldBuffer::<P>::from_values(&values);
1176
1177		let collected: Vec<F> = buffer.iter_scalars().collect();
1178		assert_eq!(collected, values);
1179
1180		// Verify it matches individual get calls
1181		for (i, &val) in collected.iter().enumerate() {
1182			assert_eq!(val, buffer.get(i));
1183		}
1184
1185		// Test with single element buffer
1186		let values = vec![F::new(42)];
1187		let buffer = FieldBuffer::<P>::from_values(&values);
1188
1189		let collected: Vec<F> = buffer.iter_scalars().collect();
1190		assert_eq!(collected, values);
1191
1192		// Test with large buffer
1193		let values: Vec<F> = (0..256).map(F::new).collect();
1194		let buffer = FieldBuffer::<P>::from_values(&values);
1195
1196		let collected: Vec<F> = buffer.iter_scalars().collect();
1197		assert_eq!(collected, values);
1198
1199		// Test that iterator is cloneable and can be used multiple times
1200		let values: Vec<F> = (0..8).map(F::new).collect();
1201		let buffer = FieldBuffer::<P>::from_values(&values);
1202
1203		let iter1 = buffer.iter_scalars();
1204		let iter2 = iter1.clone();
1205
1206		let collected1: Vec<F> = iter1.collect();
1207		let collected2: Vec<F> = iter2.collect();
1208		assert_eq!(collected1, collected2);
1209		assert_eq!(collected1, values);
1210	}
1211
1212	#[test]
1213	fn from_iter_below_packing_width() {
1214		// One element in a 4-lane word: the buffer spans 1 element, not the word's 4.
1215		let buffer: FieldBuffer<P> = std::iter::once(F::new(9)).collect();
1216		assert_eq!(buffer.log_len(), 0);
1217		assert_eq!(buffer.len(), 1);
1218		assert_eq!(buffer.get(0), F::new(9));
1219
1220		// The one word packed keeps zeros in the 3 lanes past the element.
1221		let data = buffer.into_inner();
1222		assert_eq!(data.len(), 1);
1223		assert!((1..P::WIDTH).all(|lane| get_packed_slice(&data[..], lane) == F::ZERO));
1224	}
1225
1226	#[test]
1227	#[should_panic(expected = "power of two")]
1228	fn from_iter_non_power_of_two() {
1229		// 7 elements: the count is only known once the iterator runs dry, and then it panics.
1230		let _: FieldBuffer<P> = (0..7).map(F::new).collect();
1231	}
1232
1233	#[test]
1234	#[should_panic(expected = "power of two")]
1235	fn from_iter_empty() {
1236		// Zero elements is not a power of two, so this panics rather than making a 0-length buffer.
1237		let _: FieldBuffer<P> = std::iter::empty::<F>().collect();
1238	}
1239
1240	#[test]
1241	fn iter_packed_covers_the_live_words() {
1242		// Room for 32 elements but 1 live: iteration follows the live count, not the room.
1243		let buffer = FieldBuffer::<P>::scalar_with_capacity(F::new(5), 5);
1244		let live: Vec<P> = buffer.iter_packed().copied().collect();
1245		assert_eq!(live.len(), 1);
1246		assert_eq!(get_packed_slice(&live[..], 0), F::new(5));
1247
1248		// 16 elements over 4-lane words: 4 words, holding the elements in order.
1249		let values: Vec<F> = (0..16).map(F::new).collect();
1250		let buffer = FieldBuffer::<P>::from_values(&values);
1251		let words: Vec<P> = buffer.iter_packed().copied().collect();
1252		assert_eq!(words.len(), 4);
1253		assert_eq!(P::iter_slice(&words).collect::<Vec<_>>(), values);
1254
1255		// A buffer truncated to 2 elements iterates the 1 word they sit in, not the original 4.
1256		let mut truncated = FieldBuffer::<P>::from_values(&values);
1257		truncated.truncate(1);
1258		assert_eq!(truncated.iter_packed().count(), 1);
1259	}
1260
1261	#[test]
1262	fn iter_packed_mut_writes_the_live_words() {
1263		// 8 elements over 2 words: writing both words writes every element.
1264		let mut buffer = FieldBuffer::<P>::zeros(3);
1265		assert_eq!(buffer.iter_packed_mut().count(), 2);
1266		for word in buffer.iter_packed_mut() {
1267			*word = P::broadcast(F::new(3));
1268		}
1269		assert!(buffer.iter_scalars().all(|scalar| scalar == F::new(3)));
1270
1271		// Sub-packing-width: 2 elements share one word, which is lent out whole.
1272		let mut small = FieldBuffer::<P>::zeros(1);
1273		assert_eq!(small.iter_packed_mut().count(), 1);
1274		for word in small.iter_packed_mut() {
1275			*word = P::broadcast(F::new(7));
1276		}
1277		// Only the 2 live lanes are elements, however many lanes the write touched.
1278		assert_eq!(small.iter_scalars().collect::<Vec<_>>(), vec![F::new(7); 2]);
1279	}
1280
1281	#[test]
1282	fn truncate_vec_backing() {
1283		// P::LOG_WIDTH = 2, P::WIDTH = 4.
1284		let make = || FieldBuffer::<P>::from_values(&(0..16).map(F::new).collect::<Vec<_>>());
1285
1286		// Above-width result: the backing Vec shrinks to the new packed length.
1287		let mut buffer = make();
1288		buffer.truncate(3); // 8 elements -> 2 packed words
1289		assert_eq!(buffer.log_len(), 3);
1290		assert_eq!(buffer.len(), 8);
1291		for i in 0..8 {
1292			assert_eq!(buffer.get(i), F::new(i as u128));
1293		}
1294		assert_eq!(buffer.into_inner().len(), 2);
1295
1296		// Sub-packing-width result: one word retained, live prefix kept, dead lanes zeroed.
1297		let mut buffer = make();
1298		buffer.truncate(1); // 2 elements
1299		assert_eq!(buffer.len(), 2);
1300		assert_eq!(buffer.get(0), F::new(0));
1301		assert_eq!(buffer.get(1), F::new(1));
1302		let data = buffer.into_inner();
1303		assert_eq!(data.len(), 1);
1304		assert_eq!(get_packed_slice(&data[..], 2), F::new(0));
1305		assert_eq!(get_packed_slice(&data[..], 3), F::new(0));
1306
1307		// No-op when the requested length is not smaller.
1308		let mut buffer = FieldBuffer::<P>::from_values(&(0..4).map(F::new).collect::<Vec<_>>());
1309		buffer.truncate(5);
1310		assert_eq!(buffer.log_len(), 2);
1311		assert_eq!(buffer.into_inner().len(), 1);
1312	}
1313
1314	#[test]
1315	fn truncate_slice_backing() {
1316		// Truncating a `&mut [P]` backing reslices it and zeros the sub-width dead lanes.
1317		let mut storage = vec![P::default(); 4]; // 16 elements at log_len 4
1318		let mut buffer = FieldSliceMut::from_slice(4, storage.as_mut_slice());
1319		for i in 0..16 {
1320			buffer.set(i, F::new(i as u128));
1321		}
1322
1323		buffer.truncate(1); // 2 elements, sub-width
1324		assert_eq!(buffer.len(), 2);
1325		assert_eq!(buffer.get(0), F::new(0));
1326		assert_eq!(buffer.get(1), F::new(1));
1327
1328		let data = buffer.into_inner();
1329		assert_eq!(data.len(), 1);
1330		assert_eq!(get_packed_slice(&data[..], 2), F::new(0));
1331		assert_eq!(get_packed_slice(&data[..], 3), F::new(0));
1332	}
1333
1334	#[test]
1335	fn split_half() {
1336		// Test with buffer size > P::WIDTH (multiple packed elements)
1337		let values: Vec<F> = (0..16).map(F::new).collect();
1338		let buffer = FieldBuffer::<P>::from_values(&values);
1339
1340		let (first, second) = buffer.split_half();
1341		assert_eq!(first.len(), 8);
1342		assert_eq!(second.len(), 8);
1343
1344		// Verify values
1345		for i in 0..8 {
1346			assert_eq!(first.get(i), F::new(i as u128));
1347			assert_eq!(second.get(i), F::new((i + 8) as u128));
1348		}
1349
1350		// Test with buffer size = P::WIDTH (single packed element)
1351		// P::LOG_WIDTH = 2, so P::WIDTH = 4
1352		let values: Vec<F> = (0..4).map(F::new).collect();
1353		let buffer = FieldBuffer::<P>::from_values(&values);
1354
1355		let (first, second) = buffer.split_half();
1356		assert_eq!(first.len(), 2);
1357		assert_eq!(second.len(), 2);
1358
1359		// Verify we got Single variants
1360		match &first.words {
1361			FieldSliceData::Single(_) => {}
1362			_ => panic!("Expected Single variant for first half"),
1363		}
1364		match &second.words {
1365			FieldSliceData::Single(_) => {}
1366			_ => panic!("Expected Single variant for second half"),
1367		}
1368
1369		// Verify values
1370		assert_eq!(first.get(0), F::new(0));
1371		assert_eq!(first.get(1), F::new(1));
1372		assert_eq!(second.get(0), F::new(2));
1373		assert_eq!(second.get(1), F::new(3));
1374
1375		// Test with buffer size = 2 (less than P::WIDTH)
1376		let values: Vec<F> = vec![F::new(10), F::new(20)];
1377		let buffer = FieldBuffer::<P>::from_values(&values);
1378
1379		let (first, second) = buffer.split_half();
1380		assert_eq!(first.len(), 1);
1381		assert_eq!(second.len(), 1);
1382
1383		// Verify we got Single variants
1384		match &first.words {
1385			FieldSliceData::Single(_) => {}
1386			_ => panic!("Expected Single variant for first half"),
1387		}
1388		match &second.words {
1389			FieldSliceData::Single(_) => {}
1390			_ => panic!("Expected Single variant for second half"),
1391		}
1392
1393		assert_eq!(first.get(0), F::new(10));
1394		assert_eq!(second.get(0), F::new(20));
1395	}
1396
1397	#[test]
1398	#[should_panic(expected = "precondition")]
1399	fn split_half_size_one() {
1400		let values = vec![F::new(42)];
1401		let buffer = FieldBuffer::<P>::from_values(&values);
1402		let _ = buffer.split_half();
1403	}
1404
1405	#[test]
1406	fn split_half_mut_no_closure() {
1407		// Test with buffer size > P::WIDTH (multiple packed elements)
1408		let mut buffer = FieldBuffer::<P>::zeros(4); // 16 elements
1409
1410		// Fill with test data
1411		for i in 0..16 {
1412			buffer.set(i, F::new(i as u128));
1413		}
1414
1415		{
1416			let mut split = buffer.split_half_mut();
1417			let (mut first, mut second) = split.halves();
1418
1419			assert_eq!(first.len(), 8);
1420			assert_eq!(second.len(), 8);
1421
1422			// Modify through the split halves
1423			for i in 0..8 {
1424				first.set(i, F::new((i * 10) as u128));
1425				second.set(i, F::new((i * 20) as u128));
1426			}
1427			// split drops here and writes back the changes
1428		}
1429
1430		// Verify changes were made to original buffer
1431		for i in 0..8 {
1432			assert_eq!(buffer.get(i), F::new((i * 10) as u128));
1433			assert_eq!(buffer.get(i + 8), F::new((i * 20) as u128));
1434		}
1435
1436		// Test with buffer size = P::WIDTH (single packed element)
1437		// P::LOG_WIDTH = 2, so a buffer with log_len = 2 (4 elements) can now be split
1438		let mut buffer = FieldBuffer::<P>::zeros(2); // 4 elements
1439
1440		// Fill with test data
1441		for i in 0..4 {
1442			buffer.set(i, F::new(i as u128));
1443		}
1444
1445		{
1446			let mut split = buffer.split_half_mut();
1447			let (mut first, mut second) = split.halves();
1448
1449			assert_eq!(first.len(), 2);
1450			assert_eq!(second.len(), 2);
1451
1452			// Modify values
1453			first.set(0, F::new(100));
1454			first.set(1, F::new(101));
1455			second.set(0, F::new(200));
1456			second.set(1, F::new(201));
1457			// split drops here and writes back the changes using interleave
1458		}
1459
1460		// Verify changes were written back
1461		assert_eq!(buffer.get(0), F::new(100));
1462		assert_eq!(buffer.get(1), F::new(101));
1463		assert_eq!(buffer.get(2), F::new(200));
1464		assert_eq!(buffer.get(3), F::new(201));
1465
1466		// Test with buffer size = 2
1467		let mut buffer = FieldBuffer::<P>::zeros(1); // 2 elements
1468
1469		buffer.set(0, F::new(10));
1470		buffer.set(1, F::new(20));
1471
1472		{
1473			let mut split = buffer.split_half_mut();
1474			let (mut first, mut second) = split.halves();
1475
1476			assert_eq!(first.len(), 1);
1477			assert_eq!(second.len(), 1);
1478
1479			// Modify values
1480			first.set(0, F::new(30));
1481			second.set(0, F::new(40));
1482			// split drops here and writes back the changes using interleave
1483		}
1484
1485		// Verify changes
1486		assert_eq!(buffer.get(0), F::new(30));
1487		assert_eq!(buffer.get(1), F::new(40));
1488	}
1489
1490	#[test]
1491	#[should_panic(expected = "precondition")]
1492	fn split_half_mut_size_one() {
1493		let mut buffer = FieldBuffer::<P>::zeros(0); // 1 element
1494		let _ = buffer.split_half_mut();
1495	}
1496
1497	proptest! {
1498		#[test]
1499		fn unequal_length_buffers_are_never_equal(
1500			log_len_a in 0usize..=6,
1501			log_len_b in 0usize..=6,
1502			fill in any::<u128>(),
1503		) {
1504			// Invariant: buffers are equal iff same length and same scalars.
1505			let value = F::new(fill);
1506			let buf_a = FieldBuffer::<P>::from_values(&vec![value; 1 << log_len_a]);
1507			let buf_b = FieldBuffer::<P>::from_values(&vec![value; 1 << log_len_b]);
1508
1509			// - Same length -> equal;
1510			// - Different length -> unequal despite matching scalars.
1511			if log_len_a == log_len_b {
1512				prop_assert_eq!(buf_a, buf_b);
1513			} else {
1514				prop_assert_ne!(buf_a, buf_b);
1515			}
1516		}
1517
1518		#[test]
1519		fn from_iter_matches_from_values(log_len in 0usize..=6) {
1520			// The sweep straddles the packing width: a sub-word store, then multi-word ones.
1521			let values: Vec<F> = (0..1u128 << log_len).map(F::new).collect();
1522
1523			// Collecting the scalars must rebuild what packing the slice builds, length included.
1524			let collected: FieldBuffer<P> = values.iter().copied().collect();
1525			prop_assert_eq!(collected.log_len(), log_len);
1526			prop_assert_eq!(&collected, &FieldBuffer::<P>::from_values(&values));
1527
1528			// Round trip: iterating the collected buffer hands the scalars back, in order.
1529			prop_assert_eq!(&collected.iter_scalars().collect::<Vec<_>>(), &values);
1530		}
1531	}
1532}