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