binius_field/underlier/
iteration.rs

1// Copyright 2025 Irreducible Inc.
2
3use std::marker::PhantomData;
4
5use binius_utils::{checked_arithmetics::checked_int_div, iter::IterExtensions};
6
7use super::{Divisible, NumCast, U1, U2, U4, UnderlierType, UnderlierWithBitOps};
8
9/// The iteration strategy for the given underlier type 'U' that is treated as a packed collection
10/// of 'T's.
11pub trait IterationStrategy<T, U> {
12	/// Iterate over the subvalues of the given reference.
13	fn ref_iter(value: &U) -> impl Iterator<Item = T> + Send + Clone + '_;
14
15	/// Iterate over the subvalues of the given value.
16	fn value_iter(value: U) -> impl Iterator<Item = T> + Send + Clone;
17
18	/// Iterate over the subvalues of the given slice.
19	fn slice_iter(slice: &[U]) -> impl Iterator<Item = T> + Send + Clone + '_;
20}
21
22/// The iteration strategy for the given underlier type 'U' that is treated as a packed collection
23/// of bits.
24pub struct BitIterationStrategy;
25
26impl<U> IterationStrategy<U1, U> for BitIterationStrategy
27where
28	U: Divisible<u8> + UnderlierWithBitOps,
29	U1: NumCast<U>,
30{
31	#[inline]
32	fn ref_iter(value: &U) -> impl Iterator<Item = U1> + Send + Clone + '_ {
33		(0..U::BITS).map_skippable(move |i| unsafe { value.get_subvalue(i) })
34	}
35
36	#[inline]
37	fn value_iter(value: U) -> impl Iterator<Item = U1> + Send + Clone {
38		(0..U::BITS).map_skippable(move |i| unsafe { value.get_subvalue(i) })
39	}
40
41	#[inline]
42	fn slice_iter(slice: &[U]) -> impl Iterator<Item = U1> + Send + Clone + '_ {
43		U::split_slice(slice)
44			.iter()
45			.flat_map(|val| (0..8).map(move |i| U1::new(*val >> i)))
46	}
47}
48
49/// Specialized iteration strategy for types that can be cast to an array of the elements of a
50/// smaller type 'T'.
51pub struct DivisibleStrategy;
52
53impl<U: UnderlierType + Divisible<T>, T: UnderlierType> IterationStrategy<T, U>
54	for DivisibleStrategy
55{
56	#[inline]
57	fn ref_iter(value: &U) -> impl Iterator<Item = T> + Send + Clone + '_ {
58		U::split_ref(value).iter().copied()
59	}
60
61	#[inline]
62	fn value_iter(value: U) -> impl Iterator<Item = T> + Send + Clone {
63		U::split_val(value).into_iter()
64	}
65
66	#[inline]
67	fn slice_iter(slice: &[U]) -> impl Iterator<Item = T> + Send + Clone + '_ {
68		U::split_slice(slice).iter().copied()
69	}
70}
71
72/// Fallback iteration strategy for types that do not have a specialized strategy.
73pub struct FallbackStrategy;
74
75impl<T, U> IterationStrategy<T, U> for FallbackStrategy
76where
77	T: UnderlierType + NumCast<U>,
78	U: UnderlierWithBitOps,
79{
80	#[inline]
81	fn value_iter(value: U) -> impl Iterator<Item = T> + Send + Clone {
82		(0..checked_int_div(U::BITS, T::BITS))
83			.map_skippable(move |i| unsafe { value.get_subvalue(i) })
84	}
85
86	#[inline]
87	fn ref_iter(value: &U) -> impl Iterator<Item = T> + Send + Clone + '_ {
88		(0..checked_int_div(U::BITS, T::BITS))
89			.map_skippable(move |i| unsafe { value.get_subvalue(i) })
90	}
91
92	#[inline]
93	fn slice_iter(slice: &[U]) -> impl Iterator<Item = T> + Send + Clone + '_ {
94		slice.iter().flat_map(Self::ref_iter)
95	}
96}
97
98/// `IterationMethods<T, U>` is supposed to implement an optimal strategy for the pair of types `(T,
99/// U)`.
100#[derive(Default, Copy, Clone, Eq, PartialEq, Debug)]
101pub struct IterationMethods<T, U>(PhantomData<(T, U)>);
102
103macro_rules! impl_iteration {
104	(@pairs $strategy:ident, $bigger:ty,) => {};
105	(@pairs $strategy:ident, $bigger:ty, $smaller:ty, $($tail:ty,)*) => {
106		impl $crate::underlier::IterationStrategy<$smaller, $bigger> for $crate::underlier::IterationMethods<$smaller, $bigger> {
107			fn ref_iter(value: &$bigger) -> impl Iterator<Item = $smaller> + Send + Clone + '_ {
108				$crate::underlier::$strategy::ref_iter(value)
109			}
110
111			fn value_iter(value: $bigger) -> impl Iterator<Item = $smaller> + Send + Clone {
112				$crate::underlier::$strategy::value_iter(value)
113			}
114
115			fn slice_iter(slice: &[$bigger]) -> impl Iterator<Item = $smaller> + Send + Clone + '_ {
116				$crate::underlier::$strategy::slice_iter(slice)
117			}
118		}
119
120		impl_iteration!(@pairs $strategy, $bigger, $($tail,)*);
121	};
122	($bigger:ty, $(@strategy $strategy:ident, $($smaller:ty,)*)*) => {
123		$(
124			impl_iteration!(@pairs $strategy, $bigger, $($smaller,)*);
125		)*
126	};
127}
128
129pub(crate) use impl_iteration;
130
131impl_iteration!(U1,
132	@strategy FallbackStrategy, U1,
133);
134impl_iteration!(U2,
135	@strategy FallbackStrategy, U1, U2,
136);
137impl_iteration!(U4,
138	@strategy FallbackStrategy, U1, U2, U4,
139);
140impl_iteration!(u8,
141	@strategy BitIterationStrategy, U1,
142	@strategy FallbackStrategy, U2, U4,
143	@strategy DivisibleStrategy, u8,
144);
145impl_iteration!(u16,
146	@strategy BitIterationStrategy, U1,
147	@strategy FallbackStrategy, U2, U4,
148	@strategy DivisibleStrategy, u8, u16,
149);
150impl_iteration!(u32,
151	@strategy BitIterationStrategy, U1,
152	@strategy FallbackStrategy, U2, U4,
153	@strategy DivisibleStrategy, u8, u16, u32,
154);
155impl_iteration!(u64,
156	@strategy BitIterationStrategy, U1,
157	@strategy FallbackStrategy, U2, U4,
158	@strategy DivisibleStrategy, u8, u16, u32, u64,
159);
160impl_iteration!(u128,
161	@strategy BitIterationStrategy, U1,
162	@strategy FallbackStrategy, U2, U4,
163	@strategy DivisibleStrategy, u8, u16, u32, u64, u128,
164);
binius_field/underlier/iteration.rs

binius_field/underlier/
iteration.rs
