diskann-utils 0.60.0

/*
 * Copyright (c) Microsoft Corporation.
 * Licensed under the MIT license.
 */

use std::{fmt, marker::PhantomData, num::NonZeroUsize, ptr::NonNull};
use thiserror::Error;

use crate::{
    internal,
    views::rowmajor::{self, Matrix},
    Reborrow,
};

/// The layout for [`Strided`].
///
/// This struct ensures that the [`Self::cstride`] is greater than or equal to [`Self::ncols`]
/// and that the addressable span is valid for elements of type `T`. In particular, the
/// linear length does not overflow `usize::MAX` and its size in bytes does not exceed
/// `isize::MAX`.
///
/// The linear length of [`Strided`] is given by the formula
/// ```text
/// self.nrows.saturating_sub(1) * self.cstride + self.nrows.min(1) * self.ncols
/// ```
/// This allows the last row to occupy less than a full stride.
#[derive(Debug)]
pub struct Layout<T> {
    nrows: usize,
    ncols: usize,
    cstride: usize,
    _type: PhantomData<fn() -> T>,
}

impl<T> Layout<T> {
    /// Construct a new [`Layout`].
    ///
    /// Errors if:
    ///
    /// * `cstride < ncols`.
    /// * The computation of the linear length overflows `usize::MAX`.
    /// * The number of bytes required for the full span exceeds `isize::MAX`.
    pub fn new(nrows: usize, ncols: usize, cstride: usize) -> Result<Self, LayoutError> {
        LayoutError::check::<T>(nrows, ncols, cstride)?;
        Ok(Self {
            nrows,
            ncols,
            cstride,
            _type: PhantomData,
        })
    }

    /// Return the number of rows.
    pub fn nrows(&self) -> usize {
        self.nrows
    }

    /// Return the number of columns.
    pub fn ncols(&self) -> usize {
        self.ncols
    }

    /// Return the stride between subsequent rows.
    pub fn cstride(&self) -> usize {
        self.cstride
    }

    /// Return the length of the addressable span described by this [`Layout`], including
    /// gaps between rows.
    pub fn linear_length(&self) -> usize {
        self.nrows.saturating_sub(1) * self.cstride + self.nrows.min(1) * self.ncols
    }
}

impl<T> Clone for Layout<T> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<T> Copy for Layout<T> {}

impl<T> From<rowmajor::Layout<T>> for Layout<T> {
    fn from(layout: rowmajor::Layout<T>) -> Self {
        Self {
            nrows: layout.nrows(),
            ncols: layout.ncols(),
            cstride: layout.ncols(),
            _type: PhantomData,
        }
    }
}

fn linear_length(nrows: usize, ncols: usize, cstride: usize) -> Option<usize> {
    nrows
        .saturating_sub(1)
        .checked_mul(cstride)
        .and_then(|main| main.checked_add(nrows.min(1) * ncols))
}

/// Errors for [`Layout::new`].
#[derive(Debug)]
pub struct LayoutError(LayoutErrorInner);

impl LayoutError {
    fn check<T>(nrows: usize, ncols: usize, cstride: usize) -> Result<usize, Self> {
        if cstride < ncols {
            Err(Self(LayoutErrorInner::InvalidStride { ncols, cstride }))
        } else {
            let linear_length = match linear_length(nrows, ncols, cstride) {
                Some(len) => len,
                None => {
                    return Err(Self(LayoutErrorInner::Overflow {
                        nrows,
                        cstride,
                        elsize: None,
                    }));
                }
            };

            let elsize = std::mem::size_of::<T>();
            let bytes = linear_length.saturating_mul(elsize);
            if bytes > (isize::MAX as usize) {
                Err(Self(LayoutErrorInner::Overflow {
                    nrows,
                    cstride,
                    elsize: NonZeroUsize::new(elsize),
                }))
            } else {
                Ok(linear_length)
            }
        }
    }
}

impl fmt::Display for LayoutError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl std::error::Error for LayoutError {}

#[derive(Debug)]
enum LayoutErrorInner {
    InvalidStride {
        ncols: usize,
        cstride: usize,
    },
    Overflow {
        nrows: usize,
        cstride: usize,
        elsize: Option<NonZeroUsize>,
    },
}

