Skip to main content

diskann_utils/views/
rowmajor.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6use std::{marker::PhantomData, mem::ManuallyDrop, num::NonZeroUsize, ptr::NonNull};
7
8#[cfg(feature = "rayon")]
9use rayon::prelude::{
10    IndexedParallelIterator, IntoParallelIterator, ParallelIterator, ParallelSliceMut,
11};
12use thiserror::Error;
13
14pub mod iter;
15
16use crate::{internal, Reborrow, ReborrowMut};
17
18////////////
19// Matrix //
20////////////
21
22/// Tools for treating contiguous regions of memory as a dense, row-major matrix.
23///
24/// See also: [`MatrixMut`], [`Owned`], [`Ref`], [`Mut`].
25///
26/// # Safety
27///
28/// For any shared reference `matrix: &Self`, the following must hold for the entire lifetime
29/// of that reference:
30///
31/// * Every call to `matrix.as_nonnull()` returns the same pointer.
32///
33/// * Every call to `matrix.layout()` returns the same layout.
34///
35/// * The pointer and layout describe the same valid initialized span throughout that lifetime.
36///   Given:
37///
38///   ```text
39///   let data = matrix.as_nonnull();
40///   let layout = matrix.layout();
41///   ```
42///
43///   constructing the following slice is valid:
44///
45///   ```text
46///   unsafe { std::slice::from_raw_parts(data.as_ptr().cast_const(), layout.num_elements()) };
47///   ```
48///
49///   In particular:
50///
51///   - `data` must be properly aligned, including when `layout.num_elements() == 0`.
52///
53///   - If both `layout.num_elements()` and `size_of::<Self::Element>()` are nonzero, the
54///     described span must lie within a single allocation.
55///
56///   - The span must contain `layout.num_elements()` consecutive, properly initialized
57///     values of type `Self::Element`.
58///
59///   - The referenced memory must not be mutated for the lifetime of `matrix`, except
60///     through an `UnsafeCell`.
61pub unsafe trait Matrix {
62    /// The type of the element stored in the matrix.
63    type Element;
64
65    /// Return the base pointer for the matrix.
66    fn as_nonnull(&self) -> NonNull<Self::Element>;
67
68    /// Return the [`Layout`] for the matrix.
69    fn layout(&self) -> Layout<Self::Element>;
70
71    //----------//
72    // Provided //
73    //----------//
74
75    /// Return the number of rows in the matrix.
76    fn nrows(&self) -> usize {
77        self.layout().nrows()
78    }
79
80    /// Return the number of columns in the matrix.
81    fn ncols(&self) -> usize {
82        self.layout().ncols()
83    }
84
85    /// Returns the requested row without bounds checking.
86    ///
87    /// # Safety
88    ///
89    /// The following conditions must hold to avoid undefined behavior:
90    ///
91    /// * `row < self.nrows()`.
92    unsafe fn row_unchecked(&self, row: usize) -> &[Self::Element] {
93        let layout = self.layout();
94        debug_assert!(row < layout.nrows());
95
96        // SAFETY: The caller guarantees that `row` is in-bounds. The validated layout
97        // therefore places this row entirely within the initialized span required by the
98        // `Matrix` invariant. The returned lifetime is tied to the borrow of `self`.
99        unsafe {
100            std::slice::from_raw_parts(self.as_ptr().add(layout.ncols() * row), layout.ncols())
101        }
102    }
103
104    /// Return a pointer to the base of the matrix.
105    fn as_ptr(&self) -> *const Self::Element {
106        self.as_nonnull().as_ptr().cast_const()
107    }
108
109    /// Return the underlying data as a slice.
110    fn as_slice(&self) -> &[Self::Element] {
111        // SAFETY: The `Matrix` invariant requires this pointer and length to describe a
112        // properly aligned, initialized span that remains immutable for this borrow.
113        unsafe { std::slice::from_raw_parts(self.as_ptr(), self.layout().num_elements()) }
114    }
115
116    /// Return row `row` as a slice.
117    ///
118    /// # Panic
119    ///
120    /// Panics if `row >= self.nrows()`.
121    fn row(&self, row: usize) -> &[Self::Element] {
122        assert!(
123            row < self.nrows(),
124            "tried to access row {row} of a matrix with {} rows",
125            self.nrows()
126        );
127
128        // SAFETY: `row` is in-bounds.
129        unsafe { self.row_unchecked(row) }
130    }
131
132    /// Return row `row` if `row < self.nrows()`. Otherwise, return `None`.
133    fn get_row(&self, row: usize) -> Option<&[Self::Element]> {
134        if row < self.nrows() {
135            // SAFETY: `row` is in-bounds.
136            Some(unsafe { self.row_unchecked(row) })
137        } else {
138            None
139        }
140    }
141
142    /// Return a iterator over all rows in the matrix.
143    ///
144    /// Rows are yielded sequentially beginning with row 0.
145    fn rows(&self) -> iter::Rows<'_, Self::Element> {
146        iter::Rows::new(self.as_view())
147    }
148
149    /// Returns a reference to an element without bounds checking.
150    ///
151    /// # Safety
152    ///
153    /// The following conditions must hold to avoid undefined behavior:
154    /// * `row < self.nrows()`.
155    /// * `col < self.ncols()`.
156    unsafe fn element_unchecked(&self, row: usize, col: usize) -> &Self::Element {
157        let layout = self.layout();
158        debug_assert!(row < layout.nrows());
159        debug_assert!(col < layout.ncols());
160
161        // SAFETY: The caller guarantees that both indices are in-bounds, so the validated
162        // layout places this element within the initialized span required by the `Matrix`
163        // invariant. The returned lifetime is tied to the borrow of `self`.
164        unsafe { &*self.as_ptr().add(row * layout.ncols() + col) }
165    }
166
167    /// Return the value at the specified `row` and `col`.
168    ///
169    /// If either index is out-of-bounds, return `None`.
170    fn get_element(&self, row: usize, col: usize) -> Option<&Self::Element> {
171        if row >= self.nrows() || col >= self.ncols() {
172            None
173        } else {
174            // SAFETY: We just verified that `row` and `col` are in-bounds.
175            Some(unsafe { self.element_unchecked(row, col) })
176        }
177    }
178
179    /// Return the value at the specified `row` and `col`.
180    ///
181    /// # Panics
182    ///
183    /// Panics if either `row` or `col` is out-of-bounds.
184    fn element(&self, row: usize, col: usize) -> &Self::Element {
185        assert!(
186            row < self.nrows(),
187            "row {row} is out of bounds (max: {})",
188            self.nrows()
189        );
190        assert!(
191            col < self.ncols(),
192            "col {col} is out of bounds (max: {})",
193            self.ncols()
194        );
195
196        // SAFETY: We just verified that `row` and `col` are in-bounds.
197        unsafe { self.element_unchecked(row, col) }
198    }
199
200    /// Return a view over the matrix.
201    fn as_view(&self) -> Ref<'_, Self::Element> {
202        Ref {
203            ptr: self.as_nonnull(),
204            layout: self.layout(),
205            _lifetime: PhantomData,
206        }
207    }
208
209    /// Return a view over the rows in `rows`, or `None` when the range is invalid.
210    fn subview(&self, rows: std::ops::Range<usize>) -> Option<Ref<'_, Self::Element>> {
211        if rows.start > rows.end || rows.end > self.nrows() {
212            return None;
213        }
214
215        let ncols = self.ncols();
216        // SAFETY: `rows.start <= self.nrows()`, so the validated parent layout makes the
217        // offset representable and places it within or one past the matrix span. Pointer
218        // arithmetic within that span preserves non-nullness.
219        let ptr =
220            unsafe { NonNull::new_unchecked(self.as_ptr().add(rows.start * ncols).cast_mut()) };
221        // SAFETY: The selected rows are a subset of the validated parent layout, with the
222        // same column count.
223        let layout = unsafe { Layout::new_unchecked(rows.end - rows.start, ncols) };
224        Some(Ref {
225            ptr,
226            layout,
227            _lifetime: PhantomData,
228        })
229    }
230
231    /// Return an iterator that divides the matrix into sub-matrices with (up to)
232    /// `batchsize` rows with `self.ncols()` columns.
233    ///
234    /// It is possible for yielded sub-matrices to have fewer than `batchsize` rows if the
235    /// number of rows in the parent matrix is not evenly divisible by `batchsize`.
236    fn window_iter(&self, batchsize: NonZeroUsize) -> iter::Windows<'_, Self::Element> {
237        iter::Windows::new(self.as_view(), batchsize)
238    }
239
240    /// Return an [`Owned`] with the same shape as `self` and cloned contents.
241    fn to_rowmajor_owned(&self) -> Owned<Self::Element>
242    where
243        Self::Element: Clone,
244    {
245        // SAFETY: The `Matrix` invariant makes `self.as_slice().len()` equal to
246        // `self.layout().num_elements()`.
247        unsafe { Owned::from_data_unchecked(self.as_slice().into(), self.layout()) }
248    }
249
250    /// Create a new [`Matrix`] by applying the closure `f` to each element.
251    ///
252    /// The returned matrix has the same shape as `self`.
253    fn try_map<F, R>(&self, f: F) -> Result<Owned<R>, LayoutError>
254    where
255        F: FnMut(&Self::Element) -> R,
256    {
257        let layout = self.layout().rebind::<R>()?;
258        let data: Box<[_]> = self.as_slice().iter().map(f).collect();
259
260        // SAFETY: Mapping preserves the element count, and `layout` is the source layout
261        // validated for the destination element type.
262        Ok(unsafe { Owned::from_data_unchecked(data, layout) })
263    }
264
265    /// Create a new [`Matrix`] by applying the closure `f` to each element.
266    ///
267    /// The returned matrix has the same shape as `self`.
268    ///
269    /// # Panics
270    ///
271    /// Panics if allocating space for [`Owned`] would overflow `isize::MAX`.
272    #[track_caller]
273    fn map<F, R>(&self, f: F) -> Owned<R>
274    where
275        F: FnMut(&Self::Element) -> R,
276    {
277        match self.try_map(f) {
278            Ok(owned) => owned,
279            Err(error) => panic!("`Matrix::map` failed: {error}"),
280        }
281    }
282
283    /// Transpose the elements in `self`.
284    fn transpose(&self) -> Owned<Self::Element>
285    where
286        Self::Element: Clone,
287    {
288        Owned::from_fn_with_layout(self.layout().transpose(), |RowCol { row, col }| {
289            // SAFETY: By construction, `col < self.nrows()` and `row < self.ncols()`.
290            unsafe { self.element_unchecked(col, row).clone() }
291        })
292    }
293
294    //-------//
295    // Rayon //
296    //-------//
297
298    /// Return a parallel iterator over the rows of the matrix.
299    #[cfg(feature = "rayon")]
300    fn par_rows(&self) -> impl IndexedParallelIterator<Item = &[Self::Element]>
301    where
302        Self::Element: Sync,
303    {
304        let r = self.as_view();
305
306        (0..r.nrows()).into_par_iter().map(move |row| {
307            // SAFETY: `row` comes from `0..r.nrows()`.
308            unsafe { r.into_row_unchecked(row) }
309        })
310    }
311
312    /// Return a parallel iterator that divides the matrix into sub-matrices with (up to)
313    /// `batchsize` rows with `self.ncols()` columns.
314    ///
315    /// This allows workers in parallel algorithms to work on dense subsets of the whole
316    /// matrix for better locality.
317    ///
318    /// It is possible for yielded sub-matrices to have fewer than `batchsize` rows if the
319    /// number of rows in the parent matrix is not evenly divisible by `batchsize`.
320    ///
321    /// # Panics
322    ///
323    /// Panics if `batchsize = 0`.
324    #[cfg(feature = "rayon")]
325    fn par_window_iter(
326        &self,
327        batchsize: usize,
328    ) -> impl IndexedParallelIterator<Item = Ref<'_, Self::Element>>
329    where
330        Self::Element: Sync,
331    {
332        assert!(batchsize != 0, "par_window_iter batchsize cannot be zero");
333
334        let r = self.as_view();
335        (0..r.nrows())
336            .into_par_iter()
337            .step_by(batchsize)
338            .map(move |start| {
339                let end = start.saturating_add(batchsize).min(r.nrows());
340
341                // SAFETY: `start` comes from `0..nrows` and `end` is clamped to `nrows`.
342                unsafe { r.into_subview_unchecked(start..end) }
343            })
344    }
345}
346
347///////////////
348// MatrixMut //
349///////////////
350
351/// Tools for treating contiguous regions of mutable memory as a dense, row-major matrix.
352///
353/// See also: [`Owned`], [`Ref`], [`Mut`].
354///
355/// # Safety
356///
357/// In addition to the requirements of [`Matrix`], implementations must ensure that for
358/// any **mutable** borrow `matrix: &mut Self`, it is valid to construct:
359///
360/// ```text
361/// let data = matrix.as_nonnull_mut();
362/// let layout = matrix.layout();
363///
364/// unsafe {
365///    std::slice::from_raw_parts_mut(data.as_ptr(), layout.num_elements())
366/// }
367/// ```
368///
369/// The resulting slice must have exclusive access to the described elements for its entire
370/// lifetime. No other reference may be used to access those elements during that lifetime.
371///
372/// Implementations may change their pointer or layout through exclusive access, such as when
373/// resizing or reallocating. They must not do so while any reference derived from the previous
374/// pointer and layout remains live.
375///
376/// For a given state of `matrix`, the pointers from `matrix.as_nonnull_mut()` and
377/// `matrix.as_nonnull()` must have the same address. However, the pointer obtained from
378/// `matrix.as_nonnull_mut()` must have provenance permitting mutable access to the span
379/// described by `matrix.layout()`.
380pub unsafe trait MatrixMut: Matrix {
381    /// Return the base pointer for the matrix.
382    ///
383    /// Calling this method must not change the matrix's pointer, layout, or contents.
384    ///
385    /// The returned pointer must have provenance permitting mutable access to the span
386    /// described by [`Matrix::layout`] and for a given matrix state, must have the same
387    /// address as [`Matrix::as_nonnull`].
388    fn as_nonnull_mut(&mut self) -> NonNull<Self::Element>;
389
390    //----------//
391    // Provided //
392    //----------//
393
394    /// Returns the requested row without bounds checking.
395    ///
396    /// # Safety
397    ///
398    /// The following conditions must hold to avoid undefined behavior:
399    ///
400    /// * `row < self.nrows()`.
401    unsafe fn row_unchecked_mut(&mut self, row: usize) -> &mut [Self::Element] {
402        let layout = self.layout();
403
404        debug_assert!(row < layout.nrows());
405
406        // SAFETY: The caller guarantees that `row` is in-bounds. The validated layout
407        // therefore places this row within the initialized span, and the `MatrixMut`
408        // invariant grants exclusive access for the returned borrow.
409        unsafe {
410            std::slice::from_raw_parts_mut(
411                self.as_mut_ptr().add(layout.ncols() * row),
412                layout.ncols(),
413            )
414        }
415    }
416
417    /// Return a pointer to the base of the matrix.
418    fn as_mut_ptr(&mut self) -> *mut Self::Element {
419        self.as_nonnull_mut().as_ptr()
420    }
421
422    /// Return the underlying data as a mutable slice.
423    fn as_mut_slice(&mut self) -> &mut [Self::Element] {
424        // SAFETY: The `MatrixMut` invariant requires this pointer and length to describe
425        // the initialized matrix span and grants exclusive access for this mutable borrow.
426        unsafe { std::slice::from_raw_parts_mut(self.as_mut_ptr(), self.layout().num_elements()) }
427    }
428
429    /// Return row `row` as a mutable slice.
430    ///
431    /// # Panics
432    ///
433    /// Panics if `row >= self.nrows()`.
434    fn row_mut(&mut self, row: usize) -> &mut [Self::Element] {
435        assert!(
436            row < self.nrows(),
437            "tried to access row {row} of a matrix with {} rows",
438            self.nrows()
439        );
440
441        // SAFETY: `row` is in-bounds.
442        unsafe { self.row_unchecked_mut(row) }
443    }
444
445    /// Return row `row` if `row < self.nrows()`. Otherwise, return `None`.
446    fn get_row_mut(&mut self, row: usize) -> Option<&mut [Self::Element]> {
447        if row < self.nrows() {
448            // SAFETY: `row` is in-bounds.
449            Some(unsafe { self.row_unchecked_mut(row) })
450        } else {
451            None
452        }
453    }
454
455    /// Return a mutable iterator over all rows in the matrix.
456    ///
457    /// Rows are yielded sequentially beginning with row 0.
458    fn rows_mut(&mut self) -> iter::RowsMut<'_, Self::Element> {
459        iter::RowsMut::new(self.as_view_mut())
460    }
461
462    /// Returns a mutable reference to an element without bounds checking.
463    ///
464    /// # Safety
465    ///
466    /// The following conditions must hold to avoid undefined behavior:
467    /// * `row < self.nrows()`.
468    /// * `col < self.ncols()`.
469    unsafe fn element_unchecked_mut(&mut self, row: usize, col: usize) -> &mut Self::Element {
470        let layout = self.layout();
471        debug_assert!(row < layout.nrows());
472        debug_assert!(col < layout.ncols());
473
474        // SAFETY: The caller guarantees that both indices are in-bounds, so the validated
475        // layout places this element within the initialized matrix span. The `MatrixMut`
476        // invariant grants exclusive access for the returned borrow.
477        unsafe { &mut *self.as_mut_ptr().add(row * layout.ncols() + col) }
478    }
479
480    /// Return the value at the specified `row` and `col`.
481    ///
482    /// If either index is out-of-bounds, return `None`.
483    fn get_element_mut(&mut self, row: usize, col: usize) -> Option<&mut Self::Element> {
484        if row >= self.nrows() || col >= self.ncols() {
485            None
486        } else {
487            // SAFETY: We just verified that `row` and `col` are in-bounds.
488            Some(unsafe { self.element_unchecked_mut(row, col) })
489        }
490    }
491
492    /// Return the value at the specified `row` and `col`.
493    ///
494    /// # Panics
495    ///
496    /// Panics if either `row` or `col` is out-of-bounds.
497    fn element_mut(&mut self, row: usize, col: usize) -> &mut Self::Element {
498        assert!(
499            row < self.nrows(),
500            "row {row} is out of bounds (max: {})",
501            self.nrows()
502        );
503        assert!(
504            col < self.ncols(),
505            "col {col} is out of bounds (max: {})",
506            self.ncols()
507        );
508
509        // SAFETY: We just verified that `row` and `col` are in-bounds.
510        unsafe { self.element_unchecked_mut(row, col) }
511    }
512
513    /// Return a view over the matrix.
514    fn as_view_mut(&mut self) -> Mut<'_, Self::Element> {
515        Mut {
516            ptr: self.as_nonnull_mut(),
517            layout: self.layout(),
518            _lifetime: PhantomData,
519        }
520    }
521
522    //-------//
523    // Rayon //
524    //-------//
525
526    /// Return a parallel iterator over the rows of the matrix.
527    ///
528    /// # Panics
529    ///
530    /// Panics if `self.ncols() == 0 && self.nrows() != 0`.
531    #[cfg(feature = "rayon")]
532    fn par_rows_mut(&mut self) -> impl IndexedParallelIterator<Item = &mut [Self::Element]>
533    where
534        Self::Element: Send,
535    {
536        let ncols = self.ncols();
537        assert!(
538            ncols != 0 || self.nrows() == 0,
539            "`MatrixMut::par_rows_mut` does not support matrices with rows and zero columns"
540        );
541        self.as_mut_slice().par_chunks_exact_mut(ncols.max(1))
542    }
543
544    /// Return a parallel iterator that divides the matrix into mutable sub-matrices with
545    /// (up to) `batchsize` rows with `self.ncols()` columns.
546    ///
547    /// This allows workers in parallel algorithms to work on dense subsets of the whole
548    /// matrix for better locality.
549    ///
550    /// It is possible for yielded sub-matrices to have fewer than `batchsize` rows if the
551    /// number of rows in the parent matrix is not evenly divisible by `batchsize`.
552    ///
553    /// # Panics
554    ///
555    /// Panics if `batchsize = 0` or `self.ncols() == 0 && self.nrows() != 0`.
556    #[cfg(feature = "rayon")]
557    fn par_window_iter_mut(
558        &mut self,
559        batchsize: usize,
560    ) -> impl IndexedParallelIterator<Item = Mut<'_, Self::Element>>
561    where
562        Self::Element: Send,
563    {
564        assert!(
565            batchsize != 0,
566            "par_window_iter_mut batchsize cannot be zero"
567        );
568
569        let ncols = self.ncols();
570        assert!(
571            ncols != 0 || self.nrows() == 0,
572            "`MatrixMut::par_window_iter_mut` does not support matrices with rows and zero columns"
573        );
574
575        // Ensure that `batchsize * ncols` does not overflow.
576        let batchsize = batchsize.min(self.nrows());
577        self.as_mut_slice()
578            .par_chunks_mut((ncols * batchsize).max(1))
579            .map(move |data| {
580                let blobsize = data.len();
581                let nrows = blobsize / ncols;
582                assert_eq!(blobsize % ncols, 0);
583
584                // SAFETY:
585                //
586                // * `Layout::new_unchecked` is safe because `ncols` is the parent column
587                //   count and `nrows <= self.nrows()`, so this layout cannot exceed the
588                //   validated parent layout.
589                //
590                // * `Mut::from_data_unchecked` is safe because by construction,
591                //   `data.len() == ncols * nrows`.
592                unsafe { Mut::from_data_unchecked(data, Layout::new_unchecked(nrows, ncols)) }
593            })
594    }
595}
596
597///////////////////
598// Matrix Layout //
599///////////////////
600
601/// A validated layout for [`Matrix`] or [`MatrixMut`].
602///
603/// This type guarantees the following invariants:
604///
605/// * `self.nrows() * self.ncols()` does not exceed `usize::MAX`.
606/// * `self.nrows() * self.ncols() * std::mem::size_of::<T>()` does not exceed `isize::MAX`.
607pub struct Layout<T> {
608    nrows: usize,
609    ncols: usize,
610    _type: PhantomData<fn() -> T>,
611}
612
613impl<T> Layout<T> {
614    /// Construct a new [`Layout`], validating the following:
615    ///
616    /// * `nrows * ncols` does not exceed `usize::MAX`.
617    /// * `nrows * ncols * std::mem::size_of::<T>()` does not exceed `isize::MAX` (the maximum
618    ///   addressable byte span).
619    pub const fn new(nrows: usize, ncols: usize) -> Result<Self, LayoutError> {
620        match LayoutError::check::<T>(nrows, ncols) {
621            Ok(()) => Ok(Self {
622                nrows,
623                ncols,
624                _type: PhantomData,
625            }),
626            Err(err) => Err(err),
627        }
628    }
629
630    /// Construct a layout without validating its dimensions.
631    ///
632    /// # Safety
633    ///
634    /// `LayoutError::check::<T>(nrows, ncols)` must succeed.
635    unsafe fn new_unchecked(nrows: usize, ncols: usize) -> Self {
636        debug_assert!(LayoutError::check::<T>(nrows, ncols).is_ok());
637        Self {
638            nrows,
639            ncols,
640            _type: PhantomData,
641        }
642    }
643
644    /// Return the product `self.nrows() * self.ncols()`.
645    pub fn num_elements(&self) -> usize {
646        self.nrows() * self.ncols()
647    }
648
649    /// Return the number of rows.
650    pub fn nrows(&self) -> usize {
651        self.nrows
652    }
653
654    /// Return the number of columns.
655    pub fn ncols(&self) -> usize {
656        self.ncols
657    }
658
659    /// Rebind the element type to `U`.
660    ///
661    /// # Errors
662    ///
663    /// Returns an error if the rebound layout's byte size would exceed `isize::MAX`.
664    pub fn rebind<U>(&self) -> Result<Layout<U>, LayoutError> {
665        if std::mem::size_of::<U>() <= std::mem::size_of::<T>() {
666            // This branch is mainly to communicate to the compiler situations where an
667            // erroring branch can be avoided.
668            //
669            // SAFETY: `self` is validated, and rebinding to an equally sized or smaller
670            // element type cannot increase its byte span.
671            Ok(unsafe { Layout::new_unchecked(self.nrows(), self.ncols()) })
672        } else {
673            Layout::new(self.nrows(), self.ncols())
674        }
675    }
676
677    /// Swap the rows and columns.
678    pub fn transpose(&self) -> Layout<T> {
679        // SAFETY: Swapping the dimensions preserves both the element count and byte span
680        // of this validated layout.
681        unsafe { Layout::new_unchecked(self.ncols, self.nrows) }
682    }
683}
684
685impl<T> Clone for Layout<T> {
686    fn clone(&self) -> Self {
687        *self
688    }
689}
690
691impl<T> Copy for Layout<T> {}
692
693impl<T> std::fmt::Debug for Layout<T> {
694    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
695        f.debug_struct("Layout")
696            .field("nrows", &self.nrows)
697            .field("ncols", &self.ncols)
698            .field("elsize", &std::mem::size_of::<T>())
699            .finish()
700    }
701}
702
703impl<T> PartialEq for Layout<T> {
704    fn eq(&self, other: &Self) -> bool {
705        self.nrows == other.nrows && self.ncols == other.ncols
706    }
707}
708
709impl<T> Eq for Layout<T> {}
710
711/// Errors in the invariants guaranteed by [`Layout`].
712#[derive(Debug, Clone, Copy)]
713pub struct LayoutError {
714    nrows: usize,
715    ncols: usize,
716    elsize: Option<NonZeroUsize>,
717}
718
719impl LayoutError {
720    pub(crate) const fn check<T>(nrows: usize, ncols: usize) -> Result<(), Self> {
721        // Guard the element count itself so that `num_elements()` can never overflow.
722        let elsize = std::mem::size_of::<T>();
723        let num_elements = match nrows.checked_mul(ncols) {
724            Some(num_elements) => num_elements,
725            None => {
726                return Err(Self {
727                    nrows,
728                    ncols,
729                    elsize: None,
730                })
731            }
732        };
733
734        if let Some(len) = num_elements.checked_mul(std::mem::size_of::<T>()) {
735            if len <= isize::MAX as usize {
736                return Ok(());
737            }
738        }
739
740        Err(Self {
741            nrows,
742            ncols,
743            elsize: NonZeroUsize::new(elsize),
744        })
745    }
746}
747
748impl std::fmt::Display for LayoutError {
749    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
750        match self.elsize {
751            Some(elsize) => {
752                write!(
753                    f,
754                    "a matrix of size {}x{} with elements of size {} exceeds `isize::MAX` bytes",
755                    self.nrows, self.ncols, elsize
756                )
757            }
758            None => {
759                write!(
760                    f,
761                    "a matrix of size {}x{} has a length exceeding `usize::MAX`",
762                    self.nrows, self.ncols,
763                )
764            }
765        }
766    }
767}
768
769impl std::error::Error for LayoutError {}
770
771//---------------//
772// Helper Macros //
773//---------------//
774
775macro_rules! constructors {
776    ($element:ident, $data:ty) => {
777        /// Try to construct directly from `data`.
778        ///
779        /// Returns an error if [`Layout::new`] fails for `nrows` and `ncols` or `data.len()`
780        /// is not equal to `nrows * ncols`.
781        pub fn try_from_data(
782            data: $data,
783            nrows: usize,
784            ncols: usize,
785        ) -> Result<Self, TryFromError<$data>> {
786            let layout = match Layout::<$element>::new(nrows, ncols) {
787                Ok(layout) => layout,
788                Err(err) => return Err(TryFromError::layout(data, err)),
789            };
790
791            let len = data.len();
792            if len == layout.num_elements() {
793                // SAFETY: We've checked that `data.len() == layout.num_elements()`.
794                Ok(unsafe { Self::from_data_unchecked(data, layout) })
795            } else {
796                Err(TryFromError::mismatch(
797                    data,
798                    layout.nrows(),
799                    layout.ncols(),
800                    len,
801                ))
802            }
803        }
804
805        /// Construct a row vector directly from `data`.
806        pub fn row_vector(data: $data) -> Self {
807            // SAFETY: An existing slice's byte span cannot exceed `isize::MAX`, and
808            // `1 * data.len()` cannot overflow.
809            let layout = unsafe { Layout::new_unchecked(1, data.len()) };
810
811            // SAFETY: By construction, `data.len() == layout.num_elements()`.
812            unsafe { Self::from_data_unchecked(data, layout) }
813        }
814
815        /// Construct a column vector directly from `data`.
816        pub fn column_vector(data: $data) -> Self {
817            // SAFETY: An existing slice's byte span cannot exceed `isize::MAX`, and
818            // `data.len() * 1` cannot overflow.
819            let layout = unsafe { Layout::new_unchecked(data.len(), 1) };
820
821            // SAFETY: By construction, `data.len() == layout.num_elements()`.
822            unsafe { Self::from_data_unchecked(data, layout) }
823        }
824    };
825}
826
827///////////
828// Owned //
829///////////
830
831/// An initializer argument for the closure provided to [`Owned::from_fn`] and
832/// [`Owned::try_from_fn`] to remove ambiguity of the row and column being initialiazed.
833#[derive(Debug, Clone, Copy, PartialEq, Eq)]
834pub struct RowCol {
835    pub row: usize,
836    pub col: usize,
837}
838
839/// A [`Matrix`]/[`MatrixMut`] that owns its data.
840#[derive(Debug)]
841pub struct Owned<T> {
842    ptr: NonNull<T>,
843    layout: Layout<T>,
844}
845
846impl<T> Owned<T> {
847    constructors!(T, Box<[T]>);
848
849    // NOTE: For constructors, keep `from_fn` and `from_element` first.
850    //
851    // Rust suggests methods in their declaration order, so this keeps the most common
852    // methods as top suggestions.
853
854    /// Construct a new matrix using `init`.
855    ///
856    /// Elements are initialized in memory order.
857    ///
858    /// ```
859    /// use diskann_utils::views::rowmajor::{self, Matrix};
860    ///
861    /// let mat = rowmajor::Owned::from_fn(2, 3, |rc| 3 * rc.row + rc.col);
862    ///
863    /// assert_eq!(mat.row(0), &[0, 1, 2]);
864    /// assert_eq!(mat.row(1), &[3, 4, 5]);
865    /// ```
866    ///
867    /// # Panics
868    ///
869    /// Panics if `nrows * ncols` overflows `usize::MAX`, or if the allocation size exceeds
870    /// `isize::MAX`.
871    #[track_caller]
872    pub fn from_fn<F>(nrows: usize, ncols: usize, init: F) -> Self
873    where
874        F: FnMut(RowCol) -> T,
875    {
876        match Self::try_from_fn(nrows, ncols, init) {
877            Ok(matrix) => matrix,
878            Err(error) => panic!("Owned::from_fn failed with: {error}"),
879        }
880    }
881
882    /// Construct a new matrix using `init`.
883    ///
884    /// Elements are initialized in memory order.
885    ///
886    /// ```
887    /// use diskann_utils::views::rowmajor::{self, Matrix};
888    ///
889    /// let mat = rowmajor::Owned::try_from_fn(2, 3, |rc| 3 * rc.row + rc.col).unwrap();
890    ///
891    /// assert_eq!(mat.row(0), &[0, 1, 2]);
892    /// assert_eq!(mat.row(1), &[3, 4, 5]);
893    /// ```
894    ///
895    /// # Errors
896    ///
897    /// Returns an error if `nrows * ncols` overflows `usize::MAX`, or if the allocation size
898    /// exceeds `isize::MAX`.
899    pub fn try_from_fn<F>(nrows: usize, ncols: usize, init: F) -> Result<Self, LayoutError>
900    where
901        F: FnMut(RowCol) -> T,
902    {
903        let layout = Layout::new(nrows, ncols)?;
904        Ok(Self::from_fn_with_layout(layout, init))
905    }
906
907    /// Construct a new matrix by cloning `element`.
908    ///
909    /// Elements are initialized in memory order.
910    ///
911    /// ```
912    /// use diskann_utils::views::rowmajor::{self, Matrix};
913    ///
914    /// let mat = rowmajor::Owned::from_element(2, 3, 0u32);
915    ///
916    /// assert_eq!(mat.row(0), &[0, 0, 0]);
917    /// assert_eq!(mat.row(1), &[0, 0, 0]);
918    /// ```
919    ///
920    /// # Panics
921    ///
922    /// Panics if `nrows * ncols` overflows `usize::MAX`, or if the allocation size exceeds
923    /// `isize::MAX`.
924    #[track_caller]
925    pub fn from_element(nrows: usize, ncols: usize, element: T) -> Self
926    where
927        T: Clone,
928    {
929        match Self::try_from_element(nrows, ncols, element) {
930            Ok(matrix) => matrix,
931            Err(error) => panic!("Owned::from_element failed with: {error}"),
932        }
933    }
934
935    /// Construct a new matrix by cloning `element`.
936    ///
937    /// Elements are initialized in memory order.
938    ///
939    /// ```
940    /// use diskann_utils::views::rowmajor::{self, Matrix};
941    ///
942    /// let mat = rowmajor::Owned::try_from_element(2, 3, 0u32).unwrap();
943    ///
944    /// assert_eq!(mat.row(0), &[0, 0, 0]);
945    /// assert_eq!(mat.row(1), &[0, 0, 0]);
946    /// ```
947    ///
948    /// # Errors
949    ///
950    /// Returns an error if `nrows * ncols` overflows `usize::MAX`, or if the allocation size
951    /// exceeds `isize::MAX`.
952    pub fn try_from_element(nrows: usize, ncols: usize, element: T) -> Result<Self, LayoutError>
953    where
954        T: Clone,
955    {
956        let layout = Layout::new(nrows, ncols)?;
957        Ok(Self::from_element_with_layout(layout, element))
958    }
959
960    // Less common constructors.
961
962    /// Construct a new matrix using `init`.
963    ///
964    /// Elements are initialized in memory order.
965    pub fn from_fn_with_layout<F>(layout: Layout<T>, mut init: F) -> Self
966    where
967        F: FnMut(RowCol) -> T,
968    {
969        let mut row = 0;
970        let mut col = 0;
971
972        let data: Box<[T]> = (0..layout.num_elements())
973            .map(|_| {
974                let v = (init)(RowCol { row, col });
975                col += 1;
976                if col == layout.ncols() {
977                    col = 0;
978                    row += 1;
979                }
980                v
981            })
982            .collect();
983
984        // SAFETY: We constructed `data` to have length exactly `layout.num_elements()`.
985        unsafe { Self::from_data_unchecked(data, layout) }
986    }
987
988    /// Construct a new matrix by cloning `element`.
989    ///
990    /// Elements are initialized in memory order.
991    pub fn from_element_with_layout(layout: Layout<T>, element: T) -> Self
992    where
993        T: Clone,
994    {
995        let data: Box<[T]> = std::iter::repeat_n(element, layout.num_elements()).collect();
996
997        // SAFETY: We constructed `data` to have length exactly `layout.num_elements()`.
998        unsafe { Self::from_data_unchecked(data, layout) }
999    }
1000
1001    /// # Safety
1002    ///
1003    /// `b.len()` must equal `layout.num_elements()`.
1004    unsafe fn from_data_unchecked(b: Box<[T]>, layout: Layout<T>) -> Self {
1005        debug_assert_eq!(b.len(), layout.num_elements());
1006        Self {
1007            ptr: internal::box_to_nonnull(b),
1008            layout,
1009        }
1010    }
1011
1012    /// Consume `self`, returning the unmodified contents as a boxed slice.
1013    ///
1014    /// ```
1015    /// use diskann_utils::views::rowmajor::{Matrix, Owned};
1016    ///
1017    /// let mat = Owned::from_fn(2, 3, |rc| rc.col);
1018    /// assert_eq!(mat.row(0), &[0, 1, 2]);
1019    /// assert_eq!(mat.row(1), &[0, 1, 2]);
1020    ///
1021    /// let b: Box<[usize]> = mat.into_inner();
1022    /// assert_eq!(&*b, &[0, 1, 2, 0, 1, 2]);
1023    /// ```
1024    pub fn into_inner(self) -> Box<[T]> {
1025        let me = ManuallyDrop::new(self);
1026
1027        // SAFETY: `me.ptr` came from exactly one `Box<[T]>` of this length. Suppressing
1028        // `Owned::drop` transfers that allocation back to the reconstructed box.
1029        unsafe { internal::nonnull_to_box(me.ptr, me.layout.num_elements()) }
1030    }
1031}
1032
1033// SAFETY: `Owned<T>` has the ownership semantics of `Box<[T]>`, which is `Send` when `T` is.
1034unsafe impl<T> Send for Owned<T> where T: Send {}
1035// SAFETY: Shared access to `Owned<T>` exposes only shared access to `T`.
1036unsafe impl<T> Sync for Owned<T> where T: Sync {}
1037
1038impl<T> Drop for Owned<T> {
1039    fn drop(&mut self) {
1040        // SAFETY: `self.ptr` came from exactly one `Box<[T]>` of this length, and `drop`
1041        // is the unique place that reconstructs it.
1042        let _ = unsafe { internal::nonnull_to_box(self.ptr, self.layout.num_elements()) };
1043    }
1044}
1045
1046impl<T> Clone for Owned<T>
1047where
1048    T: Clone,
1049{
1050    fn clone(&self) -> Self {
1051        // SAFETY: The constructed boxed slice has length exactly `self.layout.num_elements()`.
1052        unsafe { Owned::from_data_unchecked(self.as_slice().into(), self.layout()) }
1053    }
1054}
1055
1056// SAFETY: `Owned` keeps a stable pointer and layout to its initialized boxed slice.
1057unsafe impl<T> Matrix for Owned<T> {
1058    type Element = T;
1059
1060    fn as_nonnull(&self) -> NonNull<T> {
1061        self.ptr
1062    }
1063
1064    fn layout(&self) -> Layout<T> {
1065        self.layout
1066    }
1067}
1068
1069// SAFETY: A mutable borrow of `Owned` has exclusive access to its boxed slice.
1070unsafe impl<T> MatrixMut for Owned<T> {
1071    fn as_nonnull_mut(&mut self) -> NonNull<T> {
1072        self.ptr
1073    }
1074}
1075
1076impl<T> PartialEq for Owned<T>
1077where
1078    T: PartialEq,
1079{
1080    fn eq(&self, other: &Self) -> bool {
1081        Matrix::as_view(self).eq(&Matrix::as_view(other))
1082    }
1083}
1084
1085impl<'a, T> Reborrow<'a> for Owned<T> {
1086    type Target = Ref<'a, T>;
1087    fn reborrow(&'a self) -> Self::Target {
1088        Matrix::as_view(self)
1089    }
1090}
1091
1092impl<'a, T> ReborrowMut<'a> for Owned<T> {
1093    type Target = Mut<'a, T>;
1094    fn reborrow_mut(&'a mut self) -> Self::Target {
1095        MatrixMut::as_view_mut(self)
1096    }
1097}
1098
1099//-----//
1100// Ref //
1101//-----//
1102
1103/// A [`Matrix`] implementation that references its data.
1104#[derive(Debug)]
1105pub struct Ref<'a, T> {
1106    ptr: NonNull<T>,
1107    layout: Layout<T>,
1108    _lifetime: PhantomData<&'a [T]>,
1109}
1110
1111// SAFETY: `Ref<'_, T>` has the ownership semantics of `&[T]`, which is `Send` when `T` is
1112// `Sync`.
1113unsafe impl<T> Send for Ref<'_, T> where T: Sync {}
1114// SAFETY: `Ref<'_, T>` exposes only shared access to `T`.
1115unsafe impl<T> Sync for Ref<'_, T> where T: Sync {}
1116
1117impl<'a, T> Ref<'a, T> {
1118    constructors!(T, &'a [T]);
1119
1120    /// # Safety
1121    ///
1122    /// `b.len()` must equal `layout.num_elements()`.
1123    unsafe fn from_data_unchecked(b: &'a [T], layout: Layout<T>) -> Self {
1124        debug_assert_eq!(b.len(), layout.num_elements());
1125        Self {
1126            ptr: internal::slice_to_nonnull(b),
1127            layout,
1128            _lifetime: PhantomData,
1129        }
1130    }
1131
1132    /// Return the contents of `self` as a slice.
1133    ///
1134    /// Unlike [`Matrix::as_slice`], the returned slices inherits the lifetime of the [`Ref`].
1135    pub fn into_slice(self) -> &'a [T] {
1136        // SAFETY: `Ref` represents a valid `&'a [T]` of exactly this length.
1137        unsafe { std::slice::from_raw_parts(self.as_ptr(), self.layout().num_elements()) }
1138    }
1139
1140    /// Get the indicated `row` with the lifetime `'a`.
1141    ///
1142    /// # Safety
1143    ///
1144    /// `row < self.nrows()`.
1145    #[cfg(feature = "rayon")]
1146    unsafe fn into_row_unchecked(self, row: usize) -> &'a [T] {
1147        let layout = self.layout();
1148        debug_assert!(row < layout.nrows());
1149
1150        // SAFETY: The caller guarantees that `row` is in-bounds, so the validated layout
1151        // places this row within the `&'a [T]` represented by `self`.
1152        unsafe {
1153            std::slice::from_raw_parts(self.as_ptr().add(layout.ncols() * row), layout.ncols())
1154        }
1155    }
1156
1157    /// Return a `Ref` containing the indicated rows with the lifetime `'a`.
1158    ///
1159    /// # Safety
1160    ///
1161    /// `rows.start <= rows.end` and `rows.end <= self.nrows()`.
1162    #[cfg(feature = "rayon")]
1163    unsafe fn into_subview_unchecked(self, rows: std::ops::Range<usize>) -> Ref<'a, T> {
1164        debug_assert!(rows.start <= rows.end);
1165        debug_assert!(rows.end <= self.nrows());
1166
1167        let ncols = self.ncols();
1168        Self {
1169            // SAFETY: The validated range starts within or one past the matrix span, and
1170            // the parent layout guarantees that the offset is representable.
1171            ptr: unsafe { self.ptr.add(rows.start * ncols) },
1172            // SAFETY: This layout is no larger than `self`'s layout.
1173            layout: unsafe { Layout::new_unchecked(rows.end - rows.start, ncols) },
1174            _lifetime: PhantomData,
1175        }
1176    }
1177}
1178
1179impl<T> Clone for Ref<'_, T> {
1180    fn clone(&self) -> Self {
1181        *self
1182    }
1183}
1184
1185impl<T> Copy for Ref<'_, T> {}
1186
1187// SAFETY: `Ref` keeps a stable pointer and layout to the initialized slice it borrows.
1188unsafe impl<T> Matrix for Ref<'_, T> {
1189    type Element = T;
1190
1191    fn as_nonnull(&self) -> NonNull<T> {
1192        self.ptr
1193    }
1194
1195    fn layout(&self) -> Layout<T> {
1196        self.layout
1197    }
1198}
1199
1200impl<'a, T> Reborrow<'a> for Ref<'_, T> {
1201    type Target = Ref<'a, T>;
1202    fn reborrow(&'a self) -> Self::Target {
1203        Matrix::as_view(self)
1204    }
1205}
1206
1207impl<T> PartialEq for Ref<'_, T>
1208where
1209    T: PartialEq,
1210{
1211    fn eq(&self, other: &Self) -> bool {
1212        self.layout() == other.layout() && self.as_slice() == other.as_slice()
1213    }
1214}
1215
1216//-----//
1217// Mut //
1218//-----//
1219
1220/// A [`Matrix`]/[`MatrixMut`] implementation that mutably references its data.
1221#[derive(Debug)]
1222pub struct Mut<'a, T> {
1223    ptr: NonNull<T>,
1224    layout: Layout<T>,
1225    _lifetime: PhantomData<&'a mut [T]>,
1226}
1227
1228// SAFETY: `Mut<'_, T>` has the ownership semantics of `&mut [T]`, which is `Send` when
1229// `T` is `Send`.
1230unsafe impl<T> Send for Mut<'_, T> where T: Send {}
1231// SAFETY: Shared access to `Mut<'_, T>` exposes only shared access to `T`.
1232unsafe impl<T> Sync for Mut<'_, T> where T: Sync {}
1233
1234impl<'a, T> Mut<'a, T> {
1235    constructors!(T, &'a mut [T]);
1236
1237    /// # Safety
1238    ///
1239    /// `b.len()` must equal `layout.num_elements()`.
1240    unsafe fn from_data_unchecked(b: &'a mut [T], layout: Layout<T>) -> Self {
1241        debug_assert_eq!(b.len(), layout.num_elements());
1242        Self {
1243            ptr: internal::mut_slice_to_nonnull(b),
1244            layout,
1245            _lifetime: PhantomData,
1246        }
1247    }
1248
1249    /// Consume `self` and return the underlying data as a mutable slice.
1250    pub fn into_mut_slice(self) -> &'a mut [T] {
1251        // SAFETY: `self.ptr` and `self.layout` together describe a valid `&'a mut [T]` of
1252        // length `self.layout.num_elements()`, per the invariants of `Mut`.
1253        unsafe { std::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.layout.num_elements()) }
1254    }
1255}
1256
1257// SAFETY: `Mut` keeps a stable pointer and layout to the initialized slice it borrows.
1258unsafe impl<T> Matrix for Mut<'_, T> {
1259    type Element = T;
1260
1261    fn as_nonnull(&self) -> NonNull<T> {
1262        self.ptr
1263    }
1264
1265    fn layout(&self) -> Layout<T> {
1266        self.layout
1267    }
1268}
1269
1270// SAFETY: A mutable borrow of `Mut` has exclusive access to its borrowed slice.
1271unsafe impl<T> MatrixMut for Mut<'_, T> {
1272    fn as_nonnull_mut(&mut self) -> NonNull<T> {
1273        self.ptr
1274    }
1275}
1276
1277impl<T> PartialEq for Mut<'_, T>
1278where
1279    T: PartialEq,
1280{
1281    fn eq(&self, other: &Self) -> bool {
1282        Matrix::as_view(self).eq(&Matrix::as_view(other))
1283    }
1284}
1285
1286impl<'a, T> Reborrow<'a> for Mut<'_, T> {
1287    type Target = Ref<'a, T>;
1288    fn reborrow(&'a self) -> Self::Target {
1289        Matrix::as_view(self)
1290    }
1291}
1292
1293impl<'a, T> ReborrowMut<'a> for Mut<'_, T> {
1294    type Target = Mut<'a, T>;
1295    fn reborrow_mut(&'a mut self) -> Self::Target {
1296        MatrixMut::as_view_mut(self)
1297    }
1298}
1299
1300//--------//
1301// Errors //
1302//--------//
1303
1304/// Errors from [`Owned::try_from_data`], [`Ref::try_from_data`], and [`Mut::try_from_data`].
1305pub struct TryFromError<T> {
1306    data: T,
1307    inner: TryFromErrorInner,
1308}
1309
1310impl<T> TryFromError<T> {
1311    /// Consume the error and return the base data.
1312    pub fn into_inner(self) -> T {
1313        self.data
1314    }
1315
1316    /// Return a variation of `Self` that is guaranteed to be `'static` by removing the
1317    /// data that was passed to the original constructor.
1318    pub fn as_static(&self) -> TryFromErrorLight {
1319        TryFromErrorLight(self.inner)
1320    }
1321
1322    //--------------//
1323    // Constructors //
1324    //--------------//
1325
1326    fn layout(data: T, error: LayoutError) -> Self {
1327        Self {
1328            data,
1329            inner: TryFromErrorInner::Layout(error),
1330        }
1331    }
1332
1333    fn mismatch(data: T, nrows: usize, ncols: usize, len: usize) -> Self {
1334        Self {
1335            data,
1336            inner: TryFromErrorInner::Mismatch { nrows, ncols, len },
1337        }
1338    }
1339}
1340
1341impl<T> std::fmt::Debug for TryFromError<T> {
1342    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1343        f.debug_struct("TryFromError")
1344            .field("data", &"<hidden>")
1345            .field("inner", &self.inner)
1346            .finish()
1347    }
1348}
1349
1350impl<T> std::fmt::Display for TryFromError<T> {
1351    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1352        self.inner.fmt(f)
1353    }
1354}
1355
1356impl<T> std::error::Error for TryFromError<T> {}
1357
1358/// A guaranteed `'static` version of [`TryFromError`].
1359#[derive(Debug, Error)]
1360#[error(transparent)]
1361pub struct TryFromErrorLight(TryFromErrorInner);
1362
1363#[derive(Debug, Error, Clone, Copy)]
1364enum TryFromErrorInner {
1365    #[error(transparent)]
1366    Layout(LayoutError),
1367    #[error(
1368        "tried to construct a {}x{} matrix over a span of length {}",
1369        nrows,
1370        ncols,
1371        len
1372    )]
1373    Mismatch {
1374        nrows: usize,
1375        ncols: usize,
1376        len: usize,
1377    },
1378}
1379
1380///////////
1381// Tests //
1382///////////
1383
1384#[cfg(test)]
1385mod tests {
1386    use super::*;
1387    use crate::{assert_contains, lazy_format};
1388
1389    /// This function is only callable with copyable types.
1390    ///
1391    /// This lets us test for types we expect to be `Copy`.
1392    fn is_copyable<T: Copy>(_x: T) -> bool {
1393        true
1394    }
1395
1396    /// This function attests that `Ref` is covariant in the view lifetime.
1397    fn _matrix_view_is_covariant<'a, 'b>(m: Ref<'a, f32>) -> Ref<'b, f32>
1398    where
1399        'a: 'b,
1400    {
1401        m
1402    }
1403
1404    fn _matrix_view_is_covariant_in_t<'a, 'b, 'm>(m: Ref<'m, &'a f32>) -> Ref<'m, &'b f32>
1405    where
1406        'a: 'b,
1407    {
1408        m
1409    }
1410
1411    fn _matrix_is_covariant_in_t<'a, 'b, 'm>(m: &'m Owned<&'a f32>) -> &'m Owned<&'b f32>
1412    where
1413        'a: 'b,
1414    {
1415        m
1416    }
1417
1418    //--------//
1419    // Layout //
1420    //--------//
1421
1422    #[test]
1423    fn test_layout() {
1424        // Happy path
1425        for rows in 0..5 {
1426            for cols in 0..5 {
1427                let layout = Layout::<String>::new(rows, cols).unwrap();
1428                assert_eq!(layout.nrows(), rows);
1429                assert_eq!(layout.ncols(), cols);
1430                assert_eq!(layout.num_elements(), rows * cols);
1431
1432                let transpose = layout.transpose();
1433                assert_eq!(transpose.nrows(), cols);
1434                assert_eq!(transpose.ncols(), rows);
1435                assert_eq!(transpose.num_elements(), rows * cols);
1436
1437                let rebind = layout.rebind::<u32>().unwrap();
1438                assert_eq!(rebind.nrows(), rows);
1439                assert_eq!(rebind.ncols(), cols);
1440                assert_eq!(rebind.num_elements(), rows * cols);
1441
1442                is_copyable(layout);
1443            }
1444        }
1445
1446        #[expect(unused, reason = "we need this so the size is non-zero")]
1447        struct NotDebugOrEq(u32);
1448
1449        assert_eq!(
1450            Layout::<NotDebugOrEq>::new(10, 20).unwrap(),
1451            Layout::<NotDebugOrEq>::new(10, 20).unwrap(),
1452        );
1453
1454        assert_eq!(
1455            Layout::<NotDebugOrEq>::new(20, 0).unwrap(),
1456            Layout::<NotDebugOrEq>::new(20, 0).unwrap(),
1457        );
1458
1459        assert_ne!(
1460            Layout::<NotDebugOrEq>::new(10, 20).unwrap(),
1461            Layout::<NotDebugOrEq>::new(20, 0).unwrap(),
1462        );
1463
1464        let fmt = format!("{:?}", Layout::<NotDebugOrEq>::new(5, 6).unwrap());
1465        assert_eq!(fmt, "Layout { nrows: 5, ncols: 6, elsize: 4 }");
1466
1467        // Overflowing the element count returns an error.
1468        let error = Layout::<u8>::new(usize::MAX, 2).unwrap_err();
1469        assert_eq!(
1470            error.to_string(),
1471            format!(
1472                "a matrix of size {}x2 has a length exceeding `usize::MAX`",
1473                usize::MAX
1474            )
1475        );
1476
1477        // The largest possible byte span is valid without allocating it.
1478        let layout = Layout::<u8>::new(isize::MAX as usize, 1).unwrap();
1479        assert_eq!(layout.num_elements(), isize::MAX as usize);
1480
1481        let transpose = layout.transpose();
1482        assert_eq!(transpose.nrows(), 1);
1483        assert_eq!(transpose.ncols(), isize::MAX as usize);
1484        assert_eq!(transpose.num_elements(), layout.num_elements());
1485
1486        // One byte beyond the maximum span returns an error.
1487        let error = Layout::<u8>::new(isize::MAX as usize + 1, 1).unwrap_err();
1488        assert_eq!(
1489            error.to_string(),
1490            format!(
1491                "a matrix of size {}x1 with elements of size 1 exceeds `isize::MAX` bytes",
1492                isize::MAX as usize + 1
1493            )
1494        );
1495
1496        // Rebinding to a larger element type revalidates the byte span.
1497        let rebound = Layout::<u8>::new(3, 4).unwrap().rebind::<u16>().unwrap();
1498        assert_eq!(rebound.nrows(), 3);
1499        assert_eq!(rebound.ncols(), 4);
1500        assert_eq!(rebound.num_elements(), 12);
1501
1502        let error = layout.rebind::<u16>().unwrap_err();
1503        assert_eq!(
1504            error.to_string(),
1505            format!(
1506                "a matrix of size {}x1 with elements of size 2 exceeds `isize::MAX` bytes",
1507                isize::MAX
1508            )
1509        );
1510    }
1511
1512    ////////////
1513    // Matrix //
1514    ////////////
1515
1516    #[test]
1517    fn test_sizes() {
1518        let expected = 3 * std::mem::size_of::<usize>();
1519        assert_eq!(std::mem::size_of::<Owned<String>>(), expected);
1520        assert_eq!(std::mem::size_of::<Option<Owned<String>>>(), expected);
1521
1522        assert_eq!(std::mem::size_of::<Ref<'_, String>>(), expected);
1523        assert_eq!(std::mem::size_of::<Option<Ref<'_, String>>>(), expected);
1524
1525        assert_eq!(std::mem::size_of::<Mut<'_, String>>(), expected);
1526        assert_eq!(std::mem::size_of::<Option<Mut<'_, String>>>(), expected);
1527    }
1528
1529    #[test]
1530    fn fallible_matrix_constructors() {
1531        let err = Owned::try_from_element(usize::MAX, usize::MAX, 0u32).unwrap_err();
1532        let msg = err.to_string();
1533        assert_contains!(msg, "exceeding `usize::MAX`");
1534
1535        let err = Owned::try_from_element(isize::MAX as usize, 1, 0u32).unwrap_err();
1536        let msg = err.to_string();
1537        assert_contains!(msg, "exceeds `isize::MAX` bytes");
1538
1539        // Panicking
1540        let err = std::panic::catch_unwind(|| {
1541            Owned::from_element(usize::MAX, usize::MAX, 0u32);
1542        })
1543        .unwrap_err()
1544        .downcast::<String>()
1545        .unwrap();
1546
1547        let msg = err.to_string();
1548        assert_contains!(msg, "exceeding `usize::MAX`");
1549
1550        let err = std::panic::catch_unwind(|| {
1551            Owned::from_element(isize::MAX as usize, 1, 0u32);
1552        })
1553        .unwrap_err()
1554        .downcast::<String>()
1555        .unwrap();
1556        let msg = err.to_string();
1557        assert_contains!(msg, "exceeds `isize::MAX` bytes");
1558
1559        // Construction fails without invoking the generator.
1560        let err = Owned::try_from_fn(usize::MAX, usize::MAX, |_| panic!("boom")).unwrap_err();
1561        let msg = err.to_string();
1562        assert_contains!(msg, "exceeding `usize::MAX`");
1563
1564        let err = std::panic::catch_unwind(|| {
1565            Owned::from_fn(usize::MAX, usize::MAX, |_| {
1566                panic!("initializer must not run")
1567            });
1568        })
1569        .unwrap_err()
1570        .downcast::<String>()
1571        .unwrap();
1572        let msg = err.to_string();
1573        assert_contains!(msg, "Owned::from_fn failed");
1574        assert_contains!(msg, "exceeding `usize::MAX`");
1575    }
1576
1577    fn make_test_matrix() -> Vec<usize> {
1578        // Construct a matrix with 4 rows of length 3.
1579        // The expected layout is as follows:
1580        //
1581        // 0, 1, 2,
1582        // 1, 2, 3,
1583        // 2, 3, 4,
1584        // 3, 4, 5
1585        //
1586        vec![0, 1, 2, 1, 2, 3, 2, 3, 4, 3, 4, 5]
1587    }
1588
1589    fn striped_matrix(nrows: usize, ncols: usize) -> Owned<usize> {
1590        Owned::from_fn(nrows, ncols, |rc| rc.row * 100 + rc.col)
1591    }
1592
1593    fn assert_exact_fused<I>(mut iter: I, expected: usize)
1594    where
1595        I: ExactSizeIterator + std::iter::FusedIterator,
1596    {
1597        assert_eq!(iter.len(), expected);
1598        assert_eq!(iter.size_hint(), (expected, Some(expected)));
1599
1600        for remaining in (0..expected).rev() {
1601            assert!(iter.next().is_some());
1602            assert_eq!(iter.len(), remaining);
1603            assert_eq!(iter.size_hint(), (remaining, Some(remaining)));
1604        }
1605
1606        assert!(iter.next().is_none());
1607        assert!(iter.next().is_none());
1608        assert_eq!(iter.len(), 0);
1609        assert_eq!(iter.size_hint(), (0, Some(0)));
1610    }
1611
1612    /// Assert that the collection of rows matches the rows in `m`.
1613    fn assert_rows_match_scalar(m: Ref<'_, usize>, rows: Vec<&[usize]>) {
1614        let context = lazy_format!("nrows = {}, ncols = {}", m.nrows(), m.ncols());
1615
1616        assert_eq!(rows.len(), m.nrows(), "{context}");
1617
1618        for (row_index, row) in rows.into_iter().enumerate() {
1619            assert_eq!(row.len(), m.ncols(), "row = {row_index} -- {context}");
1620            for (col_index, value) in row.iter().enumerate() {
1621                assert_eq!(
1622                    value,
1623                    m.element(row_index, col_index),
1624                    "row = {row_index}, col = {col_index} -- {context}"
1625                );
1626            }
1627        }
1628    }
1629
1630    fn assert_windows_match_scalar(
1631        m: Ref<'_, usize>,
1632        batchsize: usize,
1633        windows: Vec<Ref<'_, usize>>,
1634    ) {
1635        let context = lazy_format!(
1636            "nrows = {}, ncols = {}, batchsize = {batchsize}",
1637            m.nrows(),
1638            m.ncols()
1639        );
1640
1641        assert_eq!(windows.len(), m.nrows().div_ceil(batchsize), "{context}");
1642
1643        for (window_index, window) in windows.into_iter().enumerate() {
1644            let window_context = lazy_format!("window = {window_index} -- {context}");
1645            let first_row = window_index * batchsize;
1646            let expected_rows = batchsize.min(m.nrows() - first_row);
1647            assert_eq!(window.nrows(), expected_rows, "{window_context}");
1648            assert_eq!(window.ncols(), m.ncols(), "{window_context}");
1649
1650            for row_index in 0..window.nrows() {
1651                for col_index in 0..window.ncols() {
1652                    assert_eq!(
1653                        window.element(row_index, col_index),
1654                        m.element(first_row + row_index, col_index),
1655                        "row = {row_index}, col = {col_index} -- {window_context}"
1656                    );
1657                }
1658            }
1659        }
1660    }
1661
1662    #[cfg(all(not(miri), feature = "rayon"))]
1663    fn assert_parallel_rows_match_scalar(m: Ref<'_, usize>) {
1664        let rows: Vec<_> = m.par_rows().collect();
1665        assert_rows_match_scalar(m, rows);
1666    }
1667
1668    #[cfg(all(not(miri), feature = "rayon"))]
1669    fn assert_parallel_windows_match_scalar(m: Ref<'_, usize>, batchsize: usize) {
1670        let windows: Vec<_> = m.par_window_iter(batchsize).collect();
1671        assert_windows_match_scalar(m, batchsize, windows);
1672    }
1673
1674    /// This expects `m` to look like the result s of [`make_test_matrix`].
1675    fn test_basic_indexing<T>(m: &T)
1676    where
1677        T: Matrix<Element = usize> + Sync,
1678    {
1679        assert_eq!(m.nrows(), 4);
1680        assert_eq!(m.ncols(), 3);
1681
1682        // Basic indexing
1683        assert_eq!(*m.element(0, 0), 0);
1684        assert_eq!(*m.element(0, 1), 1);
1685        assert_eq!(*m.element(0, 2), 2);
1686
1687        assert_eq!(*m.element(1, 0), 1);
1688        assert_eq!(*m.element(1, 1), 2);
1689        assert_eq!(*m.element(1, 2), 3);
1690
1691        assert_eq!(*m.element(2, 0), 2);
1692        assert_eq!(*m.element(2, 1), 3);
1693        assert_eq!(*m.element(2, 2), 4);
1694
1695        assert_eq!(*m.element(3, 0), 3);
1696        assert_eq!(*m.element(3, 1), 4);
1697        assert_eq!(*m.element(3, 2), 5);
1698
1699        assert_eq!(*m.get_element(0, 0).unwrap(), 0);
1700        assert_eq!(*m.get_element(0, 1).unwrap(), 1);
1701        assert_eq!(*m.get_element(0, 2).unwrap(), 2);
1702
1703        assert_eq!(*m.get_element(1, 0).unwrap(), 1);
1704        assert_eq!(*m.get_element(1, 1).unwrap(), 2);
1705        assert_eq!(*m.get_element(1, 2).unwrap(), 3);
1706
1707        assert_eq!(*m.get_element(2, 0).unwrap(), 2);
1708        assert_eq!(*m.get_element(2, 1).unwrap(), 3);
1709        assert_eq!(*m.get_element(2, 2).unwrap(), 4);
1710
1711        assert_eq!(*m.get_element(3, 0).unwrap(), 3);
1712        assert_eq!(*m.get_element(3, 1).unwrap(), 4);
1713        assert_eq!(*m.get_element(3, 2).unwrap(), 5);
1714
1715        // Row indexing.
1716        assert_eq!(m.row(0), &[0, 1, 2]);
1717        assert_eq!(m.row(1), &[1, 2, 3]);
1718        assert_eq!(m.row(2), &[2, 3, 4]);
1719        assert_eq!(m.row(3), &[3, 4, 5]);
1720
1721        let rows: Vec<Vec<usize>> = m.rows().map(|x| x.to_vec()).collect();
1722        assert_eq!(m.row(0), &rows[0]);
1723        assert_eq!(m.row(1), &rows[1]);
1724        assert_eq!(m.row(2), &rows[2]);
1725        assert_eq!(m.row(3), &rows[3]);
1726
1727        // Window Iters.
1728        let batchsize = 2;
1729        m.window_iter(NonZeroUsize::new(batchsize).unwrap())
1730            .enumerate()
1731            .for_each(|(i, submatrix)| {
1732                assert_eq!(submatrix.nrows(), batchsize);
1733                assert_eq!(submatrix.ncols(), m.ncols());
1734
1735                // Make sure we are in the correct window of the original matrix.
1736                let base = i * batchsize;
1737                assert_eq!(*submatrix.element(0, 0), base);
1738                assert_eq!(*submatrix.element(0, 1), base + 1);
1739                assert_eq!(*submatrix.element(0, 2), base + 2);
1740
1741                assert_eq!(*submatrix.element(1, 0), base + 1);
1742                assert_eq!(*submatrix.element(1, 1), base + 2);
1743                assert_eq!(*submatrix.element(1, 2), base + 3);
1744            });
1745
1746        // Try again, but with a batch size of 3 to ensure that we correctly handle cases
1747        // where the last block is under-sized.
1748        let batchsize = 3;
1749        m.window_iter(NonZeroUsize::new(batchsize).unwrap())
1750            .enumerate()
1751            .for_each(|(i, submatrix)| {
1752                if i == 0 {
1753                    assert_eq!(submatrix.nrows(), batchsize);
1754                    assert_eq!(submatrix.ncols(), m.ncols());
1755
1756                    // Check indexing
1757                    assert_eq!(*submatrix.element(0, 0), 0);
1758                    assert_eq!(*submatrix.element(0, 1), 1);
1759                    assert_eq!(*submatrix.element(0, 2), 2);
1760
1761                    assert_eq!(*submatrix.element(1, 0), 1);
1762                    assert_eq!(*submatrix.element(1, 1), 2);
1763                    assert_eq!(*submatrix.element(1, 2), 3);
1764
1765                    assert_eq!(*submatrix.element(2, 0), 2);
1766                    assert_eq!(*submatrix.element(2, 1), 3);
1767                    assert_eq!(*submatrix.element(2, 2), 4);
1768                } else {
1769                    assert_eq!(submatrix.nrows(), 1);
1770                    assert_eq!(submatrix.ncols(), m.ncols());
1771
1772                    // Check indexing
1773                    assert_eq!(*submatrix.element(0, 0), 3);
1774                    assert_eq!(*submatrix.element(0, 1), 4);
1775                    assert_eq!(*submatrix.element(0, 2), 5);
1776                }
1777            });
1778    }
1779
1780    #[test]
1781    fn matrix_happy_path() {
1782        let data = make_test_matrix();
1783        let m = Owned::try_from_data(data.into(), 4, 3).unwrap();
1784        test_basic_indexing(&m);
1785
1786        // Get the base pointer of the matrix and make sure view-conversion preserves this
1787        // value.
1788        let ptr = m.as_ptr();
1789        let view = m.as_view();
1790        assert!(is_copyable(view));
1791        assert_eq!(view.as_ptr(), ptr);
1792        assert_eq!(view.nrows(), m.nrows());
1793        assert_eq!(view.ncols(), m.ncols());
1794        test_basic_indexing(&view);
1795    }
1796
1797    #[test]
1798    fn matrix_try_from_construction_error() {
1799        let data = make_test_matrix();
1800        let ptr = data.as_ptr();
1801        let len = data.len();
1802
1803        let m = Owned::try_from_data(data.into(), 5, 4);
1804        assert!(m.is_err());
1805        let err = m.unwrap_err();
1806        assert_eq!(
1807            err.to_string(),
1808            "tried to construct a 5x4 matrix over a span of length 12"
1809        );
1810
1811        // Make sure that we can retrieve the original allocation from the interior.
1812        let data = err.into_inner();
1813        assert_eq!(data.as_ptr(), ptr);
1814        assert_eq!(data.len(), len);
1815
1816        let m = Ref::try_from_data(&data, 5, 4);
1817        assert!(m.is_err());
1818        assert_eq!(
1819            m.unwrap_err().to_string(),
1820            "tried to construct a 5x4 matrix over a span of length 12"
1821        );
1822    }
1823
1824    #[test]
1825    fn mutable_matrix_direct_construction() {
1826        let mut data = make_test_matrix();
1827
1828        {
1829            let mut m = Mut::try_from_data(data.as_mut_slice(), 4, 3).unwrap();
1830            *m.element_mut(1, 2) = 30;
1831        }
1832        assert_eq!(data[5], 30);
1833
1834        let err = Mut::try_from_data(data.as_mut_slice(), 5, 4).unwrap_err();
1835        assert_eq!(
1836            err.to_string(),
1837            "tried to construct a 5x4 matrix over a span of length 12"
1838        );
1839        let recovered = err.into_inner();
1840        recovered[0] = 10;
1841        assert_eq!(data[0], 10);
1842    }
1843
1844    #[test]
1845    fn matrix_mut_view() {
1846        let mut m = Owned::<usize>::from_element(4, 3, 0);
1847        assert_eq!(m.nrows(), 4);
1848        assert_eq!(m.ncols(), 3);
1849        assert!(m.as_slice().iter().all(|&i| i == 0));
1850        let ptr = m.as_ptr();
1851        let mut_ptr = m.as_mut_ptr();
1852        assert_eq!(ptr, mut_ptr);
1853
1854        let mut view = m.as_view_mut();
1855        assert_eq!(view.nrows(), 4);
1856        assert_eq!(view.ncols(), 3);
1857        assert_eq!(view.as_ptr(), ptr);
1858        assert_eq!(view.as_mut_ptr(), mut_ptr);
1859
1860        // Construct the test matrix manually.
1861        for i in 0..view.nrows() {
1862            for j in 0..view.ncols() {
1863                *view.element_mut(i, j) = i + j;
1864            }
1865        }
1866
1867        // Drop the view and test the original matrix.
1868        test_basic_indexing(&m);
1869
1870        // Check `PartialEq`.
1871        let mut m_clone = m.clone();
1872        assert_eq!(m.as_view_mut(), m_clone.as_view_mut());
1873
1874        let inner = m.into_inner();
1875        assert_eq!(inner.as_ptr(), ptr);
1876        assert_eq!(inner.len(), 4 * 3);
1877    }
1878
1879    #[test]
1880    fn matrix_view_zero_sizes() {
1881        let data: Vec<usize> = vec![];
1882        // Zero rows, but non-zero columns.
1883        let m = Ref::try_from_data(data.as_slice(), 0, 10).unwrap();
1884        assert_eq!(m.nrows(), 0);
1885        assert_eq!(m.ncols(), 10);
1886
1887        // Non-zero rows, but zero columns.
1888        let m = Ref::try_from_data(data.as_slice(), 3, 0).unwrap();
1889        assert_eq!(m.nrows(), 3);
1890        assert_eq!(m.ncols(), 0);
1891        let empty: &[usize] = &[];
1892        assert_eq!(m.row(0), empty);
1893        assert_eq!(m.row(1), empty);
1894        assert_eq!(m.row(2), empty);
1895
1896        // Zero rows and columns.
1897        let m = Ref::try_from_data(data.as_slice(), 0, 0).unwrap();
1898        assert_eq!(m.nrows(), 0);
1899        assert_eq!(m.ncols(), 0);
1900    }
1901
1902    #[test]
1903    fn matrix_construction_by_row() {
1904        let mut m = Owned::<usize>::from_element(4, 3, 0);
1905        assert!(m.as_slice().iter().all(|i| *i == 0));
1906
1907        let ncols = m.ncols();
1908        for i in 0..m.nrows() {
1909            let row = m.row_mut(i);
1910            assert_eq!(row.len(), ncols);
1911            row[0] = i;
1912            row[1] = i + 1;
1913            row[2] = i + 2;
1914        }
1915        test_basic_indexing(&m);
1916    }
1917
1918    // Panics
1919    #[test]
1920    #[should_panic(expected = "tried to access row 3 of a matrix with 3 rows")]
1921    fn test_get_row_panics() {
1922        let m = Owned::<usize>::from_element(3, 7, 0);
1923        m.row(3);
1924    }
1925
1926    #[test]
1927    #[should_panic(expected = "tried to access row 3 of a matrix with 3 rows")]
1928    fn test_get_row_mut_panics() {
1929        let mut m = Owned::<usize>::from_element(3, 7, 0);
1930        m.row_mut(3);
1931    }
1932
1933    #[test]
1934    #[should_panic(expected = "row 3 is out of bounds (max: 3)")]
1935    fn test_element_panics_row() {
1936        let m = Owned::<usize>::from_element(3, 7, 0);
1937        assert!(m.get_element(3, 2).is_none());
1938        let _ = m.element(3, 2);
1939    }
1940
1941    #[test]
1942    #[should_panic(expected = "col 7 is out of bounds (max: 7)")]
1943    fn test_element_panics_col() {
1944        let m = Owned::<usize>::from_element(3, 7, 0);
1945        assert!(m.get_element(2, 7).is_none());
1946        let _ = m.element(2, 7);
1947    }
1948
1949    #[test]
1950    #[should_panic(expected = "row 3 is out of bounds (max: 3)")]
1951    fn test_element_mut_panics_row() {
1952        let mut m = Owned::<usize>::from_element(3, 7, 0);
1953        assert!(m.get_element_mut(3, 2).is_none());
1954        *m.element_mut(3, 2) = 1;
1955    }
1956
1957    #[test]
1958    #[should_panic(expected = "col 7 is out of bounds (max: 7)")]
1959    fn test_element_mut_panics_col() {
1960        let mut m = Owned::<usize>::from_element(3, 7, 0);
1961        assert!(m.get_element_mut(2, 7).is_none());
1962        *m.element_mut(2, 7) = 1;
1963    }
1964
1965    #[test]
1966    #[cfg(feature = "rayon")]
1967    #[should_panic(expected = "par_window_iter batchsize cannot be zero")]
1968    fn test_par_window_iter_panics() {
1969        let m = Owned::<usize>::from_element(4, 4, 0);
1970        let _ = m.par_window_iter(0);
1971    }
1972
1973    #[test]
1974    #[cfg(feature = "rayon")]
1975    #[should_panic(expected = "par_window_iter_mut batchsize cannot be zero")]
1976    fn test_par_window_iter_mut_panics() {
1977        let mut m = Owned::<usize>::from_element(4, 4, 0);
1978        let _ = m.par_window_iter_mut(0);
1979    }
1980
1981    // Additional tests for better coverage
1982
1983    #[test]
1984    fn test_try_from_error_light() {
1985        // Incorrect slice
1986        let data = vec![1, 2, 3];
1987        let err = Ref::try_from_data(data.as_slice(), 2, 3).unwrap_err();
1988
1989        // Test `as_static` method
1990        let err_static = err.as_static();
1991        let msg = err_static.to_string();
1992        assert_contains!(
1993            msg,
1994            "tried to construct a 2x3 matrix over a span of length 3",
1995        );
1996        // Test `into_inner` method
1997        let recovered_data = err.into_inner();
1998        assert_eq!(recovered_data, data.as_slice());
1999
2000        // Invalid length.
2001        let err = Ref::try_from_data(data.as_slice(), 2, usize::MAX).unwrap_err();
2002        let msg = err.to_string();
2003        assert_contains!(msg, "usize::MAX");
2004
2005        assert_eq!(data.as_slice(), err.into_inner());
2006    }
2007
2008    #[test]
2009    fn test_map_errors() {
2010        #[derive(Debug, Clone, Copy)]
2011        struct Zst;
2012
2013        // Create a large ZST slice without taking forever on debug builds.
2014        let b = Box::<[Zst]>::new_uninit_slice((isize::MAX as usize) + 1);
2015
2016        // SAFETY: `b` has zero-sized elements, so all elements are initialized.
2017        let b = unsafe { b.assume_init() };
2018
2019        let m = Owned::column_vector(b);
2020        let err = m.try_map(|_: &Zst| 0u8).unwrap_err();
2021        let msg = err.to_string();
2022        assert!(msg.contains("isize::MAX"), "{msg}");
2023
2024        // Panicking variant.
2025        let err = std::panic::catch_unwind(|| m.map(|_: &Zst| 0u8))
2026            .unwrap_err()
2027            .downcast::<String>()
2028            .unwrap();
2029        let msg = err.to_string();
2030        assert!(msg.contains("isize::MAX"), "{msg}");
2031    }
2032
2033    #[test]
2034    fn test_get_row_optional() {
2035        let data = make_test_matrix();
2036        let mut m = Owned::try_from_data(data.into(), 4, 3).unwrap();
2037
2038        assert_eq!(m.get_row(0), Some(&[0, 1, 2][..]));
2039        assert_eq!(m.get_row(1), Some(&[1, 2, 3][..]));
2040        assert_eq!(m.get_row(3), Some(&[3, 4, 5][..]));
2041        assert_eq!(m.get_row(4), None);
2042        assert_eq!(m.get_row(100), None);
2043
2044        let row = m.get_row_mut(1).unwrap();
2045        assert_eq!(row, &[1, 2, 3]);
2046        row[0] = 10;
2047        assert_eq!(m.row(1), &[10, 2, 3]);
2048        assert!(m.get_row_mut(4).is_none());
2049        assert!(m.get_row_mut(100).is_none());
2050    }
2051
2052    #[test]
2053    fn test_unsafe_get_unchecked_methods() {
2054        let data = make_test_matrix();
2055        let mut m = Owned::try_from_data(data.into(), 4, 3).unwrap();
2056
2057        // Safety: derives from known size of matrix and access element ids
2058        unsafe {
2059            assert_eq!(*m.element_unchecked(0, 0), 0);
2060            assert_eq!(*m.element_unchecked(1, 2), 3);
2061            assert_eq!(*m.element_unchecked(3, 1), 4);
2062        }
2063
2064        // Safety: derives from known size of matrix and access element ids
2065        unsafe {
2066            *m.element_unchecked_mut(0, 0) = 100;
2067            *m.element_unchecked_mut(1, 2) = 200;
2068        }
2069
2070        assert_eq!(*m.element(0, 0), 100);
2071        assert_eq!(*m.element(1, 2), 200);
2072
2073        // Safety: derives from known size of matrix and access element ids
2074        unsafe {
2075            let row0 = m.row_unchecked(0);
2076            assert_eq!(row0[0], 100);
2077            assert_eq!(row0[1], 1);
2078            assert_eq!(row0[2], 2);
2079        }
2080
2081        // Safety: derives from known size of matrix and access element ids
2082        unsafe {
2083            let row1 = m.row_unchecked_mut(1);
2084            row1[0] = 300;
2085        }
2086
2087        assert_eq!(*m.element(1, 0), 300);
2088    }
2089
2090    #[test]
2091    fn test_to_owned() {
2092        let data = make_test_matrix();
2093        let view = Ref::try_from_data(data.as_slice(), 4, 3).unwrap();
2094
2095        // Test to_owned creates a proper clone
2096        let owned: Owned<_> = view.to_rowmajor_owned();
2097        assert_eq!(owned.nrows(), view.nrows());
2098        assert_eq!(owned.ncols(), view.ncols());
2099        assert_eq!(owned.as_slice(), view.as_slice());
2100
2101        // Verify it's actually owned (different memory location)
2102        assert_ne!(owned.as_ptr(), view.as_ptr());
2103
2104        // Test the owned matrix works properly
2105        test_basic_indexing(&owned);
2106    }
2107
2108    #[test]
2109    fn test_matrix_from_conversions() {
2110        let data = make_test_matrix();
2111        let m = Owned::try_from_data(data.into(), 4, 3).unwrap();
2112
2113        // Test Ref to slice conversion
2114        let view = m.as_view();
2115        let slice: &[usize] = view.into_slice();
2116        assert_eq!(slice.len(), 12);
2117        assert_eq!(slice[0], 0);
2118        assert_eq!(slice[11], 5);
2119
2120        // Test Mut to mutable slice conversion.
2121        let data2 = make_test_matrix();
2122        let mut m2 = Owned::try_from_data(data2.into(), 4, 3).unwrap();
2123        let mut_view = m2.as_view_mut();
2124        let slice2: &mut [usize] = mut_view.into_mut_slice();
2125        assert_eq!(slice2.len(), 12);
2126        assert_eq!(slice2[0], 0);
2127        assert_eq!(slice2[11], 5);
2128        slice2[11] = 6;
2129        assert_eq!(*m2.element(3, 2), 6);
2130    }
2131
2132    #[test]
2133    fn test_row_vector() {
2134        let data = vec![1, 2, 3];
2135        let m = Ref::row_vector(data.as_slice());
2136        assert_eq!(m.nrows(), 1);
2137        assert_eq!(m.ncols(), 3);
2138        assert_eq!(m.as_slice(), &[1, 2, 3]);
2139        assert_eq!(m.row(0), &[1, 2, 3]);
2140
2141        // Empty
2142        let empty: &[i32] = &[];
2143        let m = Ref::row_vector(empty);
2144        assert_eq!(m.nrows(), 1);
2145        assert_eq!(m.ncols(), 0);
2146
2147        // Owned
2148        let m = Owned::row_vector(vec![10u64, 20].into_boxed_slice());
2149        assert_eq!(m.nrows(), 1);
2150        assert_eq!(m.ncols(), 2);
2151        assert_eq!(*m.element(0, 0), 10);
2152        assert_eq!(*m.element(0, 1), 20);
2153    }
2154
2155    #[test]
2156    fn test_column_vector() {
2157        let data = vec![1, 2, 3];
2158        let m = Ref::column_vector(data.as_slice());
2159        assert_eq!(m.nrows(), 3);
2160        assert_eq!(m.ncols(), 1);
2161        assert_eq!(m.as_slice(), &[1, 2, 3]);
2162        assert_eq!(*m.element(0, 0), 1);
2163        assert_eq!(*m.element(1, 0), 2);
2164        assert_eq!(*m.element(2, 0), 3);
2165        assert_eq!(m.row(0), &[1]);
2166        assert_eq!(m.row(1), &[2]);
2167        assert_eq!(m.row(2), &[3]);
2168
2169        // Empty
2170        let empty: &[i32] = &[];
2171        let m = Ref::column_vector(empty);
2172        assert_eq!(m.nrows(), 0);
2173        assert_eq!(m.ncols(), 1);
2174
2175        // Owned
2176        let m = Owned::column_vector(vec![10u64, 20].into_boxed_slice());
2177        assert_eq!(m.nrows(), 2);
2178        assert_eq!(m.ncols(), 1);
2179        assert_eq!(*m.element(0, 0), 10);
2180        assert_eq!(*m.element(1, 0), 20);
2181    }
2182
2183    #[test]
2184    fn test_map() {
2185        let m = Owned::try_from_data(vec![1u32, 2, 3, 4].into(), 2, 2).unwrap();
2186        let doubled = m.map(|&x| x * 2);
2187        assert_eq!(doubled.as_slice(), &[2, 4, 6, 8]);
2188        assert_eq!(doubled.nrows(), 2);
2189        assert_eq!(doubled.ncols(), 2);
2190
2191        // Type-changing map
2192        let as_f64 = m.map(|&x| x as f64);
2193        assert_eq!(as_f64.as_slice(), &[1.0, 2.0, 3.0, 4.0]);
2194    }
2195
2196    #[test]
2197    fn test_get_element() {
2198        let mut m = Owned::try_from_data(vec![1, 2, 3, 4, 5, 6].into(), 2, 3).unwrap();
2199        assert_eq!(m.get_element(0, 0), Some(&1));
2200        assert_eq!(m.get_element(1, 2), Some(&6));
2201        assert_eq!(m.get_element(2, 0), None);
2202        assert_eq!(m.get_element(0, 3), None);
2203
2204        *m.get_element_mut(1, 2).unwrap() = 7;
2205        assert_eq!(m.get_element(1, 2), Some(&7));
2206        assert_eq!(m.get_element_mut(2, 0), None);
2207        assert_eq!(m.get_element_mut(0, 3), None);
2208    }
2209
2210    #[test]
2211    fn test_subview() {
2212        let data = make_test_matrix();
2213        let m = Owned::try_from_data(data.into(), 4, 3).unwrap();
2214
2215        // Create a subview of the first two rows
2216        {
2217            let subview = m.subview(0..4).unwrap();
2218            assert_eq!(subview.nrows(), 4);
2219            assert_eq!(subview.ncols(), 3);
2220
2221            assert_eq!(subview.row(0), &[0, 1, 2]);
2222            assert_eq!(subview.row(1), &[1, 2, 3]);
2223            assert_eq!(subview.row(2), &[2, 3, 4]);
2224            assert_eq!(subview.row(3), &[3, 4, 5]);
2225            assert!(subview.get_row(4).is_none());
2226        }
2227
2228        // Sub view over a subset that touches the end.
2229        {
2230            let subview = m.subview(1..4).unwrap();
2231            assert_eq!(subview.nrows(), 3);
2232            assert_eq!(subview.ncols(), 3);
2233
2234            assert_eq!(subview.row(0), &[1, 2, 3]);
2235            assert_eq!(subview.row(1), &[2, 3, 4]);
2236            assert_eq!(subview.row(2), &[3, 4, 5]);
2237            assert!(subview.get_row(3).is_none());
2238        }
2239
2240        // Empty subview in the middle.
2241        {
2242            let subview = m.subview(2..2).unwrap();
2243            assert_eq!(subview.nrows(), 0);
2244            assert_eq!(subview.ncols(), 3);
2245        }
2246
2247        // Empty subviews at both boundaries.
2248        {
2249            let subview = m.subview(0..0).unwrap();
2250            assert_eq!(subview.nrows(), 0);
2251            assert_eq!(subview.ncols(), 3);
2252
2253            let subview = m.subview(4..4).unwrap();
2254            assert_eq!(subview.nrows(), 0);
2255            assert_eq!(subview.ncols(), 3);
2256        }
2257
2258        // Empty out-of-bounds subview
2259        assert!(m.subview(5..5).is_none());
2260
2261        // End is out of bounds.
2262        assert!(m.subview(2..10).is_none());
2263
2264        // Reversed bounds.
2265        #[expect(
2266            clippy::reversed_empty_ranges,
2267            reason = "we want to make sure it doesn't work"
2268        )]
2269        let empty = 3..2;
2270        assert!(m.subview(empty).is_none());
2271
2272        // Extreme out-of-bounds values do not overflow internal calculations.
2273        assert!(m.subview(usize::MAX - 1..usize::MAX).is_none());
2274    }
2275
2276    #[expect(
2277        clippy::reversed_empty_ranges,
2278        reason = "we want to make sure it doesn't work"
2279    )]
2280    #[test]
2281    fn test_subview_zero_cols() {
2282        let m = Owned::from_element(10, 0, 0u32);
2283
2284        // A fully disjoint range is rejected.
2285        assert!(m.subview(100..200).is_none());
2286
2287        // A range extending one row beyond the matrix is rejected.
2288        assert!(m.subview(10..11).is_none());
2289
2290        // Reversed bounds are rejected.
2291        assert!(m.subview(5..4).is_none());
2292
2293        // An in-bounds suffix preserves its logical shape.
2294        let v = m.subview(5..10).unwrap();
2295        assert_eq!(v.nrows(), 5);
2296        assert_eq!(v.ncols(), 0);
2297
2298        // An empty range at the end is valid.
2299        let v = m.subview(10..10).unwrap();
2300        assert_eq!(v.nrows(), 0);
2301        assert_eq!(v.ncols(), 0);
2302
2303        // The full range preserves all logical rows despite having no elements.
2304        let v = m.subview(0..10).unwrap();
2305        assert_eq!(v.nrows(), 10);
2306        assert_eq!(v.ncols(), 0);
2307    }
2308
2309    #[test]
2310    #[cfg(all(not(miri), feature = "rayon"))]
2311    fn parallel_immutable_iterators_match_scalar_indexing() {
2312        for (nrows, ncols) in [(0, 0), (0, 4), (3, 0), (1, 1), (1, 4), (4, 1), (5, 3)] {
2313            let m = striped_matrix(nrows, ncols);
2314            let view = m.as_view();
2315
2316            assert_parallel_rows_match_scalar(view);
2317            for batchsize in [1, 2, 3, usize::MAX] {
2318                assert_parallel_windows_match_scalar(view, batchsize);
2319            }
2320        }
2321    }
2322
2323    #[test]
2324    #[cfg(all(not(miri), feature = "rayon"))]
2325    fn parallel_mutable_iterators_match_scalar_indexing() {
2326        use rayon::prelude::*;
2327
2328        for (nrows, ncols) in [(0, 0), (0, 4), (1, 1), (1, 4), (4, 1), (5, 3)] {
2329            let context = lazy_format!("nrows = {nrows}, ncols = {ncols}");
2330
2331            let original = striped_matrix(nrows, ncols);
2332            let mut rows = original.clone();
2333
2334            let row_views: Vec<_> = rows.par_rows_mut().collect();
2335            assert_eq!(row_views.len(), nrows, "{context}");
2336            for (row_index, row) in row_views.into_iter().enumerate() {
2337                assert_eq!(row.len(), ncols, "row = {row_index} -- {context}");
2338                for value in row {
2339                    *value = value.wrapping_add(row_index);
2340                }
2341            }
2342
2343            for row in 0..nrows {
2344                for col in 0..ncols {
2345                    assert_eq!(
2346                        *rows.element(row, col),
2347                        original.element(row, col).wrapping_add(row),
2348                        "row = {}, col = {} -- {}",
2349                        row,
2350                        col,
2351                        context,
2352                    );
2353                }
2354            }
2355
2356            for batchsize in [1, 2, 3, usize::MAX] {
2357                let context = lazy_format!("{context}, batchsize = {batchsize}");
2358
2359                let mut windows = original.clone();
2360                let window_views: Vec<_> = windows.par_window_iter_mut(batchsize).collect();
2361                assert_eq!(window_views.len(), nrows.div_ceil(batchsize), "{context}");
2362                for (window_index, mut window) in window_views.into_iter().enumerate() {
2363                    let first_row = window_index * batchsize;
2364                    assert_eq!(
2365                        window.nrows(),
2366                        batchsize.min(nrows - first_row),
2367                        "window = {window_index} -- {context}"
2368                    );
2369                    assert_eq!(
2370                        window.ncols(),
2371                        ncols,
2372                        "window = {window_index} -- {context}"
2373                    );
2374                    for value in window.as_mut_slice() {
2375                        *value = value.wrapping_add(window_index);
2376                    }
2377                }
2378
2379                for row in 0..nrows {
2380                    for col in 0..ncols {
2381                        assert_eq!(
2382                            *windows.element(row, col),
2383                            original.element(row, col).wrapping_add(row / batchsize),
2384                            "row = {}, col = {} -- {}",
2385                            row,
2386                            col,
2387                            context,
2388                        );
2389                    }
2390                }
2391            }
2392        }
2393    }
2394
2395    #[test]
2396    #[cfg(feature = "rayon")]
2397    #[should_panic(
2398        expected = "`MatrixMut::par_rows_mut` does not support matrices with rows and zero columns"
2399    )]
2400    fn par_rows_mut_rejects_nonempty_zero_column_matrix() {
2401        let mut m = striped_matrix(3, 0);
2402        let _ = m.par_rows_mut();
2403    }
2404
2405    #[test]
2406    #[cfg(feature = "rayon")]
2407    #[should_panic(
2408        expected = "`MatrixMut::par_window_iter_mut` does not support matrices with rows and zero columns"
2409    )]
2410    fn par_window_iter_mut_rejects_nonempty_zero_column_matrix() {
2411        let mut m = striped_matrix(3, 0);
2412        let _ = m.par_window_iter_mut(2);
2413    }
2414
2415    #[test]
2416    fn matrix_iterators_match_scalar_indexing() {
2417        for (nrows, ncols) in [(0, 0), (0, 4), (3, 0), (1, 1), (1, 4), (4, 1), (5, 3)] {
2418            let original = striped_matrix(nrows, ncols);
2419            let view = original.as_view();
2420
2421            let rows = view.rows().collect();
2422            assert_rows_match_scalar(view, rows);
2423            for batchsize in [1, 2, 3, usize::MAX] {
2424                let windows = view
2425                    .window_iter(NonZeroUsize::new(batchsize).unwrap())
2426                    .collect();
2427                assert_windows_match_scalar(view, batchsize, windows);
2428            }
2429
2430            let context = lazy_format!("nrows = {nrows}, ncols = {ncols}");
2431            let mut mutable = original.clone();
2432            let rows: Vec<_> = mutable.rows_mut().collect();
2433            assert_eq!(rows.len(), nrows, "{context}");
2434
2435            for (row_index, row) in rows.into_iter().enumerate() {
2436                assert_eq!(row.len(), ncols, "row = {row_index} -- {context}");
2437                for (col_index, value) in row.iter_mut().enumerate() {
2438                    assert_eq!(
2439                        *value,
2440                        *original.element(row_index, col_index),
2441                        "row = {row_index}, col = {col_index} -- {context}"
2442                    );
2443                    *value = value.wrapping_add(1);
2444                }
2445            }
2446
2447            for row in 0..nrows {
2448                for col in 0..ncols {
2449                    assert_eq!(
2450                        *mutable.element(row, col),
2451                        original.element(row, col).wrapping_add(1),
2452                        "row = {row}, col = {col} -- {context}"
2453                    );
2454                }
2455            }
2456        }
2457    }
2458
2459    #[test]
2460    fn matrix_iterators_track_exact_remaining_lengths() {
2461        for (nrows, ncols) in [(0, 0), (0, 4), (3, 0), (5, 3)] {
2462            let mut m = striped_matrix(nrows, ncols);
2463
2464            assert_exact_fused(m.rows(), nrows);
2465            assert_exact_fused(m.rows_mut(), nrows);
2466
2467            for batchsize in [1, 2, usize::MAX] {
2468                assert_exact_fused(
2469                    m.window_iter(NonZeroUsize::new(batchsize).unwrap()),
2470                    nrows.div_ceil(batchsize),
2471                );
2472            }
2473        }
2474    }
2475
2476    #[test]
2477    fn matrix_transformations_preserve_empty_shapes() {
2478        for (nrows, ncols) in [(0, 0), (0, 4), (3, 0)] {
2479            let m = striped_matrix(nrows, ncols);
2480
2481            let owned = m.to_rowmajor_owned();
2482            assert_eq!(owned.nrows(), nrows);
2483            assert_eq!(owned.ncols(), ncols);
2484
2485            let mapped = m.map(|_| -> u8 { unreachable!("empty matrix has no elements") });
2486            assert_eq!(mapped.nrows(), nrows);
2487            assert_eq!(mapped.ncols(), ncols);
2488        }
2489    }
2490
2491    #[test]
2492    fn owned_from_fn_initializes_in_memory_order() {
2493        let mut value = 0;
2494        let m = Owned::from_fn(2, 3, |_| {
2495            let result = value;
2496            value += 1;
2497            result
2498        });
2499
2500        assert_eq!(*m.element(0, 0), 0);
2501        assert_eq!(*m.element(0, 1), 1);
2502        assert_eq!(*m.element(0, 2), 2);
2503        assert_eq!(*m.element(1, 0), 3);
2504        assert_eq!(*m.element(1, 1), 4);
2505        assert_eq!(*m.element(1, 2), 5);
2506    }
2507
2508    #[test]
2509    fn test_transpose() {
2510        {
2511            let v = Owned::from_element(0, 0, 0);
2512            let t = v.transpose();
2513            assert_eq!(t.nrows(), 0);
2514            assert_eq!(t.ncols(), 0);
2515        }
2516
2517        {
2518            let v = Owned::from_element(0, 10, 0);
2519            let t = v.transpose();
2520            assert_eq!(t.nrows(), 10);
2521            assert_eq!(t.ncols(), 0);
2522        }
2523
2524        {
2525            let v = Owned::from_element(10, 0, 0);
2526            let t = v.transpose();
2527            assert_eq!(t.nrows(), 0);
2528            assert_eq!(t.ncols(), 10);
2529        }
2530
2531        {
2532            let v = Owned::<usize>::try_from_data(Box::new([1, 2, 3, 4, 5, 6]), 2, 3).unwrap();
2533            let t = v.transpose();
2534
2535            assert_eq!(t.row(0), &[1, 4]);
2536            assert_eq!(t.row(1), &[2, 5]);
2537            assert_eq!(t.row(2), &[3, 6]);
2538        }
2539    }
2540
2541    #[test]
2542    fn test_debug_error_formatting() {
2543        // Test Debug implementation for TryFromError
2544        let data = vec![1, 2, 3];
2545        let err = Owned::try_from_data(data.into(), 2, 3).unwrap_err();
2546        let debug_str = format!("{:?}", err);
2547        assert_contains!(debug_str, "TryFromError");
2548
2549        // Ensure Debug doesn't require T: Debug by using a non-Debug type
2550        #[derive(Clone)]
2551        struct NonDebug(#[expect(dead_code)] i32);
2552
2553        let non_debug_data: Box<[NonDebug]> = vec![NonDebug(1), NonDebug(2)].into();
2554        let non_debug_err = match Owned::try_from_data(non_debug_data, 1, 3) {
2555            Ok(_) => panic!("should not have succeeded!"),
2556            Err(err) => err,
2557        };
2558        let debug_str = format!("{:?}", non_debug_err);
2559        assert_contains!(debug_str, "TryFromError");
2560    }
2561}