1pub trait Divisible<T>: Sized {
22 const LOG_N: usize;
24
25 const N: usize = 1 << Self::LOG_N;
27
28 fn value_iter(value: Self) -> impl ExactSizeIterator<Item = T> + Send + Clone;
30
31 fn ref_iter(value: &Self) -> impl ExactSizeIterator<Item = T> + Send + Clone + '_;
33
34 fn slice_iter(slice: &[Self]) -> impl ExactSizeIterator<Item = T> + Send + Clone + '_;
36
37 #[inline]
43 fn get(&self, index: usize) -> T {
44 assert!(index < Self::N, "index {index} out of bounds (N = {})", Self::N);
45 unsafe { self.get_unchecked(index) }
47 }
48
49 #[inline]
55 fn set(&mut self, index: usize, val: T) {
56 assert!(index < Self::N, "index {index} out of bounds (N = {})", Self::N);
57 unsafe { self.set_unchecked(index, val) };
59 }
60
61 unsafe fn get_unchecked(&self, index: usize) -> T;
67
68 unsafe fn set_unchecked(&mut self, index: usize, val: T);
74
75 fn broadcast(val: T) -> Self;
77
78 fn from_iter(iter: impl Iterator<Item = T>) -> Self;
83}
84
85pub mod memcast {
89 use bytemuck::{Pod, Zeroable};
90
91 #[cfg(target_endian = "little")]
93 #[inline]
94 pub fn value_iter<Big, Small, const N: usize>(
95 value: Big,
96 ) -> impl ExactSizeIterator<Item = Small> + Send + Clone
97 where
98 Big: Pod,
99 Small: Pod + Send,
100 {
101 bytemuck::must_cast::<Big, [Small; N]>(value).into_iter()
102 }
103
104 #[cfg(target_endian = "big")]
106 #[inline]
107 pub fn value_iter<Big, Small, const N: usize>(
108 value: Big,
109 ) -> impl ExactSizeIterator<Item = Small> + Send + Clone
110 where
111 Big: Pod,
112 Small: Pod + Send,
113 {
114 bytemuck::must_cast::<Big, [Small; N]>(value)
115 .into_iter()
116 .rev()
117 }
118
119 #[cfg(target_endian = "little")]
121 #[inline]
122 pub fn ref_iter<Big, Small, const N: usize>(
123 value: &Big,
124 ) -> impl ExactSizeIterator<Item = Small> + Send + Clone + '_
125 where
126 Big: Pod,
127 Small: Pod + Send + Sync,
128 {
129 bytemuck::must_cast_ref::<Big, [Small; N]>(value)
130 .iter()
131 .copied()
132 }
133
134 #[cfg(target_endian = "big")]
136 #[inline]
137 pub fn ref_iter<Big, Small, const N: usize>(
138 value: &Big,
139 ) -> impl ExactSizeIterator<Item = Small> + Send + Clone + '_
140 where
141 Big: Pod,
142 Small: Pod + Send + Sync,
143 {
144 bytemuck::must_cast_ref::<Big, [Small; N]>(value)
145 .iter()
146 .rev()
147 .copied()
148 }
149
150 #[cfg(target_endian = "little")]
152 #[inline]
153 pub fn slice_iter<Big, Small>(
154 slice: &[Big],
155 ) -> impl ExactSizeIterator<Item = Small> + Send + Clone + '_
156 where
157 Big: Pod,
158 Small: Pod + Send + Sync,
159 {
160 bytemuck::must_cast_slice::<Big, Small>(slice)
161 .iter()
162 .copied()
163 }
164
165 #[cfg(target_endian = "big")]
170 #[inline]
171 pub fn slice_iter<Big, Small, const LOG_N: usize>(
172 slice: &[Big],
173 ) -> impl ExactSizeIterator<Item = Small> + Send + Clone + '_
174 where
175 Big: Pod,
176 Small: Pod + Send + Sync,
177 {
178 const N: usize = 1 << LOG_N;
179 let raw_slice = bytemuck::must_cast_slice::<Big, Small>(slice);
180 (0..raw_slice.len()).map(move |i| {
181 let element_idx = i >> LOG_N;
182 let sub_idx = i & (N - 1);
183 let reversed_sub_idx = N - 1 - sub_idx;
184 let raw_idx = element_idx * N + reversed_sub_idx;
185 raw_slice[raw_idx]
186 })
187 }
188
189 #[cfg(target_endian = "little")]
195 #[inline]
196 pub unsafe fn get<Big, Small, const N: usize>(value: &Big, index: usize) -> Small
197 where
198 Big: Pod,
199 Small: Pod,
200 {
201 unsafe { *bytemuck::must_cast_ref::<Big, [Small; N]>(value).get_unchecked(index) }
203 }
204
205 #[cfg(target_endian = "big")]
211 #[inline]
212 pub unsafe fn get<Big, Small, const N: usize>(value: &Big, index: usize) -> Small
213 where
214 Big: Pod,
215 Small: Pod,
216 {
217 unsafe { *bytemuck::must_cast_ref::<Big, [Small; N]>(value).get_unchecked(N - 1 - index) }
219 }
220
221 #[cfg(target_endian = "little")]
229 #[inline]
230 pub unsafe fn set<Big, Small, const N: usize>(value: &mut Big, index: usize, val: Small)
231 where
232 Big: Pod,
233 Small: Pod,
234 {
235 unsafe {
237 *bytemuck::must_cast_mut::<Big, [Small; N]>(value).get_unchecked_mut(index) = val;
238 }
239 }
240
241 #[cfg(target_endian = "big")]
249 #[inline]
250 pub unsafe fn set<Big, Small, const N: usize>(value: &mut Big, index: usize, val: Small)
251 where
252 Big: Pod,
253 Small: Pod,
254 {
255 unsafe {
257 *bytemuck::must_cast_mut::<Big, [Small; N]>(value).get_unchecked_mut(N - 1 - index) =
258 val;
259 }
260 }
261
262 #[inline]
264 pub fn broadcast<Big, Small, const N: usize>(val: Small) -> Big
265 where
266 Big: Pod,
267 Small: Pod + Copy,
268 {
269 bytemuck::must_cast::<[Small; N], Big>([val; N])
270 }
271
272 #[cfg(target_endian = "little")]
274 #[inline]
275 pub fn from_iter<Big, Small, const N: usize>(iter: impl Iterator<Item = Small>) -> Big
276 where
277 Big: Pod,
278 Small: Pod,
279 {
280 let mut arr: [Small; N] = Zeroable::zeroed();
281 for (i, val) in iter.take(N).enumerate() {
282 arr[i] = val;
283 }
284 bytemuck::must_cast(arr)
285 }
286
287 #[cfg(target_endian = "big")]
289 #[inline]
290 pub fn from_iter<Big, Small, const N: usize>(iter: impl Iterator<Item = Small>) -> Big
291 where
292 Big: Pod,
293 Small: Pod,
294 {
295 let mut arr: [Small; N] = Zeroable::zeroed();
296 for (i, val) in iter.take(N).enumerate() {
297 arr[N - 1 - i] = val;
298 }
299 bytemuck::must_cast(arr)
300 }
301}
302
303pub mod mapget {
308 use binius_utils::iter::IterExtensions;
309
310 use super::Divisible;
311
312 #[inline]
314 pub fn value_iter<Big, Small>(value: Big) -> impl ExactSizeIterator<Item = Small> + Send + Clone
315 where
316 Big: Divisible<Small> + Send + Clone,
317 Small: Send,
318 {
319 (0..Big::N).map_skippable(move |i| Divisible::<Small>::get(&value, i))
320 }
321
322 #[inline]
324 pub fn slice_iter<Big, Small>(
325 slice: &[Big],
326 ) -> impl ExactSizeIterator<Item = Small> + Send + Clone + '_
327 where
328 Big: Divisible<Small> + Send + Sync,
329 Small: Send,
330 {
331 let total = slice.len() * Big::N;
332 (0..total).map_skippable(move |global_idx| {
333 let elem_idx = global_idx / Big::N;
334 let sub_idx = global_idx % Big::N;
335 Divisible::<Small>::get(&slice[elem_idx], sub_idx)
336 })
337 }
338}
339
340macro_rules! impl_divisible_memcast {
354 ($big:ty, $($small:ty => |$v:ident| $broadcast:expr),+ $(,)?) => {
356 $(
357 $crate::divisible::impl_divisible_memcast!(@impl $big, $small, |$v| $broadcast);
358 )+
359 };
360 ($big:ty, $($small:ty),+ $(,)?) => {
362 $(
363 $crate::divisible::impl_divisible_memcast!(
364 @impl $big, $small,
365 |val| $crate::divisible::memcast::broadcast::<
366 $big,
367 $small,
368 { ::std::mem::size_of::<$big>() / ::std::mem::size_of::<$small>() },
369 >(val)
370 );
371 )+
372 };
373 (@impl $big:ty, $small:ty, |$v:ident| $broadcast:expr) => {
374 impl $crate::divisible::Divisible<$small> for $big {
375 const LOG_N: usize =
376 (::std::mem::size_of::<$big>() / ::std::mem::size_of::<$small>()).ilog2() as usize;
377
378 #[inline]
379 fn value_iter(value: Self) -> impl ExactSizeIterator<Item = $small> + Send + Clone {
380 const N: usize = ::std::mem::size_of::<$big>() / ::std::mem::size_of::<$small>();
381 $crate::divisible::memcast::value_iter::<$big, $small, N>(value)
382 }
383
384 #[inline]
385 fn ref_iter(value: &Self) -> impl ExactSizeIterator<Item = $small> + Send + Clone + '_ {
386 const N: usize = ::std::mem::size_of::<$big>() / ::std::mem::size_of::<$small>();
387 $crate::divisible::memcast::ref_iter::<$big, $small, N>(value)
388 }
389
390 #[inline]
391 #[cfg(target_endian = "little")]
392 fn slice_iter(slice: &[Self]) -> impl ExactSizeIterator<Item = $small> + Send + Clone + '_ {
393 $crate::divisible::memcast::slice_iter::<$big, $small>(slice)
394 }
395
396 #[inline]
397 #[cfg(target_endian = "big")]
398 fn slice_iter(slice: &[Self]) -> impl ExactSizeIterator<Item = $small> + Send + Clone + '_ {
399 const LOG_N: usize =
400 (::std::mem::size_of::<$big>() / ::std::mem::size_of::<$small>()).ilog2() as usize;
401 $crate::divisible::memcast::slice_iter::<$big, $small, LOG_N>(slice)
402 }
403
404 #[inline]
405 unsafe fn get_unchecked(&self, index: usize) -> $small {
406 const N: usize = ::std::mem::size_of::<$big>() / ::std::mem::size_of::<$small>();
407 unsafe { $crate::divisible::memcast::get::<$big, $small, N>(self, index) }
409 }
410
411 #[inline]
412 unsafe fn set_unchecked(&mut self, index: usize, val: $small) {
413 const N: usize = ::std::mem::size_of::<$big>() / ::std::mem::size_of::<$small>();
414 unsafe { $crate::divisible::memcast::set::<$big, $small, N>(self, index, val) };
416 }
417
418 #[inline]
419 fn broadcast($v: $small) -> Self {
420 $broadcast
421 }
422
423 #[inline]
424 fn from_iter(iter: impl Iterator<Item = $small>) -> Self {
425 const N: usize = ::std::mem::size_of::<$big>() / ::std::mem::size_of::<$small>();
426 $crate::divisible::memcast::from_iter::<$big, $small, N>(iter)
427 }
428 }
429 };
430}
431
432#[allow(unused)]
433pub(crate) use impl_divisible_memcast;
434
435impl_divisible_memcast!(u128, u64, u32, u16, u8);
437impl_divisible_memcast!(u64, u32, u16, u8);
438impl_divisible_memcast!(u32, u16, u8);
439impl_divisible_memcast!(u16, u8);
440
441macro_rules! impl_divisible_self {
443 ($($ty:ty),+) => {
444 $(
445 impl Divisible<$ty> for $ty {
446 const LOG_N: usize = 0;
447
448 #[inline]
449 fn value_iter(value: Self) -> impl ExactSizeIterator<Item = $ty> + Send + Clone {
450 std::iter::once(value)
451 }
452
453 #[inline]
454 fn ref_iter(value: &Self) -> impl ExactSizeIterator<Item = $ty> + Send + Clone + '_ {
455 std::iter::once(*value)
456 }
457
458 #[inline]
459 fn slice_iter(slice: &[Self]) -> impl ExactSizeIterator<Item = $ty> + Send + Clone + '_ {
460 slice.iter().copied()
461 }
462
463 #[inline]
464 unsafe fn get_unchecked(&self, _index: usize) -> $ty {
465 *self
466 }
467
468 #[inline]
469 unsafe fn set_unchecked(&mut self, _index: usize, val: $ty) {
470 *self = val;
471 }
472
473 #[inline]
474 fn broadcast(val: $ty) -> Self {
475 val
476 }
477
478 #[inline]
479 fn from_iter(mut iter: impl Iterator<Item = $ty>) -> Self {
480 iter.next().unwrap_or_else(bytemuck::Zeroable::zeroed)
481 }
482 }
483 )+
484 };
485}
486
487#[allow(unused)]
488pub(crate) use impl_divisible_self;
489
490impl_divisible_self!(u8, u16, u32, u64, u128);
491
492#[cfg(test)]
493mod tests {
494 use proptest::{arbitrary::any, proptest};
495
496 use super::*;
497
498 #[test]
499 fn test_divisible_u32_u8() {
500 let val = 0xab12cd34u32;
501
502 assert_eq!(Divisible::<u8>::get(&val, 0), 0x34u8);
504 assert_eq!(Divisible::<u8>::get(&val, 1), 0xcdu8);
505 assert_eq!(Divisible::<u8>::get(&val, 2), 0x12u8);
506 assert_eq!(Divisible::<u8>::get(&val, 3), 0xabu8);
507
508 let vals: [u32; 2] = [0x04030201, 0x08070605];
509
510 let parts: Vec<u8> = Divisible::<u8>::slice_iter(&vals).collect();
512 assert_eq!(parts.len(), 8);
513 assert_eq!(parts[0], 0x01);
515 assert_eq!(parts[1], 0x02);
516 assert_eq!(parts[2], 0x03);
517 assert_eq!(parts[3], 0x04);
518 assert_eq!(parts[4], 0x05);
519 assert_eq!(parts[5], 0x06);
520 assert_eq!(parts[6], 0x07);
521 assert_eq!(parts[7], 0x08);
522 }
523
524 #[test]
525 fn test_broadcast_u32_u8() {
526 let result: u32 = Divisible::<u8>::broadcast(0xAB);
527 assert_eq!(result, 0xABABABAB);
528 }
529
530 #[test]
531 fn test_broadcast_u64_u16() {
532 let result: u64 = Divisible::<u16>::broadcast(0x1234);
533 assert_eq!(result, 0x1234123412341234);
534 }
535
536 #[test]
537 fn test_broadcast_u128_u32() {
538 let result: u128 = Divisible::<u32>::broadcast(0xDEADBEEF);
539 assert_eq!(result, 0xDEADBEEFDEADBEEFDEADBEEFDEADBEEF);
540 }
541
542 #[test]
543 fn test_broadcast_reflexive() {
544 let result: u64 = Divisible::<u64>::broadcast(0x123456789ABCDEF0);
545 assert_eq!(result, 0x123456789ABCDEF0);
546 }
547
548 #[test]
549 fn test_from_iter_full() {
550 let result: u32 = Divisible::<u8>::from_iter([0x01, 0x02, 0x03, 0x04].into_iter());
551 assert_eq!(result, 0x04030201);
552 }
553
554 #[test]
555 fn test_from_iter_partial() {
556 let result: u32 = Divisible::<u8>::from_iter([0xAB, 0xCD].into_iter());
558 assert_eq!(result, 0x0000CDAB);
559 }
560
561 #[test]
562 fn test_from_iter_empty() {
563 let result: u32 = Divisible::<u8>::from_iter(std::iter::empty());
564 assert_eq!(result, 0);
565 }
566
567 #[test]
568 fn test_from_iter_excess() {
569 let result: u32 =
571 Divisible::<u8>::from_iter([0x01, 0x02, 0x03, 0x04, 0x05, 0x06].into_iter());
572 assert_eq!(result, 0x04030201);
573 }
574
575 #[test]
576 fn test_from_iter_u64_u16() {
577 let result: u64 = Divisible::<u16>::from_iter([0x1234, 0x5678, 0x9ABC].into_iter());
578 assert_eq!(result, 0x0000_9ABC_5678_1234);
580 }
581
582 proptest! {
583 #[test]
584 fn test_set_get_u32_u8(mut val in any::<u32>(), i in 0usize..4, elem in any::<u8>()) {
585 Divisible::<u8>::set(&mut val, i, elem);
586 assert_eq!(Divisible::<u8>::get(&val, i), elem);
587 }
588 }
589}
590
591#[cfg(test)]
592mod arch_tests {
593 use std::fmt::Debug;
594
595 use binius_utils::{SerializeBytes, bytes::BytesMut};
596 use proptest::{arbitrary::any, proptest};
597
598 use super::Divisible;
599 use crate::arch::{M128, M256, M512};
600
601 fn check_byte_anchor<Big>(value: Big)
604 where
605 Big: Divisible<u8> + SerializeBytes + Copy,
606 {
607 let mut buf = BytesMut::new();
608 value
609 .serialize(&mut buf)
610 .expect("BytesMut grows to fit the value");
611
612 assert!(Big::value_iter(value).eq(buf));
613 }
614
615 fn check_refines<Big, Small>(value: Big)
617 where
618 Big: Divisible<u8> + Divisible<Small> + Copy,
619 Small: Divisible<u8> + Copy,
620 {
621 let bytes_per_lane = <Small as Divisible<u8>>::N;
622
623 for i in 0..<Big as Divisible<Small>>::N {
624 let lane = Divisible::<Small>::get(&value, i);
625 let window =
626 (0..bytes_per_lane).map(|j| Divisible::<u8>::get(&value, i * bytes_per_lane + j));
627
628 assert!(<Small as Divisible<u8>>::value_iter(lane).eq(window), "lane {i}");
629 }
630 }
631
632 fn check_iters<Big, Small>(value: Big, other: Big)
634 where
635 Big: Divisible<Small> + Copy + Eq + Debug,
636 Small: Copy + Eq + Debug,
637 {
638 let by_index = (0..<Big as Divisible<Small>>::N)
639 .map(|i| Divisible::<Small>::get(&value, i))
640 .collect::<Vec<_>>();
641
642 assert!(Big::value_iter(value).eq(by_index.iter().copied()));
643 assert!(Big::ref_iter(&value).eq(by_index.iter().copied()));
644
645 let slice = [value, other];
647 assert!(Big::slice_iter(&slice).eq(Big::value_iter(value).chain(Big::value_iter(other))));
648
649 assert_eq!(Big::from_iter(by_index.iter().copied()), value);
650 }
651
652 fn check_broadcast<Big, Small>(source: Big)
654 where
655 Big: Divisible<Small> + Copy,
656 Small: Copy + Eq + Debug,
657 {
658 let lane = Divisible::<Small>::get(&source, 0);
660 let value = <Big as Divisible<Small>>::broadcast(lane);
661
662 for i in 0..<Big as Divisible<Small>>::N {
663 assert_eq!(Divisible::<Small>::get(&value, i), lane, "index {i}");
664 }
665 }
666
667 fn check_set<Big, Small>(value: Big, source: Big, index: usize)
669 where
670 Big: Divisible<Small> + Copy,
671 Small: Copy + Eq + Debug,
672 {
673 let index = index % <Big as Divisible<Small>>::N;
674 let lane = Divisible::<Small>::get(&source, 0);
675
676 let mut updated = value;
677 Divisible::<Small>::set(&mut updated, index, lane);
678
679 assert_eq!(Divisible::<Small>::get(&updated, index), lane);
680 for i in (0..<Big as Divisible<Small>>::N).filter(|&i| i != index) {
681 assert_eq!(
682 Divisible::<Small>::get(&updated, i),
683 Divisible::<Small>::get(&value, i),
684 "index {i}"
685 );
686 }
687 }
688
689 fn check_width<Big, Small>(a: Big, b: Big, index: usize)
691 where
692 Big: Divisible<u8> + Divisible<Small> + Copy + Eq + Debug,
693 Small: Divisible<u8> + Copy + Eq + Debug,
694 {
695 check_refines::<Big, Small>(a);
696 check_iters::<Big, Small>(a, b);
697 check_broadcast::<Big, Small>(b);
698 check_set::<Big, Small>(a, b, index);
699 }
700
701 proptest! {
702 #[test]
706 fn m128_subdivisions(a in any::<u128>(), b in any::<u128>(), index in any::<usize>()) {
707 let (a, b) = (M128::from(a), M128::from(b));
708
709 check_byte_anchor(a);
710 check_width::<M128, u128>(a, b, index);
711 check_width::<M128, u64>(a, b, index);
712 check_width::<M128, u32>(a, b, index);
713 check_width::<M128, u16>(a, b, index);
714 check_width::<M128, u8>(a, b, index);
715 }
716
717 #[test]
718 fn m256_subdivisions(
719 a in any::<[u128; 2]>(),
720 b in any::<[u128; 2]>(),
721 index in any::<usize>(),
722 ) {
723 let (a, b) = (M256::from(a), M256::from(b));
724
725 check_byte_anchor(a);
726 check_width::<M256, M128>(a, b, index);
727 check_width::<M256, u128>(a, b, index);
728 check_width::<M256, u64>(a, b, index);
729 check_width::<M256, u32>(a, b, index);
730 check_width::<M256, u16>(a, b, index);
731 check_width::<M256, u8>(a, b, index);
732 }
733
734 #[test]
735 fn m512_subdivisions(
736 a in any::<[u128; 4]>(),
737 b in any::<[u128; 4]>(),
738 index in any::<usize>(),
739 ) {
740 let (a, b) = (M512::from(a), M512::from(b));
741
742 check_byte_anchor(a);
743 check_width::<M512, M256>(a, b, index);
744 check_width::<M512, M128>(a, b, index);
745 check_width::<M512, u128>(a, b, index);
746 check_width::<M512, u64>(a, b, index);
747 check_width::<M512, u32>(a, b, index);
748 check_width::<M512, u16>(a, b, index);
749 check_width::<M512, u8>(a, b, index);
750 }
751 }
752}