1use std::borrow::Cow;
16use std::marker::PhantomData;
17use std::ops::{Index, IndexMut};
18
19use crate::element_op::{ComposableElementOp, ElementOp, ElementOpApply, Identity};
20use crate::{Result, StridedError};
21
22#[inline]
23fn empty_layout(dims: &[usize]) -> bool {
24 dims.iter().any(|&dim| dim == 0)
25}
26
27#[inline]
28fn empty_const_ptr<T>() -> *const T {
29 core::ptr::NonNull::<T>::dangling().as_ptr()
30}
31
32#[inline]
33fn empty_mut_ptr<T>() -> *mut T {
34 core::ptr::NonNull::<T>::dangling().as_ptr()
35}
36
37pub(crate) fn validate_bounds(
43 len: usize,
44 dims: &[usize],
45 strides: &[isize],
46 offset: isize,
47) -> Result<()> {
48 if dims.len() != strides.len() {
49 return Err(StridedError::StrideLengthMismatch);
50 }
51 if dims.iter().any(|&d| d == 0) {
53 return Ok(());
54 }
55 let mut min_offset = offset;
57 let mut max_offset = offset;
58 for (&dim, &stride) in dims.iter().zip(strides.iter()) {
59 if dim > 1 {
60 let last = dim
61 .checked_sub(1)
62 .and_then(|value| isize::try_from(value).ok())
63 .ok_or(StridedError::OffsetOverflow)?;
64 let end = stride
65 .checked_mul(last)
66 .ok_or(StridedError::OffsetOverflow)?;
67 if end >= 0 {
68 max_offset = max_offset
69 .checked_add(end)
70 .ok_or(StridedError::OffsetOverflow)?;
71 } else {
72 min_offset = min_offset
73 .checked_add(end)
74 .ok_or(StridedError::OffsetOverflow)?;
75 }
76 }
77 }
78 if min_offset < 0 || max_offset < 0 {
79 return Err(StridedError::OffsetOverflow);
80 }
81 if max_offset as usize >= len {
82 return Err(StridedError::OffsetOverflow);
83 }
84 Ok(())
85}
86
87#[inline]
95fn indexed_offset(dims: &[usize], strides: &[isize], indices: &[usize], base: isize) -> isize {
96 assert_eq!(indices.len(), dims.len(), "wrong number of indices");
97 let mut offset = base;
98 for (i, &index) in indices.iter().enumerate() {
99 assert!(
100 index < dims[i],
101 "index {} out of bounds for dim {}",
102 index,
103 dims[i]
104 );
105 offset = isize::try_from(index)
106 .ok()
107 .and_then(|index| index.checked_mul(strides[i]))
108 .and_then(|step| offset.checked_add(step))
109 .expect("index offset overflow");
110 }
111 offset
112}
113
114pub fn col_major_strides(dims: &[usize]) -> Vec<isize> {
116 let rank = dims.len();
117 if rank == 0 {
118 return vec![];
119 }
120 let mut strides = vec![1isize; rank];
121 for i in 1..rank {
122 strides[i] = strides[i - 1] * dims[i - 1] as isize;
123 }
124 strides
125}
126
127pub fn row_major_strides(dims: &[usize]) -> Vec<isize> {
129 let rank = dims.len();
130 if rank == 0 {
131 return vec![];
132 }
133 let mut strides = vec![1isize; rank];
134 for i in (0..rank - 1).rev() {
135 strides[i] = strides[i + 1] * dims[i + 1] as isize;
136 }
137 strides
138}
139
140pub struct StridedView<'a, T, Op = Identity> {
170 ptr: *const T,
171 data: &'a [T],
172 dims: Cow<'a, [usize]>,
173 strides: Cow<'a, [isize]>,
174 offset: isize,
175 _op: PhantomData<Op>,
176}
177
178unsafe impl<T: Sync, Op: Send> Send for StridedView<'_, T, Op> {}
179unsafe impl<T: Sync, Op: Sync> Sync for StridedView<'_, T, Op> {}
180
181impl<T, Op> Clone for StridedView<'_, T, Op> {
182 fn clone(&self) -> Self {
183 Self {
184 ptr: self.ptr,
185 data: self.data,
186 dims: self.dims.clone(),
187 strides: self.strides.clone(),
188 offset: self.offset,
189 _op: PhantomData,
190 }
191 }
192}
193
194impl<T: std::fmt::Debug, Op> std::fmt::Debug for StridedView<'_, T, Op> {
195 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
196 f.debug_struct("StridedView")
197 .field("dims", &self.dims)
198 .field("strides", &self.strides)
199 .field("offset", &self.offset)
200 .finish()
201 }
202}
203
204#[inline]
209fn base_ptr<T>(data: &[T], dims: &[usize], offset: isize) -> *const T {
210 if empty_layout(dims) {
211 empty_const_ptr()
212 } else {
213 unsafe { data.as_ptr().offset(offset) }
214 }
215}
216
217#[inline]
219fn from_meta<'a, T, Op>(
220 data: &'a [T],
221 dims: Cow<'a, [usize]>,
222 strides: Cow<'a, [isize]>,
223 offset: isize,
224) -> StridedView<'a, T, Op> {
225 StridedView {
226 ptr: base_ptr(data, &dims, offset),
227 data,
228 dims,
229 strides,
230 offset,
231 _op: PhantomData,
232 }
233}
234
235impl<'a, T, Op> StridedView<'a, T, Op> {
236 pub fn new(
241 data: &'a [T],
242 dims: &'a [usize],
243 strides: &'a [isize],
244 offset: isize,
245 ) -> Result<Self> {
246 validate_bounds(data.len(), dims, strides, offset)?;
247 Ok(from_meta(
248 data,
249 Cow::Borrowed(dims),
250 Cow::Borrowed(strides),
251 offset,
252 ))
253 }
254
255 pub fn new_owned(
260 data: &'a [T],
261 dims: Vec<usize>,
262 strides: Vec<isize>,
263 offset: isize,
264 ) -> Result<Self> {
265 validate_bounds(data.len(), &dims, &strides, offset)?;
266 Ok(from_meta(
267 data,
268 Cow::Owned(dims),
269 Cow::Owned(strides),
270 offset,
271 ))
272 }
273
274 pub unsafe fn new_unchecked(
279 data: &'a [T],
280 dims: &'a [usize],
281 strides: &'a [isize],
282 offset: isize,
283 ) -> Self {
284 from_meta(data, Cow::Borrowed(dims), Cow::Borrowed(strides), offset)
285 }
286
287 pub unsafe fn new_owned_unchecked(
292 data: &'a [T],
293 dims: Vec<usize>,
294 strides: Vec<isize>,
295 offset: isize,
296 ) -> Self {
297 from_meta(data, Cow::Owned(dims), Cow::Owned(strides), offset)
298 }
299
300 #[inline]
302 pub fn dims(&self) -> &[usize] {
303 &self.dims
304 }
305
306 #[inline]
308 pub fn strides(&self) -> &[isize] {
309 &self.strides
310 }
311
312 #[inline]
314 pub fn offset(&self) -> isize {
315 self.offset
316 }
317
318 #[inline]
320 pub fn ndim(&self) -> usize {
321 self.dims.len()
322 }
323
324 #[inline]
326 pub fn len(&self) -> usize {
327 self.dims.iter().product()
328 }
329
330 #[inline]
332 pub fn is_empty(&self) -> bool {
333 self.dims.iter().any(|&d| d == 0)
334 }
335
336 #[inline]
338 pub fn data(&self) -> &'a [T] {
339 self.data
340 }
341
342 #[inline]
344 pub fn ptr(&self) -> *const T {
345 self.ptr
346 }
347
348 pub fn permute(&self, perm: &[usize]) -> Result<StridedView<'a, T, Op>> {
350 let rank = self.dims.len();
351 if perm.len() != rank {
352 return Err(StridedError::RankMismatch(perm.len(), rank));
353 }
354 let mut seen = vec![false; rank];
355 for &p in perm {
356 if p >= rank {
357 return Err(StridedError::InvalidAxis { axis: p, rank });
358 }
359 if seen[p] {
360 return Err(StridedError::InvalidAxis { axis: p, rank });
361 }
362 seen[p] = true;
363 }
364 let new_dims: Vec<usize> = perm.iter().map(|&p| self.dims[p]).collect();
365 let new_strides: Vec<isize> = perm.iter().map(|&p| self.strides[p]).collect();
366 Ok(from_meta(
367 self.data,
368 Cow::Owned(new_dims),
369 Cow::Owned(new_strides),
370 self.offset,
371 ))
372 }
373
374 pub fn diagonal_view(&self, axis_pairs: &[(usize, usize)]) -> Result<StridedView<'a, T, Op>> {
395 let ndim = self.ndim();
396 let mut dims: Vec<usize> = self.dims().to_vec();
397 let mut strides: Vec<isize> = self.strides().to_vec();
398
399 let mut axes_to_remove = Vec::new();
400 for &(a, b) in axis_pairs {
401 let (lo, hi) = if a < b { (a, b) } else { (b, a) };
402 if hi >= ndim {
403 return Err(StridedError::InvalidAxis {
404 axis: hi,
405 rank: ndim,
406 });
407 }
408 if lo == hi {
409 return Err(StridedError::InvalidAxis {
410 axis: lo,
411 rank: ndim,
412 });
413 }
414 strides[lo] = strides[lo]
415 .checked_add(strides[hi])
416 .ok_or(StridedError::OffsetOverflow)?;
417 dims[lo] = dims[lo].min(dims[hi]);
418 axes_to_remove.push(hi);
419 }
420
421 axes_to_remove.sort_unstable();
422 axes_to_remove.dedup();
423 for &ax in axes_to_remove.iter().rev() {
424 dims.remove(ax);
425 strides.remove(ax);
426 }
427
428 validate_bounds(self.data().len(), &dims, &strides, self.offset())?;
432
433 Ok(from_meta(
434 self.data(),
435 Cow::Owned(dims),
436 Cow::Owned(strides),
437 self.offset(),
438 ))
439 }
440
441 pub fn broadcast(&self, target_dims: &[usize]) -> Result<StridedView<'a, T, Op>> {
445 if self.dims.len() != target_dims.len() {
446 return Err(StridedError::RankMismatch(
447 self.dims.len(),
448 target_dims.len(),
449 ));
450 }
451 let mut new_strides = Vec::with_capacity(self.dims.len());
452 for i in 0..self.dims.len() {
453 if self.dims[i] == target_dims[i] {
454 new_strides.push(self.strides[i]);
455 } else if self.dims[i] == 1 {
456 new_strides.push(0);
457 } else {
458 return Err(StridedError::ShapeMismatch(
459 self.dims.to_vec(),
460 target_dims.to_vec(),
461 ));
462 }
463 }
464 Ok(from_meta(
465 self.data,
466 Cow::Owned(target_dims.to_vec()),
467 Cow::Owned(new_strides),
468 self.offset,
469 ))
470 }
471}
472
473impl<'a, T: Copy + ElementOpApply, Op: ComposableElementOp<T>> StridedView<'a, T, Op> {
475 pub fn transpose_2d(&self) -> Result<StridedView<'a, T, Op::ComposeTranspose>> {
479 if self.dims.len() != 2 {
480 return Err(StridedError::RankMismatch(self.dims.len(), 2));
481 }
482 Ok(from_meta(
483 self.data,
484 Cow::Owned(vec![self.dims[1], self.dims[0]]),
485 Cow::Owned(vec![self.strides[1], self.strides[0]]),
486 self.offset,
487 ))
488 }
489
490 pub fn adjoint_2d(&self) -> Result<StridedView<'a, T, Op::ComposeAdjoint>> {
494 if self.dims.len() != 2 {
495 return Err(StridedError::RankMismatch(self.dims.len(), 2));
496 }
497 Ok(from_meta(
498 self.data,
499 Cow::Owned(vec![self.dims[1], self.dims[0]]),
500 Cow::Owned(vec![self.strides[1], self.strides[0]]),
501 self.offset,
502 ))
503 }
504
505 pub fn conj(&self) -> StridedView<'a, T, Op::ComposeConj> {
509 StridedView {
510 ptr: self.ptr,
511 data: self.data,
512 dims: self.dims.clone(),
513 strides: self.strides.clone(),
514 offset: self.offset,
515 _op: PhantomData,
516 }
517 }
518}
519
520impl<'a, T: Copy, Op: ElementOp<T>> StridedView<'a, T, Op> {
522 pub fn get(&self, indices: &[usize]) -> T {
529 let offset = indexed_offset(self.dims(), self.strides(), indices, 0);
530 Op::apply(unsafe { *self.ptr.offset(offset) })
531 }
532
533 #[inline]
538 pub unsafe fn get_unchecked(&self, indices: &[usize]) -> T {
539 let mut idx = 0isize;
540 for (i, &index) in indices.iter().enumerate() {
541 idx += index as isize * self.strides[i];
542 }
543 Op::apply(*self.ptr.offset(idx))
544 }
545}
546
547pub struct StridedViewMut<'a, T> {
556 ptr: *mut T,
557 data: &'a mut [T],
558 dims: Cow<'a, [usize]>,
559 strides: Cow<'a, [isize]>,
560 offset: isize,
561}
562
563unsafe impl<T: Send> Send for StridedViewMut<'_, T> {}
564
565impl<T: std::fmt::Debug> std::fmt::Debug for StridedViewMut<'_, T> {
566 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
567 f.debug_struct("StridedViewMut")
568 .field("dims", &self.dims)
569 .field("strides", &self.strides)
570 .field("offset", &self.offset)
571 .finish()
572 }
573}
574
575impl<'a, T> StridedViewMut<'a, T> {
576 #[inline]
578 fn from_meta(
579 data: &'a mut [T],
580 dims: Cow<'a, [usize]>,
581 strides: Cow<'a, [isize]>,
582 offset: isize,
583 ) -> Self {
584 let ptr = if empty_layout(&dims) {
585 empty_mut_ptr()
586 } else {
587 unsafe { data.as_mut_ptr().offset(offset) }
588 };
589 Self {
590 ptr,
591 data,
592 dims,
593 strides,
594 offset,
595 }
596 }
597
598 pub fn new(
600 data: &'a mut [T],
601 dims: &'a [usize],
602 strides: &'a [isize],
603 offset: isize,
604 ) -> Result<Self> {
605 validate_bounds(data.len(), dims, strides, offset)?;
606 Ok(Self::from_meta(
607 data,
608 Cow::Borrowed(dims),
609 Cow::Borrowed(strides),
610 offset,
611 ))
612 }
613
614 pub unsafe fn new_unchecked(
619 data: &'a mut [T],
620 dims: &'a [usize],
621 strides: &'a [isize],
622 offset: isize,
623 ) -> Self {
624 Self::from_meta(data, Cow::Borrowed(dims), Cow::Borrowed(strides), offset)
625 }
626
627 #[inline]
629 pub fn dims(&self) -> &[usize] {
630 &self.dims
631 }
632
633 #[inline]
635 pub fn strides(&self) -> &[isize] {
636 &self.strides
637 }
638
639 #[inline]
641 pub fn offset(&self) -> isize {
642 self.offset
643 }
644
645 #[inline]
647 pub fn ndim(&self) -> usize {
648 self.dims.len()
649 }
650
651 #[inline]
653 pub fn len(&self) -> usize {
654 self.dims.iter().product()
655 }
656
657 #[inline]
659 pub fn is_empty(&self) -> bool {
660 self.dims.iter().any(|&d| d == 0)
661 }
662
663 #[inline]
665 pub fn ptr(&self) -> *const T {
666 self.ptr as *const T
667 }
668
669 #[inline]
671 pub fn as_mut_ptr(&self) -> *mut T {
672 self.ptr
673 }
674
675 #[inline]
677 pub fn data(&self) -> &[T] {
678 &self.data
679 }
680
681 #[inline]
683 pub fn data_mut(&mut self) -> &mut [T] {
684 &mut self.data
685 }
686
687 pub fn permute(self, perm: &[usize]) -> Result<StridedViewMut<'a, T>> {
692 let rank = self.dims.len();
693 if perm.len() != rank {
694 return Err(StridedError::RankMismatch(perm.len(), rank));
695 }
696 let mut seen = vec![false; rank];
697 for &p in perm {
698 if p >= rank {
699 return Err(StridedError::InvalidAxis { axis: p, rank });
700 }
701 if seen[p] {
702 return Err(StridedError::InvalidAxis { axis: p, rank });
703 }
704 seen[p] = true;
705 }
706 let new_dims: Vec<usize> = perm.iter().map(|&p| self.dims[p]).collect();
707 let new_strides: Vec<isize> = perm.iter().map(|&p| self.strides[p]).collect();
708 Ok(Self::from_meta(
709 self.data,
710 Cow::Owned(new_dims),
711 Cow::Owned(new_strides),
712 self.offset,
713 ))
714 }
715
716 pub fn as_view(&self) -> StridedView<'_, T, Identity> {
721 StridedView {
722 ptr: self.ptr as *const T,
723 data: unsafe { std::slice::from_raw_parts(self.data.as_ptr(), self.data.len()) },
724 dims: Cow::Borrowed(self.dims.as_ref()),
725 strides: Cow::Borrowed(self.strides.as_ref()),
726 offset: self.offset,
727 _op: PhantomData,
728 }
729 }
730}
731
732impl<'a, T: Copy> StridedViewMut<'a, T> {
733 pub fn get(&self, indices: &[usize]) -> T {
740 let offset = indexed_offset(self.dims(), self.strides(), indices, 0);
741 unsafe { *self.ptr.offset(offset) }
742 }
743
744 pub fn set(&mut self, indices: &[usize], value: T) {
751 let offset = indexed_offset(self.dims(), self.strides(), indices, 0);
752 unsafe {
753 *self.ptr.offset(offset) = value;
754 }
755 }
756}
757
758pub struct StridedArray<T> {
766 data: Vec<T>,
767 dims: Vec<usize>,
768 strides: Vec<isize>,
769 offset: isize,
770}
771
772impl<T: std::fmt::Debug> std::fmt::Debug for StridedArray<T> {
773 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
774 f.debug_struct("StridedArray")
775 .field("dims", &self.dims)
776 .field("strides", &self.strides)
777 .field("offset", &self.offset)
778 .finish()
779 }
780}
781
782impl<T: Clone> Clone for StridedArray<T> {
783 fn clone(&self) -> Self {
784 Self {
785 data: self.data.clone(),
786 dims: self.dims.clone(),
787 strides: self.strides.clone(),
788 offset: self.offset,
789 }
790 }
791}
792
793impl<T: Clone + Default> StridedArray<T> {
794 pub fn col_major(dims: &[usize]) -> Self {
796 let total: usize = dims.iter().product();
797 let data = vec![T::default(); total];
798 let strides = col_major_strides(dims);
799 Self {
800 data,
801 dims: dims.to_vec(),
802 strides,
803 offset: 0,
804 }
805 }
806
807 pub fn row_major(dims: &[usize]) -> Self {
809 let total: usize = dims.iter().product();
810 let data = vec![T::default(); total];
811 let strides = row_major_strides(dims);
812 Self {
813 data,
814 dims: dims.to_vec(),
815 strides,
816 offset: 0,
817 }
818 }
819
820 pub fn from_fn_col_major(dims: &[usize], mut f: impl FnMut(&[usize]) -> T) -> Self {
824 let total: usize = dims.iter().product();
825 let strides = col_major_strides(dims);
826 let rank = dims.len();
827 let mut data = Vec::with_capacity(total);
828 let mut idx = vec![0usize; rank];
829 for _ in 0..total {
830 data.push(f(&idx));
831 for d in 0..rank {
832 idx[d] += 1;
833 if idx[d] < dims[d] {
834 break;
835 }
836 idx[d] = 0;
837 }
838 }
839 Self {
840 data,
841 dims: dims.to_vec(),
842 strides,
843 offset: 0,
844 }
845 }
846
847 pub fn from_fn_row_major(dims: &[usize], mut f: impl FnMut(&[usize]) -> T) -> Self {
851 let total: usize = dims.iter().product();
852 let strides = row_major_strides(dims);
853 let rank = dims.len();
854 let mut data = Vec::with_capacity(total);
855 let mut idx = vec![0usize; rank];
856 for _ in 0..total {
857 data.push(f(&idx));
858 for d in (0..rank).rev() {
859 idx[d] += 1;
860 if idx[d] < dims[d] {
861 break;
862 }
863 idx[d] = 0;
864 }
865 }
866 Self {
867 data,
868 dims: dims.to_vec(),
869 strides,
870 offset: 0,
871 }
872 }
873}
874
875impl<T> StridedArray<T> {
876 pub fn from_parts(
878 data: Vec<T>,
879 dims: &[usize],
880 strides: &[isize],
881 offset: isize,
882 ) -> Result<Self> {
883 validate_bounds(data.len(), dims, strides, offset)?;
884 Ok(Self {
885 data,
886 dims: dims.to_vec(),
887 strides: strides.to_vec(),
888 offset,
889 })
890 }
891
892 #[inline]
894 pub fn dims(&self) -> &[usize] {
895 &self.dims
896 }
897
898 #[inline]
900 pub fn strides(&self) -> &[isize] {
901 &self.strides
902 }
903
904 #[inline]
906 pub fn ndim(&self) -> usize {
907 self.dims.len()
908 }
909
910 #[inline]
912 pub fn len(&self) -> usize {
913 self.dims.iter().product()
914 }
915
916 #[inline]
918 pub fn is_empty(&self) -> bool {
919 self.dims.iter().any(|&d| d == 0)
920 }
921
922 #[inline]
924 pub fn data(&self) -> &[T] {
925 &self.data
926 }
927
928 #[inline]
930 pub fn data_mut(&mut self) -> &mut [T] {
931 &mut self.data
932 }
933
934 pub fn view(&self) -> StridedView<'_, T> {
936 let ptr = if self.is_empty() {
937 empty_const_ptr()
938 } else {
939 unsafe { self.data.as_ptr().offset(self.offset) }
940 };
941 StridedView {
942 ptr,
943 data: &self.data,
944 dims: Cow::Borrowed(&self.dims),
945 strides: Cow::Borrowed(&self.strides),
946 offset: self.offset,
947 _op: PhantomData,
948 }
949 }
950
951 pub fn view_mut(&mut self) -> StridedViewMut<'_, T> {
953 let ptr = if self.is_empty() {
954 empty_mut_ptr()
955 } else {
956 unsafe { self.data.as_mut_ptr().offset(self.offset) }
957 };
958 StridedViewMut {
959 ptr,
960 data: &mut self.data,
961 dims: Cow::Borrowed(&self.dims),
962 strides: Cow::Borrowed(&self.strides),
963 offset: self.offset,
964 }
965 }
966
967 pub fn permuted(self, perm: &[usize]) -> Result<Self> {
972 let rank = self.dims.len();
973 if perm.len() != rank {
974 return Err(StridedError::RankMismatch(perm.len(), rank));
975 }
976 let mut seen = vec![false; rank];
977 for &p in perm {
978 if p >= rank {
979 return Err(StridedError::InvalidAxis { axis: p, rank });
980 }
981 if seen[p] {
982 return Err(StridedError::InvalidAxis { axis: p, rank });
983 }
984 seen[p] = true;
985 }
986 let new_dims: Vec<usize> = perm.iter().map(|&p| self.dims[p]).collect();
987 let new_strides: Vec<isize> = perm.iter().map(|&p| self.strides[p]).collect();
988 Ok(Self {
989 data: self.data,
990 dims: new_dims,
991 strides: new_strides,
992 offset: self.offset,
993 })
994 }
995
996 pub fn into_data(self) -> Vec<T> {
998 self.data
999 }
1000
1001 pub fn iter(&self) -> std::slice::Iter<'_, T> {
1003 self.data.iter()
1004 }
1005
1006 pub fn iter_mut(&mut self) -> std::slice::IterMut<'_, T> {
1008 self.data.iter_mut()
1009 }
1010}
1011
1012impl<T: Default> StridedArray<T> {
1013 pub fn col_major_from_buffer(mut buf: Vec<T>, dims: &[usize]) -> Self {
1018 let total: usize = dims.iter().product();
1019 if buf.len() >= total {
1020 buf.truncate(total);
1021 } else {
1022 buf.resize_with(total, T::default);
1023 }
1024 for v in buf.iter_mut() {
1026 *v = T::default();
1027 }
1028 let strides = col_major_strides(dims);
1029 Self {
1030 data: buf,
1031 dims: dims.to_vec(),
1032 strides,
1033 offset: 0,
1034 }
1035 }
1036}
1037
1038impl<T: Copy> StridedArray<T> {
1039 pub unsafe fn col_major_uninit(dims: &[usize]) -> Self {
1044 let total: usize = dims.iter().product();
1045 let mut data = Vec::with_capacity(total);
1046 data.set_len(total);
1047 let strides = col_major_strides(dims);
1048 Self {
1049 data,
1050 dims: dims.to_vec(),
1051 strides,
1052 offset: 0,
1053 }
1054 }
1055
1056 pub unsafe fn col_major_from_buffer_uninit(mut buf: Vec<T>, dims: &[usize]) -> Self {
1061 let total: usize = dims.iter().product();
1062 if buf.capacity() < total {
1063 buf.reserve(total - buf.len());
1064 }
1065 buf.set_len(total);
1066 let strides = col_major_strides(dims);
1067 Self {
1068 data: buf,
1069 dims: dims.to_vec(),
1070 strides,
1071 offset: 0,
1072 }
1073 }
1074}
1075
1076impl<T: Copy> StridedArray<T> {
1077 pub fn get(&self, indices: &[usize]) -> T {
1084 self.view().get(indices)
1085 }
1086
1087 pub fn set(&mut self, indices: &[usize], value: T) {
1094 let offset = indexed_offset(&self.dims, &self.strides, indices, self.offset);
1095 self.data[offset as usize] = value;
1096 }
1097}
1098
1099impl<T: Copy> Index<&[usize]> for StridedArray<T> {
1100 type Output = T;
1101
1102 fn index(&self, indices: &[usize]) -> &T {
1103 let offset = indexed_offset(&self.dims, &self.strides, indices, self.offset);
1104 &self.data[offset as usize]
1105 }
1106}
1107
1108impl<T: Copy> IndexMut<&[usize]> for StridedArray<T> {
1109 fn index_mut(&mut self, indices: &[usize]) -> &mut T {
1110 let offset = indexed_offset(&self.dims, &self.strides, indices, self.offset);
1111 &mut self.data[offset as usize]
1112 }
1113}
1114
1115#[cfg(test)]
1120mod tests {
1121 use super::*;
1122
1123 #[test]
1124 fn empty_views_use_dangling_base_pointers_for_extreme_offsets() {
1125 let dims = [0usize];
1126 let strides = [1isize];
1127 let data: [f64; 0] = [];
1128 let view = StridedView::<f64>::new(&data, &dims, &strides, isize::MAX).unwrap();
1129 assert_eq!(view.ptr(), empty_const_ptr());
1130
1131 let mut data: [f64; 0] = [];
1132 let view = StridedViewMut::new(&mut data, &dims, &strides, isize::MAX).unwrap();
1133 assert_eq!(view.ptr(), empty_const_ptr());
1134 assert_eq!(view.as_mut_ptr(), empty_mut_ptr());
1135
1136 let array =
1137 StridedArray::<f64>::from_parts(Vec::new(), &dims, &strides, isize::MAX).unwrap();
1138 assert_eq!(array.view().ptr(), empty_const_ptr());
1139
1140 let mut array =
1141 StridedArray::<f64>::from_parts(Vec::new(), &dims, &strides, isize::MAX).unwrap();
1142 let view = array.view_mut();
1143 assert_eq!(view.ptr(), empty_const_ptr());
1144 assert_eq!(view.as_mut_ptr(), empty_mut_ptr());
1145 }
1146 use num_complex::Complex64;
1147
1148 #[test]
1149 fn test_col_major_strides() {
1150 assert_eq!(col_major_strides(&[3, 4]), vec![1, 3]);
1151 assert_eq!(col_major_strides(&[2, 3, 4]), vec![1, 2, 6]);
1152 }
1153
1154 #[test]
1155 fn test_row_major_strides() {
1156 assert_eq!(row_major_strides(&[3, 4]), vec![4, 1]);
1157 assert_eq!(row_major_strides(&[2, 3, 4]), vec![12, 4, 1]);
1158 }
1159
1160 #[test]
1161 fn test_strided_view_new() {
1162 let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
1163 let view = StridedView::<f64>::new(&data, &[2, 3], &[3, 1], 0).unwrap();
1164 assert_eq!(view.ndim(), 2);
1165 assert_eq!(view.dims(), &[2, 3]);
1166 assert_eq!(view.strides(), &[3, 1]);
1167 assert_eq!(view.len(), 6);
1168 }
1169
1170 #[test]
1171 fn test_strided_view_get() {
1172 let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
1173 let view = StridedView::<f64>::new(&data, &[2, 3], &[3, 1], 0).unwrap();
1174 assert_eq!(view.get(&[0, 0]), 1.0);
1175 assert_eq!(view.get(&[0, 1]), 2.0);
1176 assert_eq!(view.get(&[0, 2]), 3.0);
1177 assert_eq!(view.get(&[1, 0]), 4.0);
1178 assert_eq!(view.get(&[1, 2]), 6.0);
1179 }
1180
1181 #[test]
1182 fn test_strided_view_col_major() {
1183 let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
1185 let view = StridedView::<f64>::new(&data, &[2, 3], &[1, 2], 0).unwrap();
1186 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); }
1193
1194 #[test]
1195 fn test_strided_view_permute() {
1196 let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
1197 let view = StridedView::<f64>::new(&data, &[2, 3], &[3, 1], 0).unwrap();
1198 let perm = view.permute(&[1, 0]).unwrap();
1199 assert_eq!(perm.dims(), &[3, 2]);
1200 assert_eq!(perm.strides(), &[1, 3]);
1201 assert_eq!(perm.get(&[0, 0]), 1.0);
1202 assert_eq!(perm.get(&[1, 0]), 2.0);
1203 assert_eq!(perm.get(&[0, 1]), 4.0);
1204 }
1205
1206 #[test]
1207 fn test_strided_view_transpose_2d() {
1208 let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
1209 let view = StridedView::<f64>::new(&data, &[2, 3], &[3, 1], 0).unwrap();
1210 let t = view.transpose_2d().unwrap();
1211 assert_eq!(t.dims(), &[3, 2]);
1212 assert_eq!(t.get(&[0, 0]), 1.0);
1213 assert_eq!(t.get(&[1, 0]), 2.0);
1214 assert_eq!(t.get(&[0, 1]), 4.0);
1215 }
1216
1217 #[test]
1218 fn test_strided_view_conj() {
1219 let data = vec![Complex64::new(1.0, 2.0), Complex64::new(3.0, 4.0)];
1220 let view = StridedView::<Complex64>::new(&data, &[2], &[1], 0).unwrap();
1221 let c = view.conj();
1222 assert_eq!(c.get(&[0]), Complex64::new(1.0, -2.0));
1223 assert_eq!(c.get(&[1]), Complex64::new(3.0, -4.0));
1224 }
1225
1226 #[test]
1227 fn test_strided_view_adjoint_2d() {
1228 let data = vec![
1229 Complex64::new(1.0, 2.0),
1230 Complex64::new(3.0, 4.0),
1231 Complex64::new(5.0, 6.0),
1232 Complex64::new(7.0, 8.0),
1233 ];
1234 let view = StridedView::<Complex64>::new(&data, &[2, 2], &[2, 1], 0).unwrap();
1236 let adj = view.adjoint_2d().unwrap();
1237 assert_eq!(adj.dims(), &[2, 2]);
1238 assert_eq!(adj.get(&[0, 0]), Complex64::new(1.0, -2.0));
1240 assert_eq!(adj.get(&[1, 0]), Complex64::new(3.0, -4.0));
1241 assert_eq!(adj.get(&[0, 1]), Complex64::new(5.0, -6.0));
1242 }
1243
1244 #[test]
1245 fn test_strided_view_broadcast() {
1246 let data = vec![1.0, 2.0, 3.0];
1247 let view = StridedView::<f64>::new(&data, &[1, 3], &[3, 1], 0).unwrap();
1248 let broad = view.broadcast(&[4, 3]).unwrap();
1249 assert_eq!(broad.dims(), &[4, 3]);
1250 for i in 0..4 {
1251 assert_eq!(broad.get(&[i, 0]), 1.0);
1252 assert_eq!(broad.get(&[i, 1]), 2.0);
1253 assert_eq!(broad.get(&[i, 2]), 3.0);
1254 }
1255 }
1256
1257 #[test]
1258 fn test_strided_view_mut() {
1259 let mut data = vec![0.0; 6];
1260 {
1261 let mut view = StridedViewMut::<f64>::new(&mut data, &[2, 3], &[3, 1], 0).unwrap();
1262 view.set(&[0, 0], 1.0);
1263 view.set(&[1, 2], 6.0);
1264 }
1265 assert_eq!(data[0], 1.0);
1266 assert_eq!(data[5], 6.0);
1267 }
1268
1269 #[test]
1270 fn test_strided_view_mut_as_view() {
1271 let mut data = vec![1.0, 2.0, 3.0];
1272 let vm = StridedViewMut::<f64>::new(&mut data, &[3], &[1], 0).unwrap();
1273 let v = vm.as_view();
1274 assert_eq!(v.get(&[0]), 1.0);
1275 assert_eq!(v.get(&[2]), 3.0);
1276 }
1277
1278 #[test]
1279 fn test_strided_tensor_col_major() {
1280 let t = StridedArray::<f64>::from_fn_col_major(&[2, 3], |idx| (idx[0] * 3 + idx[1]) as f64);
1281 assert_eq!(t.dims(), &[2, 3]);
1282 assert_eq!(t.strides(), &[1, 2]); assert_eq!(t.get(&[0, 0]), 0.0);
1284 assert_eq!(t.get(&[1, 0]), 3.0);
1285 assert_eq!(t.get(&[0, 1]), 1.0);
1286 assert_eq!(t.get(&[1, 2]), 5.0);
1287 }
1288
1289 #[test]
1290 fn test_strided_tensor_row_major() {
1291 let t = StridedArray::<f64>::from_fn_row_major(&[2, 3], |idx| (idx[0] * 3 + idx[1]) as f64);
1292 assert_eq!(t.dims(), &[2, 3]);
1293 assert_eq!(t.strides(), &[3, 1]); assert_eq!(t.get(&[0, 0]), 0.0);
1295 assert_eq!(t.get(&[0, 1]), 1.0);
1296 assert_eq!(t.get(&[1, 0]), 3.0);
1297 assert_eq!(t.get(&[1, 2]), 5.0);
1298 }
1299
1300 #[test]
1301 fn test_strided_tensor_view() {
1302 let t =
1303 StridedArray::<f64>::from_fn_col_major(&[2, 3], |idx| (idx[0] * 10 + idx[1]) as f64);
1304 let v = t.view();
1305 assert_eq!(v.get(&[0, 0]), 0.0);
1306 assert_eq!(v.get(&[1, 0]), 10.0);
1307 assert_eq!(v.get(&[0, 2]), 2.0);
1308 }
1309
1310 #[test]
1311 fn test_strided_tensor_view_mut() {
1312 let mut t = StridedArray::<f64>::col_major(&[2, 3]);
1313 {
1314 let mut vm = t.view_mut();
1315 vm.set(&[1, 2], 42.0);
1316 }
1317 assert_eq!(t.get(&[1, 2]), 42.0);
1318 }
1319
1320 #[test]
1321 fn test_strided_tensor_index() {
1322 let t =
1323 StridedArray::<f64>::from_fn_row_major(&[3, 4], |idx| (idx[0] * 10 + idx[1]) as f64);
1324 assert_eq!(t[&[0usize, 0] as &[usize]], 0.0);
1325 assert_eq!(t[&[2usize, 3] as &[usize]], 23.0);
1326 }
1327
1328 #[test]
1329 fn test_strided_tensor_index_mut() {
1330 let mut t = StridedArray::<f64>::row_major(&[2, 3]);
1331 t[&[1usize, 2] as &[usize]] = 99.0;
1332 assert_eq!(t.get(&[1, 2]), 99.0);
1333 }
1334
1335 #[test]
1336 fn test_validate_bounds_ok() {
1337 assert!(validate_bounds(6, &[2, 3], &[3, 1], 0).is_ok());
1338 assert!(validate_bounds(6, &[2, 3], &[1, 2], 0).is_ok());
1339 }
1340
1341 #[test]
1342 fn test_validate_bounds_out_of_range() {
1343 assert!(validate_bounds(5, &[2, 3], &[3, 1], 0).is_err());
1344 }
1345
1346 #[test]
1347 fn test_validate_bounds_empty() {
1348 assert!(validate_bounds(0, &[0, 3], &[3, 1], 0).is_ok());
1349 }
1350
1351 #[test]
1352 fn test_validate_bounds_with_offset() {
1353 assert!(validate_bounds(7, &[2, 3], &[3, 1], 1).is_ok());
1354 assert!(validate_bounds(6, &[2, 3], &[3, 1], 1).is_err());
1355 }
1356
1357 #[test]
1358 fn test_strided_tensor_3d() {
1359 let t = StridedArray::<f64>::from_fn_col_major(&[2, 3, 4], |idx| {
1360 (idx[0] * 100 + idx[1] * 10 + idx[2]) as f64
1361 });
1362 assert_eq!(t.ndim(), 3);
1363 assert_eq!(t.strides(), &[1, 2, 6]); assert_eq!(t.get(&[0, 0, 0]), 0.0);
1365 assert_eq!(t.get(&[1, 0, 0]), 100.0);
1366 assert_eq!(t.get(&[0, 1, 0]), 10.0);
1367 assert_eq!(t.get(&[0, 0, 1]), 1.0);
1368 assert_eq!(t.get(&[1, 2, 3]), 123.0);
1369 }
1370
1371 #[test]
1372 fn test_diagonal_view_2d() {
1373 let data: Vec<f64> = (0..9).map(|x| x as f64).collect();
1376 let view = StridedView::<f64>::new(&data, &[3, 3], &[3, 1], 0).unwrap();
1377 let diag = view.diagonal_view(&[(0, 1)]).unwrap();
1378 assert_eq!(diag.dims(), &[3]);
1379 assert_eq!(diag.strides(), &[4]);
1380 assert_eq!(diag.get(&[0]), 0.0); assert_eq!(diag.get(&[1]), 4.0); assert_eq!(diag.get(&[2]), 8.0); }
1384
1385 #[test]
1386 fn test_diagonal_view_3d_adjacent() {
1387 let data: Vec<f64> = (0..12).map(|x| x as f64).collect();
1390 let view = StridedView::<f64>::new(&data, &[2, 2, 3], &[6, 3, 1], 0).unwrap();
1391 let diag = view.diagonal_view(&[(0, 1)]).unwrap();
1392 assert_eq!(diag.dims(), &[2, 3]);
1393 assert_eq!(diag.strides(), &[9, 1]);
1394 assert_eq!(diag.get(&[0, 0]), 0.0);
1395 assert_eq!(diag.get(&[0, 2]), 2.0);
1396 assert_eq!(diag.get(&[1, 0]), 9.0);
1397 }
1398
1399 #[test]
1400 fn test_diagonal_view_3d_non_adjacent() {
1401 let data: Vec<f64> = (0..12).map(|x| x as f64).collect();
1404 let view = StridedView::<f64>::new(&data, &[2, 3, 2], &[6, 2, 1], 0).unwrap();
1405 let diag = view.diagonal_view(&[(0, 2)]).unwrap();
1406 assert_eq!(diag.dims(), &[2, 3]);
1407 assert_eq!(diag.strides(), &[7, 2]);
1408 assert_eq!(diag.get(&[0, 0]), 0.0);
1409 assert_eq!(diag.get(&[0, 1]), 2.0);
1410 assert_eq!(diag.get(&[1, 0]), 7.0);
1411 assert_eq!(diag.get(&[1, 2]), 11.0);
1412 }
1413
1414 #[test]
1415 fn test_diagonal_view_two_pairs() {
1416 let data: Vec<f64> = (0..36).map(|x| x as f64).collect();
1418 let view = StridedView::<f64>::new(&data, &[2, 3, 2, 3], &[18, 6, 3, 1], 0).unwrap();
1419 let diag = view.diagonal_view(&[(0, 2), (1, 3)]).unwrap();
1420 assert_eq!(diag.dims(), &[2, 3]);
1421 assert_eq!(diag.strides(), &[21, 7]);
1422 assert_eq!(diag.get(&[0, 0]), 0.0);
1423 assert_eq!(diag.get(&[1, 1]), 28.0);
1424 }
1425
1426 #[test]
1427 fn test_diagonal_view_rejects_pairs_that_leave_the_backing() {
1428 let data: Vec<f64> = (0..8).map(|x| x as f64).collect();
1429 let view = StridedView::<f64>::new(&data, &[2, 2, 2], &[4, 2, 1], 0).unwrap();
1430 assert!(view.diagonal_view(&[(0, 2), (1, 2)]).is_err());
1433 assert!(view.diagonal_view(&[(0, 1), (0, 1)]).is_err());
1435 }
1436
1437 #[test]
1438 fn test_diagonal_view_fuses_a_shared_low_axis_in_bounds() {
1439 let data: Vec<f64> = (0..8).map(|x| x as f64).collect();
1442 let view = StridedView::<f64>::new(&data, &[2, 2, 2], &[4, 2, 1], 0).unwrap();
1443 let diag = view.diagonal_view(&[(0, 1), (0, 2)]).unwrap();
1444 assert_eq!(diag.dims(), &[2]);
1445 assert_eq!(diag.strides(), &[7]);
1446 assert_eq!(diag.get(&[0]), 0.0);
1447 assert_eq!(diag.get(&[1]), 7.0);
1448 }
1449
1450 #[test]
1451 fn test_custom_type_without_element_op_apply() {
1452 #[derive(Debug, Clone, Copy, PartialEq)]
1455 struct MyScalar(f64);
1456
1457 impl Default for MyScalar {
1458 fn default() -> Self {
1459 MyScalar(0.0)
1460 }
1461 }
1462
1463 let data = vec![MyScalar(1.0), MyScalar(2.0), MyScalar(3.0), MyScalar(4.0)];
1465 let arr = StridedArray::from_parts(data, &[2, 2], &[2, 1], 0).unwrap();
1466
1467 assert_eq!(arr.get(&[0, 0]), MyScalar(1.0));
1469 assert_eq!(arr.get(&[1, 0]), MyScalar(3.0));
1470
1471 let view: StridedView<MyScalar> = arr.view();
1473 assert_eq!(view.get(&[0, 1]), MyScalar(2.0));
1474
1475 let perm = view.permute(&[1, 0]).unwrap();
1477 assert_eq!(perm.get(&[0, 0]), MyScalar(1.0));
1478 assert_eq!(perm.get(&[0, 1]), MyScalar(3.0)); }
1480
1481 #[test]
1482 #[should_panic(expected = "wrong number of indices")]
1483 fn test_index_rejects_a_short_index() {
1484 let arr = StridedArray::from_parts(vec![0.0f64; 4], &[2, 2], &[2, 1], 0).unwrap();
1485 let _ = arr[&[1][..]];
1487 }
1488
1489 #[test]
1490 #[should_panic(expected = "wrong number of indices")]
1491 fn test_view_get_rejects_a_short_index() {
1492 let data = vec![0.0f64; 4];
1493 let view = StridedView::<f64>::new(&data, &[2, 2], &[2, 1], 0).unwrap();
1494 let _ = view.get(&[1]);
1495 }
1496
1497 #[test]
1498 #[should_panic(expected = "out of bounds")]
1499 fn test_view_set_rejects_an_out_of_bounds_index() {
1500 let mut data = vec![0.0f64; 4];
1501 let mut view = StridedViewMut::<f64>::new(&mut data, &[2, 2], &[2, 1], 0).unwrap();
1502 view.set(&[2, 0], 1.0);
1503 }
1504}