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