impl fmt::Display for LayoutErrorInner {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::InvalidStride { ncols, cstride } => write!(
                f,
                "column stride {} must be greater than or equal to number of columns {}",
                cstride, ncols
            ),
            Self::Overflow {
                nrows,
                cstride,
                elsize,
            } => match elsize {
                Some(elsize) => write!(
                    f,
                    "a {}x{} strided matrix with element size {} exceeds isize::MAX bytes",
                    nrows, cstride, elsize
                ),
                None => write!(
                    f,
                    "a {}x{} strided matrix has a length exceeding usize::MAX",
                    nrows, cstride
                ),
            },
        }
    }
}

/// A row-major strided matrix.
///
/// This is a generalization of the [`Matrix`] trait as it does not mandate a dense
/// layout in memory.
///
/// ```text
///            |<------ cstride ----->|
///                     |<-- ncols -->|
///                     +-------------+
/// slice 0 -> a0 a1 a2 | a3 a4 a5 a6 |    ^
/// slice 1 -> b0 b1 b2 | b3 b4 b4 b5 |    |
/// slice 2 -> c0 c1 c2 | c3 c4 c5 c6 |  nrows
/// slice 3 -> d0 d1 d2 | d3 d4 d5 d6 |    |
/// slice 4 -> e0 e1 e2 | e3 e4 e5 e6 |    |
/// slice 5 -> f0 f1 f2 | f3 f4 f5 f6 |    v
///                     +-------------+
///                           ^
///                           |
///                        Strided
/// ```
///
/// This abstraction is useful when performing PQ related operations such as training or
/// compression as it provides a convenient abstraction for working with columnar subsets
/// of dense data in-place.
#[derive(Debug)]
pub struct Strided<'a, T> {
    ptr: NonNull<T>,
    layout: Layout<T>,
    _lifetime: PhantomData<&'a [T]>,
}

impl<'a, T> Strided<'a, T> {
    /// Construct a strided view over data slice, shrinking the slice as needed.
    ///
    /// Returns an error if the layout is invalid or `data` is not at least
    /// [`Layout::linear_length`].
    pub fn try_from_data(
        data: &'a [T],
        nrows: usize,
        ncols: usize,
        cstride: usize,
    ) -> Result<Self, TryFromError> {
        let layout = Layout::new(nrows, ncols, cstride).map_err(TryFromError::LayoutError)?;
        let expected = layout.linear_length();
        if data.len() < expected {
            Err(TryFromError::InvalidLength {
                got: data.len(),
                expected,
            })
        } else {
            // SAFETY: `data.len() >= layout.linear_length()`.
            Ok(unsafe { Self::from_data_unchecked(data, layout) })
        }
    }

    /// Construct a strided view over `data`, shrinking the slice as needed, without
    /// verifying that `data.len() >= layout.linear_length()`.
    ///
    /// # Safety
    ///
    /// `data.len()` must be greater-than or equal to `layout.linear_length()`.
    unsafe fn from_data_unchecked(data: &'a [T], layout: Layout<T>) -> Self {
        debug_assert!(data.len() >= layout.linear_length());
        Self {
            ptr: internal::slice_to_nonnull(data),
            layout,
            _lifetime: PhantomData,
        }
    }

    fn as_nonnull(&self) -> NonNull<T> {
        self.ptr
    }

    /// Return the [`Layout`] for the matrix.
    pub fn layout(&self) -> Layout<T> {
        self.layout
    }

    /// Return a pointer to the base of the matrix.
    pub fn as_ptr(&self) -> *const T {
        self.as_nonnull().as_ptr().cast_const()
    }

    /// Return the number of columns in the matrix.
    pub fn ncols(&self) -> usize {
        self.layout().ncols()
    }

    /// Return the number of rows in the matrix.
    pub fn nrows(&self) -> usize {
        self.layout().nrows()
    }

    /// Return the count of elements between the start of each row.
    pub fn cstride(&self) -> usize {
        self.layout().cstride()
    }

    /// Return the underlying data as a slice.
    ///
    /// # Note
    ///
    /// The underlying representation for a strided matrix is not necessarily dense.
    pub fn as_slice(&self) -> &[T] {
        let layout = self.layout();

        // SAFETY: Constructors verify that the backing memory has a length at least
        // `layout.linear_length()`.
        unsafe { std::slice::from_raw_parts(self.as_ptr(), layout.linear_length()) }
    }

    // Element Access

