binius_field/underlier/
divisible.rs

1// Copyright 2024-2025 Irreducible Inc.
2
3use core::slice;
4use std::{
5	mem::{align_of, size_of},
6	slice::{from_raw_parts, from_raw_parts_mut},
7};
8
9/// Underlier value that can be split into a slice of smaller `U` values.
10/// This trait is unsafe because it allows to reinterpret the memory of a type as a slice of another type.
11///
12/// # Safety
13/// Implementors must ensure that `&Self` can be safely bit-cast to `&[U; Self::WIDTH]` and
14/// `&mut Self` can be safely bit-cast to `&mut [U; Self::WIDTH]`.
15pub unsafe trait Divisible<U: UnderlierType>: UnderlierType {
16	const WIDTH: usize = {
17		assert!(size_of::<Self>() % size_of::<U>() == 0);
18		assert!(align_of::<Self>() >= align_of::<U>());
19		size_of::<Self>() / size_of::<U>()
20	};
21
22	/// This is actually `[U; Self::WIDTH]` but we can't use it as the default value in the trait definition
23	/// without `generic_const_exprs` feature enabled.
24	type Array: IntoIterator<Item = U, IntoIter: Send + Clone>;
25
26	fn split_val(self) -> Self::Array;
27	fn split_ref(&self) -> &[U];
28	fn split_mut(&mut self) -> &mut [U];
29
30	fn split_slice(values: &[Self]) -> &[U] {
31		let ptr = values.as_ptr() as *const U;
32		// Safety: if `&Self` can be reinterpreted as a sequence of `Self::WIDTH` elements of `U` then
33		// `&[Self]` can be reinterpreted as a sequence of `Self::Width * values.len()` elements of `U`.
34		unsafe { from_raw_parts(ptr, values.len() * Self::WIDTH) }
35	}
36
37	fn split_slice_mut(values: &mut [Self]) -> &mut [U] {
38		let ptr = values.as_mut_ptr() as *mut U;
39		// Safety: if `&mut Self` can be reinterpreted as a sequence of `Self::WIDTH` elements of `U` then
40		// `&mut [Self]` can be reinterpreted as a sequence of `Self::Width * values.len()` elements of `U`.
41		unsafe { from_raw_parts_mut(ptr, values.len() * Self::WIDTH) }
42	}
43}
44
45unsafe impl<U: UnderlierType> Divisible<U> for U {
46	type Array = [U; 1];
47
48	fn split_val(self) -> Self::Array {
49		[self]
50	}
51
52	fn split_ref(&self) -> &[U] {
53		slice::from_ref(self)
54	}
55
56	fn split_mut(&mut self) -> &mut [U] {
57		slice::from_mut(self)
58	}
59}
60
61macro_rules! impl_divisible {
62    (@pairs $name:ty,?) => {};
63    (@pairs $bigger:ty, $smaller:ty) => {
64        unsafe impl $crate::underlier::Divisible<$smaller> for $bigger {
65            type Array = [$smaller; {size_of::<Self>() / size_of::<$smaller>()}];
66
67            fn split_val(self) -> Self::Array {
68                bytemuck::must_cast::<_, Self::Array>(self)
69            }
70
71            fn split_ref(&self) -> &[$smaller] {
72                bytemuck::must_cast_ref::<_, [$smaller;{(<$bigger>::BITS as usize / <$smaller>::BITS as usize ) }]>(self)
73            }
74
75            fn split_mut(&mut self) -> &mut [$smaller] {
76                bytemuck::must_cast_mut::<_, [$smaller;{(<$bigger>::BITS as usize / <$smaller>::BITS as usize ) }]>(self)
77            }
78        }
79
80		unsafe impl $crate::underlier::Divisible<$smaller> for $crate::underlier::ScaledUnderlier<$bigger, 2> {
81            type Array = [$smaller; {2 * size_of::<$bigger>() / size_of::<$smaller>()}];
82
83            fn split_val(self) -> Self::Array {
84                bytemuck::must_cast::<_, Self::Array>(self)
85            }
86
87            fn split_ref(&self) -> &[$smaller] {
88                bytemuck::must_cast_ref::<_, [$smaller;{(2 * <$bigger>::BITS as usize / <$smaller>::BITS as usize ) }]>(&self.0)
89            }
90
91            fn split_mut(&mut self) -> &mut [$smaller] {
92                bytemuck::must_cast_mut::<_, [$smaller;{(2 * <$bigger>::BITS as usize / <$smaller>::BITS as usize ) }]>(&mut self.0)
93            }
94        }
95
96		unsafe impl $crate::underlier::Divisible<$smaller> for $crate::underlier::ScaledUnderlier<$crate::underlier::ScaledUnderlier<$bigger, 2>, 2> {
97            type Array = [$smaller; {4 * size_of::<$bigger>() / size_of::<$smaller>()}];
98
99            fn split_val(self) -> Self::Array {
100                bytemuck::must_cast::<_, Self::Array>(self)
101            }
102
103            fn split_ref(&self) -> &[$smaller] {
104                bytemuck::must_cast_ref::<_, [$smaller;{(4 * <$bigger>::BITS as usize / <$smaller>::BITS as usize ) }]>(&self.0)
105            }
106
107            fn split_mut(&mut self) -> &mut [$smaller] {
108                bytemuck::must_cast_mut::<_, [$smaller;{(4 * <$bigger>::BITS as usize / <$smaller>::BITS as usize ) }]>(&mut self.0)
109            }
110        }
111    };
112    (@pairs $first:ty, $second:ty, $($tail:ty),*) => {
113        impl_divisible!(@pairs $first, $second);
114        impl_divisible!(@pairs $first, $($tail),*);
115    };
116    ($_:ty) => {};
117    ($head:ty, $($tail:ty),*) => {
118        impl_divisible!(@pairs $head, $($tail),*);
119        impl_divisible!($($tail),*);
120    }
121}
122
123#[allow(unused)]
124pub(crate) use impl_divisible;
125
126use super::UnderlierType;
127
128impl_divisible!(u128, u64, u32, u16, u8);