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