    /// Return the specified element.
    ///
    /// # Safety
    ///
    /// * `row < self.nrows()`.
    /// * `col < self.ncols()`.
    pub unsafe fn element_unchecked(&self, row: usize, col: usize) -> &T {
        let layout = self.layout();
        debug_assert!(row < layout.nrows());
        debug_assert!(col < layout.ncols());

        // SAFETY: Constructors verify that the backing memory has a length of at least
        // `layout.linear_length`. Since `row` and `col` are in-bounds, the pointer offset
        // is valid and it is safe to return a reference.
        unsafe { &*self.as_ptr().add(layout.cstride() * row + col) }
    }

    /// Return the specified element if `row < self.nrows()` and `col < self.ncols()`.
    pub fn get_element(&self, row: usize, col: usize) -> Option<&T> {
        if row < self.nrows() && col < self.ncols() {
            // SAFETY: `row` and `col` are in-bounds.
            Some(unsafe { self.element_unchecked(row, col) })
        } else {
            None
        }
    }

    /// Return the specified element.
    ///
    /// # Panics
    ///
    /// Panics if `row >= self.nrows()` or `col >= self.ncols()`.
    pub fn element(&self, row: usize, col: usize) -> &T {
        assert!(
            row < self.nrows(),
            "row {} is out of bounds for a matrix with {} rows",
            row,
            self.nrows()
        );
        assert!(
            col < self.ncols(),
            "col {} is out of bounds for a matrix with {} cols",
            col,
            self.ncols()
        );

        // SAFETY: `row` and `col` are in-bounds.
        unsafe { self.element_unchecked(row, col) }
    }

    // Row Access

    /// Returns the requested row without boundschecking.
    ///
    /// # Safety
    ///
    /// Caller must ensure `row < self.nrows()`.
    pub unsafe fn row_unchecked(&self, row: usize) -> &[T] {
        let layout = self.layout();
        debug_assert!(row < layout.nrows());

        // SAFETY: Constructors verify that the backing memory has a length of at least
        // `layout.linear_length`. Since `row` is in-bounds, the pointer offset
        // is valid and it is safe to form a slice of length `layout.ncols()`.
        unsafe {
            std::slice::from_raw_parts(self.as_ptr().add(layout.cstride() * row), layout.ncols())
        }
    }

    /// Return the requested row if `row < self.nrows()`.
    pub fn get_row(&self, row: usize) -> Option<&[T]> {
        if row < self.nrows() {
            // SAFETY: `row` is in-bounds.
            Some(unsafe { self.row_unchecked(row) })
        } else {
            None
        }
    }

    /// Return row `row` as a slice.
    ///
    /// # Panics
    ///
    /// Panics if `row >= self.nrows()`.
    pub fn row(&self, row: usize) -> &[T] {
        assert!(
            row < self.nrows(),
            "row {} is out of bounds for a matrix with {} rows",
            row,
            self.nrows()
        );

        // SAFETY: `row` is in-bounds.
        unsafe { self.row_unchecked(row) }
    }

    /// Return a iterator over all rows in the matrix.
    ///
    /// Rows are yielded sequentially beginning with row 0.
    pub fn rows(&self) -> Rows<'_, T> {
        Rows::new(*self)
    }
}

/// Errors for [`Strided::try_from_data`].
#[derive(Debug, Error)]
pub enum TryFromError {
    #[error(transparent)]
    LayoutError(LayoutError),
    #[error(
        "argument of length {} is shorter than the expected length {}",
        got,
        expected
    )]
    InvalidLength { got: usize, expected: usize },
}

impl<'a, T> From<rowmajor::Ref<'a, T>> for Strided<'a, T> {
    fn from(matrix: rowmajor::Ref<'a, T>) -> Self {
        let layout = Layout::from(matrix.layout());

        // SAFETY: `rowmajor::Ref` guarantees that the length of the base slice for `matrix`
        // is exactly `layout.linear_length()`.
        unsafe { Self::from_data_unchecked(matrix.into_slice(), layout) }
    }
}

// SAFETY: `Strided` only exposes shared access to its `T` elements (via `&T`/`&[T]`), so
// sharing or transferring a `Strided<T>` across threads is sound whenever `T: Sync`.
unsafe impl<T> Send for Strided<'_, T> where T: Sync {}
// SAFETY: See above.
unsafe impl<T> Sync for Strided<'_, T> where T: Sync {}

impl<T> Clone for Strided<'_, T> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<T> Copy for Strided<'_, T> {}

