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

binius_field/underlier/
iteration.rs
