strided_view/view.rs
1//! Julia-like dynamic-rank strided view types.
2//!
3//! This module provides the canonical view types for strided operations,
4//! matching Julia's StridedViews.jl data model:
5//!
6//! - [`StridedView`]: Immutable dynamic-rank strided view with lazy element operations
7//! - [`StridedViewMut`]: Mutable dynamic-rank strided view (Identity op only)
8//! - [`StridedArray`]: Owned strided multidimensional array
9//!
10//! A view's shape and strides are a small inline-or-heap container with
11//! capacity [`crate::INLINE_RANK`]: every constructor, clone, and metadata
12//! transform allocates nothing up to that rank and spills gracefully above
13//! it. The view itself only borrows the element data.
14//!
15//! The metadata container is `tinyvec::ArrayVec`, unified across tensor4all
16//! repositories (tensor4all-agent-rules#22, strided-rs#312). `ArrayVec` is
17//! `Copy` for `Copy` items, so a clone is one fixed-size bitwise copy, and the
18//! heap fallback for rank above [`INLINE_RANK`] lives in a separate
19//! `Option<Box<SpilledMeta>>` so the inline value has no drop glue.
20
21use 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/// Build an `ArrayVec` from at most [`INLINE_RANK`] elements.
30///
31/// The value is written with literal-count arms: a runtime-length
32/// `copy_from_slice` or loop is recognized by LLVM and lowered to a `memcpy`
33/// PLT call, whose overhead dominates a handful of words.
34#[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/// Heap-owned metadata used only above [`INLINE_RANK`].
96#[derive(Clone, Debug)]
97struct SpilledMeta {
98 dims: Box<[usize]>,
99 strides: Box<[isize]>,
100}
101
102/// Shape and strides of a view or array.
103///
104/// Rank `<= INLINE_RANK` is stored inline in two `ArrayVec`s and the
105/// `spilled` slot is `None`, so cloning the whole metadata is a fixed-size
106/// bitwise copy with one predictable `None` check. A higher rank keeps the
107/// inline buffers empty and boxes both vectors.
108#[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// ============================================================================
181// Validation helpers
182// ============================================================================
183
184/// Validate that all accessed offsets stay within `[0, len)`.
185#[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 // Empty array - no access needed
196 if dims.iter().any(|&d| d == 0) {
197 return Ok(());
198 }
199 // Compute min and max offsets
200 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/// Offset of `indices` inside `dims`/`strides`, relative to `base`.
232///
233/// # Panics
234///
235/// Panics when `indices` is not the full index for `dims`, an index is out of
236/// bounds, or an indexed offset does not fit `isize`. Every panicking element
237/// accessor below reports those conditions through this helper.
238#[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/// Number of elements of `dims`: the checked product of the extents.
259///
260/// A zero extent wins before the product, so a zero-sized shape is empty even
261/// when its other extents would overflow. Returns
262/// [`StridedError::OffsetOverflow`] when the product does not fit `usize`
263/// (tensor4all/strided-rs#205).
264#[inline]
265fn checked_element_count(dims: &[usize]) -> Result<usize> {
266 if dims.contains(&0) {
267 return Ok(0);
268 }
269 dims.iter()
270 .try_fold(1usize, |acc, &dim| acc.checked_mul(dim))
271 .ok_or(StridedError::OffsetOverflow)
272}
273
274/// [`checked_element_count`] for the infallible `len` accessors.
275///
276/// # Panics
277///
278/// Panics when the element count does not fit `usize`, which only a
279/// broadcast (stride-0) layout can reach.
280#[inline]
281fn element_count(dims: &[usize]) -> usize {
282 checked_element_count(dims).expect("element count overflows usize")
283}
284
285/// Dense strides of `dims` with the axes visited fastest first in `order`.
286///
287/// Each stride is the product of the extents of the faster axes and must fit
288/// `isize`; otherwise [`StridedError::OffsetOverflow`]. A zero-sized shape is
289/// never dereferenced, so its strides saturate instead of failing: an empty
290/// array with huge extents stays constructible.
291fn checked_dense_strides(dims: &[usize], order: impl Iterator<Item = usize>) -> Result<StrideVec> {
292 let empty = dims.contains(&0);
293 let mut strides = StrideVec::new();
294 strides.resize(dims.len(), 0isize);
295 let mut step = 1isize;
296 for axis in order {
297 strides[axis] = step;
298 let dim = isize::try_from(dims[axis]).ok();
299 step = match dim.and_then(|dim| step.checked_mul(dim)) {
300 Some(next) => next,
301 None if empty => isize::MAX,
302 // The last axis's step is never stored, but it is the element
303 // count, which must also fit for the allocation and the offsets.
304 None => return Err(StridedError::OffsetOverflow),
305 };
306 }
307 Ok(strides)
308}
309
310fn checked_col_major_strides(dims: &[usize]) -> Result<StrideVec> {
311 checked_dense_strides(dims, 0..dims.len())
312}
313
314fn checked_row_major_strides(dims: &[usize]) -> Result<StrideVec> {
315 checked_dense_strides(dims, (0..dims.len()).rev())
316}
317
318/// Compute column-major strides (Julia default: first index varies fastest).
319///
320/// Returns an inline [`StrideVec`]: the result does not allocate at rank
321/// `<= INLINE_RANK` and spills above it.
322///
323/// # Panics
324///
325/// Panics when a stride or the element count of a non-empty shape does not
326/// fit `isize`. A shape with a zero extent never panics.
327///
328/// # Examples
329///
330/// ```
331/// use strided_view::col_major_strides;
332///
333/// // Julia default: first index varies fastest.
334/// assert_eq!(col_major_strides(&[2, 3]).as_slice(), &[1, 2]);
335/// ```
336#[inline]
337pub fn col_major_strides(dims: &[usize]) -> StrideVec {
338 checked_col_major_strides(dims).expect("column-major strides overflow isize")
339}
340
341/// Compute row-major strides (C default: last index varies fastest).
342///
343/// Returns an inline [`StrideVec`]: the result does not allocate at rank
344/// `<= INLINE_RANK` and spills above it.
345///
346/// # Panics
347///
348/// Panics when a stride or the element count of a non-empty shape does not
349/// fit `isize`. A shape with a zero extent never panics.
350///
351/// # Examples
352///
353/// ```
354/// use strided_view::row_major_strides;
355///
356/// // C default: last index varies fastest.
357/// assert_eq!(row_major_strides(&[2, 3]).as_slice(), &[3, 1]);
358/// ```
359#[inline]
360pub fn row_major_strides(dims: &[usize]) -> StrideVec {
361 checked_row_major_strides(dims).expect("row-major strides overflow isize")
362}
363
364// ============================================================================
365// StridedView
366// ============================================================================
367
368/// Dynamic-rank immutable strided view with lazy element operations.
369///
370/// This is the Julia-equivalent `StridedView` type with:
371/// - Dynamic rank, with metadata stored inline for up to [`INLINE_RANK`]
372/// dimensions ([`StridedView::new`] copies the caller's layout, so the view
373/// borrows only the element data)
374/// - Lazy element operations via the `Op` type parameter
375/// - Zero-copy transformations (permute, transpose, adjoint, conj)
376///
377/// # Type Parameters
378/// - `'a`: Lifetime of the underlying data
379/// - `T`: Element type
380/// - `Op`: Element operation applied lazily on access (default: `Identity`)
381///
382/// # Size
383///
384/// The inline metadata makes the view 184 bytes on 64-bit targets for
385/// `usize`/`isize` strides.
386///
387/// # Auto traits
388///
389/// The view borrows its elements, so it is `Send` only when `T: Sync` (a
390/// shared slice is). A view of non-`Sync` data is not `Send`:
391///
392/// ```compile_fail
393/// use std::cell::Cell;
394/// use strided_view::StridedView;
395///
396/// fn assert_send<T: Send>() {}
397/// assert_send::<StridedView<'static, Cell<i32>>>();
398/// ```
399pub struct StridedView<'a, T, Op = Identity> {
400 ptr: *const T,
401 data: &'a [T],
402 meta: ViewMeta,
403 offset: isize,
404 _op: PhantomData<Op>,
405}
406
407// SAFETY: `StridedView` owns no `T`. It holds `&'a [T]` and a `*const T`
408// derived from it, so moving the view moves a shared borrow, which is `Send`
409// exactly when `T: Sync`. `Op` appears only in `PhantomData`, and `ViewMeta` is
410// owned metadata (an inline `tinyvec::ArrayVec` plus an optional box).
411unsafe impl<T: Sync, Op: Send> Send for StridedView<'_, T, Op> {}
412// SAFETY: a shared `&StridedView` hands out only `&T` (`get`, `ptr`, the
413// `ElementOp` read path) and never `&mut T`, so sharing is sound when `T: Sync`.
414// `Op` is `Sync` and appears only in `PhantomData`.
415unsafe impl<T: Sync, Op: Sync> Sync for StridedView<'_, T, Op> {}
416
417impl<T, Op> Clone for StridedView<'_, T, Op> {
418 #[inline]
419 fn clone(&self) -> Self {
420 Self {
421 ptr: self.ptr,
422 data: self.data,
423 meta: self.meta.clone(),
424 offset: self.offset,
425 _op: PhantomData,
426 }
427 }
428}
429
430impl<T: std::fmt::Debug, Op> std::fmt::Debug for StridedView<'_, T, Op> {
431 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
432 f.debug_struct("StridedView")
433 .field("dims", &self.meta.dims())
434 .field("strides", &self.meta.strides())
435 .field("offset", &self.offset)
436 .finish()
437 }
438}
439
440/// The cached base pointer of a view over `data` at `offset`.
441///
442/// An empty layout gets the dangling pointer rather than `data + offset`: a
443/// zero-sized extent may sit at an offset no element of `data` reaches.
444#[inline]
445fn base_ptr<T>(data: &[T], dims: &[usize], offset: isize) -> *const T {
446 if empty_layout(dims) {
447 empty_const_ptr()
448 } else {
449 // SAFETY: every caller reaches `base_ptr` from `from_meta`, whose
450 // `dims`/`strides`/`offset` were either accepted by `validate_bounds`
451 // (view.rs:186) — `StridedView::new`, `StridedViewMut::new`, the
452 // re-validated diagonal path, or a metadata transform that only reorders
453 // axes or zeroes strides and therefore cannot enlarge the reachable
454 // offset set — or were supplied under the `new_unchecked` contract. The
455 // empty-layout branch above excludes the case where `offset` may point
456 // past `data`.
457 unsafe { data.as_ptr().offset(offset) }
458 }
459}
460
461/// A view from already-validated metadata, with the element op chosen here.
462#[inline]
463fn from_meta<'a, T, Op>(
464 data: &'a [T],
465 dims: &[usize],
466 strides: &[isize],
467 offset: isize,
468) -> StridedView<'a, T, Op> {
469 let meta = ViewMeta::from_slices(dims, strides);
470 let ptr = base_ptr(data, meta.dims(), offset);
471 StridedView {
472 ptr,
473 data,
474 meta,
475 offset,
476 _op: PhantomData,
477 }
478}
479
480impl<'a, T, Op> StridedView<'a, T, Op> {
481 /// Create an immutable strided view, copying the caller's metadata inline.
482 ///
483 /// The layout is copied into the view's inline storage, so it does not
484 /// need to outlive the view and no heap allocation happens at rank
485 /// `<= INLINE_RANK`. The view's lifetime is tied to `data` only.
486 ///
487 /// # Errors
488 ///
489 /// Returns [`StridedError::StrideLengthMismatch`] when `dims` and
490 /// `strides` have different lengths, and
491 /// [`StridedError::OffsetOverflow`] when a reachable offset is not
492 /// representable: `dim - 1` of a non-unit extent does not fit `isize`, a
493 /// stride times that last index or the accumulated offsets overflow
494 /// `isize` or go negative, or the highest reachable offset is not below
495 /// the buffer length. A layout with a zero extent skips this check.
496 ///
497 /// # Examples
498 ///
499 /// ```
500 /// use strided_view::StridedView;
501 ///
502 /// let data = [1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0];
503 /// let view = StridedView::<f64>::new(&data, &[2, 3], &[1, 2], 0).unwrap();
504 /// assert_eq!(view.dims(), &[2, 3]);
505 /// assert_eq!(view.get(&[1, 2]), 6.0);
506 /// ```
507 #[inline]
508 pub fn new(data: &'a [T], dims: &[usize], strides: &[isize], offset: isize) -> Result<Self> {
509 validate_bounds(data.len(), dims, strides, offset)?;
510 Ok(from_meta(data, dims, strides, offset))
511 }
512
513 /// Create a view without bounds checking, copying the caller's metadata.
514 ///
515 /// # Safety
516 /// The caller must ensure all index combinations stay within bounds.
517 ///
518 /// # Examples
519 ///
520 /// ```
521 /// use strided_view::StridedView;
522 ///
523 /// let data = [1.0f64, 2.0, 3.0, 4.0];
524 /// let view = unsafe { StridedView::<f64>::new_unchecked(&data, &[2, 2], &[1, 2], 0) };
525 /// assert_eq!(view.ndim(), 2);
526 /// ```
527 #[inline]
528 pub unsafe fn new_unchecked(
529 data: &'a [T],
530 dims: &[usize],
531 strides: &[isize],
532 offset: isize,
533 ) -> Self {
534 from_meta(data, dims, strides, offset)
535 }
536
537 /// Returns the shape (dimension sizes) of this view.
538 #[inline]
539 pub fn dims(&self) -> &[usize] {
540 self.meta.dims()
541 }
542
543 /// Returns the strides (in units of `T`) for each dimension.
544 #[inline]
545 pub fn strides(&self) -> &[isize] {
546 self.meta.strides()
547 }
548
549 /// Returns the byte offset into the backing data.
550 #[inline]
551 pub fn offset(&self) -> isize {
552 self.offset
553 }
554
555 /// Returns the number of dimensions (rank).
556 #[inline]
557 pub fn ndim(&self) -> usize {
558 self.dims().len()
559 }
560
561 /// Returns the total number of elements.
562 ///
563 /// # Panics
564 ///
565 /// Panics when the element count does not fit `usize`, which only a
566 /// broadcast (stride-0) layout can reach.
567 #[inline]
568 pub fn len(&self) -> usize {
569 element_count(self.dims())
570 }
571
572 /// Returns `true` if any dimension is zero.
573 #[inline]
574 pub fn is_empty(&self) -> bool {
575 self.dims().iter().any(|&d| d == 0)
576 }
577
578 /// Returns a reference to the backing data slice.
579 #[inline]
580 pub fn data(&self) -> &'a [T] {
581 self.data
582 }
583
584 /// Raw const pointer to element at the view's base offset.
585 #[inline]
586 pub fn ptr(&self) -> *const T {
587 self.ptr
588 }
589
590 /// Permute dimensions.
591 ///
592 /// # Errors
593 ///
594 /// Returns [`StridedError::RankMismatch`] when `perm` does not have one
595 /// entry per axis, and [`StridedError::InvalidAxis`] when an entry
596 /// names an axis outside the rank or repeats an axis.
597 pub fn permute(&self, perm: &[usize]) -> Result<StridedView<'a, T, Op>> {
598 let rank = self.dims().len();
599 if perm.len() != rank {
600 return Err(StridedError::RankMismatch(perm.len(), rank));
601 }
602 let mut seen: tinyvec::TinyVec<[bool; INLINE_RANK]> = tinyvec::TinyVec::new();
603 seen.resize(rank, false);
604 for &p in perm {
605 if p >= rank {
606 return Err(StridedError::InvalidAxis { axis: p, rank });
607 }
608 if seen[p] {
609 return Err(StridedError::InvalidAxis { axis: p, rank });
610 }
611 seen[p] = true;
612 }
613 let new_dims: AxisVec = perm.iter().map(|&p| self.dims()[p]).collect();
614 let new_strides: StrideVec = perm.iter().map(|&p| self.strides()[p]).collect();
615 Ok(from_meta(self.data, &new_dims, &new_strides, self.offset))
616 }
617
618 /// Create a diagonal view by fusing repeated axis pairs via stride trick (zero-copy).
619 ///
620 /// Each pair `(a, b)` ties axes `a` and `b` to one index. Pairs that share
621 /// an axis chain into one class (`[(0, 1), (1, 2)]` and `[(0, 1), (0, 2)]`
622 /// both tie axes 0, 1 and 2), and each class becomes one axis:
623 /// - New stride = the sum of the strides of the axes in the class
624 /// - New dim = the minimum of their dims
625 /// - The class sits at its lowest-numbered axis; the others are removed
626 /// - Pairs use **original** axis numbering; their order does not matter,
627 /// and repeating a pair has no further effect
628 ///
629 /// # Errors
630 ///
631 /// Returns [`StridedError::InvalidAxis`] when a pair names an axis outside
632 /// the rank or names the same axis twice, and
633 /// [`StridedError::OffsetOverflow`] when a fused stride does not fit
634 /// `isize` or the fused layout reaches past the backing allocation.
635 ///
636 /// # Example
637 /// `A[i,i,j]` shape=`[n,n,m]` strides=`[s0,s1,s2]` -> shape=`[n,m]` strides=`[s0+s1, s2]`
638 pub fn diagonal_view(&self, axis_pairs: &[(usize, usize)]) -> Result<StridedView<'a, T, Op>> {
639 let ndim = self.ndim();
640 // Each pair ties two axes to one index, so the pairs partition the
641 // axes into classes. Track each class's root as its lowest axis; the
642 // result is then independent of pair order and of how pairs chain
643 // (tensor4all/strided-rs#329).
644 let mut root: tinyvec::TinyVec<[usize; INLINE_RANK]> = (0..ndim).collect();
645 fn find(root: &mut [usize], mut ax: usize) -> usize {
646 while root[ax] != ax {
647 root[ax] = root[root[ax]];
648 ax = root[ax];
649 }
650 ax
651 }
652 for &(a, b) in axis_pairs {
653 let hi = a.max(b);
654 if hi >= ndim {
655 return Err(StridedError::InvalidAxis {
656 axis: hi,
657 rank: ndim,
658 });
659 }
660 if a == b {
661 return Err(StridedError::InvalidAxis {
662 axis: a,
663 rank: ndim,
664 });
665 }
666 let (ra, rb) = (find(&mut root, a), find(&mut root, b));
667 root[ra.max(rb)] = ra.min(rb);
668 }
669
670 // Every class collapses onto its root: the root's stride sums the
671 // strides of all axes in the class and its dim is their minimum. A
672 // root is lower than every other axis of its class, so it has already
673 // been kept when a later member folds into it.
674 let mut dims: AxisVec = AxisVec::new();
675 let mut strides: StrideVec = StrideVec::new();
676 let mut position: tinyvec::TinyVec<[usize; INLINE_RANK]> =
677 tinyvec::TinyVec::with_capacity(ndim);
678 for ax in 0..ndim {
679 let r = find(&mut root, ax);
680 if r == ax {
681 position.push(dims.len());
682 dims.push(self.dims()[ax]);
683 strides.push(self.strides()[ax]);
684 } else {
685 position.push(usize::MAX);
686 let p = position[r];
687 strides[p] = strides[p]
688 .checked_add(self.strides()[ax])
689 .ok_or(StridedError::OffsetOverflow)?;
690 dims[p] = dims[p].min(self.dims()[ax]);
691 }
692 }
693
694 // A fused stride sums several input strides, so the result carries no
695 // bound proof of its own and is validated here instead
696 // (tensor4all/strided-rs#204).
697 validate_bounds(self.data().len(), &dims, &strides, self.offset())?;
698
699 Ok(from_meta(self.data(), &dims, &strides, self.offset()))
700 }
701
702 /// Broadcast this view to a target shape.
703 ///
704 /// Size-1 dimensions are expanded (stride set to 0) to match target.
705 ///
706 /// # Errors
707 ///
708 /// Returns [`StridedError::RankMismatch`] when `target_dims` does not
709 /// have one entry per axis, and [`StridedError::ShapeMismatch`] when an
710 /// axis is neither already the target extent nor of size one.
711 pub fn broadcast(&self, target_dims: &[usize]) -> Result<StridedView<'a, T, Op>> {
712 if self.dims().len() != target_dims.len() {
713 return Err(StridedError::RankMismatch(
714 self.dims().len(),
715 target_dims.len(),
716 ));
717 }
718 let mut new_strides = StrideVec::with_capacity(self.dims().len());
719 for i in 0..self.dims().len() {
720 if self.dims()[i] == target_dims[i] {
721 new_strides.push(self.strides()[i]);
722 } else if self.dims()[i] == 1 {
723 new_strides.push(0);
724 } else {
725 return Err(StridedError::ShapeMismatch(
726 self.dims().to_vec(),
727 target_dims.to_vec(),
728 ));
729 }
730 }
731 Ok(from_meta(self.data, target_dims, &new_strides, self.offset))
732 }
733}
734
735/// Composition methods: require `T: ElementOpApply` and `Op: ComposableElementOp<T>`.
736impl<'a, T: Copy + ElementOpApply, Op: ComposableElementOp<T>> StridedView<'a, T, Op> {
737 /// Transpose a 2D view: reverses dimensions and composes the Transpose element op.
738 ///
739 /// Julia equivalent: `Base.transpose(a::AbstractStridedView{<:Any, 2})`
740 ///
741 /// # Errors
742 ///
743 /// Returns [`StridedError::RankMismatch`] when the view is not 2-D.
744 pub fn transpose_2d(&self) -> Result<StridedView<'a, T, Op::ComposeTranspose>> {
745 if self.dims().len() != 2 {
746 return Err(StridedError::RankMismatch(self.dims().len(), 2));
747 }
748 Ok(from_meta(
749 self.data,
750 &[self.dims()[1], self.dims()[0]],
751 &[self.strides()[1], self.strides()[0]],
752 self.offset,
753 ))
754 }
755
756 /// Adjoint (conjugate transpose) of a 2D view.
757 ///
758 /// Julia equivalent: `Base.adjoint(a::AbstractStridedView{<:Any, 2})`
759 ///
760 /// # Errors
761 ///
762 /// Returns [`StridedError::RankMismatch`] when the view is not 2-D.
763 pub fn adjoint_2d(&self) -> Result<StridedView<'a, T, Op::ComposeAdjoint>> {
764 if self.dims().len() != 2 {
765 return Err(StridedError::RankMismatch(self.dims().len(), 2));
766 }
767 Ok(from_meta(
768 self.data,
769 &[self.dims()[1], self.dims()[0]],
770 &[self.strides()[1], self.strides()[0]],
771 self.offset,
772 ))
773 }
774
775 /// Complex conjugate (compose Conj without changing dims/strides).
776 ///
777 /// Julia equivalent: `Base.conj(a::AbstractStridedView)`
778 #[inline]
779 pub fn conj(&self) -> StridedView<'a, T, Op::ComposeConj> {
780 StridedView {
781 ptr: self.ptr,
782 data: self.data,
783 meta: self.meta.clone(),
784 offset: self.offset,
785 _op: PhantomData,
786 }
787 }
788}
789
790/// Element access: requires `Op: ElementOp<T>`.
791impl<'a, T: Copy, Op: ElementOp<T>> StridedView<'a, T, Op> {
792 /// Get an element with the element operation applied.
793 ///
794 /// # Panics
795 ///
796 /// Panics if `indices.len()` is not the view's rank, an index is out of
797 /// bounds, or the indexed offset does not fit `isize`.
798 pub fn get(&self, indices: &[usize]) -> T {
799 let offset = indexed_offset(self.dims(), self.strides(), indices, 0);
800 // SAFETY: `indexed_offset` (view.rs:239) asserts the index count and
801 // `index < dims[i]`, so the offset it returns lies within the reachable
802 // extent that `validate_bounds` (view.rs:186) proved stays inside
803 // `data` for this view's construction.
804 Op::apply(unsafe { *self.ptr.offset(offset) })
805 }
806
807 /// Get an element without bounds checking.
808 ///
809 /// # Safety
810 /// Caller must ensure indices are within bounds.
811 #[inline]
812 pub unsafe fn get_unchecked(&self, indices: &[usize]) -> T {
813 let mut idx = 0isize;
814 for (i, &index) in indices.iter().enumerate() {
815 idx += index as isize * self.strides()[i];
816 }
817 Op::apply(*self.ptr.offset(idx))
818 }
819}
820
821// ============================================================================
822// StridedViewMut
823// ============================================================================
824
825/// Dynamic-rank mutable strided view.
826///
827/// Always uses `Identity` element operation for write simplicity.
828/// Julia typically applies ops on the read side.
829pub struct StridedViewMut<'a, T> {
830 /// Element base: `base` plus `offset`, or a dangling pointer for an empty
831 /// layout. `get` and `set` go through it, so it stays cached.
832 ///
833 /// Nothing here reborrows the caller's `&mut [T]` any more: a stored slice
834 /// field was retagged at construction and by `data_mut`, which invalidated
835 /// this pointer and made every later `get`, `set`, and `as_view` access
836 /// undefined behavior.
837 ptr: *mut T,
838 /// Base of the whole backing slice.
839 ///
840 /// `base` with `src_len` is the handle `data` and `data_mut` re-derive the
841 /// slice from, so the backing extent stays independent of the shape (a
842 /// `[0]` shape still has a backing slice).
843 base: *mut T,
844 /// Length of the whole backing slice.
845 src_len: usize,
846 meta: ViewMeta,
847 offset: isize,
848 /// The view borrows the caller's slice for `'a` without holding it.
849 _marker: PhantomData<&'a mut [T]>,
850}
851
852// SAFETY: `StridedViewMut` uniquely owns its `&'a mut [T]` (held as
853// `PhantomData<&'a mut [T]>` plus the two raw pointers `base` and `ptr` derived
854// from it), and there is no `Sync` impl for it, so sending it moves that
855// exclusive borrow; `T: Send` is the bound `&mut [T]: Send` requires.
856unsafe impl<T: Send> Send for StridedViewMut<'_, T> {}
857
858impl<T: std::fmt::Debug> std::fmt::Debug for StridedViewMut<'_, T> {
859 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
860 f.debug_struct("StridedViewMut")
861 .field("dims", &self.meta.dims())
862 .field("strides", &self.meta.strides())
863 .field("offset", &self.offset)
864 .finish()
865 }
866}
867
868impl<'a, T> StridedViewMut<'a, T> {
869 /// A view from already-validated metadata.
870 #[inline]
871 fn from_meta(data: &'a mut [T], dims: &[usize], strides: &[isize], offset: isize) -> Self {
872 let meta = ViewMeta::from_slices(dims, strides);
873 let src_len = data.len();
874 let base = data.as_mut_ptr();
875 let ptr = if empty_layout(meta.dims()) {
876 empty_mut_ptr()
877 } else {
878 // SAFETY: as `base_ptr` above — `offset` was accepted by
879 // `validate_bounds` (view.rs:186) in `StridedViewMut::new`, or by a
880 // metadata transform that only reorders axes or zeroes strides, or
881 // came from the `new_unchecked` contract. `base` is this view's own
882 // live handle on `data` and the empty-layout branch above excludes
883 // the case where `offset` may point past it.
884 unsafe { base.offset(offset) }
885 };
886 Self {
887 ptr,
888 base,
889 src_len,
890 meta,
891 offset,
892 _marker: PhantomData,
893 }
894 }
895
896 /// Create a mutable view, copying the caller's metadata inline.
897 ///
898 /// The layout is copied, so it does not need to outlive the view and no
899 /// heap allocation happens at rank `<= INLINE_RANK`.
900 ///
901 /// # Errors
902 ///
903 /// Returns [`StridedError::StrideLengthMismatch`] when `dims` and
904 /// `strides` have different lengths, and
905 /// [`StridedError::OffsetOverflow`] when a reachable offset is not
906 /// representable: `dim - 1` of a non-unit extent does not fit `isize`, a
907 /// stride times that last index or the accumulated offsets overflow
908 /// `isize` or go negative, or the highest reachable offset is not below
909 /// the buffer length. A layout with a zero extent skips this check.
910 ///
911 /// # Examples
912 ///
913 /// ```
914 /// use strided_view::StridedViewMut;
915 ///
916 /// let mut data = [0.0f64; 4];
917 /// let mut view = StridedViewMut::<f64>::new(&mut data, &[2, 2], &[1, 2], 0).unwrap();
918 /// view.set(&[1, 1], 5.0);
919 /// assert_eq!(data[3], 5.0);
920 /// ```
921 #[inline]
922 pub fn new(
923 data: &'a mut [T],
924 dims: &[usize],
925 strides: &[isize],
926 offset: isize,
927 ) -> Result<Self> {
928 validate_bounds(data.len(), dims, strides, offset)?;
929 Ok(Self::from_meta(data, dims, strides, offset))
930 }
931
932 /// Create without bounds checking, copying the caller's metadata.
933 ///
934 /// # Safety
935 /// Caller must ensure all index combinations stay within bounds.
936 ///
937 /// # Examples
938 ///
939 /// ```
940 /// use strided_view::StridedViewMut;
941 ///
942 /// let mut data = [0.0f64; 4];
943 /// let view = unsafe { StridedViewMut::<f64>::new_unchecked(&mut data, &[2, 2], &[1, 2], 0) };
944 /// assert_eq!(view.ndim(), 2);
945 /// ```
946 #[inline]
947 pub unsafe fn new_unchecked(
948 data: &'a mut [T],
949 dims: &[usize],
950 strides: &[isize],
951 offset: isize,
952 ) -> Self {
953 Self::from_meta(data, dims, strides, offset)
954 }
955
956 /// Returns the shape (dimension sizes) of this view.
957 #[inline]
958 pub fn dims(&self) -> &[usize] {
959 self.meta.dims()
960 }
961
962 /// Returns the strides (in units of `T`) for each dimension.
963 #[inline]
964 pub fn strides(&self) -> &[isize] {
965 self.meta.strides()
966 }
967
968 /// Returns the byte offset into the backing data.
969 #[inline]
970 pub fn offset(&self) -> isize {
971 self.offset
972 }
973
974 /// Returns the number of dimensions (rank).
975 #[inline]
976 pub fn ndim(&self) -> usize {
977 self.dims().len()
978 }
979
980 /// Returns the total number of elements.
981 ///
982 /// # Panics
983 ///
984 /// Panics when the element count does not fit `usize`, which only a
985 /// broadcast (stride-0) layout can reach.
986 #[inline]
987 pub fn len(&self) -> usize {
988 element_count(self.dims())
989 }
990
991 /// Returns `true` if any dimension is zero.
992 #[inline]
993 pub fn is_empty(&self) -> bool {
994 self.dims().iter().any(|&d| d == 0)
995 }
996
997 /// Raw const pointer to element at the view's base offset.
998 #[inline]
999 pub fn ptr(&self) -> *const T {
1000 self.ptr as *const T
1001 }
1002
1003 /// Raw mutable pointer to element at the view's base offset.
1004 #[inline]
1005 pub fn as_mut_ptr(&self) -> *mut T {
1006 self.ptr
1007 }
1008
1009 /// Returns a reference to the backing data slice.
1010 ///
1011 /// The length is the backing slice's, not the shape's: a view whose shape
1012 /// has a zero extent still reports the whole buffer.
1013 #[inline]
1014 pub fn data(&self) -> &[T] {
1015 // SAFETY: `base` is the caller's `data.as_mut_ptr()` and `src_len` is
1016 // that slice's length, both recorded at construction, so this rebuilds
1017 // exactly the borrowed slice. No length is derived from the shape, which
1018 // is why a zero-extent layout still reports the whole buffer. `T` is
1019 // initialized across it, and `&self` restricts this to a shared borrow.
1020 unsafe { std::slice::from_raw_parts(self.base, self.src_len) }
1021 }
1022
1023 /// Returns a mutable reference to the backing data slice.
1024 ///
1025 /// The slice is re-derived from the view's storage handle, so `ptr`, `get`,
1026 /// and `set` stay valid after the returned borrow ends.
1027 #[inline]
1028 pub fn data_mut(&mut self) -> &mut [T] {
1029 // SAFETY: as `data` above — the same `base`/`src_len` pair, so this
1030 // rebuilds the caller's slice rather than a shape-sized one. `&mut self`
1031 // keeps the returned slice the only reference into the view while it
1032 // lives. In a non-empty layout `ptr` is `base + offset`, derived without
1033 // a retag, so re-deriving the slice here does not invalidate it — the
1034 // property `get`/`set` rely on; an empty layout caches a dangling `ptr`
1035 // that no element access dereferences.
1036 unsafe { std::slice::from_raw_parts_mut(self.base, self.src_len) }
1037 }
1038
1039 /// Permute dimensions, consuming the mutable view.
1040 ///
1041 /// Returns a new mutable view with reordered dimensions and strides.
1042 /// Takes `self` by value to prevent aliasing of mutable views.
1043 ///
1044 /// # Errors
1045 ///
1046 /// Returns [`StridedError::RankMismatch`] when `perm` does not have one
1047 /// entry per axis, and [`StridedError::InvalidAxis`] when an entry
1048 /// names an axis outside the rank or repeats an axis.
1049 pub fn permute(self, perm: &[usize]) -> Result<StridedViewMut<'a, T>> {
1050 let rank = self.dims().len();
1051 if perm.len() != rank {
1052 return Err(StridedError::RankMismatch(perm.len(), rank));
1053 }
1054 let mut seen: tinyvec::TinyVec<[bool; INLINE_RANK]> = tinyvec::TinyVec::new();
1055 seen.resize(rank, false);
1056 for &p in perm {
1057 if p >= rank {
1058 return Err(StridedError::InvalidAxis { axis: p, rank });
1059 }
1060 if seen[p] {
1061 return Err(StridedError::InvalidAxis { axis: p, rank });
1062 }
1063 seen[p] = true;
1064 }
1065 let new_dims: AxisVec = perm.iter().map(|&p| self.dims()[p]).collect();
1066 let new_strides: StrideVec = perm.iter().map(|&p| self.strides()[p]).collect();
1067 Ok(Self {
1068 meta: ViewMeta::from_slices(&new_dims, &new_strides),
1069 ..self
1070 })
1071 }
1072
1073 /// Reborrow as an immutable view.
1074 ///
1075 /// The reborrow copies this view's metadata into the immutable view's
1076 /// inline storage, which allocates nothing up to [`INLINE_RANK`] and
1077 /// spills above it (the mutable owner always keeps its own copy).
1078 ///
1079 /// # Examples
1080 ///
1081 /// ```
1082 /// use strided_view::StridedViewMut;
1083 ///
1084 /// let mut data = [1.0f64, 2.0, 3.0];
1085 /// let view = StridedViewMut::<f64>::new(&mut data, &[3], &[1], 0).unwrap();
1086 /// let read = view.as_view();
1087 /// assert_eq!(read.get(&[2]), 3.0);
1088 /// ```
1089 #[inline]
1090 pub fn as_view(&self) -> StridedView<'_, T, Identity> {
1091 // SAFETY: `data()` re-derives the slice from this view's own handle
1092 // (`base` + `src_len`), and `ptr` is `base + offset` for the same
1093 // `offset`, so the immutable view ends up pointing at the same storage
1094 // the mutable view owns. A zero-extent layout yields a dangling pointer
1095 // on both sides (`empty_mut_ptr` here, the `empty_layout` branch of
1096 // `base_ptr` in the immutable view), so they agree there too. `&self`
1097 // keeps that storage exclusively borrowed for the returned view's
1098 // lifetime, so the shared view can only read it, and `T` is initialized
1099 // across the whole slice.
1100 unsafe { StridedView::new_unchecked(self.data(), self.dims(), self.strides(), self.offset) }
1101 }
1102}
1103
1104impl<'a, T: Copy> StridedViewMut<'a, T> {
1105 /// Get an element.
1106 ///
1107 /// # Panics
1108 ///
1109 /// Panics if `indices.len()` is not the view's rank, an index is out of
1110 /// bounds, or the indexed offset does not fit `isize`.
1111 pub fn get(&self, indices: &[usize]) -> T {
1112 let offset = indexed_offset(self.dims(), self.strides(), indices, 0);
1113 // SAFETY: in a non-empty layout `ptr` is `base + offset`, derived from
1114 // the caller's `&'a mut [T]` through raw pointers only, so a reborrow of
1115 // the backing slice does not invalidate it: `data_mut` re-derives that
1116 // slice from `base` and never rewrites `ptr`, and offsetting a raw
1117 // pointer keeps `base`'s provenance. (An empty layout caches a dangling
1118 // `ptr`, but no index satisfies `index < dims[i]` there, so this access
1119 // is unreachable.) `indexed_offset` (view.rs:239) asserts the index
1120 // count and `index < dims[i]`, so the relative offset it returns keeps
1121 // `ptr + offset` inside the extent `validate_bounds` (view.rs:186)
1122 // proved reachable, i.e. inside `base..base + src_len`, where every
1123 // element is initialized because the view borrows an initialized
1124 // `&'a mut [T]` and no method shortens that borrow.
1125 unsafe { *self.ptr.offset(offset) }
1126 }
1127
1128 /// Set an element.
1129 ///
1130 /// # Panics
1131 ///
1132 /// Panics if `indices.len()` is not the view's rank, an index is out of
1133 /// bounds, or the indexed offset does not fit `isize`.
1134 pub fn set(&mut self, indices: &[usize], value: T) {
1135 let offset = indexed_offset(self.dims(), self.strides(), indices, 0);
1136 // SAFETY: as `get` above — the same `ptr` survives an earlier `data_mut`
1137 // reborrow, and `indexed_offset` (view.rs:239) bounds-checks every index,
1138 // so the store lands inside `base..base + src_len`. `&mut self` makes
1139 // this the only access through the view for the call, and the view holds
1140 // the sole `&'a mut [T]` borrow of that storage, so nothing else can
1141 // alias the written place.
1142 unsafe {
1143 *self.ptr.offset(offset) = value;
1144 }
1145 }
1146}
1147
1148// ============================================================================
1149// StridedArray
1150// ============================================================================
1151
1152/// Owned strided multidimensional array.
1153///
1154/// Supports both column-major (Julia default) and row-major (C default) layouts.
1155pub struct StridedArray<T> {
1156 data: Vec<T>,
1157 meta: ViewMeta,
1158 offset: isize,
1159}
1160
1161impl<T: std::fmt::Debug> std::fmt::Debug for StridedArray<T> {
1162 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1163 f.debug_struct("StridedArray")
1164 .field("dims", &self.meta.dims())
1165 .field("strides", &self.meta.strides())
1166 .field("offset", &self.offset)
1167 .finish()
1168 }
1169}
1170
1171impl<T: Clone> Clone for StridedArray<T> {
1172 fn clone(&self) -> Self {
1173 Self {
1174 data: self.data.clone(),
1175 meta: self.meta.clone(),
1176 offset: self.offset,
1177 }
1178 }
1179}
1180
1181impl<T: Clone + Default> StridedArray<T> {
1182 /// Create a column-major (Julia default) tensor filled with Default values.
1183 ///
1184 /// The shape and strides are stored inline, so a rank `<= INLINE_RANK`
1185 /// tensor allocates only its element buffer.
1186 ///
1187 /// # Examples
1188 ///
1189 /// ```
1190 /// use strided_view::StridedArray;
1191 ///
1192 /// let a = StridedArray::<f64>::col_major(&[2, 3]);
1193 /// assert_eq!(a.dims(), &[2, 3]);
1194 /// assert_eq!(a.strides(), &[1, 2]);
1195 /// ```
1196 ///
1197 /// # Panics
1198 ///
1199 /// Panics when the element count or a stride does not fit `isize`; use
1200 /// [`Self::try_col_major`] to get [`StridedError::OffsetOverflow`] instead.
1201 pub fn col_major(dims: &[usize]) -> Self {
1202 Self::try_col_major(dims).expect("col_major: shape overflows isize")
1203 }
1204
1205 /// Fallible [`Self::col_major`].
1206 ///
1207 /// # Errors
1208 ///
1209 /// Returns [`StridedError::OffsetOverflow`] when the element count or a
1210 /// stride does not fit `isize`.
1211 ///
1212 /// # Examples
1213 ///
1214 /// ```
1215 /// use strided_view::{StridedArray, StridedError};
1216 ///
1217 /// assert_eq!(StridedArray::<f64>::try_col_major(&[2, 3]).unwrap().strides(), &[1, 2]);
1218 /// assert!(matches!(
1219 /// StridedArray::<f64>::try_col_major(&[1 << 62, 4]),
1220 /// Err(StridedError::OffsetOverflow)
1221 /// ));
1222 /// ```
1223 pub fn try_col_major(dims: &[usize]) -> Result<Self> {
1224 let strides = checked_col_major_strides(dims)?;
1225 let total = checked_element_count(dims)?;
1226 Ok(Self {
1227 data: vec![T::default(); total],
1228 meta: ViewMeta::from_slices(dims, &strides),
1229 offset: 0,
1230 })
1231 }
1232
1233 /// Create a row-major (C default) tensor filled with Default values.
1234 ///
1235 /// The shape and strides are stored inline, so a rank `<= INLINE_RANK`
1236 /// tensor allocates only its element buffer.
1237 ///
1238 /// # Examples
1239 ///
1240 /// ```
1241 /// use strided_view::StridedArray;
1242 ///
1243 /// let a = StridedArray::<f64>::row_major(&[2, 3]);
1244 /// assert_eq!(a.dims(), &[2, 3]);
1245 /// assert_eq!(a.strides(), &[3, 1]);
1246 /// ```
1247 ///
1248 /// # Panics
1249 ///
1250 /// Panics when the element count or a stride does not fit `isize`; use
1251 /// [`Self::try_row_major`] to get [`StridedError::OffsetOverflow`] instead.
1252 pub fn row_major(dims: &[usize]) -> Self {
1253 Self::try_row_major(dims).expect("row_major: shape overflows isize")
1254 }
1255
1256 /// Fallible [`Self::row_major`].
1257 ///
1258 /// # Errors
1259 ///
1260 /// Returns [`StridedError::OffsetOverflow`] when the element count or a
1261 /// stride does not fit `isize`.
1262 pub fn try_row_major(dims: &[usize]) -> Result<Self> {
1263 let strides = checked_row_major_strides(dims)?;
1264 let total = checked_element_count(dims)?;
1265 Ok(Self {
1266 data: vec![T::default(); total],
1267 meta: ViewMeta::from_slices(dims, &strides),
1268 offset: 0,
1269 })
1270 }
1271
1272 /// Create a column-major tensor with values produced by a function.
1273 ///
1274 /// The function is called with indices in column-major iteration order.
1275 ///
1276 /// # Panics
1277 ///
1278 /// Panics when the element count or a stride does not fit `isize`; use
1279 /// [`Self::try_from_fn_col_major`] to get [`StridedError::OffsetOverflow`] instead.
1280 pub fn from_fn_col_major(dims: &[usize], f: impl FnMut(&[usize]) -> T) -> Self {
1281 Self::try_from_fn_col_major(dims, f).expect("from_fn_col_major: shape overflows isize")
1282 }
1283
1284 /// Fallible [`Self::from_fn_col_major`]; `f` is not called on error.
1285 ///
1286 /// # Errors
1287 ///
1288 /// Returns [`StridedError::OffsetOverflow`] when the element count or a
1289 /// stride does not fit `isize`.
1290 pub fn try_from_fn_col_major(dims: &[usize], mut f: impl FnMut(&[usize]) -> T) -> Result<Self> {
1291 let strides = checked_col_major_strides(dims)?;
1292 let total = checked_element_count(dims)?;
1293 let rank = dims.len();
1294 let mut data = Vec::with_capacity(total);
1295 let mut idx: AxisVec = AxisVec::new();
1296 idx.resize(rank, 0usize);
1297 // Borrow the cursor once; indexing the `TinyVec` per access would pay
1298 // its inline/heap match in the element-count loop.
1299 let idx: &mut [usize] = &mut idx[..];
1300 for _ in 0..total {
1301 data.push(f(idx));
1302 for d in 0..rank {
1303 idx[d] += 1;
1304 if idx[d] < dims[d] {
1305 break;
1306 }
1307 idx[d] = 0;
1308 }
1309 }
1310 Ok(Self {
1311 data,
1312 meta: ViewMeta::from_slices(dims, &strides),
1313 offset: 0,
1314 })
1315 }
1316
1317 /// Create a row-major tensor with values produced by a function.
1318 ///
1319 /// The function is called with indices in row-major iteration order.
1320 ///
1321 /// # Panics
1322 ///
1323 /// Panics when the element count or a stride does not fit `isize`; use
1324 /// [`Self::try_from_fn_row_major`] to get [`StridedError::OffsetOverflow`] instead.
1325 pub fn from_fn_row_major(dims: &[usize], f: impl FnMut(&[usize]) -> T) -> Self {
1326 Self::try_from_fn_row_major(dims, f).expect("from_fn_row_major: shape overflows isize")
1327 }
1328
1329 /// Fallible [`Self::from_fn_row_major`]; `f` is not called on error.
1330 ///
1331 /// # Errors
1332 ///
1333 /// Returns [`StridedError::OffsetOverflow`] when the element count or a
1334 /// stride does not fit `isize`.
1335 pub fn try_from_fn_row_major(dims: &[usize], mut f: impl FnMut(&[usize]) -> T) -> Result<Self> {
1336 let strides = checked_row_major_strides(dims)?;
1337 let total = checked_element_count(dims)?;
1338 let rank = dims.len();
1339 let mut data = Vec::with_capacity(total);
1340 let mut idx: AxisVec = AxisVec::new();
1341 idx.resize(rank, 0usize);
1342 // Borrow the cursor once; see `from_fn_col_major`.
1343 let idx: &mut [usize] = &mut idx[..];
1344 for _ in 0..total {
1345 data.push(f(idx));
1346 for d in (0..rank).rev() {
1347 idx[d] += 1;
1348 if idx[d] < dims[d] {
1349 break;
1350 }
1351 idx[d] = 0;
1352 }
1353 }
1354 Ok(Self {
1355 data,
1356 meta: ViewMeta::from_slices(dims, &strides),
1357 offset: 0,
1358 })
1359 }
1360}
1361
1362impl<T> StridedArray<T> {
1363 /// Create from raw parts.
1364 ///
1365 /// # Errors
1366 ///
1367 /// Returns [`StridedError::StrideLengthMismatch`] when `dims` and
1368 /// `strides` have different lengths, and
1369 /// [`StridedError::OffsetOverflow`] when a reachable offset is not
1370 /// representable: `dim - 1` of a non-unit extent does not fit `isize`, a
1371 /// stride times that last index or the accumulated offsets overflow
1372 /// `isize` or go negative, or the highest reachable offset is not below
1373 /// the buffer length. A layout with a zero extent skips this check.
1374 pub fn from_parts(
1375 data: Vec<T>,
1376 dims: &[usize],
1377 strides: &[isize],
1378 offset: isize,
1379 ) -> Result<Self> {
1380 validate_bounds(data.len(), dims, strides, offset)?;
1381 Ok(Self {
1382 data,
1383 meta: ViewMeta::from_slices(dims, strides),
1384 offset,
1385 })
1386 }
1387
1388 /// Returns the shape (dimension sizes) of this array.
1389 #[inline]
1390 pub fn dims(&self) -> &[usize] {
1391 self.meta.dims()
1392 }
1393
1394 /// Returns the strides (in units of `T`) for each dimension.
1395 #[inline]
1396 pub fn strides(&self) -> &[isize] {
1397 self.meta.strides()
1398 }
1399
1400 /// Returns the number of dimensions (rank).
1401 #[inline]
1402 pub fn ndim(&self) -> usize {
1403 self.dims().len()
1404 }
1405
1406 /// Returns the total number of elements.
1407 ///
1408 /// # Panics
1409 ///
1410 /// Panics when the element count does not fit `usize`, which only a
1411 /// broadcast (stride-0) layout can reach.
1412 #[inline]
1413 pub fn len(&self) -> usize {
1414 element_count(self.dims())
1415 }
1416
1417 /// Returns `true` if any dimension is zero.
1418 #[inline]
1419 pub fn is_empty(&self) -> bool {
1420 self.dims().iter().any(|&d| d == 0)
1421 }
1422
1423 /// Returns a reference to the backing data slice.
1424 #[inline]
1425 pub fn data(&self) -> &[T] {
1426 &self.data
1427 }
1428
1429 /// Returns a mutable reference to the backing data slice.
1430 #[inline]
1431 pub fn data_mut(&mut self) -> &mut [T] {
1432 &mut self.data
1433 }
1434
1435 /// Create an immutable view over this tensor.
1436 #[inline]
1437 pub fn view(&self) -> StridedView<'_, T> {
1438 let ptr = if self.is_empty() {
1439 empty_const_ptr()
1440 } else {
1441 // SAFETY: `offset` is either 0 from a constructor or was accepted by
1442 // `validate_bounds` in `StridedArray::from_parts` (view.rs:1374), so
1443 // `data.as_ptr() + offset` stays inside `data`. The empty branch
1444 // above excludes a zero-sized layout, whose offset need not be
1445 // reachable.
1446 unsafe { self.data.as_ptr().offset(self.offset) }
1447 };
1448 StridedView {
1449 ptr,
1450 data: &self.data,
1451 meta: self.meta.clone(),
1452 offset: self.offset,
1453 _op: PhantomData,
1454 }
1455 }
1456
1457 /// Create a mutable view over this tensor.
1458 #[inline]
1459 pub fn view_mut(&mut self) -> StridedViewMut<'_, T> {
1460 let src_len = self.data.len();
1461 let base = self.data.as_mut_ptr();
1462 let ptr = if self.is_empty() {
1463 empty_mut_ptr()
1464 } else {
1465 // SAFETY: as `view` above — `self.offset` is either 0 from a
1466 // constructor or was accepted by `validate_bounds` in
1467 // `StridedArray::from_parts` (view.rs:1374), and `&mut self` keeps
1468 // this reborrow exclusive. The empty branch above excludes a
1469 // zero-sized layout, whose offset need not be reachable.
1470 unsafe { base.offset(self.offset) }
1471 };
1472 StridedViewMut {
1473 ptr,
1474 base,
1475 src_len,
1476 meta: self.meta.clone(),
1477 offset: self.offset,
1478 _marker: PhantomData,
1479 }
1480 }
1481
1482 /// Permute dimensions (metadata-only reorder, no data copy).
1483 ///
1484 /// Returns a new array with reordered dims and strides.
1485 /// The underlying data buffer is not touched.
1486 ///
1487 /// # Errors
1488 ///
1489 /// Returns [`StridedError::RankMismatch`] when `perm` does not have one
1490 /// entry per axis, and [`StridedError::InvalidAxis`] when an entry
1491 /// names an axis outside the rank or repeats an axis.
1492 pub fn permuted(self, perm: &[usize]) -> Result<Self> {
1493 let rank = self.dims().len();
1494 if perm.len() != rank {
1495 return Err(StridedError::RankMismatch(perm.len(), rank));
1496 }
1497 let mut seen: tinyvec::TinyVec<[bool; INLINE_RANK]> = tinyvec::TinyVec::new();
1498 seen.resize(rank, false);
1499 for &p in perm {
1500 if p >= rank {
1501 return Err(StridedError::InvalidAxis { axis: p, rank });
1502 }
1503 if seen[p] {
1504 return Err(StridedError::InvalidAxis { axis: p, rank });
1505 }
1506 seen[p] = true;
1507 }
1508 let new_dims: AxisVec = perm.iter().map(|&p| self.dims()[p]).collect();
1509 let new_strides: StrideVec = perm.iter().map(|&p| self.strides()[p]).collect();
1510 Ok(Self {
1511 data: self.data,
1512 meta: ViewMeta::from_slices(&new_dims, &new_strides),
1513 offset: self.offset,
1514 })
1515 }
1516
1517 /// Consume the array and return the backing `Vec<T>`.
1518 pub fn into_data(self) -> Vec<T> {
1519 self.data
1520 }
1521
1522 /// Iterate over all elements in memory order.
1523 pub fn iter(&self) -> std::slice::Iter<'_, T> {
1524 self.data.iter()
1525 }
1526
1527 /// Mutable iteration over all elements in memory order.
1528 pub fn iter_mut(&mut self) -> std::slice::IterMut<'_, T> {
1529 self.data.iter_mut()
1530 }
1531}
1532
1533impl<T: Default> StridedArray<T> {
1534 /// Create a column-major tensor reusing an existing buffer.
1535 ///
1536 /// If `buf` has at least `product(dims)` elements, it is truncated and
1537 /// zeroed. Otherwise a fresh buffer is allocated.
1538 ///
1539 /// # Panics
1540 ///
1541 /// Panics when the element count or a stride does not fit `isize`; use
1542 /// [`Self::try_col_major_from_buffer`] to get
1543 /// [`StridedError::OffsetOverflow`] instead.
1544 pub fn col_major_from_buffer(buf: Vec<T>, dims: &[usize]) -> Self {
1545 Self::try_col_major_from_buffer(buf, dims)
1546 .expect("col_major_from_buffer: shape overflows isize")
1547 }
1548
1549 /// Fallible [`Self::col_major_from_buffer`]; `buf` is dropped on error.
1550 ///
1551 /// # Errors
1552 ///
1553 /// Returns [`StridedError::OffsetOverflow`] when the element count or a
1554 /// stride does not fit `isize`.
1555 pub fn try_col_major_from_buffer(mut buf: Vec<T>, dims: &[usize]) -> Result<Self> {
1556 let strides = checked_col_major_strides(dims)?;
1557 let total = checked_element_count(dims)?;
1558 if buf.len() >= total {
1559 buf.truncate(total);
1560 } else {
1561 buf.resize_with(total, T::default);
1562 }
1563 // Zero the reused region
1564 for v in buf.iter_mut() {
1565 *v = T::default();
1566 }
1567 Ok(Self {
1568 data: buf,
1569 meta: ViewMeta::from_slices(dims, &strides),
1570 offset: 0,
1571 })
1572 }
1573}
1574
1575impl<T: Copy> StridedArray<T> {
1576 /// Create a column-major tensor with **uninitialized** data.
1577 ///
1578 /// # Safety
1579 /// Caller must write every element before reading.
1580 ///
1581 /// # Panics
1582 ///
1583 /// Panics when the element count or a stride does not fit `isize`; use
1584 /// [`Self::try_col_major_uninit`] to get [`StridedError::OffsetOverflow`]
1585 /// instead.
1586 pub unsafe fn col_major_uninit(dims: &[usize]) -> Self {
1587 Self::try_col_major_uninit(dims).expect("col_major_uninit: shape overflows isize")
1588 }
1589
1590 /// Fallible [`Self::col_major_uninit`].
1591 ///
1592 /// # Safety
1593 /// Caller must write every element before reading.
1594 ///
1595 /// # Errors
1596 ///
1597 /// Returns [`StridedError::OffsetOverflow`] when the element count or a
1598 /// stride does not fit `isize`.
1599 pub unsafe fn try_col_major_uninit(dims: &[usize]) -> Result<Self> {
1600 let strides = checked_col_major_strides(dims)?;
1601 let total = checked_element_count(dims)?;
1602 let mut data = Vec::with_capacity(total);
1603 data.set_len(total);
1604 Ok(Self {
1605 data,
1606 meta: ViewMeta::from_slices(dims, &strides),
1607 offset: 0,
1608 })
1609 }
1610
1611 /// Reuse an existing buffer as a column-major tensor **without zeroing**.
1612 ///
1613 /// # Safety
1614 /// Caller must write every element before reading.
1615 ///
1616 /// # Panics
1617 ///
1618 /// Panics when the element count or a stride does not fit `isize`; use
1619 /// [`Self::try_col_major_from_buffer_uninit`] to get
1620 /// [`StridedError::OffsetOverflow`] instead.
1621 pub unsafe fn col_major_from_buffer_uninit(buf: Vec<T>, dims: &[usize]) -> Self {
1622 Self::try_col_major_from_buffer_uninit(buf, dims)
1623 .expect("col_major_from_buffer_uninit: shape overflows isize")
1624 }
1625
1626 /// Fallible [`Self::col_major_from_buffer_uninit`]; `buf` is dropped on
1627 /// error.
1628 ///
1629 /// # Safety
1630 /// Caller must write every element before reading.
1631 ///
1632 /// # Errors
1633 ///
1634 /// Returns [`StridedError::OffsetOverflow`] when the element count or a
1635 /// stride does not fit `isize`.
1636 pub unsafe fn try_col_major_from_buffer_uninit(
1637 mut buf: Vec<T>,
1638 dims: &[usize],
1639 ) -> Result<Self> {
1640 let strides = checked_col_major_strides(dims)?;
1641 let total = checked_element_count(dims)?;
1642 if buf.capacity() < total {
1643 buf.reserve(total - buf.len());
1644 }
1645 buf.set_len(total);
1646 Ok(Self {
1647 data: buf,
1648 meta: ViewMeta::from_slices(dims, &strides),
1649 offset: 0,
1650 })
1651 }
1652}
1653
1654impl<T: Copy> StridedArray<T> {
1655 /// Get an element by multi-dimensional index.
1656 ///
1657 /// # Panics
1658 ///
1659 /// Panics if `indices.len()` is not the array's rank, an index is out of
1660 /// bounds, or the indexed offset does not fit `isize`.
1661 pub fn get(&self, indices: &[usize]) -> T {
1662 self.view().get(indices)
1663 }
1664
1665 /// Set an element by multi-dimensional index.
1666 ///
1667 /// # Panics
1668 ///
1669 /// Panics if `indices.len()` is not the array's rank, an index is out of
1670 /// bounds, or the indexed offset does not fit `isize`.
1671 pub fn set(&mut self, indices: &[usize], value: T) {
1672 let offset = indexed_offset(self.dims(), self.strides(), indices, self.offset);
1673 self.data[offset as usize] = value;
1674 }
1675}
1676
1677impl<T: Copy> Index<&[usize]> for StridedArray<T> {
1678 type Output = T;
1679
1680 fn index(&self, indices: &[usize]) -> &T {
1681 let offset = indexed_offset(self.dims(), self.strides(), indices, self.offset);
1682 &self.data[offset as usize]
1683 }
1684}
1685
1686impl<T: Copy> IndexMut<&[usize]> for StridedArray<T> {
1687 fn index_mut(&mut self, indices: &[usize]) -> &mut T {
1688 let offset = indexed_offset(self.dims(), self.strides(), indices, self.offset);
1689 &mut self.data[offset as usize]
1690 }
1691}
1692
1693// ============================================================================
1694// Tests
1695// ============================================================================
1696
1697#[cfg(test)]
1698mod tests;