impl<'a, T> Reborrow<'a> for Strided<'_, T> {
    type Target = Strided<'a, T>;
    fn reborrow(&'a self) -> Self::Target {
        *self
    }
}

/// Iterator for [`Strided::rows`].
#[derive(Debug)]
pub struct Rows<'a, T> {
    ptr: NonNull<T>,
    remaining: usize,
    ncols: usize,
    cstride: usize,
    _lifetime: PhantomData<&'a T>,
}

impl<'a, T> Rows<'a, T> {
    fn new(strided: Strided<'a, T>) -> Self {
        let layout = strided.layout();
        Self {
            ptr: strided.as_nonnull(),
            remaining: layout.nrows(),
            ncols: layout.ncols(),
            cstride: layout.cstride(),
            _lifetime: PhantomData,
        }
    }
}

// SAFETY: `Rows` only yields shared `&[T]` slices borrowed from a `Strided`, so sharing or
// transferring a `Rows<T>` across threads is sound whenever `T: Sync`.
unsafe impl<T> Send for Rows<'_, T> where T: Sync {}
// SAFETY: See above.
unsafe impl<T> Sync for Rows<'_, T> where T: Sync {}

impl<'a, T> Iterator for Rows<'a, T> {
    type Item = &'a [T];
    fn next(&mut self) -> Option<&'a [T]> {
        self.remaining.checked_sub(1).map(|remaining| {
            // SAFETY: The originating `Strided` guarantees `self.remaining` rows of
            // `self.ncols` elements each are readable starting from `self.ptr`, so the
            // current row is valid for `self.ncols` elements.
            let item =
                unsafe { std::slice::from_raw_parts(self.ptr.as_ptr().cast_const(), self.ncols) };
            self.remaining = remaining;
            if remaining != 0 {
                // SAFETY: There is at least one more row remaining, so advancing by
                // `self.cstride` stays within (or one-past-the-end of) the originating
                // allocation, per `Layout::linear_length`.
                self.ptr = unsafe { self.ptr.add(self.cstride) };
            }
            item
        })
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.remaining, Some(self.remaining))
    }
}

impl<T> ExactSizeIterator for Rows<'_, T> {}
impl<T> std::iter::FusedIterator for Rows<'_, T> {}

///////////
// Tests //
///////////

#[cfg(test)]
mod tests {
    use super::*;

    use crate::views::rowmajor::MatrixMut;

    #[test]
    fn test_linear_length() {
        // If the number of rows is zero - the output should always be zero.
        assert_eq!(linear_length(0, 1, 1).unwrap(), 0);
        assert_eq!(linear_length(0, 2, 2).unwrap(), 0);
        assert_eq!(linear_length(0, 2, 3).unwrap(), 0);
        assert_eq!(linear_length(0, 2, 4).unwrap(), 0);

        // If `cstride == ncols`, then the computation should be trivial.
        for row in 1..10 {
            for col in 1..10 {
                assert_eq!(linear_length(row, col, col).unwrap(), row * col);
            }
        }

        // If there is only one row, then `cstride` should be ignored.
        assert_eq!(linear_length(1, 5, 10).unwrap(), 5);
        assert_eq!(linear_length(1, 7, 99).unwrap(), 7);

        // Otherwise, the computation is a block of `nrows - 1` chunks of `cstride` and then
        // `ncols`. Yes - this runs a bunch of computations.
        for row in 2..10 {
            for col in 0..10 {
                for cstride in col..12 {
                    assert_eq!(
                        linear_length(row, col, cstride).unwrap(),
                        (row - 1) * cstride + col
                    );
                }
            }
        }

        // Check OOB detection.
        assert!(linear_length(usize::MAX, 2, 2).is_none());
        assert!(linear_length(2, usize::MAX, 2).is_none());
        assert!(linear_length(2, 2, usize::MAX).is_none());
    }

