1use std::marker::PhantomData;
22use std::ops::{Index, IndexMut};
23
24use tinyvec::ArrayVec;
25
26use crate::element_op::{ComposableElementOp, ElementOp, ElementOpApply, Identity};
27use crate::{AxisVec, Result, StrideVec, StridedError, INLINE_RANK};
28
29#[inline(always)]
35fn arrayvec_from_slice<T: Copy + Default>(src: &[T]) -> ArrayVec<[T; INLINE_RANK]> {
36 debug_assert!(src.len() <= INLINE_RANK);
37 let n = src.len();
38 let mut buf = [T::default(); INLINE_RANK];
39 match n {
40 0 => {}
41 1 => buf[0] = src[0],
42 2 => {
43 buf[0] = src[0];
44 buf[1] = src[1];
45 }
46 3 => {
47 buf[0] = src[0];
48 buf[1] = src[1];
49 buf[2] = src[2];
50 }
51 4 => {
52 buf[0] = src[0];
53 buf[1] = src[1];
54 buf[2] = src[2];
55 buf[3] = src[3];
56 }
57 5 => {
58 buf[0] = src[0];
59 buf[1] = src[1];
60 buf[2] = src[2];
61 buf[3] = src[3];
62 buf[4] = src[4];
63 }
64 6 => {
65 buf[0] = src[0];
66 buf[1] = src[1];
67 buf[2] = src[2];
68 buf[3] = src[3];
69 buf[4] = src[4];
70 buf[5] = src[5];
71 }
72 7 => {
73 buf[0] = src[0];
74 buf[1] = src[1];
75 buf[2] = src[2];
76 buf[3] = src[3];
77 buf[4] = src[4];
78 buf[5] = src[5];
79 buf[6] = src[6];
80 }
81 _ => {
82 buf[0] = src[0];
83 buf[1] = src[1];
84 buf[2] = src[2];
85 buf[3] = src[3];
86 buf[4] = src[4];
87 buf[5] = src[5];
88 buf[6] = src[6];
89 buf[7] = src[7];
90 }
91 }
92 ArrayVec::from_array_len(buf, n)
93}
94
95#[derive(Clone, Debug)]
97struct SpilledMeta {
98 dims: Box<[usize]>,
99 strides: Box<[isize]>,
100}
101
102#[derive(Debug)]
109struct ViewMeta {
110 inline_dims: ArrayVec<[usize; INLINE_RANK]>,
111 inline_strides: ArrayVec<[isize; INLINE_RANK]>,
112 spilled: Option<Box<SpilledMeta>>,
113}
114
115impl Clone for ViewMeta {
116 #[inline]
117 fn clone(&self) -> Self {
118 Self {
119 inline_dims: self.inline_dims,
120 inline_strides: self.inline_strides,
121 spilled: self.spilled.clone(),
122 }
123 }
124}
125
126impl ViewMeta {
127 #[inline(always)]
128 fn from_slices(dims: &[usize], strides: &[isize]) -> Self {
129 debug_assert_eq!(dims.len(), strides.len());
130 if dims.len() <= INLINE_RANK {
131 Self {
132 inline_dims: arrayvec_from_slice(dims),
133 inline_strides: arrayvec_from_slice(strides),
134 spilled: None,
135 }
136 } else {
137 Self {
138 inline_dims: ArrayVec::new(),
139 inline_strides: ArrayVec::new(),
140 spilled: Some(Box::new(SpilledMeta {
141 dims: dims.into(),
142 strides: strides.into(),
143 })),
144 }
145 }
146 }
147
148 #[inline]
149 fn dims(&self) -> &[usize] {
150 match &self.spilled {
151 Some(spilled) => &spilled.dims,
152 None => self.inline_dims.as_slice(),
153 }
154 }
155
156 #[inline]
157 fn strides(&self) -> &[isize] {
158 match &self.spilled {
159 Some(spilled) => &spilled.strides,
160 None => self.inline_strides.as_slice(),
161 }
162 }
163}
164
165#[inline]
166fn empty_layout(dims: &[usize]) -> bool {
167 dims.iter().any(|&dim| dim == 0)
168}
169
170#[inline]
171fn empty_const_ptr<T>() -> *const T {
172 core::ptr::NonNull::<T>::dangling().as_ptr()
173}
174
175#[inline]
176fn empty_mut_ptr<T>() -> *mut T {
177 core::ptr::NonNull::<T>::dangling().as_ptr()
178}
179
180#[inline]
186pub(crate) fn validate_bounds(
187 len: usize,
188 dims: &[usize],
189 strides: &[isize],
190 offset: isize,
191) -> Result<()> {
192 if dims.len() != strides.len() {
193 return Err(StridedError::StrideLengthMismatch);
194 }
195 if dims.iter().any(|&d| d == 0) {
197 return Ok(());
198 }
199 let mut min_offset = offset;
201 let mut max_offset = offset;
202 for (&dim, &stride) in dims.iter().zip(strides.iter()) {
203 if dim > 1 {
204 let last = dim
205 .checked_sub(1)
206 .and_then(|value| isize::try_from(value).ok())
207 .ok_or(StridedError::OffsetOverflow)?;
208 let end = stride
209 .checked_mul(last)
210 .ok_or(StridedError::OffsetOverflow)?;
211 if end >= 0 {
212 max_offset = max_offset
213 .checked_add(end)
214 .ok_or(StridedError::OffsetOverflow)?;
215 } else {
216 min_offset = min_offset
217 .checked_add(end)
218 .ok_or(StridedError::OffsetOverflow)?;
219 }
220 }
221 }
222 if min_offset < 0 || max_offset < 0 {
223 return Err(StridedError::OffsetOverflow);
224 }
225 if max_offset as usize >= len {
226 return Err(StridedError::OffsetOverflow);
227 }
228 Ok(())
229}
230
231#[inline]
239fn indexed_offset(dims: &[usize], strides: &[isize], indices: &[usize], base: isize) -> isize {
240 assert_eq!(indices.len(), dims.len(), "wrong number of indices");
241 let mut offset = base;
242 for (i, &index) in indices.iter().enumerate() {
243 assert!(
244 index < dims[i],
245 "index {} out of bounds for dim {}",
246 index,
247 dims[i]
248 );
249 offset = isize::try_from(index)
250 .ok()
251 .and_then(|index| index.checked_mul(strides[i]))
252 .and_then(|step| offset.checked_add(step))
253 .expect("index offset overflow");
254 }
255 offset
256}
257
258#[inline]
272pub fn col_major_strides(dims: &[usize]) -> StrideVec {
273 let rank = dims.len();
274 if rank == 0 {
275 return StrideVec::new();
276 }
277 let mut strides = StrideVec::new();
278 strides.resize(rank, 1isize);
279 for i in 1..rank {
280 strides[i] = strides[i - 1] * dims[i - 1] as isize;
281 }
282 strides
283}
284
285#[inline]
299pub fn row_major_strides(dims: &[usize]) -> StrideVec {
300 let rank = dims.len();
301 if rank == 0 {
302 return StrideVec::new();
303 }
304 let mut strides = StrideVec::new();
305 strides.resize(rank, 1isize);
306 for i in (0..rank - 1).rev() {
307 strides[i] = strides[i + 1] * dims[i + 1] as isize;
308 }
309 strides
310}
311
312pub struct StridedView<'a, T, Op = Identity> {
348 ptr: *const T,
349 data: &'a [T],
350 meta: ViewMeta,
351 offset: isize,
352 _op: PhantomData<Op>,
353}
354
355unsafe impl<T: Sync, Op: Send> Send for StridedView<'_, T, Op> {}
356unsafe impl<T: Sync, Op: Sync> Sync for StridedView<'_, T, Op> {}
357
358impl<T, Op> Clone for StridedView<'_, T, Op> {
359 #[inline]
360 fn clone(&self) -> Self {
361 Self {
362 ptr: self.ptr,
363 data: self.data,
364 meta: self.meta.clone(),
365 offset: self.offset,
366 _op: PhantomData,
367 }
368 }
369}
370
371impl<T: std::fmt::Debug, Op> std::fmt::Debug for StridedView<'_, T, Op> {
372 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
373 f.debug_struct("StridedView")
374 .field("dims", &self.meta.dims())
375 .field("strides", &self.meta.strides())
376 .field("offset", &self.offset)
377 .finish()
378 }
379}
380
381#[inline]
386fn base_ptr<T>(data: &[T], dims: &[usize], offset: isize) -> *const T {
387 if empty_layout(dims) {
388 empty_const_ptr()
389 } else {
390 unsafe { data.as_ptr().offset(offset) }
391 }
392}
393
394#[inline]
396fn from_meta<'a, T, Op>(
397 data: &'a [T],
398 dims: &[usize],
399 strides: &[isize],
400 offset: isize,
401) -> StridedView<'a, T, Op> {
402 let meta = ViewMeta::from_slices(dims, strides);
403 let ptr = base_ptr(data, meta.dims(), offset);
404 StridedView {
405 ptr,
406 data,
407 meta,
408 offset,
409 _op: PhantomData,
410 }
411}
412
413impl<'a, T, Op> StridedView<'a, T, Op> {
414 #[inline]
431 pub fn new(data: &'a [T], dims: &[usize], strides: &[isize], offset: isize) -> Result<Self> {
432 validate_bounds(data.len(), dims, strides, offset)?;
433 Ok(from_meta(data, dims, strides, offset))
434 }
435
436 #[inline]
451 pub unsafe fn new_unchecked(
452 data: &'a [T],
453 dims: &[usize],
454 strides: &[isize],
455 offset: isize,
456 ) -> Self {
457 from_meta(data, dims, strides, offset)
458 }
459
460 #[inline]
462 pub fn dims(&self) -> &[usize] {
463 self.meta.dims()
464 }
465
466 #[inline]
468 pub fn strides(&self) -> &[isize] {
469 self.meta.strides()
470 }
471
472 #[inline]
474 pub fn offset(&self) -> isize {
475 self.offset
476 }
477
478 #[inline]
480 pub fn ndim(&self) -> usize {
481 self.dims().len()
482 }
483
484 #[inline]
486 pub fn len(&self) -> usize {
487 self.dims().iter().product()
488 }
489
490 #[inline]
492 pub fn is_empty(&self) -> bool {
493 self.dims().iter().any(|&d| d == 0)
494 }
495
496 #[inline]
498 pub fn data(&self) -> &'a [T] {
499 self.data
500 }
501
502 #[inline]
504 pub fn ptr(&self) -> *const T {
505 self.ptr
506 }
507
508 pub fn permute(&self, perm: &[usize]) -> Result<StridedView<'a, T, Op>> {
510 let rank = self.dims().len();
511 if perm.len() != rank {
512 return Err(StridedError::RankMismatch(perm.len(), rank));
513 }
514 let mut seen: tinyvec::TinyVec<[bool; INLINE_RANK]> = tinyvec::TinyVec::new();
515 seen.resize(rank, false);
516 for &p in perm {
517 if p >= rank {
518 return Err(StridedError::InvalidAxis { axis: p, rank });
519 }
520 if seen[p] {
521 return Err(StridedError::InvalidAxis { axis: p, rank });
522 }
523 seen[p] = true;
524 }
525 let new_dims: AxisVec = perm.iter().map(|&p| self.dims()[p]).collect();
526 let new_strides: StrideVec = perm.iter().map(|&p| self.strides()[p]).collect();
527 Ok(from_meta(self.data, &new_dims, &new_strides, self.offset))
528 }
529
530 pub fn diagonal_view(&self, axis_pairs: &[(usize, usize)]) -> Result<StridedView<'a, T, Op>> {
551 let ndim = self.ndim();
552 let mut dims: AxisVec = AxisVec::new();
553 dims.extend_from_slice(self.dims());
554 let mut strides: StrideVec = StrideVec::new();
555 strides.extend_from_slice(self.strides());
556
557 let mut axes_to_remove: tinyvec::TinyVec<[usize; INLINE_RANK]> = tinyvec::TinyVec::new();
558 for &(a, b) in axis_pairs {
559 let (lo, hi) = if a < b { (a, b) } else { (b, a) };
560 if hi >= ndim {
561 return Err(StridedError::InvalidAxis {
562 axis: hi,
563 rank: ndim,
564 });
565 }
566 if lo == hi {
567 return Err(StridedError::InvalidAxis {
568 axis: lo,
569 rank: ndim,
570 });
571 }
572 strides[lo] = strides[lo]
573 .checked_add(strides[hi])
574 .ok_or(StridedError::OffsetOverflow)?;
575 dims[lo] = dims[lo].min(dims[hi]);
576 if !axes_to_remove.contains(&hi) {
580 axes_to_remove.push(hi);
581 }
582 }
583
584 axes_to_remove.sort_unstable();
587 for &ax in axes_to_remove.iter().rev() {
588 dims.remove(ax);
589 strides.remove(ax);
590 }
591
592 validate_bounds(self.data().len(), &dims, &strides, self.offset())?;
596
597 Ok(from_meta(self.data(), &dims, &strides, self.offset()))
598 }
599
600 pub fn broadcast(&self, target_dims: &[usize]) -> Result<StridedView<'a, T, Op>> {
604 if self.dims().len() != target_dims.len() {
605 return Err(StridedError::RankMismatch(
606 self.dims().len(),
607 target_dims.len(),
608 ));
609 }
610 let mut new_strides = StrideVec::with_capacity(self.dims().len());
611 for i in 0..self.dims().len() {
612 if self.dims()[i] == target_dims[i] {
613 new_strides.push(self.strides()[i]);
614 } else if self.dims()[i] == 1 {
615 new_strides.push(0);
616 } else {
617 return Err(StridedError::ShapeMismatch(
618 self.dims().to_vec(),
619 target_dims.to_vec(),
620 ));
621 }
622 }
623 Ok(from_meta(self.data, target_dims, &new_strides, self.offset))
624 }
625}
626
627impl<'a, T: Copy + ElementOpApply, Op: ComposableElementOp<T>> StridedView<'a, T, Op> {
629 pub fn transpose_2d(&self) -> Result<StridedView<'a, T, Op::ComposeTranspose>> {
633 if self.dims().len() != 2 {
634 return Err(StridedError::RankMismatch(self.dims().len(), 2));
635 }
636 Ok(from_meta(
637 self.data,
638 &[self.dims()[1], self.dims()[0]],
639 &[self.strides()[1], self.strides()[0]],
640 self.offset,
641 ))
642 }
643
644 pub fn adjoint_2d(&self) -> Result<StridedView<'a, T, Op::ComposeAdjoint>> {
648 if self.dims().len() != 2 {
649 return Err(StridedError::RankMismatch(self.dims().len(), 2));
650 }
651 Ok(from_meta(
652 self.data,
653 &[self.dims()[1], self.dims()[0]],
654 &[self.strides()[1], self.strides()[0]],
655 self.offset,
656 ))
657 }
658
659 #[inline]
663 pub fn conj(&self) -> StridedView<'a, T, Op::ComposeConj> {
664 StridedView {
665 ptr: self.ptr,
666 data: self.data,
667 meta: self.meta.clone(),
668 offset: self.offset,
669 _op: PhantomData,
670 }
671 }
672}
673
674impl<'a, T: Copy, Op: ElementOp<T>> StridedView<'a, T, Op> {
676 pub fn get(&self, indices: &[usize]) -> T {
683 let offset = indexed_offset(self.dims(), self.strides(), indices, 0);
684 Op::apply(unsafe { *self.ptr.offset(offset) })
685 }
686
687 #[inline]
692 pub unsafe fn get_unchecked(&self, indices: &[usize]) -> T {
693 let mut idx = 0isize;
694 for (i, &index) in indices.iter().enumerate() {
695 idx += index as isize * self.strides()[i];
696 }
697 Op::apply(*self.ptr.offset(idx))
698 }
699}
700
701pub struct StridedViewMut<'a, T> {
710 ptr: *mut T,
711 data: &'a mut [T],
712 meta: ViewMeta,
713 offset: isize,
714}
715
716unsafe impl<T: Send> Send for StridedViewMut<'_, T> {}
717
718impl<T: std::fmt::Debug> std::fmt::Debug for StridedViewMut<'_, T> {
719 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
720 f.debug_struct("StridedViewMut")
721 .field("dims", &self.meta.dims())
722 .field("strides", &self.meta.strides())
723 .field("offset", &self.offset)
724 .finish()
725 }
726}
727
728impl<'a, T> StridedViewMut<'a, T> {
729 #[inline]
731 fn from_meta(data: &'a mut [T], dims: &[usize], strides: &[isize], offset: isize) -> Self {
732 let meta = ViewMeta::from_slices(dims, strides);
733 let ptr = if empty_layout(meta.dims()) {
734 empty_mut_ptr()
735 } else {
736 unsafe { data.as_mut_ptr().offset(offset) }
737 };
738 Self {
739 ptr,
740 data,
741 meta,
742 offset,
743 }
744 }
745
746 #[inline]
762 pub fn new(
763 data: &'a mut [T],
764 dims: &[usize],
765 strides: &[isize],
766 offset: isize,
767 ) -> Result<Self> {
768 validate_bounds(data.len(), dims, strides, offset)?;
769 Ok(Self::from_meta(data, dims, strides, offset))
770 }
771
772 #[inline]
787 pub unsafe fn new_unchecked(
788 data: &'a mut [T],
789 dims: &[usize],
790 strides: &[isize],
791 offset: isize,
792 ) -> Self {
793 Self::from_meta(data, dims, strides, offset)
794 }
795
796 #[inline]
798 pub fn dims(&self) -> &[usize] {
799 self.meta.dims()
800 }
801
802 #[inline]
804 pub fn strides(&self) -> &[isize] {
805 self.meta.strides()
806 }
807
808 #[inline]
810 pub fn offset(&self) -> isize {
811 self.offset
812 }
813
814 #[inline]
816 pub fn ndim(&self) -> usize {
817 self.dims().len()
818 }
819
820 #[inline]
822 pub fn len(&self) -> usize {
823 self.dims().iter().product()
824 }
825
826 #[inline]
828 pub fn is_empty(&self) -> bool {
829 self.dims().iter().any(|&d| d == 0)
830 }
831
832 #[inline]
834 pub fn ptr(&self) -> *const T {
835 self.ptr as *const T
836 }
837
838 #[inline]
840 pub fn as_mut_ptr(&self) -> *mut T {
841 self.ptr
842 }
843
844 #[inline]
846 pub fn data(&self) -> &[T] {
847 &self.data
848 }
849
850 #[inline]
852 pub fn data_mut(&mut self) -> &mut [T] {
853 &mut self.data
854 }
855
856 pub fn permute(self, perm: &[usize]) -> Result<StridedViewMut<'a, T>> {
861 let rank = self.dims().len();
862 if perm.len() != rank {
863 return Err(StridedError::RankMismatch(perm.len(), rank));
864 }
865 let mut seen: tinyvec::TinyVec<[bool; INLINE_RANK]> = tinyvec::TinyVec::new();
866 seen.resize(rank, false);
867 for &p in perm {
868 if p >= rank {
869 return Err(StridedError::InvalidAxis { axis: p, rank });
870 }
871 if seen[p] {
872 return Err(StridedError::InvalidAxis { axis: p, rank });
873 }
874 seen[p] = true;
875 }
876 let new_dims: AxisVec = perm.iter().map(|&p| self.dims()[p]).collect();
877 let new_strides: StrideVec = perm.iter().map(|&p| self.strides()[p]).collect();
878 Ok(Self::from_meta(
879 self.data,
880 &new_dims,
881 &new_strides,
882 self.offset,
883 ))
884 }
885
886 #[inline]
903 pub fn as_view(&self) -> StridedView<'_, T, Identity> {
904 StridedView {
905 ptr: self.ptr as *const T,
906 data: unsafe { std::slice::from_raw_parts(self.data.as_ptr(), self.data.len()) },
907 meta: self.meta.clone(),
908 offset: self.offset,
909 _op: PhantomData,
910 }
911 }
912}
913
914impl<'a, T: Copy> StridedViewMut<'a, T> {
915 pub fn get(&self, indices: &[usize]) -> T {
922 let offset = indexed_offset(self.dims(), self.strides(), indices, 0);
923 unsafe { *self.ptr.offset(offset) }
924 }
925
926 pub fn set(&mut self, indices: &[usize], value: T) {
933 let offset = indexed_offset(self.dims(), self.strides(), indices, 0);
934 unsafe {
935 *self.ptr.offset(offset) = value;
936 }
937 }
938}
939
940pub struct StridedArray<T> {
948 data: Vec<T>,
949 meta: ViewMeta,
950 offset: isize,
951}
952
953impl<T: std::fmt::Debug> std::fmt::Debug for StridedArray<T> {
954 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
955 f.debug_struct("StridedArray")
956 .field("dims", &self.meta.dims())
957 .field("strides", &self.meta.strides())
958 .field("offset", &self.offset)
959 .finish()
960 }
961}
962
963impl<T: Clone> Clone for StridedArray<T> {
964 fn clone(&self) -> Self {
965 Self {
966 data: self.data.clone(),
967 meta: self.meta.clone(),
968 offset: self.offset,
969 }
970 }
971}
972
973impl<T: Clone + Default> StridedArray<T> {
974 pub fn col_major(dims: &[usize]) -> Self {
989 let total: usize = dims.iter().product();
990 let data = vec![T::default(); total];
991 let strides = col_major_strides(dims);
992 Self {
993 data,
994 meta: ViewMeta::from_slices(dims, &strides),
995 offset: 0,
996 }
997 }
998
999 pub fn row_major(dims: &[usize]) -> Self {
1014 let total: usize = dims.iter().product();
1015 let data = vec![T::default(); total];
1016 let strides = row_major_strides(dims);
1017 Self {
1018 data,
1019 meta: ViewMeta::from_slices(dims, &strides),
1020 offset: 0,
1021 }
1022 }
1023
1024 pub fn from_fn_col_major(dims: &[usize], mut f: impl FnMut(&[usize]) -> T) -> Self {
1028 let total: usize = dims.iter().product();
1029 let strides = col_major_strides(dims);
1030 let rank = dims.len();
1031 let mut data = Vec::with_capacity(total);
1032 let mut idx: AxisVec = AxisVec::new();
1033 idx.resize(rank, 0usize);
1034 let idx: &mut [usize] = &mut idx[..];
1037 for _ in 0..total {
1038 data.push(f(idx));
1039 for d in 0..rank {
1040 idx[d] += 1;
1041 if idx[d] < dims[d] {
1042 break;
1043 }
1044 idx[d] = 0;
1045 }
1046 }
1047 Self {
1048 data,
1049 meta: ViewMeta::from_slices(dims, &strides),
1050 offset: 0,
1051 }
1052 }
1053
1054 pub fn from_fn_row_major(dims: &[usize], mut f: impl FnMut(&[usize]) -> T) -> Self {
1058 let total: usize = dims.iter().product();
1059 let strides = row_major_strides(dims);
1060 let rank = dims.len();
1061 let mut data = Vec::with_capacity(total);
1062 let mut idx: AxisVec = AxisVec::new();
1063 idx.resize(rank, 0usize);
1064 let idx: &mut [usize] = &mut idx[..];
1066 for _ in 0..total {
1067 data.push(f(idx));
1068 for d in (0..rank).rev() {
1069 idx[d] += 1;
1070 if idx[d] < dims[d] {
1071 break;
1072 }
1073 idx[d] = 0;
1074 }
1075 }
1076 Self {
1077 data,
1078 meta: ViewMeta::from_slices(dims, &strides),
1079 offset: 0,
1080 }
1081 }
1082}
1083
1084impl<T> StridedArray<T> {
1085 pub fn from_parts(
1087 data: Vec<T>,
1088 dims: &[usize],
1089 strides: &[isize],
1090 offset: isize,
1091 ) -> Result<Self> {
1092 validate_bounds(data.len(), dims, strides, offset)?;
1093 Ok(Self {
1094 data,
1095 meta: ViewMeta::from_slices(dims, strides),
1096 offset,
1097 })
1098 }
1099
1100 #[inline]
1102 pub fn dims(&self) -> &[usize] {
1103 self.meta.dims()
1104 }
1105
1106 #[inline]
1108 pub fn strides(&self) -> &[isize] {
1109 self.meta.strides()
1110 }
1111
1112 #[inline]
1114 pub fn ndim(&self) -> usize {
1115 self.dims().len()
1116 }
1117
1118 #[inline]
1120 pub fn len(&self) -> usize {
1121 self.dims().iter().product()
1122 }
1123
1124 #[inline]
1126 pub fn is_empty(&self) -> bool {
1127 self.dims().iter().any(|&d| d == 0)
1128 }
1129
1130 #[inline]
1132 pub fn data(&self) -> &[T] {
1133 &self.data
1134 }
1135
1136 #[inline]
1138 pub fn data_mut(&mut self) -> &mut [T] {
1139 &mut self.data
1140 }
1141
1142 #[inline]
1144 pub fn view(&self) -> StridedView<'_, T> {
1145 let ptr = if self.is_empty() {
1146 empty_const_ptr()
1147 } else {
1148 unsafe { self.data.as_ptr().offset(self.offset) }
1149 };
1150 StridedView {
1151 ptr,
1152 data: &self.data,
1153 meta: self.meta.clone(),
1154 offset: self.offset,
1155 _op: PhantomData,
1156 }
1157 }
1158
1159 #[inline]
1161 pub fn view_mut(&mut self) -> StridedViewMut<'_, T> {
1162 let ptr = if self.is_empty() {
1163 empty_mut_ptr()
1164 } else {
1165 unsafe { self.data.as_mut_ptr().offset(self.offset) }
1166 };
1167 StridedViewMut {
1168 ptr,
1169 data: &mut self.data,
1170 meta: self.meta.clone(),
1171 offset: self.offset,
1172 }
1173 }
1174
1175 pub fn permuted(self, perm: &[usize]) -> Result<Self> {
1180 let rank = self.dims().len();
1181 if perm.len() != rank {
1182 return Err(StridedError::RankMismatch(perm.len(), rank));
1183 }
1184 let mut seen: tinyvec::TinyVec<[bool; INLINE_RANK]> = tinyvec::TinyVec::new();
1185 seen.resize(rank, false);
1186 for &p in perm {
1187 if p >= rank {
1188 return Err(StridedError::InvalidAxis { axis: p, rank });
1189 }
1190 if seen[p] {
1191 return Err(StridedError::InvalidAxis { axis: p, rank });
1192 }
1193 seen[p] = true;
1194 }
1195 let new_dims: AxisVec = perm.iter().map(|&p| self.dims()[p]).collect();
1196 let new_strides: StrideVec = perm.iter().map(|&p| self.strides()[p]).collect();
1197 Ok(Self {
1198 data: self.data,
1199 meta: ViewMeta::from_slices(&new_dims, &new_strides),
1200 offset: self.offset,
1201 })
1202 }
1203
1204 pub fn into_data(self) -> Vec<T> {
1206 self.data
1207 }
1208
1209 pub fn iter(&self) -> std::slice::Iter<'_, T> {
1211 self.data.iter()
1212 }
1213
1214 pub fn iter_mut(&mut self) -> std::slice::IterMut<'_, T> {
1216 self.data.iter_mut()
1217 }
1218}
1219
1220impl<T: Default> StridedArray<T> {
1221 pub fn col_major_from_buffer(mut buf: Vec<T>, dims: &[usize]) -> Self {
1226 let total: usize = dims.iter().product();
1227 if buf.len() >= total {
1228 buf.truncate(total);
1229 } else {
1230 buf.resize_with(total, T::default);
1231 }
1232 for v in buf.iter_mut() {
1234 *v = T::default();
1235 }
1236 let strides = col_major_strides(dims);
1237 Self {
1238 data: buf,
1239 meta: ViewMeta::from_slices(dims, &strides),
1240 offset: 0,
1241 }
1242 }
1243}
1244
1245impl<T: Copy> StridedArray<T> {
1246 pub unsafe fn col_major_uninit(dims: &[usize]) -> Self {
1251 let total: usize = dims.iter().product();
1252 let mut data = Vec::with_capacity(total);
1253 data.set_len(total);
1254 let strides = col_major_strides(dims);
1255 Self {
1256 data,
1257 meta: ViewMeta::from_slices(dims, &strides),
1258 offset: 0,
1259 }
1260 }
1261
1262 pub unsafe fn col_major_from_buffer_uninit(mut buf: Vec<T>, dims: &[usize]) -> Self {
1267 let total: usize = dims.iter().product();
1268 if buf.capacity() < total {
1269 buf.reserve(total - buf.len());
1270 }
1271 buf.set_len(total);
1272 let strides = col_major_strides(dims);
1273 Self {
1274 data: buf,
1275 meta: ViewMeta::from_slices(dims, &strides),
1276 offset: 0,
1277 }
1278 }
1279}
1280
1281impl<T: Copy> StridedArray<T> {
1282 pub fn get(&self, indices: &[usize]) -> T {
1289 self.view().get(indices)
1290 }
1291
1292 pub fn set(&mut self, indices: &[usize], value: T) {
1299 let offset = indexed_offset(self.dims(), self.strides(), indices, self.offset);
1300 self.data[offset as usize] = value;
1301 }
1302}
1303
1304impl<T: Copy> Index<&[usize]> for StridedArray<T> {
1305 type Output = T;
1306
1307 fn index(&self, indices: &[usize]) -> &T {
1308 let offset = indexed_offset(self.dims(), self.strides(), indices, self.offset);
1309 &self.data[offset as usize]
1310 }
1311}
1312
1313impl<T: Copy> IndexMut<&[usize]> for StridedArray<T> {
1314 fn index_mut(&mut self, indices: &[usize]) -> &mut T {
1315 let offset = indexed_offset(self.dims(), self.strides(), indices, self.offset);
1316 &mut self.data[offset as usize]
1317 }
1318}
1319
1320#[cfg(test)]
1325mod tests {
1326 use super::*;
1327
1328 #[test]
1329 fn empty_views_use_dangling_base_pointers_for_extreme_offsets() {
1330 let dims = [0usize];
1331 let strides = [1isize];
1332 let data: [f64; 0] = [];
1333 let view = StridedView::<f64>::new(&data, &dims, &strides, isize::MAX).unwrap();
1334 assert_eq!(view.ptr(), empty_const_ptr());
1335
1336 let mut data: [f64; 0] = [];
1337 let view = StridedViewMut::new(&mut data, &dims, &strides, isize::MAX).unwrap();
1338 assert_eq!(view.ptr(), empty_const_ptr());
1339 assert_eq!(view.as_mut_ptr(), empty_mut_ptr());
1340
1341 let array =
1342 StridedArray::<f64>::from_parts(Vec::new(), &dims, &strides, isize::MAX).unwrap();
1343 assert_eq!(array.view().ptr(), empty_const_ptr());
1344
1345 let mut array =
1346 StridedArray::<f64>::from_parts(Vec::new(), &dims, &strides, isize::MAX).unwrap();
1347 let view = array.view_mut();
1348 assert_eq!(view.ptr(), empty_const_ptr());
1349 assert_eq!(view.as_mut_ptr(), empty_mut_ptr());
1350 }
1351 use num_complex::Complex64;
1352
1353 #[test]
1354 fn test_col_major_strides() {
1355 assert_eq!(col_major_strides(&[3, 4]).as_slice(), &[1, 3]);
1356 assert_eq!(col_major_strides(&[2, 3, 4]).as_slice(), &[1, 2, 6]);
1357 }
1358
1359 #[test]
1360 fn test_row_major_strides() {
1361 assert_eq!(row_major_strides(&[3, 4]).as_slice(), &[4, 1]);
1362 assert_eq!(row_major_strides(&[2, 3, 4]).as_slice(), &[12, 4, 1]);
1363 }
1364
1365 #[test]
1366 fn test_strided_view_new() {
1367 let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
1368 let view = StridedView::<f64>::new(&data, &[2, 3], &[3, 1], 0).unwrap();
1369 assert_eq!(view.ndim(), 2);
1370 assert_eq!(view.dims(), &[2, 3]);
1371 assert_eq!(view.strides(), &[3, 1]);
1372 assert_eq!(view.len(), 6);
1373 }
1374
1375 #[test]
1376 fn test_strided_view_get() {
1377 let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
1378 let view = StridedView::<f64>::new(&data, &[2, 3], &[3, 1], 0).unwrap();
1379 assert_eq!(view.get(&[0, 0]), 1.0);
1380 assert_eq!(view.get(&[0, 1]), 2.0);
1381 assert_eq!(view.get(&[0, 2]), 3.0);
1382 assert_eq!(view.get(&[1, 0]), 4.0);
1383 assert_eq!(view.get(&[1, 2]), 6.0);
1384 }
1385
1386 #[test]
1387 fn test_strided_view_col_major() {
1388 let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
1390 let view = StridedView::<f64>::new(&data, &[2, 3], &[1, 2], 0).unwrap();
1391 assert_eq!(view.get(&[0, 0]), 1.0); assert_eq!(view.get(&[1, 0]), 2.0); assert_eq!(view.get(&[0, 1]), 3.0); assert_eq!(view.get(&[1, 1]), 4.0); assert_eq!(view.get(&[0, 2]), 5.0); assert_eq!(view.get(&[1, 2]), 6.0); }
1398
1399 #[test]
1400 fn test_strided_view_permute() {
1401 let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
1402 let view = StridedView::<f64>::new(&data, &[2, 3], &[3, 1], 0).unwrap();
1403 let perm = view.permute(&[1, 0]).unwrap();
1404 assert_eq!(perm.dims(), &[3, 2]);
1405 assert_eq!(perm.strides(), &[1, 3]);
1406 assert_eq!(perm.get(&[0, 0]), 1.0);
1407 assert_eq!(perm.get(&[1, 0]), 2.0);
1408 assert_eq!(perm.get(&[0, 1]), 4.0);
1409 }
1410
1411 #[test]
1412 fn test_strided_view_transpose_2d() {
1413 let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
1414 let view = StridedView::<f64>::new(&data, &[2, 3], &[3, 1], 0).unwrap();
1415 let t = view.transpose_2d().unwrap();
1416 assert_eq!(t.dims(), &[3, 2]);
1417 assert_eq!(t.get(&[0, 0]), 1.0);
1418 assert_eq!(t.get(&[1, 0]), 2.0);
1419 assert_eq!(t.get(&[0, 1]), 4.0);
1420 }
1421
1422 #[test]
1423 fn test_strided_view_conj() {
1424 let data = vec![Complex64::new(1.0, 2.0), Complex64::new(3.0, 4.0)];
1425 let view = StridedView::<Complex64>::new(&data, &[2], &[1], 0).unwrap();
1426 let c = view.conj();
1427 assert_eq!(c.get(&[0]), Complex64::new(1.0, -2.0));
1428 assert_eq!(c.get(&[1]), Complex64::new(3.0, -4.0));
1429 }
1430
1431 #[test]
1432 fn test_strided_view_adjoint_2d() {
1433 let data = vec![
1434 Complex64::new(1.0, 2.0),
1435 Complex64::new(3.0, 4.0),
1436 Complex64::new(5.0, 6.0),
1437 Complex64::new(7.0, 8.0),
1438 ];
1439 let view = StridedView::<Complex64>::new(&data, &[2, 2], &[2, 1], 0).unwrap();
1441 let adj = view.adjoint_2d().unwrap();
1442 assert_eq!(adj.dims(), &[2, 2]);
1443 assert_eq!(adj.get(&[0, 0]), Complex64::new(1.0, -2.0));
1445 assert_eq!(adj.get(&[1, 0]), Complex64::new(3.0, -4.0));
1446 assert_eq!(adj.get(&[0, 1]), Complex64::new(5.0, -6.0));
1447 }
1448
1449 #[test]
1450 fn test_strided_view_broadcast() {
1451 let data = vec![1.0, 2.0, 3.0];
1452 let view = StridedView::<f64>::new(&data, &[1, 3], &[3, 1], 0).unwrap();
1453 let broad = view.broadcast(&[4, 3]).unwrap();
1454 assert_eq!(broad.dims(), &[4, 3]);
1455 for i in 0..4 {
1456 assert_eq!(broad.get(&[i, 0]), 1.0);
1457 assert_eq!(broad.get(&[i, 1]), 2.0);
1458 assert_eq!(broad.get(&[i, 2]), 3.0);
1459 }
1460 }
1461
1462 #[test]
1463 fn test_strided_view_mut() {
1464 let mut data = vec![0.0; 6];
1465 {
1466 let mut view = StridedViewMut::<f64>::new(&mut data, &[2, 3], &[3, 1], 0).unwrap();
1467 view.set(&[0, 0], 1.0);
1468 view.set(&[1, 2], 6.0);
1469 }
1470 assert_eq!(data[0], 1.0);
1471 assert_eq!(data[5], 6.0);
1472 }
1473
1474 #[test]
1475 fn test_strided_view_mut_as_view() {
1476 let mut data = vec![1.0, 2.0, 3.0];
1477 let vm = StridedViewMut::<f64>::new(&mut data, &[3], &[1], 0).unwrap();
1478 let v = vm.as_view();
1479 assert_eq!(v.get(&[0]), 1.0);
1480 assert_eq!(v.get(&[2]), 3.0);
1481 }
1482
1483 #[test]
1484 fn test_strided_tensor_col_major() {
1485 let t = StridedArray::<f64>::from_fn_col_major(&[2, 3], |idx| (idx[0] * 3 + idx[1]) as f64);
1486 assert_eq!(t.dims(), &[2, 3]);
1487 assert_eq!(t.strides(), &[1, 2]); assert_eq!(t.get(&[0, 0]), 0.0);
1489 assert_eq!(t.get(&[1, 0]), 3.0);
1490 assert_eq!(t.get(&[0, 1]), 1.0);
1491 assert_eq!(t.get(&[1, 2]), 5.0);
1492 }
1493
1494 #[test]
1495 fn test_strided_tensor_row_major() {
1496 let t = StridedArray::<f64>::from_fn_row_major(&[2, 3], |idx| (idx[0] * 3 + idx[1]) as f64);
1497 assert_eq!(t.dims(), &[2, 3]);
1498 assert_eq!(t.strides(), &[3, 1]); assert_eq!(t.get(&[0, 0]), 0.0);
1500 assert_eq!(t.get(&[0, 1]), 1.0);
1501 assert_eq!(t.get(&[1, 0]), 3.0);
1502 assert_eq!(t.get(&[1, 2]), 5.0);
1503 }
1504
1505 #[test]
1506 fn test_strided_tensor_view() {
1507 let t =
1508 StridedArray::<f64>::from_fn_col_major(&[2, 3], |idx| (idx[0] * 10 + idx[1]) as f64);
1509 let v = t.view();
1510 assert_eq!(v.get(&[0, 0]), 0.0);
1511 assert_eq!(v.get(&[1, 0]), 10.0);
1512 assert_eq!(v.get(&[0, 2]), 2.0);
1513 }
1514
1515 #[test]
1516 fn test_strided_tensor_view_mut() {
1517 let mut t = StridedArray::<f64>::col_major(&[2, 3]);
1518 {
1519 let mut vm = t.view_mut();
1520 vm.set(&[1, 2], 42.0);
1521 }
1522 assert_eq!(t.get(&[1, 2]), 42.0);
1523 }
1524
1525 #[test]
1526 fn test_strided_tensor_index() {
1527 let t =
1528 StridedArray::<f64>::from_fn_row_major(&[3, 4], |idx| (idx[0] * 10 + idx[1]) as f64);
1529 assert_eq!(t[&[0usize, 0] as &[usize]], 0.0);
1530 assert_eq!(t[&[2usize, 3] as &[usize]], 23.0);
1531 }
1532
1533 #[test]
1534 fn test_strided_tensor_index_mut() {
1535 let mut t = StridedArray::<f64>::row_major(&[2, 3]);
1536 t[&[1usize, 2] as &[usize]] = 99.0;
1537 assert_eq!(t.get(&[1, 2]), 99.0);
1538 }
1539
1540 #[test]
1541 fn test_validate_bounds_ok() {
1542 assert!(validate_bounds(6, &[2, 3], &[3, 1], 0).is_ok());
1543 assert!(validate_bounds(6, &[2, 3], &[1, 2], 0).is_ok());
1544 }
1545
1546 #[test]
1547 fn test_validate_bounds_out_of_range() {
1548 assert!(validate_bounds(5, &[2, 3], &[3, 1], 0).is_err());
1549 }
1550
1551 #[test]
1552 fn test_validate_bounds_empty() {
1553 assert!(validate_bounds(0, &[0, 3], &[3, 1], 0).is_ok());
1554 }
1555
1556 #[test]
1557 fn test_validate_bounds_with_offset() {
1558 assert!(validate_bounds(7, &[2, 3], &[3, 1], 1).is_ok());
1559 assert!(validate_bounds(6, &[2, 3], &[3, 1], 1).is_err());
1560 }
1561
1562 #[test]
1563 fn test_strided_tensor_3d() {
1564 let t = StridedArray::<f64>::from_fn_col_major(&[2, 3, 4], |idx| {
1565 (idx[0] * 100 + idx[1] * 10 + idx[2]) as f64
1566 });
1567 assert_eq!(t.ndim(), 3);
1568 assert_eq!(t.strides(), &[1, 2, 6]); assert_eq!(t.get(&[0, 0, 0]), 0.0);
1570 assert_eq!(t.get(&[1, 0, 0]), 100.0);
1571 assert_eq!(t.get(&[0, 1, 0]), 10.0);
1572 assert_eq!(t.get(&[0, 0, 1]), 1.0);
1573 assert_eq!(t.get(&[1, 2, 3]), 123.0);
1574 }
1575
1576 #[test]
1577 fn test_diagonal_view_2d() {
1578 let data: Vec<f64> = (0..9).map(|x| x as f64).collect();
1581 let view = StridedView::<f64>::new(&data, &[3, 3], &[3, 1], 0).unwrap();
1582 let diag = view.diagonal_view(&[(0, 1)]).unwrap();
1583 assert_eq!(diag.dims(), &[3]);
1584 assert_eq!(diag.strides(), &[4]);
1585 assert_eq!(diag.get(&[0]), 0.0); assert_eq!(diag.get(&[1]), 4.0); assert_eq!(diag.get(&[2]), 8.0); }
1589
1590 #[test]
1591 fn test_diagonal_view_3d_adjacent() {
1592 let data: Vec<f64> = (0..12).map(|x| x as f64).collect();
1595 let view = StridedView::<f64>::new(&data, &[2, 2, 3], &[6, 3, 1], 0).unwrap();
1596 let diag = view.diagonal_view(&[(0, 1)]).unwrap();
1597 assert_eq!(diag.dims(), &[2, 3]);
1598 assert_eq!(diag.strides(), &[9, 1]);
1599 assert_eq!(diag.get(&[0, 0]), 0.0);
1600 assert_eq!(diag.get(&[0, 2]), 2.0);
1601 assert_eq!(diag.get(&[1, 0]), 9.0);
1602 }
1603
1604 #[test]
1605 fn test_diagonal_view_3d_non_adjacent() {
1606 let data: Vec<f64> = (0..12).map(|x| x as f64).collect();
1609 let view = StridedView::<f64>::new(&data, &[2, 3, 2], &[6, 2, 1], 0).unwrap();
1610 let diag = view.diagonal_view(&[(0, 2)]).unwrap();
1611 assert_eq!(diag.dims(), &[2, 3]);
1612 assert_eq!(diag.strides(), &[7, 2]);
1613 assert_eq!(diag.get(&[0, 0]), 0.0);
1614 assert_eq!(diag.get(&[0, 1]), 2.0);
1615 assert_eq!(diag.get(&[1, 0]), 7.0);
1616 assert_eq!(diag.get(&[1, 2]), 11.0);
1617 }
1618
1619 #[test]
1620 fn test_diagonal_view_two_pairs() {
1621 let data: Vec<f64> = (0..36).map(|x| x as f64).collect();
1623 let view = StridedView::<f64>::new(&data, &[2, 3, 2, 3], &[18, 6, 3, 1], 0).unwrap();
1624 let diag = view.diagonal_view(&[(0, 2), (1, 3)]).unwrap();
1625 assert_eq!(diag.dims(), &[2, 3]);
1626 assert_eq!(diag.strides(), &[21, 7]);
1627 assert_eq!(diag.get(&[0, 0]), 0.0);
1628 assert_eq!(diag.get(&[1, 1]), 28.0);
1629 }
1630
1631 #[test]
1632 fn test_diagonal_view_rejects_pairs_that_leave_the_backing() {
1633 let data: Vec<f64> = (0..8).map(|x| x as f64).collect();
1634 let view = StridedView::<f64>::new(&data, &[2, 2, 2], &[4, 2, 1], 0).unwrap();
1635 assert!(view.diagonal_view(&[(0, 2), (1, 2)]).is_err());
1638 assert!(view.diagonal_view(&[(0, 1), (0, 1)]).is_err());
1640 }
1641
1642 #[test]
1643 fn test_diagonal_view_fuses_a_shared_low_axis_in_bounds() {
1644 let data: Vec<f64> = (0..8).map(|x| x as f64).collect();
1647 let view = StridedView::<f64>::new(&data, &[2, 2, 2], &[4, 2, 1], 0).unwrap();
1648 let diag = view.diagonal_view(&[(0, 1), (0, 2)]).unwrap();
1649 assert_eq!(diag.dims(), &[2]);
1650 assert_eq!(diag.strides(), &[7]);
1651 assert_eq!(diag.get(&[0]), 0.0);
1652 assert_eq!(diag.get(&[1]), 7.0);
1653 }
1654
1655 #[test]
1656 fn test_custom_type_without_element_op_apply() {
1657 #[derive(Debug, Clone, Copy, PartialEq)]
1660 struct MyScalar(f64);
1661
1662 impl Default for MyScalar {
1663 fn default() -> Self {
1664 MyScalar(0.0)
1665 }
1666 }
1667
1668 let data = vec![MyScalar(1.0), MyScalar(2.0), MyScalar(3.0), MyScalar(4.0)];
1670 let arr = StridedArray::from_parts(data, &[2, 2], &[2, 1], 0).unwrap();
1671
1672 assert_eq!(arr.get(&[0, 0]), MyScalar(1.0));
1674 assert_eq!(arr.get(&[1, 0]), MyScalar(3.0));
1675
1676 let view: StridedView<MyScalar> = arr.view();
1678 assert_eq!(view.get(&[0, 1]), MyScalar(2.0));
1679
1680 let perm = view.permute(&[1, 0]).unwrap();
1682 assert_eq!(perm.get(&[0, 0]), MyScalar(1.0));
1683 assert_eq!(perm.get(&[0, 1]), MyScalar(3.0)); }
1685
1686 #[test]
1687 #[should_panic(expected = "wrong number of indices")]
1688 fn test_index_rejects_a_short_index() {
1689 let arr = StridedArray::from_parts(vec![0.0f64; 4], &[2, 2], &[2, 1], 0).unwrap();
1690 let _ = arr[&[1][..]];
1692 }
1693
1694 #[test]
1695 #[should_panic(expected = "wrong number of indices")]
1696 fn test_view_get_rejects_a_short_index() {
1697 let data = vec![0.0f64; 4];
1698 let view = StridedView::<f64>::new(&data, &[2, 2], &[2, 1], 0).unwrap();
1699 let _ = view.get(&[1]);
1700 }
1701
1702 #[test]
1703 #[should_panic(expected = "out of bounds")]
1704 fn test_view_set_rejects_an_out_of_bounds_index() {
1705 let mut data = vec![0.0f64; 4];
1706 let mut view = StridedViewMut::<f64>::new(&mut data, &[2, 2], &[2, 1], 0).unwrap();
1707 view.set(&[2, 0], 1.0);
1708 }
1709}