diff --git a/diskann-benchmark-core/src/recall.rs b/diskann-benchmark-core/src/recall.rs index 12c9f11ba4..7b26e2b2a7 100644 --- a/diskann-benchmark-core/src/recall.rs +++ b/diskann-benchmark-core/src/recall.rs @@ -9,7 +9,7 @@ use std::{ }; use diskann_utils::{ - strided::StridedView, + strided::Strided, views::{Matrix, MatrixView}, }; use thiserror::Error; @@ -145,7 +145,7 @@ pub enum GroundTruthMode { /// than `recall_k` candidates. pub fn knn( groundtruth: &dyn Rows, - groundtruth_distances: Option>, + groundtruth_distances: Option>, results: &dyn Rows, recall_k: usize, recall_n: usize, diff --git a/diskann-providers/src/model/pq/strided.rs b/diskann-providers/src/model/pq/strided.rs index 4ff3c4d2de..7cdab00c25 100644 --- a/diskann-providers/src/model/pq/strided.rs +++ b/diskann-providers/src/model/pq/strided.rs @@ -4,14 +4,14 @@ */ use diskann::ANNError; -use diskann_utils::{strided, views}; +use diskann_utils::strided; use crate::utils::Bridge; -impl From>> for ANNError { +impl From> for ANNError { #[track_caller] - fn from(value: Bridge>) -> Self { - ANNError::new(value.into_inner().as_static()) + fn from(value: Bridge) -> Self { + ANNError::new(value.into_inner()) } } @@ -32,7 +32,7 @@ mod tests { let x = vec![u8::default(); nrows * ncols]; // Provided the incorrect dimensions. - let err = strided::StridedView::try_from(&x, nrows, ncols + 1, ncols + 1) + let err = strided::Strided::try_from_data(&x, nrows, ncols + 1, ncols + 1) .bridge_err() .unwrap_err(); let message = format!("{}", err); diff --git a/diskann-quantization/src/algorithms/kmeans/lloyds.rs b/diskann-quantization/src/algorithms/kmeans/lloyds.rs index f81894f1e8..86bc08e57f 100644 --- a/diskann-quantization/src/algorithms/kmeans/lloyds.rs +++ b/diskann-quantization/src/algorithms/kmeans/lloyds.rs @@ -8,7 +8,7 @@ use diskann_wide::{SIMDMask, SIMDMulAdd, SIMDPartialOrd, SIMDSelect, SIMDSumTree use super::common::square_norm; use crate::multi_vector::{BlockTransposed, BlockTransposedRef}; use diskann_utils::{ - strided::StridedView, + strided::Strided, views::{Matrix, MatrixView, MutMatrixView}, }; @@ -342,10 +342,10 @@ fn update((d0, i0): (f32s, u32s), (d1, i1): (f32s, u32s)) -> (f32s, u32s) { // Update Step // ///////////////// -fn update_centroids(mut centers: MutMatrixView<'_, f32>, data: StridedView<'_, f32>, map: &[u32]) { +fn update_centroids(mut centers: MutMatrixView<'_, f32>, data: Strided<'_, f32>, map: &[u32]) { let mut sums = Matrix::::new(0.0, centers.nrows(), centers.ncols()); let mut counts: Vec = vec![0; centers.nrows()]; - data.row_iter().zip(map.iter()).for_each(|(row, ¢er)| { + data.rows().zip(map.iter()).for_each(|(row, ¢er)| { counts[center as usize] += 1; let sum = sums.row_mut(center as usize); std::iter::zip(sum.iter_mut(), row.iter()).for_each(|(s, r)| { @@ -370,7 +370,7 @@ fn update_centroids(mut centers: MutMatrixView<'_, f32>, data: StridedView<'_, f //////////// pub(crate) fn lloyds_inner( - data: StridedView<'_, f32>, + data: Strided<'_, f32>, square_norms: &[f32], transpose: BlockTransposedRef<'_, f32, 16>, mut centers: MutMatrixView<'_, f32>, diff --git a/diskann-quantization/src/algorithms/kmeans/plusplus.rs b/diskann-quantization/src/algorithms/kmeans/plusplus.rs index 79abf0e947..084b64a8a1 100644 --- a/diskann-quantization/src/algorithms/kmeans/plusplus.rs +++ b/diskann-quantization/src/algorithms/kmeans/plusplus.rs @@ -6,7 +6,7 @@ use std::{collections::HashSet, fmt}; use diskann_utils::{ - strided::StridedView, + strided::Strided, views::{MatrixView, MutMatrixView}, }; use diskann_wide::{SIMDMulAdd, SIMDPartialOrd, SIMDSelect, SIMDVector}; @@ -380,7 +380,7 @@ impl KMeansPlusPlusError { pub(crate) fn kmeans_plusplus_into_inner( mut points: MutMatrixView<'_, f32>, - data: StridedView<'_, f32>, + data: Strided<'_, f32>, transpose: BlockTransposedRef<'_, f32, N>, norms: &[f32], rng: &mut dyn RngCore, diff --git a/diskann-quantization/src/multi_vector/block_transposed.rs b/diskann-quantization/src/multi_vector/block_transposed.rs index 159dc4bbba..39024cecac 100644 --- a/diskann-quantization/src/multi_vector/block_transposed.rs +++ b/diskann-quantization/src/multi_vector/block_transposed.rs @@ -81,7 +81,7 @@ use std::{alloc::Layout, marker::PhantomData, ptr::NonNull}; use diskann_utils::{ Reborrow, ReborrowMut, - strided::StridedView, + strided::Strided, views::{MatrixView, MutMatrixView}, }; @@ -1142,7 +1142,7 @@ impl BlockTransposed BlockTransposed) -> Self { + pub fn from_strided(v: Strided<'_, T>) -> Self { let nrows = v.nrows(); let ncols = v.ncols(); let mut mat = Self::new(nrows, ncols); @@ -1175,7 +1175,7 @@ impl BlockTransposed::from_strided(strided); diff --git a/diskann-quantization/src/product/tables/transposed/pivots.rs b/diskann-quantization/src/product/tables/transposed/pivots.rs index c76c5c5197..e5f2333795 100644 --- a/diskann-quantization/src/product/tables/transposed/pivots.rs +++ b/diskann-quantization/src/product/tables/transposed/pivots.rs @@ -5,7 +5,7 @@ use std::fmt; -use diskann_utils::strided; +use diskann_utils::strided::Strided; use diskann_wide::{SIMDMask, SIMDMulAdd, SIMDPartialOrd, SIMDSelect, SIMDVector}; use crate::{ @@ -21,9 +21,9 @@ diskann_wide::alias!(u32s = u32x8); /// Error types returned by Chunk construction. #[derive(Debug, Clone)] pub enum ChunkConstructionError { - /// A `StridedView` was provided with a dimension of zero. + /// A `Strided` was provided with a dimension of zero. DimensionCannotBeZero, - /// A `StridedView` was provided with a length of zero. + /// A `Strided` was provided with a length of zero. LengthCannotBeZero, } @@ -220,7 +220,7 @@ impl Chunk { /// /// 1. `data.ncols() == 0` /// 2. `data.nrows() == 0` - pub(super) fn new(data: strided::StridedView<'_, f32>) -> Result { + pub(super) fn new(data: Strided<'_, f32>) -> Result { // Error handling. if data.ncols() == 0 { return Err(ChunkConstructionError::DimensionCannotBeZero); @@ -229,7 +229,7 @@ impl Chunk { return Err(ChunkConstructionError::LengthCannotBeZero); } - let square_norms = data.row_iter().map(kmeans::square_norm).collect(); + let square_norms = data.rows().map(kmeans::square_norm).collect(); let data = BlockTransposed::::from_strided(data); Ok(Self { data, square_norms }) } @@ -346,9 +346,9 @@ impl Chunk { /// /// This method is generally more efficient to call. /// - /// **IMPORTANT**: The provided `StridedView` must have a length of `Self::batchsize()`. + /// **IMPORTANT**: The provided `Strided` must have a length of `Self::batchsize()`. /// - /// Providing a `StridedView` as an argument allows the compiler to infer that each + /// Providing a `Strided` as an argument allows the compiler to infer that each /// row in `x` has a strided offset from the base, allowing for better code generation. /// /// If the distances between a row in `x` and all centers is not finite, return @@ -385,7 +385,7 @@ impl Chunk { /// inner products. pub(super) fn find_closest_batch( &self, - x: strided::StridedView<'_, T>, + x: Strided<'_, T>, ) -> [CompressionResult; Self::batchsize()] where T: Copy + Into, @@ -394,7 +394,7 @@ impl Chunk { assert_eq!( x.nrows(), Self::batchsize(), - "argument StridedView must have a length of {}", + "argument Strided must have a length of {}", Self::batchsize() ); @@ -437,7 +437,7 @@ impl Chunk { debug_assert!(k < x.ncols()); // SAFETY: - // * `StridedView` indexing: We have checked in Assertion 1 that it is safe + // * `Strided` indexing: We have checked in Assertion 1 that it is safe // to index `x` at indices `0, 1, 2, and 3`. // * Inner slice indexing: It is the caller's responsibility to ensure that // `k < x.dim()` so that indexing the slices return by `x.get_unchecked` are @@ -449,19 +449,19 @@ impl Chunk { ( f32s::splat( diskann_wide::ARCH, - >::into(*x.get_row_unchecked(0).get_unchecked(k)), + >::into(*x.row_unchecked(0).get_unchecked(k)), ), f32s::splat( diskann_wide::ARCH, - >::into(*x.get_row_unchecked(1).get_unchecked(k)), + >::into(*x.row_unchecked(1).get_unchecked(k)), ), f32s::splat( diskann_wide::ARCH, - >::into(*x.get_row_unchecked(2).get_unchecked(k)), + >::into(*x.row_unchecked(2).get_unchecked(k)), ), f32s::splat( diskann_wide::ARCH, - >::into(*x.get_row_unchecked(3).get_unchecked(k)), + >::into(*x.row_unchecked(3).get_unchecked(k)), ), ) } @@ -1198,19 +1198,13 @@ mod tests { let query_batch = flatten(&[copy_query(0), copy_query(1), copy_query(2), copy_query(3)]); - let view = strided::StridedView::try_from( - query_batch.as_slice(), - Chunk::batchsize(), - dim, - dim, - ) - .unwrap(); + let view = Strided::try_from_data(query_batch.as_slice(), Chunk::batchsize(), dim, dim) + .unwrap(); // Make sure that the query batch was constructed correctly. assert_eq!(view.nrows(), Chunk::batchsize()); assert_eq!(view.ncols(), dim); - for k in 0..view.nrows() { - let row = view.row(k); + for (k, row) in view.rows().enumerate() { if k == j { assert_eq!( row, query, @@ -1261,13 +1255,8 @@ mod tests { maybe_broadcast(3, values[3]), ]); - let view = strided::StridedView::try_from( - query_batch.as_slice(), - Chunk::batchsize(), - dim, - dim, - ) - .unwrap(); + let view = Strided::try_from_data(query_batch.as_slice(), Chunk::batchsize(), dim, dim) + .unwrap(); let closest = chunk.find_closest_batch(view); // Lane `j` should not be okay. All other lanes should return the correct @@ -1299,7 +1288,7 @@ mod tests { let mut data_aggregate = create_test_pattern(dim, total); let data = flatten(&data_aggregate); - let sliced = strided::StridedView::try_from(data.as_slice(), total, dim, dim).unwrap(); + let sliced = Strided::try_from_data(data.as_slice(), total, dim, dim).unwrap(); let chunk = Chunk::new(sliced).unwrap(); assert_eq!(chunk.num_centers(), total); @@ -1309,7 +1298,7 @@ mod tests { for row in 0..sliced.nrows() { for col in 0..sliced.ncols() { assert_eq!( - sliced[(row, col)], + *sliced.element(row, col), chunk.get(row, col), "failed on row {} and col {}", row, @@ -1383,7 +1372,7 @@ mod tests { let last = data_aggregate.last().unwrap().clone(); data_aggregate[0].clone_from(&last); let data = flatten(&data_aggregate); - let sliced = strided::StridedView::try_from(data.as_slice(), total, dim, dim).unwrap(); + let sliced = Strided::try_from_data(data.as_slice(), total, dim, dim).unwrap(); let chunk = Chunk::new(sliced).unwrap(); assert_eq!(chunk.num_centers(), total); @@ -1447,7 +1436,7 @@ mod tests { #[test] fn test_chunk_construction_error() { // No dimensions - let chunk = Chunk::new(strided::StridedView::try_from(&[], 3, 0, 0).unwrap()); + let chunk = Chunk::new(Strided::try_from_data(&[], 3, 0, 0).unwrap()); let err = chunk.unwrap_err(); assert!( err.to_string() @@ -1455,7 +1444,7 @@ mod tests { ); // No length - let chunk = Chunk::new(strided::StridedView::try_from(&[], 0, 10, 10).unwrap()); + let chunk = Chunk::new(Strided::try_from_data(&[], 0, 10, 10).unwrap()); let err = chunk.unwrap_err(); assert!( err.to_string() @@ -1470,7 +1459,7 @@ mod tests { let dim = 10; let total = 13; let data = flatten(&create_test_pattern(dim, total)); - let sliced = strided::StridedView::try_from(data.as_slice(), total, dim, dim).unwrap(); + let sliced = Strided::try_from_data(data.as_slice(), total, dim, dim).unwrap(); let chunk = Chunk::new(sliced).unwrap(); let query: Vec = vec![0.0; total]; @@ -1486,13 +1475,12 @@ mod tests { let dim = 10; let total = 13; let data = flatten(&create_test_pattern(dim, total)); - let sliced = strided::StridedView::try_from(data.as_slice(), total, dim, dim).unwrap(); + let sliced = Strided::try_from_data(data.as_slice(), total, dim, dim).unwrap(); let chunk = Chunk::new(sliced).unwrap(); let query: Vec = vec![0.0; 4 * total]; let query_view = - strided::StridedView::try_from(query.as_slice(), Chunk::batchsize(), total, total) - .unwrap(); + Strided::try_from_data(query.as_slice(), Chunk::batchsize(), total, total).unwrap(); // PANICS chunk.find_closest_batch(query_view); @@ -1500,18 +1488,17 @@ mod tests { // Make sure `find_closest_batch` panics for an incorrect length #[test] - #[should_panic(expected = "argument StridedView must have a length of")] + #[should_panic(expected = "argument Strided must have a length of")] fn test_find_closest_batch_panics_on_non_batch_length() { let dim = 10; let total = 13; let data = flatten(&create_test_pattern(dim, total)); - let sliced = strided::StridedView::try_from(data.as_slice(), total, dim, dim).unwrap(); + let sliced = Strided::try_from_data(data.as_slice(), total, dim, dim).unwrap(); let chunk = Chunk::new(sliced).unwrap(); let query: Vec = vec![0.0; (Chunk::batchsize() + 1) * dim]; let query_view = - strided::StridedView::try_from(query.as_slice(), Chunk::batchsize() + 1, dim, dim) - .unwrap(); + Strided::try_from_data(query.as_slice(), Chunk::batchsize() + 1, dim, dim).unwrap(); // PANICS chunk.find_closest_batch(query_view); diff --git a/diskann-quantization/src/product/tables/transposed/table.rs b/diskann-quantization/src/product/tables/transposed/table.rs index bc527fd8a2..6e497f5e5c 100644 --- a/diskann-quantization/src/product/tables/transposed/table.rs +++ b/diskann-quantization/src/product/tables/transposed/table.rs @@ -12,7 +12,7 @@ use crate::{ views::{ChunkOffsets, ChunkOffsetsView}, }; use diskann_utils::{ - strided, + strided::Strided, views::{self, MatrixView, MutMatrixView}, }; use thiserror::Error; @@ -89,7 +89,7 @@ impl TransposedTable { .map(|i| { let range = offsets.at(i); largest = largest.max(range.len()); - let view = strided::StridedView::try_shrink_from( + let view = Strided::try_from_data( &(pivots.as_slice()[range.start..]), pivots.nrows(), range.len(), diff --git a/diskann-quantization/src/product/train.rs b/diskann-quantization/src/product/train.rs index 3d9a1a47d1..48605c4a87 100644 --- a/diskann-quantization/src/product/train.rs +++ b/diskann-quantization/src/product/train.rs @@ -4,7 +4,7 @@ */ use diskann_utils::{ - strided::StridedView, + strided::Strided, views::{self, Matrix}, }; #[cfg(feature = "rayon")] @@ -153,7 +153,7 @@ impl TrainQuantizer for LightPQTrainingParameters { // the remaining tasks to exit early. exit_if_canceled()?; - let view = StridedView::try_shrink_from( + let view = Strided::try_from_data( &(data.as_slice()[range.start..]), data.nrows(), range.len(), @@ -163,11 +163,11 @@ impl TrainQuantizer for LightPQTrainingParameters { chunk: i, of: schema.len(), dim: range.len(), - kind: PQTrainingErrorKind::InternalError(Box::new(err.as_static())), + kind: PQTrainingErrorKind::InternalError(Box::new(err)), })?; // Allocate scratch data structures. - let norms: Vec = view.row_iter().map(square_norm).collect(); + let norms: Vec = view.rows().map(square_norm).collect(); let transpose = BlockTransposed::::from_strided(view); let mut centers = Matrix::new(0.0, trainer.ncenters, range.len()); diff --git a/diskann-utils/src/internal.rs b/diskann-utils/src/internal.rs new file mode 100644 index 0000000000..b4a30ef5f5 --- /dev/null +++ b/diskann-utils/src/internal.rs @@ -0,0 +1,9 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +pub(crate) fn slice_to_nonnull(s: &[T]) -> std::ptr::NonNull { + // SAFETY: slices are guaranteed to have non-null base pointers. + unsafe { std::ptr::NonNull::new_unchecked(s.as_ptr().cast_mut()) } +} diff --git a/diskann-utils/src/lib.rs b/diskann-utils/src/lib.rs index 8b16d5acf4..764ee6ad49 100644 --- a/diskann-utils/src/lib.rs +++ b/diskann-utils/src/lib.rs @@ -24,6 +24,8 @@ pub mod views; mod lazystring; pub use lazystring::LazyString; +mod internal; + #[cfg(feature = "testing")] #[doc(hidden)] pub fn workspace_root() -> std::path::PathBuf { diff --git a/diskann-utils/src/strided.rs b/diskann-utils/src/strided.rs index efbfd621ef..b74ee2e0fb 100644 --- a/diskann-utils/src/strided.rs +++ b/diskann-utils/src/strided.rs @@ -3,244 +3,207 @@ * Licensed under the MIT license. */ -use std::{ - fmt, - ops::{Index, IndexMut}, -}; +use std::{fmt, marker::PhantomData, ptr::NonNull}; use thiserror::Error; -use crate::views::{self, DenseData, MutDenseData}; +use crate::{ + internal, + views::{self}, + Reborrow, +}; -/// A row-major strided matrix. +/// The layout for [`Strided`]. /// -/// This is a generalization of the `MatrixBase` class as it does not mandate a dense -/// layout in memory. +/// This struct ensures that the [`Self::cstride`] is greater than or equal to [`Self::ncols`] +/// and that the linear length of the representation does not overflow `usize::MAX`. /// +/// The linear length of [`Strided`] is given by the forumula /// ```text -/// |<------ cstride ----->| -/// |<-- ncols -->| -/// +-------------+ -/// slice 0 -> | a0 a1 a2 a3 | a4 a5 a6 ^ -/// slice 1 -> | b0 b1 b2 b3 | b4 b5 b6 | -/// 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 | -/// slicf 5 -> | f0 f1 f2 f3 | f4 f5 f6 v -/// +-------------+ -/// ^ -/// | -/// StridedView +/// self.nrows.saturating_sub(1) * self.cstride + self.nrows.min(1) * self.ncols /// ``` -/// -/// 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. +/// This allows the last row to occupy less than a full stride. #[derive(Debug, Clone, Copy)] -pub struct StridedBase -where - T: DenseData, -{ - data: T, +pub struct Layout { nrows: usize, ncols: usize, - // The stride along the columns. This must be greater than or equal to `ncols`. cstride: usize, } -/// Return the linear length of a slice underlying a `StridedBase` with the given parameters. -pub fn linear_length(nrows: usize, ncols: usize, cstride: usize) -> usize { - (nrows.max(1) - 1) * cstride + nrows.min(1) * ncols +impl Layout { + /// Construct a new [`Layout`]. + /// + /// Errors if: + /// + /// * `cstride < ncols`. + /// * The computation of the linear length overflows `usize::MAX`. + pub fn new(nrows: usize, ncols: usize, cstride: usize) -> Result { + LayoutError::check(nrows, ncols, cstride)?; + Ok(Self { + nrows, + ncols, + cstride, + }) + } + + /// 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 linear length of the [`Strided`] described by this [`Layout`]. + pub fn linear_length(&self) -> usize { + self.nrows.saturating_sub(1) * self.cstride + self.nrows.min(1) * self.ncols + } } -#[derive(Debug, Error)] -#[non_exhaustive] -#[error( - "tried to construct a strided matrix with {nrows} rows and {ncols} cols and \ - column stride {cstride} over a slice of length {} (expected {})", - len, - linear_length(self.nrows, self.ncols, self.cstride) -)] -pub struct TryFromErrorLight { - len: usize, - nrows: usize, - ncols: usize, - cstride: usize, +fn linear_length(nrows: usize, ncols: usize, cstride: usize) -> Option { + nrows + .saturating_sub(1) + .checked_mul(cstride) + .and_then(|main| main.checked_add(nrows.min(1) * ncols)) } -#[derive(Error)] -#[non_exhaustive] -#[error( - "tried to construct a strided matrix with {nrows} rows and {ncols} cols and \ - column stride {cstride} over a slice of length {} (expected {})", - data.as_slice().len(), - linear_length(self.nrows, self.ncols, self.cstride) -)] -pub struct TryFromError { - data: T, - nrows: usize, - ncols: usize, - cstride: usize, +/// Errors for [`Layout::new`]. +#[derive(Debug)] +pub struct LayoutError(LayoutErrorInner); + +impl LayoutError { + fn check(nrows: usize, ncols: usize, cstride: usize) -> Result { + if cstride < ncols { + Err(Self(LayoutErrorInner::InvalidStride { ncols, cstride })) + } else { + linear_length(nrows, ncols, cstride) + .ok_or(Self(LayoutErrorInner::Overflow { nrows, cstride })) + } + } } -// Manually implement `fmt::Debug` so we don't require `T::Debug`. -impl fmt::Debug for TryFromError { +impl fmt::Display for LayoutError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("TryFromError") - .field("data_len", &self.data.as_slice().len()) - .field("nrows", &self.nrows) - .field("ncols", &self.ncols) - .field("cstride", &self.cstride) - .finish() + self.0.fmt(f) } } -impl TryFromError { - /// Consume the error and return the base data. - pub fn into_inner(self) -> T { - self.data - } +impl std::error::Error for LayoutError {} - /// Drop the data portion of the error and return an equivalent error that is guaranteed - /// to be `'static`. - pub fn as_static(&self) -> TryFromErrorLight { - TryFromErrorLight { - len: self.data.as_slice().len(), - nrows: self.nrows, - ncols: self.ncols, - cstride: self.cstride, - } - } +#[derive(Debug)] +enum LayoutErrorInner { + InvalidStride { ncols: usize, cstride: usize }, + Overflow { nrows: usize, cstride: usize }, } -impl<'a, T> StridedBase<&'a [T]> { - /// Construct a strided view over data slice, shrinking the slice as needed. - /// - /// Returns an error if `data` is shorter than the value returned by `linear_length`. - /// - /// # Panics - /// - /// * Panics if `cstride < ncols`. - pub fn try_shrink_from( - data: &'a [T], - nrows: usize, - ncols: usize, - cstride: usize, - ) -> Result> { - assert!( - cstride >= ncols, - "cstride must be greater than or equal to ncols" - ); - let required_length = linear_length(nrows, ncols, cstride); - match data.get(..required_length) { - Some(data) => Ok(Self { - data, - nrows, - ncols, - cstride, - }), - None => Err(TryFromError { - data, - nrows, - ncols, - cstride, - }), +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 } => write!( + f, + "a {}x{} strided matrix has a length exceeding usize::MAX", + nrows, cstride + ), } } } -impl<'a, T> StridedBase<&'a mut [T]> { - /// Construct a strided view over data slice, shrinking the slice as needed. - /// - /// Returns an error if `data` is shorter than the value returned by `linear_length`. - /// - /// # Panics - /// - /// * Panics if `cstride < ncols`. - pub fn try_shrink_from_mut( - data: &'a mut [T], - nrows: usize, - ncols: usize, - cstride: usize, - ) -> Result> { - assert!( - cstride >= ncols, - "cstride must be greater than or equal to ncols" - ); - let required_length = linear_length(nrows, ncols, cstride); - if data.as_slice().len() >= required_length { - Ok(Self { - data: &mut data[..required_length], - nrows, - ncols, - cstride, - }) - } else { - Err(TryFromError { - data, - nrows, - ncols, - cstride, - }) - } - } +/// A row-major strided matrix. +/// +/// This is a generalization of the `MatrixBase` class 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, + layout: Layout, + _lifetime: PhantomData<&'a [T]>, } -impl StridedBase -where - T: DenseData, -{ +impl<'a, T> Strided<'a, T> { /// Construct a strided view over data slice, shrinking the slice as needed. /// - /// Returns an error if `data` is not equal to the expected length as determined - /// by `linear_length`. - /// - /// # Panics - /// - /// * Panics if `cstride < ncols`. - pub fn try_from( - data: T, + /// 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> { - assert!( - cstride >= ncols, - "cstride must be greater than or equal to ncols" - ); - // This computation needs to be set up such that: - // 1. When `nrows == 0`, the expected length is 0. - // 2. We make a tight upper-bound on the expected length for the last row. - let required_length = linear_length(nrows, ncols, cstride); - if data.as_slice().len() == required_length { - Ok(Self { - data, - nrows, - ncols, - cstride, + ) -> Result { + 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 { - Err(TryFromError { - data, - nrows, - ncols, - cstride, + Ok(Self { + ptr: internal::slice_to_nonnull(data), + layout, + _lifetime: PhantomData, }) } } + fn as_nonnull(&self) -> NonNull { + self.ptr + } + + /// Return the [`Layout`] for the matrix. + pub fn layout(&self) -> Layout { + 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.ncols + self.layout().ncols() } /// Return the number of rows in the matrix. pub fn nrows(&self) -> usize { - self.nrows + self.layout().nrows() } /// Return the count of elements between the start of each row. pub fn cstride(&self) -> usize { - self.cstride + self.layout().cstride() } /// Return the underlying data as a slice. @@ -248,273 +211,231 @@ where /// # Note /// /// The underlying representation for a strided matrix is not necessarily dense. - pub fn as_slice(&self) -> &[T::Elem] { - self.data.as_slice() - } + pub fn as_slice(&self) -> &[T] { + let layout = self.layout(); - /// Return the underlying data as a slice. - /// - /// # Note - /// - /// The underlying representation for a strided matrix is not necessarily dense. - pub fn as_mut_slice(&mut self) -> &mut [T::Elem] - where - T: MutDenseData, - { - self.data.as_mut_slice() + // 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()) } } - /// Return row `row` as a slice. - /// - /// # Panic - /// - /// Panics if `row >= self.nrows()`. - pub fn row(&self, row: usize) -> &[T::Elem] { - assert!( - row < self.nrows(), - "tried to access row {row} of a matrix with {} rows", - self.nrows() - ); - - // SAFETY: `row` is in-bounds. - unsafe { self.get_row_unchecked(row) } - } + // Element Access - /// Returns the requested row without boundschecking. + /// Return the specified element. /// /// # Safety /// - /// The following conditions must hold to avoid undefined behavior: /// * `row < self.nrows()`. - pub unsafe fn get_row_unchecked(&self, row: usize) -> &[T::Elem] { - debug_assert!(row < self.nrows); - let cstride = self.cstride; - let ncols = self.ncols; - let start = row * cstride; - - debug_assert!(start + ncols <= self.as_slice().len()); - // SAFETY: The idempotency requirement of `as_slice` and our audited constructors - // mean that `self.as_slice()` has a length of `self.nrows * self.ncols`. - // - // Therefore, this access is in-bounds. - unsafe { self.as_slice().get_unchecked(start..start + ncols) } + /// * `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 row `row` as a mutable slice. + /// Return the specified element. /// /// # Panics /// - /// Panics if `row >= self.nrows()`. - pub fn row_mut(&mut self, row: usize) -> &mut [T::Elem] - where - T: MutDenseData, - { + /// Panics if `row >= self.nrows()` or `col >= self.ncols()`. + pub fn element(&self, row: usize, col: usize) -> &T { assert!( row < self.nrows(), - "tried to access row {row} of a matrix with {} rows", + "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` is in-bounds. - unsafe { self.get_row_unchecked_mut(row) } + // SAFETY: `row` and `col` are in-bounds. + unsafe { self.element_unchecked(row, col) } } + // Row Access + /// Returns the requested row without boundschecking. /// /// # Safety /// - /// The following conditions must hold to avoid undefined behavior: - /// * `row < self.nrows()`. - pub unsafe fn get_row_unchecked_mut(&mut self, row: usize) -> &mut [T::Elem] - where - T: MutDenseData, - { - debug_assert!(row < self.nrows); - let cstride = self.cstride; - let ncols = self.ncols; - let start = row * cstride; - - debug_assert!(start + ncols <= self.as_slice().len()); - // SAFETY: The idempotency requirement of `as_mut_slice` and our audited constructors - // mean that `self.as_mut_slice()` has a length of `self.nrows * self.ncols`. - // - // Therefore, this access is in-bounds. + /// 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 { - self.data - .as_mut_slice() - .get_unchecked_mut(start..start + ncols) + std::slice::from_raw_parts(self.as_ptr().add(layout.cstride() * row), layout.ncols()) } } - /// Return a iterator over all rows in the matrix. - /// - /// Rows are yielded sequentially beginning with row 0. + /// 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 + /// # Panic /// - /// Panics if `self.ncols() == 0` (because the implementation does not work correctly - /// in this case and it's too corner-case to bother fixing). This restriction may - /// be lifted in the future. - pub fn row_iter(&self) -> impl Iterator { - assert!(self.ncols() != 0); - let ncols = self.ncols; + /// 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() + ); - self.data - .as_slice() - .chunks(self.cstride()) - .map(move |i| &i[..ncols]) + // SAFETY: `row` is in-bounds. + unsafe { self.row_unchecked(row) } } - /// Return a mutable iterator over all rows in the matrix. + /// Return a iterator over all rows in the matrix. /// /// Rows are yielded sequentially beginning with row 0. - /// - /// # Panics - /// - /// Panics if `self.ncols() == 0` (because the implementation does not work correctly - /// in this case and it's too corner-case to bother fixing). This restriction may - /// be lifted in the future. - pub fn row_iter_mut(&mut self) -> impl Iterator - where - T: MutDenseData, - { - assert!(self.ncols() != 0); - - let ncols = self.ncols(); - let cstride = self.cstride(); - self.data - .as_mut_slice() - .chunks_mut(cstride) - .map(move |i| &mut i[..ncols]) + pub fn rows(&self) -> Rows<'_, T> { + Rows::new(*self) } +} - /// Return a pointer to the base of the matrix. - pub fn as_ptr(&self) -> *const T::Elem { - self.as_slice().as_ptr() - } +/// 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 }, +} - /// Return a pointer to the base of the matrix. - pub fn as_mut_ptr(&mut self) -> *mut T::Elem - where - T: MutDenseData, - { - self.as_mut_slice().as_mut_ptr() +impl<'a, T> From> for Strided<'a, T> { + fn from(matrix: views::MatrixView<'a, T>) -> Self { + // FIXME: `views::MatrixView` doesn't quite ensure that `nrows * ncols` doesn't + // overflow. So for now, we need to be pessimistic. + // + // This will be fixed when `MatrixView` gets the same treatment. + Self::try_from_data( + matrix.into_inner(), + matrix.nrows(), + matrix.ncols(), + matrix.ncols(), + ) + .expect("this will be made infallible in the future") } +} - /// Returns a reference to an element without boundschecking. - /// - /// # Safety - /// - /// The following conditions must hold to avoid undefined behavior: - /// * `row < self.nrows()`. - /// * `col < self.ncols()`. - pub unsafe fn get_unchecked(&self, row: usize, col: usize) -> &T::Elem { - debug_assert!(row < self.nrows); - debug_assert!(col < self.ncols); - self.as_slice().get_unchecked(row * self.cstride + col) - } +// SAFETY: `Strided` only exposes shared access to its `T` elements (via `&T`/`&[T]`), so +// sharing or transferring a `Strided` across threads is sound whenever `T: Sync`. +unsafe impl Send for Strided<'_, T> where T: Sync {} +// SAFETY: See above. +unsafe impl Sync for Strided<'_, T> where T: Sync {} - /// Returns a mutable reference to an element without boundschecking. - /// - /// # Safety - /// - /// The following conditions must hold to avoid undefined behavior: - /// * `row < self.nrows()`. - /// * `col < self.ncols()`. - pub unsafe fn get_unchecked_mut(&mut self, row: usize, col: usize) -> &mut T::Elem - where - T: MutDenseData, - { - let cstride = self.cstride; - debug_assert!(row < self.nrows); - debug_assert!(col < self.ncols); - self.as_mut_slice().get_unchecked_mut(row * cstride + col) - } - - /// Return a view over the matrix. - pub fn as_view(&self) -> StridedView<'_, T::Elem> { - StridedView { - data: self.as_slice(), - nrows: self.nrows, - ncols: self.ncols, - cstride: self.cstride, - } +impl Clone for Strided<'_, T> { + fn clone(&self) -> Self { + *self } } -pub type StridedView<'a, T> = StridedBase<&'a [T]>; -pub type MutStridedView<'a, T> = StridedBase<&'a mut [T]>; +impl Copy for Strided<'_, T> {} -/// Return a reference to the item at entry `(row, col)` in the matrix. -/// -/// # Panics -/// -/// Panics if `row >= self.nrows()` or `col >= self.ncols()`. -impl Index<(usize, usize)> for StridedBase -where - T: DenseData, -{ - type Output = T::Elem; - - fn index(&self, (row, col): (usize, usize)) -> &Self::Output { - assert!( - row < self.nrows(), - "row {row} is out of bounds (max: {})", - self.nrows() - ); - assert!( - col < self.ncols(), - "col {col} is out of bounds (max: {})", - self.ncols() - ); - // SAFETY: We have checked that `row` and `col` are in-bounds. - unsafe { self.get_unchecked(row, col) } +impl<'a, T> Reborrow<'a> for Strided<'_, T> { + type Target = Strided<'a, T>; + fn reborrow(&'a self) -> Self::Target { + *self } } -/// Return a mutable reference to the item at entry `(row, col)` in the matrix. -/// -/// # Panics -/// -/// Panics if `row >= self.nrows()` or `col >= self.ncols()`. -impl IndexMut<(usize, usize)> for StridedBase -where - T: MutDenseData, -{ - fn index_mut(&mut self, (row, col): (usize, usize)) -> &mut Self::Output { - assert!( - row < self.nrows(), - "row {row} is out of bounds (max: {})", - self.nrows() - ); - assert!( - col < self.ncols(), - "col {col} is out of bounds (max: {})", - self.ncols() - ); - // SAFETY: We have checked that `row` and `col` are in-bounds. - unsafe { self.get_unchecked_mut(row, col) } - } +/// Iterator for [`Strided::rows`]. +#[derive(Debug)] +pub struct Rows<'a, T> { + ptr: NonNull, + remaining: usize, + ncols: usize, + cstride: usize, + _lifetime: PhantomData<&'a T>, } -impl From> for StridedBase -where - T: DenseData, - U: DenseData, - T: Into, -{ - fn from(matrix: views::MatrixBase) -> Self { - let nrows = matrix.nrows(); - let ncols = matrix.ncols(); +impl<'a, T> Rows<'a, T> { + fn new(strided: Strided<'a, T>) -> Self { + let layout = strided.layout(); Self { - data: matrix.into_inner().into(), - nrows, - ncols, - cstride: ncols, + 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` across threads is sound whenever `T: Sync`. +unsafe impl Send for Rows<'_, T> where T: Sync {} +// SAFETY: See above. +unsafe impl 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) { + (self.remaining, Some(self.remaining)) + } +} + +impl ExactSizeIterator for Rows<'_, T> {} +impl std::iter::FusedIterator for Rows<'_, T> {} + +/////////// +// Tests // +/////////// + #[cfg(test)] mod tests { use super::*; @@ -522,73 +443,231 @@ mod tests { #[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), 0); - assert_eq!(linear_length(0, 2, 2), 0); - assert_eq!(linear_length(0, 2, 3), 0); - assert_eq!(linear_length(0, 2, 4), 0); + 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), row * col); + 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), 5); - assert_eq!(linear_length(1, 7, 99), 7); + 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), (row - 1) * cstride + col); + 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()); } - fn assert_is_static(_x: &T) {} + #[test] + fn test_layout_new() { + // Valid layouts. + let layout = Layout::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::new(3, 4, 6).unwrap(); + assert_eq!(layout.linear_length(), 2 * 6 + 4); + + // `cstride == ncols` is fine, even at zero. + assert!(Layout::new(0, 0, 0).is_ok()); + + // Invalid stride: `cstride < ncols`. + let err = Layout::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::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 + ) + ); + } #[test] - fn try_from_error_misc() { - let x = TryFromError::<&[f32]> { - data: &[], - nrows: 1, - ncols: 2, - cstride: 3, - }; - - let display = format!("{}", x); - let debug = format!("{:?}", x); - println!("debug = {}", debug); - assert!(debug.contains("TryFromError")); - assert!(debug.contains("data_len: 0")); - assert!(debug.contains("nrows: 1")); - assert!(debug.contains("ncols: 2")); - assert!(debug.contains("cstride: 3")); - - let x = x.as_static(); - assert_is_static(&x); + fn test_try_from_data_errors() { + let m = views::Matrix::::new(0, 10, 10); + 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!( - display, - format!("{}", x), - "static version of the error must hav ethe same message" + 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", ); } - fn expected_error(len: usize, nrows: usize, ncols: usize, cstride: usize) -> String { - format!( - "tried to construct a strided matrix with {nrows} rows and {ncols} cols and \ - column stride {cstride} over a slice of length {} (expected {})", - len, - linear_length(nrows, ncols, cstride) - ) + #[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() {} + assert_send_sync::>(); + assert_send_sync::>(); + } + + #[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: StridedView<'_, usize>, expected: views::MatrixView<'_, usize>) { + fn test_indexing(dut: Strided<'_, usize>, expected: views::MatrixView<'_, usize>) { assert_eq!(dut.nrows(), expected.nrows()); assert_eq!(dut.ncols(), expected.ncols()); @@ -602,9 +681,19 @@ mod tests { // Compare via linear indexing. for row in 0..dut.nrows() { for col in 0..dut.ncols() { + let e = expected[(row, col)]; + + assert_eq!( + *dut.element(row, col), + e, + "failed on (row, col) = ({}, {})", + row, + col + ); + assert_eq!( - dut[(row, col)], - expected[(row, col)], + *dut.get_element(row, col).unwrap(), + e, "failed on (row, col) = ({}, {})", row, col @@ -615,10 +704,17 @@ mod tests { // 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.row_iter().eq(expected.row_iter())); + assert!(dut.rows().eq(expected.row_iter())); } // Create a base Matrix with the following pattern: @@ -645,9 +741,9 @@ mod tests { fn test_basic_indexing() { let m = create_test_matrix(5, 3); - // First - test a dense StridedView over the entire matrix. + // First - test a dense Strided view over the entire matrix. let ptr = m.as_ptr(); - let v = StridedView::try_from(m.as_slice(), m.nrows(), m.ncols(), m.ncols()).unwrap(); + 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()); @@ -656,7 +752,7 @@ mod tests { test_indexing(v, m.as_view()); // Now - create a truly strided view over the first two columns. - let v = StridedView::try_from( + let v = Strided::try_from_data( &(m.as_slice()[..(4 * m.ncols() + 2)]), m.nrows(), 2, @@ -675,7 +771,7 @@ mod tests { test_indexing(v, expected.as_view()); // Create a strided view over the last two columns. - let v = StridedView::try_from(&(m.as_slice()[1..]), m.nrows(), 2, m.ncols()).unwrap(); + let v = Strided::try_from_data(&(m.as_slice()[1..]), m.nrows(), 2, m.ncols()).unwrap(); let mut expected = views::Matrix::new(0, 5, 2); for row in 0..expected.nrows() { for col in 0..expected.ncols() { @@ -685,109 +781,11 @@ mod tests { test_indexing(v, expected.as_view()); } - #[test] - fn test_mutable_indexing() { - // The source matrix. - let src = create_test_matrix(5, 4); - - // Initialize using 2d indexing. - { - let mut dst = views::Matrix::::new(0, 5, 10); - - let ptr = dst.as_ptr(); - - let ncols = src.ncols(); - let nrows = src.nrows(); - let cstride = dst.ncols(); - let mut dst_view = - MutStridedView::try_shrink_from_mut(dst.as_mut_slice(), nrows, ncols, cstride) - .unwrap(); - - assert_eq!(dst_view.as_ptr(), ptr); - assert_eq!(dst_view.as_mut_ptr().cast_const(), ptr); - assert_eq!(dst_view.nrows(), nrows); - assert_eq!(dst_view.ncols(), ncols); - assert_eq!(dst_view.cstride(), cstride); - - // Initialize using linear indexing. - for row in 0..dst_view.nrows() { - for col in 0..dst_view.ncols() { - dst_view[(row, col)] = src[(row, col)] - } - } - - // Check equality. - test_indexing(dst_view.as_view(), src.as_view()); - } - - // Initialize using row-wise indexing. - { - let mut dst = views::Matrix::::new(0, 5, 10); - - let ptr = dst.as_ptr(); - - let ncols = src.ncols(); - let nrows = src.nrows(); - let cstride = dst.ncols(); - let mut dst_view = - MutStridedView::try_shrink_from_mut(dst.as_mut_slice(), nrows, ncols, cstride) - .unwrap(); - - assert_eq!(dst_view.as_ptr(), ptr); - assert_eq!(dst_view.as_mut_ptr().cast_const(), ptr); - assert_eq!(dst_view.nrows(), nrows); - assert_eq!(dst_view.ncols(), ncols); - assert_eq!(dst_view.cstride(), cstride); - - // Initialize by looping over rows. - for row in 0..dst_view.nrows() { - dst_view.row_mut(row).copy_from_slice(src.row(row)) - } - - // Check equality. - test_indexing(dst_view.as_view(), src.as_view()); - } - - // Initialize using row-iterator indexing. - { - let mut dst = views::Matrix::::new(0, 5, 10); - - let offset = 2; - // SAFETY: The underlying allocation is valid for much more than 2 elements. - let ptr = unsafe { dst.as_ptr().add(offset) }; - - let ncols = src.ncols(); - let nrows = src.nrows(); - let cstride = dst.ncols(); - let mut dst_view = MutStridedView::try_shrink_from_mut( - &mut dst.as_mut_slice()[2..], - nrows, - ncols, - cstride, - ) - .unwrap(); - - assert_eq!(dst_view.as_ptr(), ptr); - assert_eq!(dst_view.as_mut_ptr().cast_const(), ptr); - assert_eq!(dst_view.nrows(), nrows); - assert_eq!(dst_view.ncols(), ncols); - assert_eq!(dst_view.cstride(), cstride); - - // Initialize using row iterators. - for (d, s) in std::iter::zip(dst_view.row_iter_mut(), src.row_iter()) { - d.copy_from_slice(s) - } - - // Check equality. - test_indexing(dst_view.as_view(), src.as_view()); - } - } - #[test] fn matrix_conversion() { let m = create_test_matrix(3, 4); let ptr = m.as_ptr(); - let v: StridedView<_> = m.as_view().into(); + let v: Strided<_> = m.as_view().into(); assert_eq!(v.as_ptr(), ptr); test_indexing(v, m.as_view()); } @@ -795,186 +793,47 @@ mod tests { #[test] fn test_zero_sized() { let m = create_test_matrix(5, 5); - let v = StridedView::try_shrink_from(m.as_slice(), 0, 4, 5).unwrap(); + 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 = StridedView::try_shrink_from(m.as_slice(), 5, 0, 5).unwrap(); + 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.row(row), empty); + assert_eq!(v.get_row(row).unwrap(), empty); } } - #[test] - #[should_panic] - fn test_row_iter_panics() { - let m = create_test_matrix(5, 5); - let v = StridedView::try_shrink_from(m.as_slice(), 5, 0, 5).unwrap(); - let _ = v.row_iter(); - } - - #[test] - #[should_panic] - fn test_row_iter_mut_panics() { - let mut m = create_test_matrix(5, 5); - let mut v = MutStridedView::try_shrink_from_mut(m.as_mut_slice(), 5, 0, 5).unwrap(); - let _ = v.row_iter_mut(); - } - #[test] fn test_try_shrink_from() { // Exact is okay. let m = views::Matrix::::new(0, 10, 10); let nrows = m.nrows(); let ncols = m.ncols(); - let s = StridedView::try_shrink_from(m.as_slice(), nrows, ncols, ncols).unwrap(); + 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 = StridedView::try_shrink_from(m.as_slice(), nrows, 5, ncols).unwrap(); + 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. - let s = StridedView::try_shrink_from(m.as_slice(), nrows, ncols, ncols + 1); + // 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()); - let err = s.unwrap_err(); - assert_eq!( - err.to_string(), - expected_error(m.as_slice().len(), nrows, ncols, ncols + 1) - ); - assert_eq!(err.into_inner(), m.as_slice()); } #[test] - #[should_panic(expected = "cstride must be greater than or equal to ncols")] - fn test_try_shink_from_panics() { + 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 = views::Matrix::::new(0, 4, 4); - let _ = StridedView::try_shrink_from(m.as_slice(), 2, 2, 1); - } - - #[test] - fn test_try_shrink_from_mut() { - // Exact is okay. - let mut m = views::Matrix::::new(0, 10, 10); - - let nrows = m.nrows(); - let ncols = m.ncols(); - let ptr = m.as_ptr(); - let len = m.as_slice().len(); - - let s = MutStridedView::try_shrink_from_mut(m.as_mut_slice(), nrows, ncols, ncols).unwrap(); - assert_eq!(s.as_ptr(), ptr); - assert_eq!(s.as_slice().len(), len); - - // Giving a slice that is too large is okay. - let s = MutStridedView::try_shrink_from_mut(m.as_mut_slice(), nrows, 5, ncols).unwrap(); - assert_eq!(s.as_ptr(), ptr); - - // Too small is a problem. - let s = MutStridedView::try_shrink_from_mut(m.as_mut_slice(), nrows, ncols, ncols + 1); - assert!(s.is_err()); - let err = s.unwrap_err(); - assert_eq!( - err.to_string(), - expected_error(len, nrows, ncols, ncols + 1) - ); - } - - #[test] - #[should_panic(expected = "cstride must be greater than or equal to ncols")] - fn test_try_shink_from_mut_panics() { - let mut m = views::Matrix::::new(0, 4, 4); - let _ = MutStridedView::try_shrink_from_mut(m.as_mut_slice(), 2, 2, 1); - } - - #[test] - fn test_try_from() { - // Exact is okay. - let m = views::Matrix::::new(0, 10, 10); - let nrows = m.nrows(); - let ncols = m.ncols(); - let s = StridedView::try_from(m.as_slice(), nrows, ncols, ncols).unwrap(); - assert_eq!(s.as_slice(), m.as_slice()); - - // Giving a slice that is too large is a problem. - let s = StridedView::try_from(m.as_slice(), nrows, 5, ncols); - assert!(s.is_err()); - let err = s.unwrap_err(); - assert_eq!( - err.to_string(), - expected_error(m.as_slice().len(), nrows, 5, ncols) - ); - - // Too small is a problem. - let s = StridedView::try_from(m.as_slice(), nrows, ncols, ncols + 1); - assert!(s.is_err()); - let err = s.unwrap_err(); - assert_eq!( - err.to_string(), - expected_error(m.as_slice().len(), nrows, ncols, ncols + 1) - ); - assert_eq!(err.into_inner(), m.as_slice()); - } - - #[test] - #[should_panic(expected = "cstride must be greater than or equal to ncols")] - fn test_try_frompanics() { - let mut m = views::Matrix::::new(0, 4, 4); - let _ = MutStridedView::try_from(m.as_mut_slice(), 2, 2, 1); - } - - #[test] - #[should_panic(expected = "tried to access row 3 of a matrix with 3 rows")] - fn test_get_row_panics() { - let m = views::Matrix::::new(0, 3, 7); - let v: StridedView<_> = m.as_view().into(); - v.row(3); - } - - #[test] - #[should_panic(expected = "tried to access row 3 of a matrix with 3 rows")] - fn test_get_row_mut_panics() { - let mut m = views::Matrix::::new(0, 3, 7); - let mut v: MutStridedView<_> = m.as_mut_view().into(); - v.row_mut(3); - } - - #[test] - #[should_panic(expected = "row 3 is out of bounds (max: 3)")] - fn test_index_panics_row() { - let m = views::Matrix::::new(0, 3, 7); - let v: StridedView<_> = m.as_view().into(); - let _ = v[(3, 2)]; - } - - #[test] - #[should_panic(expected = "col 7 is out of bounds (max: 7)")] - fn test_index_panics_col() { - let m = views::Matrix::::new(0, 3, 7); - let v: StridedView<_> = m.as_view().into(); - let _ = v[(2, 7)]; - } - - #[test] - #[should_panic(expected = "row 3 is out of bounds (max: 3)")] - fn test_index_mut_panics_row() { - let mut m = views::Matrix::::new(0, 3, 7); - let mut v: MutStridedView<_> = m.as_mut_view().into(); - v[(3, 2)] = 1; - } - - #[test] - #[should_panic(expected = "col 7 is out of bounds (max: 7)")] - fn test_index_mut_panics_col() { - let mut m = views::Matrix::::new(0, 3, 7); - let mut v: MutStridedView<_> = m.as_mut_view().into(); - v[(2, 7)] = 1; + let err = Strided::try_from_data(m.as_slice(), 2, 2, 1).unwrap_err(); + assert!(matches!(err, TryFromError::LayoutError(_))); } }