    #[test]
    fn test_layout_new() {
        // Valid layouts.
        let layout = Layout::<usize>::new(3, 4, 4).unwrap();
        assert_eq!(layout.nrows(), 3);
        assert_eq!(layout.ncols(), 4);
        assert_eq!(layout.cstride(), 4);
        assert_eq!(layout.linear_length(), 12);

        let layout = Layout::<usize>::new(3, 4, 6).unwrap();
        assert_eq!(layout.linear_length(), 2 * 6 + 4);

        // `cstride == ncols` is fine, even at zero.
        assert!(Layout::<usize>::new(0, 0, 0).is_ok());

        // Invalid stride: `cstride < ncols`.
        let err = Layout::<usize>::new(3, 4, 3).unwrap_err();
        assert_eq!(
            err.to_string(),
            "column stride 3 must be greater than or equal to number of columns 4"
        );

        // Overflow: linear length exceeds `usize::MAX`.
        let err = Layout::<usize>::new(usize::MAX, usize::MAX, usize::MAX).unwrap_err();
        assert_eq!(
            err.to_string(),
            format!(
                "a {}x{} strided matrix has a length exceeding usize::MAX",
                usize::MAX,
                usize::MAX
            )
        );

        // Overflow: bytes exceeds `isize::MAX`.
        let err = Layout::<usize>::new(isize::MAX as usize, 1, 1).unwrap_err();
        assert_eq!(
            err.to_string(),
            format!(
                "a {}x{} strided matrix with element size {} exceeds isize::MAX bytes",
                isize::MAX,
                1,
                std::mem::size_of::<usize>(),
            )
        );

        // The element type affects whether the addressable span is valid.
        let length = isize::MAX as usize;
        let layout = Layout::<u8>::new(length, 1, 1).unwrap();
        assert_eq!(layout.linear_length(), length);
        assert!(Layout::<u8>::new(length + 1, 1, 1).is_err());
        assert!(Layout::<u16>::new(length, 1, 1).is_err());
    }

    #[test]
    fn test_try_from_data_errors() {
        let m = rowmajor::Owned::<usize>::from_element(10, 10, 0);
        let nrows = m.nrows();
        let ncols = m.ncols();

        // An invalid layout (`cstride < ncols`) is reported as a `LayoutError`, not a panic.
        let err = Strided::try_from_data(m.as_slice(), 2, 2, 1).unwrap_err();
        assert_eq!(
            err.to_string(),
            "column stride 1 must be greater than or equal to number of columns 2"
        );

        // A slice shorter than `Layout::linear_length` is reported as `InvalidLength`.
        let err = Strided::try_from_data(m.as_slice(), nrows, ncols, ncols + 1).unwrap_err();

        assert_eq!(
            err.to_string(),
            "argument of length 100 is shorter than the expected length 109",
        );
    }

    #[test]
    fn test_element_and_row_out_of_bounds() {
        let m = create_test_matrix(3, 4);
        let v = Strided::try_from_data(m.as_slice(), m.nrows(), m.ncols(), m.ncols()).unwrap();

        // In-bounds accesses succeed.
        assert!(v.get_element(2, 3).is_some());
        assert!(v.get_row(2).is_some());

        // Out-of-bounds row and/or col return `None` rather than panicking.
        assert!(v.get_element(3, 0).is_none(), "row out-of-bounds");
        assert!(v.get_element(0, 4).is_none(), "col out-of-bounds");
        assert!(v.get_element(3, 4).is_none(), "both out-of-bounds");
        assert!(v.get_row(3).is_none());
    }

    #[test]
    #[should_panic(expected = "row 3 is out of bounds for a matrix with 3 rows")]
    fn test_element_panics_on_row() {
        let m = create_test_matrix(3, 4);
        let v = Strided::try_from_data(m.as_slice(), m.nrows(), m.ncols(), m.ncols()).unwrap();
        v.element(3, 0);
    }

    #[test]
    #[should_panic(expected = "col 4 is out of bounds for a matrix with 4 cols")]
    fn test_element_panics_on_col() {
        let m = create_test_matrix(3, 4);
        let v = Strided::try_from_data(m.as_slice(), m.nrows(), m.ncols(), m.ncols()).unwrap();
        v.element(0, 4);
    }

    #[test]
    #[should_panic(expected = "row 3 is out of bounds for a matrix with 3 rows")]
    fn test_row_panics() {
        let m = create_test_matrix(3, 4);
        let v = Strided::try_from_data(m.as_slice(), m.nrows(), m.ncols(), m.ncols()).unwrap();
        v.row(3);
    }

    #[test]
    fn test_clone_copy_reborrow() {
        let m = create_test_matrix(3, 4);
        let v = Strided::try_from_data(m.as_slice(), m.nrows(), m.ncols(), m.ncols()).unwrap();

        // `Copy`/`Clone` produce an independent handle to the same data.
        let copied = v;
        let cloned = Clone::clone(&v);
        assert_eq!(v.as_ptr(), copied.as_ptr());
        assert_eq!(v.as_ptr(), cloned.as_ptr());

        // `Reborrow` should yield an equivalent view.
        let reborrowed = v.reborrow();
        assert_eq!(reborrowed.as_ptr(), v.as_ptr());
        assert_eq!(reborrowed.nrows(), v.nrows());
        assert_eq!(reborrowed.ncols(), v.ncols());
    }

    #[test]
    fn test_send_sync() {
        fn assert_send_sync<T: Send + Sync>() {}
        assert_send_sync::<Strided<'_, u8>>();
        assert_send_sync::<Rows<'_, u8>>();
    }

    #[test]
    fn test_rows_iterator_properties() {
        let m = create_test_matrix(4, 3);
        let v = Strided::try_from_data(&m.as_slice()[1..], m.nrows(), m.ncols() - 1, m.ncols())
            .unwrap();

        let mut rows = v.rows();
        assert_eq!(rows.len(), 4);
        assert_eq!(rows.size_hint(), (4, Some(4)));

        for expected_row in 0..4 {
            let row = rows.next().unwrap();
            assert_eq!(row, &m.row(expected_row)[1..]);
        }

        // Exhausted iterators keep returning `None` (`FusedIterator`).
        assert_eq!(rows.next(), None);
        assert_eq!(rows.next(), None);
        assert_eq!(rows.len(), 0);
    }

    #[test]
    fn test_rows_iterator_zero_rows() {
        let m = create_test_matrix(5, 5);
        let v = Strided::try_from_data(m.as_slice(), 0, 4, 5).unwrap();

        let mut rows = v.rows();
        assert_eq!(rows.len(), 0);
        assert_eq!(rows.next(), None);
    }

    #[test]
    fn test_rows_iterator_zero_cols() {
        let m = create_test_matrix(5, 5);
        let v = Strided::try_from_data(m.as_slice(), 5, 0, 5).unwrap();

        let rows = v.rows();
        assert_eq!(rows.len(), 5);
        assert_eq!(rows.size_hint(), (5, Some(5)));

        let mut count = 0;
        for r in rows {
            assert!(r.is_empty());
            count += 1;
        }

        assert_eq!(count, 5);
    }

    #[test]
    fn test_rows_iterator_zero_cstride() {
        let m = create_test_matrix(5, 5);
        let v = Strided::try_from_data(m.as_slice(), 5, 0, 0).unwrap();

        let rows = v.rows();
        assert_eq!(rows.len(), 5);
        assert_eq!(rows.size_hint(), (5, Some(5)));

        let mut count = 0;
        for r in rows {
            assert!(r.is_empty());
            count += 1;
        }

        assert_eq!(count, 5);
    }

    // Test that the contents of `dut` match those in the dense 2d matrix.
    fn test_indexing(dut: Strided<'_, usize>, expected: rowmajor::Ref<'_, usize>) {
        assert_eq!(dut.nrows(), expected.nrows());
        assert_eq!(dut.ncols(), expected.ncols());

        // Check the underlying data.
        if dut.cstride() == dut.ncols() {
            assert_eq!(dut.as_slice(), expected.as_slice());
        } else {
            assert_ne!(dut.as_slice(), expected.as_slice());
        }

        // Compare via linear indexing.
        for row in 0..dut.nrows() {
            for col in 0..dut.ncols() {
                let e = *expected.element(row, col);

                assert_eq!(
                    *dut.element(row, col),
                    e,
                    "failed on (row, col) = ({}, {})",
                    row,
                    col
                );

                assert_eq!(
                    *dut.get_element(row, col).unwrap(),
                    e,
                    "failed on (row, col) = ({}, {})",
                    row,
                    col
                );
            }
        }

        // Compare via row.
        for row in 0..dut.nrows() {
            assert_eq!(dut.row(row), expected.row(row), "failed on row {}", row);

            assert_eq!(
                dut.get_row(row).unwrap(),
                expected.row(row),
                "failed on row {}",
                row
            );
        }

        // Compare via row iterators.
        assert!(dut.rows().eq(expected.rows()));
    }

    // Create a base Matrix with the following pattern:
    // ```text
    //       0         1         2 ...   ncols-1
    //   ncols   ncols+1   ncols+2 ... 2*ncols-1
    // 2*ncols 2*ncols+1 2*ncols+2 ... 3*ncols-1
    // ...
    // ```
    fn create_test_matrix(nrows: usize, ncols: usize) -> rowmajor::Owned<usize> {
        let mut i = 0;
        rowmajor::Owned::from_fn(nrows, ncols, |_| {
            let v = i;
            i += 1;
            v
        })
    }

    #[test]
    fn test_basic_indexing() {
        let m = create_test_matrix(5, 3);

        // First - test a dense Strided view over the entire matrix.
        let ptr = m.as_ptr();
        let v = Strided::try_from_data(m.as_slice(), m.nrows(), m.ncols(), m.ncols()).unwrap();
        assert_eq!(v.as_ptr(), ptr, "base pointer was not preserved");

        assert_eq!(v.nrows(), m.nrows());
        assert_eq!(v.ncols(), m.ncols());
        assert_eq!(v.cstride(), m.ncols());
        test_indexing(v, m.as_view());

        // Now - create a truly strided view over the first two columns.
        let v = Strided::try_from_data(
            &(m.as_slice()[..(4 * m.ncols() + 2)]),
            m.nrows(),
            2,
            m.ncols(),
        )
        .unwrap();
        assert_eq!(v.as_ptr(), ptr, "base pointer was not preserved");

        // Create the expected matrix.
        let mut expected = rowmajor::Owned::from_element(5, 2, 0);
        for row in 0..expected.nrows() {
            for col in 0..expected.ncols() {
                *expected.element_mut(row, col) = *m.element(row, col);
            }
        }
        test_indexing(v, expected.as_view());

        // Create a strided view over the last two columns.
        let v = Strided::try_from_data(&(m.as_slice()[1..]), m.nrows(), 2, m.ncols()).unwrap();
        let mut expected = rowmajor::Owned::from_element(5, 2, 0);
        for row in 0..expected.nrows() {
            for col in 0..expected.ncols() {
                *expected.element_mut(row, col) = *m.element(row, col + 1);
            }
        }
        test_indexing(v, expected.as_view());
    }

    #[test]
    fn matrix_conversion() {
        let m = create_test_matrix(3, 4);
        let ptr = m.as_ptr();
        let v: Strided<_> = m.as_view().into();
        assert_eq!(v.as_ptr(), ptr);
        assert_eq!(v.cstride(), m.ncols());
        assert_eq!(v.layout().linear_length(), m.layout().num_elements());
        test_indexing(v, m.as_view());
    }

    #[test]
    fn test_zero_sized() {
        let m = create_test_matrix(5, 5);
        let v = Strided::try_from_data(m.as_slice(), 0, 4, 5).unwrap();

        assert_eq!(v.nrows(), 0);
        assert_eq!(v.ncols(), 4);
        assert_eq!(v.cstride(), 5);

        let v = Strided::try_from_data(m.as_slice(), 5, 0, 5).unwrap();
        assert_eq!(v.nrows(), 5);
        assert_eq!(v.ncols(), 0);
        assert_eq!(v.cstride(), 5);

        for row in 0..v.nrows() {
            let empty: &[usize] = &[];
            assert_eq!(v.get_row(row).unwrap(), empty);
        }
    }

    #[test]
    fn test_try_shrink_from() {
        // Exact is okay.
        let m = rowmajor::Owned::<usize>::from_element(10, 10, 0);
        let nrows = m.nrows();
        let ncols = m.ncols();
        let s = Strided::try_from_data(m.as_slice(), nrows, ncols, ncols).unwrap();
        assert_eq!(s.as_slice(), m.as_slice());

        // Giving a slice that is too large is okay.
        let s = Strided::try_from_data(m.as_slice(), nrows, 5, ncols).unwrap();
        assert_eq!(s.as_ptr(), m.as_ptr());

        // Too small is a problem, and is reported as an `Err`, not a panic.
        let s = Strided::try_from_data(m.as_slice(), nrows, ncols, ncols + 1);
        assert!(s.is_err());
    }

    #[test]
    fn test_invalid_stride_is_an_error_not_a_panic() {
        // Constructing a `Strided` with an invalid layout (`cstride < ncols`) returns an
        // `Err` rather than panicking - only unwrapping the result panics.
        let m = rowmajor::Owned::<usize>::from_element(4, 4, 0);
        let err = Strided::try_from_data(m.as_slice(), 2, 2, 1).unwrap_err();
        assert!(matches!(err, TryFromError::LayoutError(_)));
    }
}