diff --git a/Cargo.lock b/Cargo.lock index 20d6ad95d..2bfc99171 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -766,6 +766,7 @@ dependencies = [ "rayon", "rstest", "thiserror 2.0.17", + "trybuild", ] [[package]] diff --git a/diskann-benchmark/src/multi_vector/driver.rs b/diskann-benchmark/src/multi_vector/driver.rs index e69c70845..e91d23b18 100644 --- a/diskann-benchmark/src/multi_vector/driver.rs +++ b/diskann-benchmark/src/multi_vector/driver.rs @@ -14,7 +14,8 @@ use diskann_benchmark_runner::{ }, Checker, Input, }; -use diskann_quantization::multi_vector::{Mat, MatRef, MaxSimKernel, Overflow, Standard}; +use diskann_quantization::multi_vector::{Mat, MatRef, MaxSimKernel, RowMajor}; +use diskann_utils::views::Init; use rand::{ distr::{Distribution, StandardUniform}, rngs::StdRng, @@ -68,25 +69,27 @@ impl Input for MultiVectorTolerance { /// Random query / doc fixture for a single benchmark run. pub(super) struct Data { - pub(super) queries: Mat>, - pub(super) docs: Mat>, + pub(super) queries: Mat>, + pub(super) docs: Mat>, } impl Data where StandardUniform: Distribution, { - pub(super) fn new(run: &Run) -> Result { + pub(super) fn new(run: &Run) -> Self { let mut rng = StdRng::seed_from_u64(0x12345); - let queries = Mat::from_fn( - Standard::new(run.num_query_vectors.get(), run.dim.get())?, - || StandardUniform.sample(&mut rng), + let queries = Mat::new( + Init(|| StandardUniform.sample(&mut rng)), + run.num_query_vectors.get(), + run.dim.get(), ); - let docs = Mat::from_fn( - Standard::new(run.num_doc_vectors.get(), run.dim.get())?, - || StandardUniform.sample(&mut rng), + let docs = Mat::new( + Init(|| StandardUniform.sample(&mut rng)), + run.num_doc_vectors.get(), + run.dim.get(), ); - Ok(Self { queries, docs }) + Self { queries, docs } } } @@ -96,7 +99,7 @@ where pub(super) fn run_with_kernel( run: &Run, - doc: MatRef<'_, Standard>, + doc: MatRef<'_, RowMajor>, kernel: &dyn MaxSimKernel, ) -> RunResult { let mut scores = vec![0.0f32; run.num_query_vectors.get()]; diff --git a/diskann-benchmark/src/multi_vector/kernels.rs b/diskann-benchmark/src/multi_vector/kernels.rs index d45ac646b..12759ea11 100644 --- a/diskann-benchmark/src/multi_vector/kernels.rs +++ b/diskann-benchmark/src/multi_vector/kernels.rs @@ -66,7 +66,7 @@ where writeln!(output, "{}", input)?; let mut results = Vec::with_capacity(input.runs.len()); for run in input.runs.iter() { - let data = Data::::new(run)?; + let data = Data::::new(run); let kernel = build_max_sim::(input.isa.into(), data.queries.as_view(), BoxErase)?; results.push(run_with_kernel(run, data.docs.as_view(), &*kernel)); } diff --git a/diskann-disk/src/search/pq/quantizer_preprocess.rs b/diskann-disk/src/search/pq/quantizer_preprocess.rs index c5a202630..344570ba5 100644 --- a/diskann-disk/src/search/pq/quantizer_preprocess.rs +++ b/diskann-disk/src/search/pq/quantizer_preprocess.rs @@ -22,7 +22,7 @@ pub fn quantizer_preprocess( ) -> ANNResult<()> { let table = pq_data.pq_table(); let expected_len = table.ncenters() * table.nchunks(); - let dst = diskann_utils::views::MutMatrixView::try_from( + let dst = diskann_utils::views::MatrixViewMut::try_from( &mut (*pq_scratch.aligned_pqtable_dist_scratch)[..expected_len], table.nchunks(), table.ncenters(), diff --git a/diskann-disk/src/storage/quant/compressor.rs b/diskann-disk/src/storage/quant/compressor.rs index acddd8994..dfd82b3a3 100644 --- a/diskann-disk/src/storage/quant/compressor.rs +++ b/diskann-disk/src/storage/quant/compressor.rs @@ -4,7 +4,7 @@ */ use diskann::{utils::VectorRepr, ANNResult}; -use diskann_utils::views::{MatrixView, MutMatrixView}; +use diskann_utils::views::{MatrixView, MatrixViewMut}; /// [`QuantCompressor`] defines the interface for quantizer with [`QuantDataGenerator`] /// @@ -30,6 +30,6 @@ where type CompressorContext; fn new(context: &Self::CompressorContext) -> ANNResult; - fn compress(&self, vector: MatrixView, output: MutMatrixView) -> ANNResult<()>; + fn compress(&self, vector: MatrixView, output: MatrixViewMut) -> ANNResult<()>; fn compressed_bytes(&self) -> usize; } diff --git a/diskann-disk/src/storage/quant/generator.rs b/diskann-disk/src/storage/quant/generator.rs index 7c09c1182..45ed4f6bb 100644 --- a/diskann-disk/src/storage/quant/generator.rs +++ b/diskann-disk/src/storage/quant/generator.rs @@ -147,7 +147,7 @@ where // Wrap the data in `MatrixViews` so we do not need to manually construct view // in the compression loop. - let mut compressed_block = views::MutMatrixView::try_from( + let mut compressed_block = views::MatrixViewMut::try_from( block_compressed_base, cur_block_size, compressed_size, @@ -221,7 +221,7 @@ mod generator_tests { fn compress( &self, _vector: views::MatrixView, - mut output: views::MutMatrixView, + mut output: views::MatrixViewMut, ) -> ANNResult<()> { output .row_iter_mut() diff --git a/diskann-disk/src/storage/quant/pq/pq_generation.rs b/diskann-disk/src/storage/quant/pq/pq_generation.rs index b1eae30c3..e317d8e3c 100644 --- a/diskann-disk/src/storage/quant/pq/pq_generation.rs +++ b/diskann-disk/src/storage/quant/pq/pq_generation.rs @@ -16,7 +16,7 @@ use diskann_providers::{ utils::{BridgeErr, RayonThreadPoolRef}, }; use diskann_quantization::{product::TransposedTable, CompressInto}; -use diskann_utils::views::MatrixBase; +use diskann_utils::views::{MatrixView, MatrixViewMut}; use diskann_vector::distance::Metric; use tracing::info; @@ -119,7 +119,7 @@ where context.storage_provider, )?; - let mut full_pivot_data_mat = diskann_utils::views::MutMatrixView::try_from( + let mut full_pivot_data_mat = diskann_utils::views::MatrixViewMut::try_from( full_pivot_data.as_mut_slice(), context.num_centers, full_dim, @@ -146,8 +146,8 @@ where fn compress( &self, - vector: MatrixBase<&[f32]>, - output: MatrixBase<&mut [u8]>, + vector: MatrixView<'_, f32>, + output: MatrixViewMut<'_, u8>, ) -> Result<(), diskann::ANNError> { self.table .compress_into(vector, output) @@ -175,7 +175,7 @@ mod pq_generation_tests { use diskann_utils::{ io::{read_bin, write_bin}, test_data_root, - views::{MatrixView, MutMatrixView}, + views::{MatrixView, MatrixViewMut}, }; use diskann_vector::distance::Metric; use rstest::rstest; @@ -329,7 +329,7 @@ mod pq_generation_tests { let mut compressed_mat = vec![0_u8; num_chunks * npts]; let result = compressor.unwrap().compress( data_matrix.as_view(), - MutMatrixView::try_from(&mut compressed_mat, npts, num_chunks).unwrap(), + MatrixViewMut::try_from(&mut compressed_mat, npts, num_chunks).unwrap(), ); assert!(result.is_ok()); diff --git a/diskann-garnet/src/provider.rs b/diskann-garnet/src/provider.rs index 5024b7b9d..b4c608334 100644 --- a/diskann-garnet/src/provider.rs +++ b/diskann-garnet/src/provider.rs @@ -499,7 +499,7 @@ impl GarnetProvider { Ok(v) => v, Err(_) => return false, }; - let view = match MatrixView::try_from(&*converted, view.nrows(), view.ncols()) { + let view = match MatrixView::try_from(&converted, view.nrows(), view.ncols()) { Ok(v) => v, Err(_) => return false, }; diff --git a/diskann-providers/src/model/graph/provider/determinant_diversity.rs b/diskann-providers/src/model/graph/provider/determinant_diversity.rs index 5e35c8708..a553a5cb6 100644 --- a/diskann-providers/src/model/graph/provider/determinant_diversity.rs +++ b/diskann-providers/src/model/graph/provider/determinant_diversity.rs @@ -53,7 +53,7 @@ use std::fmt; -use diskann_utils::views::MutMatrixView; +use diskann_utils::views::MatrixViewMut; use diskann_vector::{MathematicalValue, PureDistanceFunction, distance::InnerProduct}; /// Parameters for Determinant-Diversity post-processor with validation. @@ -199,7 +199,7 @@ struct DistanceRange { /// An empty candidate set, a `k` of zero, or zero-dimensional vectors yield an /// empty result. pub fn determinant_diversity( - candidates: MutMatrixView<'_, f32>, + candidates: MatrixViewMut<'_, f32>, distances: &[f32], query: &[f32], k: usize, @@ -338,7 +338,7 @@ pub fn determinant_diversity( /// O(n * k * dim) -- for each of k pivots we touch all n residual rows of /// length `dim`. Memory is O(n * dim) for the contiguous residual matrix. fn greedy_orthogonal_select( - mut candidates: MutMatrixView<'_, f32>, + mut candidates: MatrixViewMut<'_, f32>, distances: &[f32], k: usize, power: f32, diff --git a/diskann-providers/src/model/pq/fixed_chunk_pq_table.rs b/diskann-providers/src/model/pq/fixed_chunk_pq_table.rs index 2973db8fd..2ef019f73 100644 --- a/diskann-providers/src/model/pq/fixed_chunk_pq_table.rs +++ b/diskann-providers/src/model/pq/fixed_chunk_pq_table.rs @@ -9,7 +9,7 @@ use diskann_quantization::{ product::{self, BasicTable}, views::ChunkOffsetsBase, }; -use diskann_utils::views::{self, MatrixBase, MatrixView}; +use diskann_utils::views::{self, Matrix, MatrixView}; use diskann_vector::{PureDistanceFunction, distance}; use diskann_wide::ARCH; @@ -133,7 +133,7 @@ impl FixedChunkPQTable { pub fn new(dim: usize, pq_table: Box<[f32]>, chunk_offsets: Box<[usize]>) -> ANNResult { let len = pq_table.len(); let table = BasicTable::new( - MatrixBase::try_from(pq_table, len / dim, dim).bridge_err()?, + Matrix::try_from(pq_table, len / dim, dim).bridge_err()?, ChunkOffsetsBase::new(chunk_offsets).bridge_err()?, ) .map_err(|err| ANNError::log_pq_error(diskann_quantization::error::format(&err)))?; diff --git a/diskann-providers/src/model/pq/pq_construction.rs b/diskann-providers/src/model/pq/pq_construction.rs index 8ca3f11a4..91fa6ca5e 100644 --- a/diskann-providers/src/model/pq/pq_construction.rs +++ b/diskann-providers/src/model/pq/pq_construction.rs @@ -25,7 +25,7 @@ use diskann_quantization::{ }; use diskann_utils::{ io::Metadata, - views::{MatrixView, MutMatrixView}, + views::{MatrixView, MatrixViewMut}, }; use rand::{Rng, distr::Distribution}; use rayon::prelude::*; @@ -293,7 +293,7 @@ pub fn move_train_data_by_centroid( /// # Panics /// /// Panics if `y.len() != x.ncols()`. -pub fn accum_row_inplace(mut x: MutMatrixView, y: &[T]) +pub fn accum_row_inplace(mut x: MatrixViewMut, y: &[T]) where T: Copy + std::ops::AddAssign, { @@ -362,7 +362,7 @@ where // Instead of subtracting the center from each data set component, we instead // add it to each center. let mut full_pivot_data_mat = - MutMatrixView::try_from(full_pivot_data.as_mut_slice(), num_centers, full_dim) + MatrixViewMut::try_from(full_pivot_data.as_mut_slice(), num_centers, full_dim) .bridge_err()?; accum_row_inplace(full_pivot_data_mat.as_mut_view(), centroid.as_slice()); @@ -430,7 +430,7 @@ where // Wrap the data in `MatrixViews` so we do not need to manually construct view // in the compression loop. let mut compressed_block = - MutMatrixView::try_from(&mut block_compressed_base, cur_block_size, num_pq_chunks) + MatrixViewMut::try_from(&mut block_compressed_base, cur_block_size, num_pq_chunks) .bridge_err()?; let base_block = MatrixView::try_from(block_data, cur_block_size, full_dim).bridge_err()?; diff --git a/diskann-providers/src/model/pq/views.rs b/diskann-providers/src/model/pq/views.rs index 17028de3c..b6dc8652b 100644 --- a/diskann-providers/src/model/pq/views.rs +++ b/diskann-providers/src/model/pq/views.rs @@ -106,7 +106,7 @@ mod tests { let data = vec![0; ncols * nrows]; test_error(|| { - views::MatrixView::try_from(&*data, nrows, ncols + 1) + views::MatrixView::try_from(&data, nrows, ncols + 1) .bridge_err() .unwrap_err() }); diff --git a/diskann-providers/src/storage/pq_storage.rs b/diskann-providers/src/storage/pq_storage.rs index a2d49391b..fcae33c24 100644 --- a/diskann-providers/src/storage/pq_storage.rs +++ b/diskann-providers/src/storage/pq_storage.rs @@ -523,7 +523,7 @@ mod pq_storage_tests { #[test] fn write_read_roundtrip_with_legacy_centroid() { use crate::model::pq::accum_row_inplace; - use diskann_utils::views::MutMatrixView; + use diskann_utils::views::MatrixViewMut; let storage_provider = VirtualStorageProvider::new_memory(); let pivot_path = "/roundtrip_legacy_centroid_pivots.bin"; @@ -563,7 +563,7 @@ mod pq_storage_tests { // Fold the centroid into the pivots — this is what production callers do. let mut pivot_mat = - MutMatrixView::try_from(loaded_pivots.as_mut_slice(), num_centers, dim).unwrap(); + MatrixViewMut::try_from(loaded_pivots.as_mut_slice(), num_centers, dim).unwrap(); accum_row_inplace(pivot_mat.as_mut_view(), &loaded_centroid); // Each pivot row should have the centroid added element-wise. diff --git a/diskann-quantization/src/algorithms/kmeans/lloyds.rs b/diskann-quantization/src/algorithms/kmeans/lloyds.rs index f81894f1e..5681b48f1 100644 --- a/diskann-quantization/src/algorithms/kmeans/lloyds.rs +++ b/diskann-quantization/src/algorithms/kmeans/lloyds.rs @@ -9,7 +9,7 @@ use super::common::square_norm; use crate::multi_vector::{BlockTransposed, BlockTransposedRef}; use diskann_utils::{ strided::StridedView, - views::{Matrix, MatrixView, MutMatrixView}, + views::{Matrix, MatrixView, MatrixViewMut}, }; //////////////////////////////// @@ -342,7 +342,7 @@ 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: MatrixViewMut<'_, f32>, data: StridedView<'_, 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)| { @@ -373,7 +373,7 @@ pub(crate) fn lloyds_inner( data: StridedView<'_, f32>, square_norms: &[f32], transpose: BlockTransposedRef<'_, f32, 16>, - mut centers: MutMatrixView<'_, f32>, + mut centers: MatrixViewMut<'_, f32>, max_reps: usize, ) -> (Vec, f32) { // Check our requirements. @@ -440,7 +440,7 @@ pub(crate) fn lloyds_inner( /// dimension. pub fn lloyds( data: MatrixView<'_, f32>, - centers: MutMatrixView<'_, f32>, + centers: MatrixViewMut<'_, f32>, max_reps: usize, ) -> (Vec, f32) { assert_eq!( diff --git a/diskann-quantization/src/algorithms/kmeans/plusplus.rs b/diskann-quantization/src/algorithms/kmeans/plusplus.rs index 79abf0e94..652139335 100644 --- a/diskann-quantization/src/algorithms/kmeans/plusplus.rs +++ b/diskann-quantization/src/algorithms/kmeans/plusplus.rs @@ -7,7 +7,7 @@ use std::{collections::HashSet, fmt}; use diskann_utils::{ strided::StridedView, - views::{MatrixView, MutMatrixView}, + views::{MatrixView, MatrixViewMut}, }; use diskann_wide::{SIMDMulAdd, SIMDPartialOrd, SIMDSelect, SIMDVector}; use rand::{ @@ -379,7 +379,7 @@ impl KMeansPlusPlusError { } pub(crate) fn kmeans_plusplus_into_inner( - mut points: MutMatrixView<'_, f32>, + mut points: MatrixViewMut<'_, f32>, data: StridedView<'_, f32>, transpose: BlockTransposedRef<'_, f32, N>, norms: &[f32], @@ -498,7 +498,7 @@ where } pub fn kmeans_plusplus_into( - centers: MutMatrixView<'_, f32>, + centers: MatrixViewMut<'_, f32>, data: MatrixView<'_, f32>, rng: &mut dyn RngCore, ) -> Result<(), KMeansPlusPlusError> { @@ -581,7 +581,7 @@ mod tests { /// ... /// K-1, K, K+1, K+3 ... N+K-2 /// ``` - fn set_default_values(mut x: MutMatrixView<'_, f32>) { + fn set_default_values(mut x: MatrixViewMut<'_, f32>) { for (i, row) in x.row_iter_mut().enumerate() { for (j, r) in row.iter_mut().enumerate() { *r = (i + j) as f32; diff --git a/diskann-quantization/src/lib.rs b/diskann-quantization/src/lib.rs index 248d02f23..5446f9773 100644 --- a/diskann-quantization/src/lib.rs +++ b/diskann-quantization/src/lib.rs @@ -216,7 +216,6 @@ mod tests { t.pass("tests/compile-fail/bootstrap/bootstrap.rs"); t.compile_fail("tests/compile-fail/*.rs"); t.compile_fail("tests/compile-fail/error/*.rs"); - t.compile_fail("tests/compile-fail/multi/*.rs"); } } diff --git a/diskann-quantization/src/minmax/multi/max_sim.rs b/diskann-quantization/src/minmax/multi/max_sim.rs index 5bc0aa406..247add67a 100644 --- a/diskann-quantization/src/minmax/multi/max_sim.rs +++ b/diskann-quantization/src/minmax/multi/max_sim.rs @@ -144,7 +144,7 @@ mod tests { use crate::algorithms::transforms::NullTransform; use crate::bits::{Representation, Unsigned}; use crate::minmax::{Data, MinMaxQuantizer}; - use crate::multi_vector::{Defaulted, Mat, Standard}; + use crate::multi_vector::{Mat, RowMajor}; use crate::num::Positive; use diskann_utils::ReborrowMut; use std::num::NonZeroUsize; @@ -202,9 +202,9 @@ mod tests { where Unsigned: Representation, { - let input_mat = MatRef::new(Standard::::new(n, dim).unwrap(), input).unwrap(); + let input_mat = MatRef::from_repr(RowMajor::::new(n, dim).unwrap(), input).unwrap(); let mut output: Mat> = - Mat::new(MinMaxMeta::new(n, dim), Defaulted).unwrap(); + Mat::from_default(MinMaxMeta::new(n, dim)).unwrap(); quantizer .compress_into(input_mat, output.reborrow_mut()) .unwrap(); @@ -301,10 +301,10 @@ mod tests { let query_data = vec![0u8; 2 * row_bytes]; let doc_data = vec![0u8; 3 * row_bytes]; - let query: QueryMatRef<_> = MatRef::new(MinMaxMeta::<8>::new(2, dim), &query_data) + let query: QueryMatRef<_> = MatRef::from_repr(MinMaxMeta::<8>::new(2, dim), &query_data) .unwrap() .into(); - let doc = MatRef::new(MinMaxMeta::<8>::new(3, dim), &doc_data).unwrap(); + let doc = MatRef::from_repr(MinMaxMeta::<8>::new(3, dim), &doc_data).unwrap(); let mut scores = vec![0.0f32; 5]; // Wrong size MaxSim::new(&mut scores).evaluate(query, doc); diff --git a/diskann-quantization/src/minmax/multi/meta.rs b/diskann-quantization/src/minmax/multi/meta.rs index 678cbd529..cf136c594 100644 --- a/diskann-quantization/src/minmax/multi/meta.rs +++ b/diskann-quantization/src/minmax/multi/meta.rs @@ -10,10 +10,10 @@ use super::super::vectors::DataMutRef; use crate::CompressInto; use crate::bits::{Representation, Unsigned}; use crate::minmax::{self, Data}; -use crate::multi_vector::matrix::{ - Defaulted, NewMut, NewOwned, NewRef, Repr, ReprMut, ReprOwned, SliceError, +use crate::multi_vector::views::{ + NewDefault, NewMut, NewRef, Repr, ReprMut, ReprOwned, SliceError, }; -use crate::multi_vector::{LayoutError, Mat, MatMut, MatRef, Standard}; +use crate::multi_vector::{LayoutError, Mat, MatMut, MatRef, RowMajor}; use crate::scalar::InputContainsNaN; use crate::utils; @@ -165,13 +165,13 @@ where // SAFETY: `ptr` points to a properly sized slice that is compatible with the drop // logic in `Self as ReprOwned`. Box guarantees that the initial construction // will be non-null. -unsafe impl NewOwned for MinMaxMeta +unsafe impl NewDefault for MinMaxMeta where Unsigned: Representation, { type Error = crate::error::Infallible; - fn new_owned(self, _: Defaulted) -> Result, Self::Error> { - let b: Box<[u8]> = (0..self.bytes()).map(|_| u8::default()).collect(); + fn new_default(self) -> Result, Self::Error> { + let b: Box<[u8]> = vec![0u8; self.bytes()].into_boxed_slice(); let ptr = utils::box_into_nonnull(b).cast::(); @@ -231,7 +231,7 @@ where ////////////////// impl<'a, 'b, const NBITS: usize, T> - CompressInto>, MatMut<'b, MinMaxMeta>> for MinMaxQuantizer + CompressInto>, MatMut<'b, MinMaxMeta>> for MinMaxQuantizer where T: Copy + Into, Unsigned: Representation, @@ -257,7 +257,7 @@ where /// * The output intrinsic dimension doesn't match `self.output_dim()`. fn compress_into( &self, - from: MatRef<'a, Standard>, + from: MatRef<'a, RowMajor>, mut to: MatMut<'b, MinMaxMeta>, ) -> Result<(), Self::Error> { assert_eq!( @@ -372,7 +372,7 @@ mod tests { mod construction { use super::*; - /// Tests NewOwned (Mat::new with Defaulted) for various dimensions and sizes. + /// Tests NewDefault (Mat::from_default) for various dimensions and sizes. fn test_new_owned() where Unsigned: Representation, @@ -381,7 +381,7 @@ mod tests { for &num_vectors in TEST_NVECS { let meta = MinMaxMeta::::new(num_vectors, dim); let mat: Mat> = - Mat::new(meta, Defaulted).expect("NewOwned should succeed"); + Mat::from_default(meta).expect("NewOwned should succeed"); // Verify num_vectors and intrinsic_dim assert_eq!(mat.num_vectors(), num_vectors); @@ -412,7 +412,7 @@ mod tests { expand_to_bitrates!(new_owned, test_new_owned); - /// Tests NewRef (MatRef::new) for valid slices. + /// Tests NewRef (MatRef::from_repr) for valid slices. fn test_new_ref() where Unsigned: Representation, @@ -425,7 +425,7 @@ mod tests { // Valid slice let data = vec![0u8; expected_bytes]; - let mat_ref = MatRef::new(meta, &data); + let mat_ref = MatRef::from_repr(meta, &data); assert!(mat_ref.is_ok(), "NewRef should succeed for correct size"); let mat_ref = mat_ref.unwrap(); @@ -446,7 +446,7 @@ mod tests { expand_to_bitrates!(new_ref, test_new_ref); - /// Tests NewMut (MatMut::new) for valid slices. + /// Tests NewMut (MatMut::from_repr) for valid slices. fn test_new_mut() where Unsigned: Representation, @@ -458,7 +458,7 @@ mod tests { let expected_bytes = expected_row_bytes * num_vectors; let mut data = vec![0u8; expected_bytes]; - let mat_mut = MatMut::new(meta, &mut data); + let mat_mut = MatMut::from_repr(meta, &mut data); assert!(mat_mut.is_ok(), "NewMut should succeed for correct size"); let mat_mut = mat_mut.unwrap(); @@ -488,7 +488,7 @@ mod tests { // Too short let short_data = vec![0u8; expected_bytes - 1]; - let result = MatRef::new(meta, &short_data); + let result = MatRef::from_repr(meta, &short_data); assert!( matches!(result, Err(SliceError::LengthMismatch { .. })), "should fail for too-short slice" @@ -496,7 +496,7 @@ mod tests { // Too long let long_data = vec![0u8; expected_bytes + 1]; - let result = MatRef::new(meta, &long_data); + let result = MatRef::from_repr(meta, &long_data); assert!( matches!(result, Err(SliceError::LengthMismatch { .. })), "should fail for too-long slice" @@ -504,7 +504,7 @@ mod tests { // Mutable version - too short let mut short_mut = vec![0u8; expected_bytes - 1]; - let result = MatMut::new(meta, &mut short_mut); + let result = MatMut::from_repr(meta, &mut short_mut); assert!( matches!(result, Err(SliceError::LengthMismatch { .. })), "MatMut should fail for too-short slice" @@ -532,11 +532,11 @@ mod tests { // Multi-vector compression let input_view = - MatRef::new(Standard::new(num_vectors, dim).unwrap(), &input_data) + MatRef::from_repr(RowMajor::new(num_vectors, dim).unwrap(), &input_data) .expect("input view creation"); let mut multi_mat: Mat> = - Mat::new(MinMaxMeta::new(num_vectors, dim), Defaulted) + Mat::from_default(MinMaxMeta::new(num_vectors, dim)) .expect("output mat creation"); quantizer @@ -587,11 +587,12 @@ mod tests { let quantizer = make_quantizer(dim); let input_data = generate_test_data(num_vectors, dim); - let input_view = MatRef::new(Standard::new(num_vectors, dim).unwrap(), &input_data) - .expect("input view"); + let input_view = + MatRef::from_repr(RowMajor::new(num_vectors, dim).unwrap(), &input_data) + .expect("input view"); let mut mat: Mat> = - Mat::new(MinMaxMeta::new(num_vectors, dim), Defaulted).expect("mat creation"); + Mat::from_default(MinMaxMeta::new(num_vectors, dim)).expect("mat creation"); quantizer .compress_into(input_view, mat.reborrow_mut()) @@ -627,11 +628,11 @@ mod tests { // Input has 3 vectors let input_data = generate_test_data(3, dim); let input_view = - MatRef::new(Standard::new(3, dim).unwrap(), &input_data).expect("input view"); + MatRef::from_repr(RowMajor::new(3, dim).unwrap(), &input_data).expect("input view"); // Output has 2 vectors (mismatch) let mut mat: Mat> = - Mat::new(MinMaxMeta::new(2, dim), Defaulted).expect("mat creation"); + Mat::from_default(MinMaxMeta::new(2, dim)).expect("mat creation"); let _ = quantizer.compress_into(input_view, mat.reborrow_mut()); } @@ -645,11 +646,11 @@ mod tests { // Input has dim=8 (mismatch) let input_data = generate_test_data(2, 8); let input_view = - MatRef::new(Standard::new(2, 8).unwrap(), &input_data).expect("input view"); + MatRef::from_repr(RowMajor::new(2, 8).unwrap(), &input_data).expect("input view"); // Output correctly has dim=4 let mut mat: Mat> = - Mat::new(MinMaxMeta::new(2, 4), Defaulted).expect("mat creation"); + Mat::from_default(MinMaxMeta::new(2, 4)).expect("mat creation"); let _ = quantizer.compress_into(input_view, mat.reborrow_mut()); } @@ -665,13 +666,13 @@ mod tests { // Input correctly has dim=4 let input_data = generate_test_data(2, 4); let input_view = - MatRef::new(Standard::new(2, 4).unwrap(), &input_data).expect("input view"); + MatRef::from_repr(RowMajor::new(2, 4).unwrap(), &input_data).expect("input view"); // Output has intrinsic_dim=8 (mismatch) let row_bytes = Data::::canonical_bytes(8); let mut output_data = vec![0u8; 2 * row_bytes]; - let output_view = - MatMut::new(MinMaxMeta::::new(2, 8), &mut output_data).expect("output view"); + let output_view = MatMut::from_repr(MinMaxMeta::::new(2, 8), &mut output_data) + .expect("output view"); let _ = quantizer.compress_into(input_view, output_view); } diff --git a/diskann-quantization/src/minmax/multi/mod.rs b/diskann-quantization/src/minmax/multi/mod.rs index 0fb4e57e6..b1d61bbcf 100644 --- a/diskann-quantization/src/minmax/multi/mod.rs +++ b/diskann-quantization/src/minmax/multi/mod.rs @@ -20,7 +20,7 @@ //! minmax::{MinMaxMeta, MinMaxQuantizer}, //! multi_vector::{ //! distance::{Chamfer, MaxSim, QueryMatRef}, -//! Defaulted, Mat, MatRef, Standard, +//! Mat, MatRef, RowMajor, //! }, //! num::Positive, //! CompressInto, @@ -44,8 +44,8 @@ //! 1.0, 0.0, 0.0, 0.0, // query vector 0 //! 0.0, 1.0, 0.0, 0.0, // query vector 1 //! ]; -//! let query_input = MatRef::new( -//! Standard::new(num_query_vectors, dim).unwrap(), &query_data +//! let query_input = MatRef::from_repr( +//! RowMajor::new(num_query_vectors, dim).unwrap(), &query_data //! ).unwrap(); //! //! // Full-precision document multi-vector (3 vectors × 4 dimensions) @@ -54,15 +54,15 @@ //! 1.0, 0.0, 0.0, 0.0, // doc vector 1 //! 0.0, 0.0, 1.0, 0.0, // doc vector 2 //! ]; -//! let doc_input = MatRef::new( -//! Standard::new(num_doc_vectors, dim).unwrap(), &doc_data +//! let doc_input = MatRef::from_repr( +//! RowMajor::new(num_doc_vectors, dim).unwrap(), &doc_data //! ).unwrap(); //! -//! // Create owned matrices for quantized output using Mat::new +//! // Create owned matrices for the quantized output //! let mut query_out: Mat> = -//! Mat::new(MinMaxMeta::new(num_query_vectors, dim), Defaulted).unwrap(); +//! Mat::from_default(MinMaxMeta::new(num_query_vectors, dim)).unwrap(); //! let mut doc_out: Mat> = -//! Mat::new(MinMaxMeta::new(num_doc_vectors, dim), Defaulted).unwrap(); +//! Mat::from_default(MinMaxMeta::new(num_doc_vectors, dim)).unwrap(); //! //! // Quantize both multi-vectors //! quantizer.compress_into(query_input, query_out.reborrow_mut()).unwrap(); diff --git a/diskann-quantization/src/multi_vector/block_transposed.rs b/diskann-quantization/src/multi_vector/block_transposed.rs index 6fe9a1315..acc56f739 100644 --- a/diskann-quantization/src/multi_vector/block_transposed.rs +++ b/diskann-quantization/src/multi_vector/block_transposed.rs @@ -82,12 +82,12 @@ use std::{alloc::Layout, marker::PhantomData, ptr::NonNull}; use diskann_utils::{ Reborrow, ReborrowMut, strided::StridedView, - views::{MatrixView, MutMatrixView}, + views::{MatrixView, MatrixViewMut}, }; -use super::matrix::{ - Defaulted, LayoutError, Mat, MatMut, MatRef, NewMut, NewOwned, NewRef, Overflow, Repr, ReprMut, - ReprOwned, SliceError, +use super::views::{ + LayoutError, Mat, MatMut, MatRef, NewDefault, NewMut, NewOwned, NewRef, Overflow, Repr, + ReprMut, ReprOwned, SliceError, }; use crate::bits::{AsMutPtr, AsPtr, MutSlicePtr, SlicePtr}; use crate::utils; @@ -591,14 +591,17 @@ unsafe impl NewOwned } } -// SAFETY: This safely re-uses `>`. -unsafe impl NewOwned +// SAFETY: The returned `Mat` contains a `Box` with exactly `self.storage_len()` elements. +unsafe impl NewDefault for BlockTransposedRepr { type Error = crate::error::Infallible; - fn new_owned(self, _: Defaulted) -> Result, Self::Error> { - self.new_owned(T::default()) + fn new_default(self) -> Result, Self::Error> { + let b: Box<[T]> = vec![T::default(); self.storage_len()].into_boxed_slice(); + + // SAFETY: By construction, `b.len() == self.storage_len()`. + Ok(unsafe { self.box_to_mat(b) }) } } @@ -944,12 +947,12 @@ impl<'a, T: Copy, const GROUP: usize, const PACK: usize> BlockTransposedMut<'a, /// /// Panics if `block >= self.full_blocks()`. #[allow(clippy::expect_used)] - pub fn block_mut(&mut self, block: usize) -> MutMatrixView<'_, T> { + pub fn block_mut(&mut self, block: usize) -> MatrixViewMut<'_, T> { self.reborrow_mut().block_mut_inner(block) } #[allow(clippy::expect_used)] - fn block_mut_inner(mut self, block: usize) -> MutMatrixView<'a, T> { + fn block_mut_inner(mut self, block: usize) -> MatrixViewMut<'a, T> { let repr = *self.data.repr(); assert!(block < repr.full_blocks()); let offset = repr.block_offset(block); @@ -962,19 +965,19 @@ impl<'a, T: Copy, const GROUP: usize, const PACK: usize> BlockTransposedMut<'a, stride, ) }; - MutMatrixView::try_from(data, pncols / PACK, GROUP * PACK) + MatrixViewMut::try_from(data, pncols / PACK, GROUP * PACK) .expect("base data should have been sized correctly") } /// Return a mutable view over the remainder block, or `None` if there is no /// remainder. #[allow(clippy::expect_used)] - pub fn remainder_block_mut(&mut self) -> Option> { + pub fn remainder_block_mut(&mut self) -> Option> { self.reborrow_mut().remainder_block_mut_inner() } #[allow(clippy::expect_used)] - fn remainder_block_mut_inner(mut self) -> Option> { + fn remainder_block_mut_inner(mut self) -> Option> { let repr = *self.data.repr(); if repr.remainder() == 0 { None @@ -990,7 +993,7 @@ impl<'a, T: Copy, const GROUP: usize, const PACK: usize> BlockTransposedMut<'a, ) }; Some( - MutMatrixView::try_from(data, pncols / PACK, GROUP * PACK) + MatrixViewMut::try_from(data, pncols / PACK, GROUP * PACK) .expect("base data should have been sized correctly"), ) } @@ -1075,13 +1078,13 @@ impl BlockTransposed MutMatrixView<'_, T> { + pub fn block_mut(&mut self, block: usize) -> MatrixViewMut<'_, T> { self.as_view_mut().block_mut_inner(block) } /// See [`BlockTransposedMut::remainder_block_mut`]. #[allow(clippy::expect_used)] - pub fn remainder_block_mut(&mut self) -> Option> { + pub fn remainder_block_mut(&mut self) -> Option> { self.as_view_mut().remainder_block_mut_inner() } @@ -1118,7 +1121,7 @@ impl BlockTransposed::new(nrows, ncols) .expect("dimensions should not overflow"); Self { - data: Mat::new(repr, Defaulted).expect("infallible"), + data: Mat::from_default(repr).expect("infallible"), } } @@ -1126,7 +1129,7 @@ impl BlockTransposed Result { let repr = BlockTransposedRepr::::new(nrows, ncols)?; Ok(Self { - data: Mat::new(repr, Defaulted).expect("infallible"), + data: Mat::from_default(repr).expect("infallible"), }) } diff --git a/diskann-quantization/src/multi_vector/distance/factory.rs b/diskann-quantization/src/multi_vector/distance/factory.rs index 5dcd4b8cd..817d64a87 100644 --- a/diskann-quantization/src/multi_vector/distance/factory.rs +++ b/diskann-quantization/src/multi_vector/distance/factory.rs @@ -20,7 +20,7 @@ use super::kernels::f16::F16Entry; use super::kernels::f32::F32Kernel; use super::max_sim::{MaxSim, MaxSimError}; use crate::multi_vector::distance::QueryMatRef; -use crate::multi_vector::{BlockTransposed, BlockTransposedRef, Mat, MatRef, Standard}; +use crate::multi_vector::{BlockTransposed, BlockTransposedRef, Mat, MatRef, RowMajor}; // ───────────────────────────────────────────────────────────────────────── // Prepared — concrete kernel for the arch-dispatched paths. @@ -39,7 +39,7 @@ where A, (), BlockTransposedRef<'a, f32, GROUP>, - MatRef<'a, Standard>, + MatRef<'a, RowMajor>, &'a mut [f32], >, { @@ -49,7 +49,7 @@ where fn compute_max_sim( &self, - doc: MatRef<'_, Standard>, + doc: MatRef<'_, RowMajor>, scores: &mut [f32], ) -> Result<(), MaxSimError> { if scores.len() != self.nrows() { @@ -81,7 +81,7 @@ where A, (), BlockTransposedRef<'a, half::f16, GROUP>, - MatRef<'a, Standard>, + MatRef<'a, RowMajor>, &'a mut [f32], >, { @@ -91,7 +91,7 @@ where fn compute_max_sim( &self, - doc: MatRef<'_, Standard>, + doc: MatRef<'_, RowMajor>, scores: &mut [f32], ) -> Result<(), MaxSimError> { if scores.len() != self.nrows() { @@ -120,7 +120,7 @@ where // ───────────────────────────────────────────────────────────────────────── struct ReferenceKernel { - query: Mat>, + query: Mat>, } impl std::fmt::Debug for ReferenceKernel { @@ -132,7 +132,7 @@ impl std::fmt::Debug for ReferenceKernel { } impl ReferenceKernel { - fn new(query: MatRef<'_, Standard>) -> Self { + fn new(query: MatRef<'_, RowMajor>) -> Self { Self { query: query.to_owned(), } @@ -150,7 +150,7 @@ where fn compute_max_sim( &self, - doc: MatRef<'_, Standard>, + doc: MatRef<'_, RowMajor>, scores: &mut [f32], ) -> Result<(), MaxSimError> { if scores.len() != self.nrows() { @@ -160,7 +160,7 @@ where scores.fill(f32::MAX); return Ok(()); } - let query: QueryMatRef<'_, Standard> = self.query.as_view().into(); + let query: QueryMatRef<'_, RowMajor> = self.query.as_view().into(); let mut max_sim = MaxSim::new(scores); max_sim.evaluate(query, doc) } @@ -174,30 +174,30 @@ struct BuildAndErase(E); // ───── f32 Target1 impls ───── -impl> diskann_wide::arch::Target1>> +impl> diskann_wide::arch::Target1>> for BuildAndErase { - fn run(self, arch: Scalar, query: MatRef<'_, Standard>) -> E::Output { + fn run(self, arch: Scalar, query: MatRef<'_, RowMajor>) -> E::Output { let prepared = BlockTransposed::::from_matrix_view(query.as_matrix_view()); self.0.erase(Prepared { arch, prepared }) } } #[cfg(target_arch = "x86_64")] -impl> diskann_wide::arch::Target1>> +impl> diskann_wide::arch::Target1>> for BuildAndErase { - fn run(self, arch: V3, query: MatRef<'_, Standard>) -> E::Output { + fn run(self, arch: V3, query: MatRef<'_, RowMajor>) -> E::Output { let prepared = BlockTransposed::::from_matrix_view(query.as_matrix_view()); self.0.erase(Prepared { arch, prepared }) } } #[cfg(target_arch = "x86_64")] -impl> diskann_wide::arch::Target1>> +impl> diskann_wide::arch::Target1>> for BuildAndErase { - fn run(self, arch: V4, query: MatRef<'_, Standard>) -> E::Output { + fn run(self, arch: V4, query: MatRef<'_, RowMajor>) -> E::Output { // V4 dispatches to V3 (no V4-specific kernel). let arch = arch.retarget(); let prepared = BlockTransposed::::from_matrix_view(query.as_matrix_view()); @@ -206,10 +206,10 @@ impl> diskann_wide::arch::Target1> diskann_wide::arch::Target1>> +impl> diskann_wide::arch::Target1>> for BuildAndErase { - fn run(self, arch: Neon, query: MatRef<'_, Standard>) -> E::Output { + fn run(self, arch: Neon, query: MatRef<'_, RowMajor>) -> E::Output { // Neon dispatches to Scalar (no Neon-specific kernel). let arch = arch.retarget(); let prepared = BlockTransposed::::from_matrix_view(query.as_matrix_view()); @@ -220,10 +220,10 @@ impl> diskann_wide::arch::Target1> - diskann_wide::arch::Target1>> + diskann_wide::arch::Target1>> for BuildAndErase { - fn run(self, arch: Scalar, query: MatRef<'_, Standard>) -> E::Output { + fn run(self, arch: Scalar, query: MatRef<'_, RowMajor>) -> E::Output { let prepared = BlockTransposed::::from_matrix_view(query.as_matrix_view()); self.0.erase(Prepared { arch, prepared }) } @@ -231,10 +231,10 @@ impl> #[cfg(target_arch = "x86_64")] impl> - diskann_wide::arch::Target1>> + diskann_wide::arch::Target1>> for BuildAndErase { - fn run(self, arch: V3, query: MatRef<'_, Standard>) -> E::Output { + fn run(self, arch: V3, query: MatRef<'_, RowMajor>) -> E::Output { let prepared = BlockTransposed::::from_matrix_view(query.as_matrix_view()); self.0.erase(Prepared { arch, prepared }) } @@ -242,10 +242,10 @@ impl> #[cfg(target_arch = "x86_64")] impl> - diskann_wide::arch::Target1>> + diskann_wide::arch::Target1>> for BuildAndErase { - fn run(self, arch: V4, query: MatRef<'_, Standard>) -> E::Output { + fn run(self, arch: V4, query: MatRef<'_, RowMajor>) -> E::Output { // V4 dispatches to V3 (no V4-specific kernel). let arch = arch.retarget(); let prepared = BlockTransposed::::from_matrix_view(query.as_matrix_view()); @@ -255,10 +255,10 @@ impl> #[cfg(target_arch = "aarch64")] impl> - diskann_wide::arch::Target1>> + diskann_wide::arch::Target1>> for BuildAndErase { - fn run(self, arch: Neon, query: MatRef<'_, Standard>) -> E::Output { + fn run(self, arch: Neon, query: MatRef<'_, RowMajor>) -> E::Output { // Neon dispatches to Scalar (no Neon-specific kernel). let arch = arch.retarget(); let prepared = BlockTransposed::::from_matrix_view(query.as_matrix_view()); @@ -278,7 +278,7 @@ mod sealed { /// /// Sealed: external crates cannot add impls. Quantized representations /// (PQ, SQ, packed sub-byte) are intentionally excluded — they need -/// codebook/scale state that [`MatRef<'_, Standard>`] can't carry. +/// codebook/scale state that [`MatRef<'_, RowMajor>`] can't carry. pub trait MaxSimElement: sealed::Sealed + Sized + Copy + Send + Sync + 'static { /// Build the concrete kernel for this element type and hand it to /// `erase.erase(...)`. @@ -289,7 +289,7 @@ pub trait MaxSimElement: sealed::Sealed + Sized + Copy + Send + Sync + 'static { /// build (e.g. AVX-512 unavailable; aarch64 on x86_64). fn build>( isa: MaxSimIsa, - query: MatRef<'_, Standard>, + query: MatRef<'_, RowMajor>, erase: E, ) -> Result; } @@ -300,7 +300,7 @@ impl sealed::Sealed for half::f16 {} impl MaxSimElement for f32 { fn build>( isa: MaxSimIsa, - query: MatRef<'_, Standard>, + query: MatRef<'_, RowMajor>, erase: E, ) -> Result { match isa { @@ -351,7 +351,7 @@ impl MaxSimElement for f32 { impl MaxSimElement for half::f16 { fn build>( isa: MaxSimIsa, - query: MatRef<'_, Standard>, + query: MatRef<'_, RowMajor>, erase: E, ) -> Result { match isa { @@ -413,7 +413,7 @@ impl MaxSimElement for half::f16 { /// Returns [`NotSupported`] when the requested ISA cannot run on this build. pub fn build_max_sim>( isa: MaxSimIsa, - query: MatRef<'_, Standard>, + query: MatRef<'_, RowMajor>, erase: E, ) -> Result { T::build(isa, query, erase) @@ -443,8 +443,8 @@ mod tests { } } - fn make_mat(data: &[T], nrows: usize, ncols: usize) -> MatRef<'_, Standard> { - MatRef::new(Standard::new(nrows, ncols).unwrap(), data).unwrap() + fn make_mat(data: &[T], nrows: usize, ncols: usize) -> MatRef<'_, RowMajor> { + MatRef::from_repr(RowMajor::new(nrows, ncols).unwrap(), data).unwrap() } fn make_test_data(len: usize, ceil: usize, shift: usize) -> Vec { diff --git a/diskann-quantization/src/multi_vector/distance/fallback.rs b/diskann-quantization/src/multi_vector/distance/fallback.rs index ed8da7a3e..86319a831 100644 --- a/diskann-quantization/src/multi_vector/distance/fallback.rs +++ b/diskann-quantization/src/multi_vector/distance/fallback.rs @@ -10,7 +10,7 @@ use diskann_vector::{DistanceFunctionMut, PureDistanceFunction}; use super::max_sim::{Chamfer, MaxSim}; use super::projected_eigen::ProjectedEigen; -use crate::multi_vector::{MatRef, MaxSimError, Repr, Standard}; +use crate::multi_vector::{MatRef, MaxSimError, Repr, RowMajor}; ///////////////// // QueryMatRef // @@ -25,11 +25,11 @@ use crate::multi_vector::{MatRef, MaxSimError, Repr, Standard}; /// # Example /// /// ``` -/// use diskann_quantization::multi_vector::{MatRef, Standard}; +/// use diskann_quantization::multi_vector::{MatRef, RowMajor}; /// use diskann_quantization::multi_vector::distance::QueryMatRef; /// /// let data = [1.0f32, 2.0, 3.0, 4.0]; -/// let view = MatRef::new(Standard::new(2, 2).unwrap(), &data).unwrap(); +/// let view = MatRef::from_repr(RowMajor::new(2, 2).unwrap(), &data).unwrap(); /// let query: QueryMatRef<_> = view.into(); /// ``` #[derive(Debug, Clone, Copy)] @@ -77,8 +77,8 @@ impl FallbackKernel { /// * `f` - Callback invoked with `(query_index, similarity)` for each query vector #[inline] pub(crate) fn max_sim_kernel( - query: QueryMatRef<'_, Standard>, - doc: MatRef<'_, Standard>, + query: QueryMatRef<'_, RowMajor>, + doc: MatRef<'_, RowMajor>, mut f: F, ) where F: FnMut(usize, f32), @@ -119,8 +119,8 @@ impl FallbackKernel { /// * `f` - Callback invoked with `(query_index, score)` for each query vector #[inline] pub(crate) fn projected_eigen_kernel( - query: QueryMatRef<'_, Standard>, - doc: MatRef<'_, Standard>, + query: QueryMatRef<'_, RowMajor>, + doc: MatRef<'_, RowMajor>, mut f: F, ) where F: FnMut(usize, f32), @@ -153,8 +153,8 @@ impl FallbackKernel { impl DistanceFunctionMut< - QueryMatRef<'_, Standard>, - MatRef<'_, Standard>, + QueryMatRef<'_, RowMajor>, + MatRef<'_, RowMajor>, Result<(), MaxSimError>, > for MaxSim<'_> where @@ -163,8 +163,8 @@ where #[inline(always)] fn evaluate( &mut self, - query: QueryMatRef<'_, Standard>, - doc: MatRef<'_, Standard>, + query: QueryMatRef<'_, RowMajor>, + doc: MatRef<'_, RowMajor>, ) -> Result<(), MaxSimError> { let size = self.size(); let n_queries = query.num_vectors(); @@ -187,13 +187,13 @@ where // Chamfer // ///////////// -impl PureDistanceFunction>, MatRef<'_, Standard>, f32> +impl PureDistanceFunction>, MatRef<'_, RowMajor>, f32> for Chamfer where InnerProduct: for<'a, 'b> PureDistanceFunction<&'a [T], &'b [T], f32>, { #[inline(always)] - fn evaluate(query: QueryMatRef<'_, Standard>, doc: MatRef<'_, Standard>) -> f32 { + fn evaluate(query: QueryMatRef<'_, RowMajor>, doc: MatRef<'_, RowMajor>) -> f32 { let mut sum = 0.0f32; FallbackKernel::max_sim_kernel(query, doc, |_i, score| { @@ -208,13 +208,13 @@ where // ProjectedEigen // ///////////////////// -impl PureDistanceFunction>, MatRef<'_, Standard>, f32> +impl PureDistanceFunction>, MatRef<'_, RowMajor>, f32> for ProjectedEigen where InnerProduct: for<'a, 'b> PureDistanceFunction<&'a [T], &'b [T], f32>, { #[inline(always)] - fn evaluate(query: QueryMatRef<'_, Standard>, doc: MatRef<'_, Standard>) -> f32 { + fn evaluate(query: QueryMatRef<'_, RowMajor>, doc: MatRef<'_, RowMajor>) -> f32 { let mut sum = 0.0f32; FallbackKernel::projected_eigen_kernel(query, doc, |_i, score| { @@ -230,19 +230,19 @@ mod tests { use super::*; /// Helper to create a QueryMatRef from raw data - fn make_query(data: &[f32], nrows: usize, ncols: usize) -> QueryMatRef<'_, Standard> { - MatRef::new(Standard::new(nrows, ncols).unwrap(), data) + fn make_query(data: &[f32], nrows: usize, ncols: usize) -> QueryMatRef<'_, RowMajor> { + MatRef::from_repr(RowMajor::new(nrows, ncols).unwrap(), data) .unwrap() .into() } /// Helper to create a MatRef from raw data - fn make_doc(data: &[f32], nrows: usize, ncols: usize) -> MatRef<'_, Standard> { - MatRef::new(Standard::new(nrows, ncols).unwrap(), data).unwrap() + fn make_doc(data: &[f32], nrows: usize, ncols: usize) -> MatRef<'_, RowMajor> { + MatRef::from_repr(RowMajor::new(nrows, ncols).unwrap(), data).unwrap() } /// Naive implementation of max-sim for a single query vector against all doc vectors. - fn naive_max_sim_single(query_vec: &[f32], doc: &MatRef<'_, Standard>) -> f32 { + fn naive_max_sim_single(query_vec: &[f32], doc: &MatRef<'_, RowMajor>) -> f32 { doc.rows() .map(|d_vec| { let ip: f32 = query_vec.iter().zip(d_vec.iter()).map(|(a, b)| a * b).sum(); @@ -253,7 +253,7 @@ mod tests { /// Naive implementation of projected-eigen for a single query vector /// against all doc vectors: `\sum_{j} -IP(q, d_{j})^2`. - fn naive_projected_eigen_single(query_vec: &[f32], doc: &MatRef<'_, Standard>) -> f32 { + fn naive_projected_eigen_single(query_vec: &[f32], doc: &MatRef<'_, RowMajor>) -> f32 { doc.rows() .map(|d_vec| { let ip: f32 = query_vec.iter().zip(d_vec.iter()).map(|(a, b)| a * b).sum(); @@ -273,7 +273,7 @@ mod tests { #[test] fn from_mat_ref_and_deref() { let data = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0]; - let view = MatRef::new(Standard::new(2, 3).unwrap(), &data).unwrap(); + let view = MatRef::from_repr(RowMajor::new(2, 3).unwrap(), &data).unwrap(); let query: QueryMatRef<_> = view.into(); // Deref access works diff --git a/diskann-quantization/src/multi_vector/distance/kernel.rs b/diskann-quantization/src/multi_vector/distance/kernel.rs index b292def54..f663b115e 100644 --- a/diskann-quantization/src/multi_vector/distance/kernel.rs +++ b/diskann-quantization/src/multi_vector/distance/kernel.rs @@ -3,7 +3,7 @@ //! Object-safe kernel boundary trait plus BYOTE visitor trait. -use crate::multi_vector::{MatRef, MaxSimError, Standard}; +use crate::multi_vector::{MatRef, MaxSimError, RowMajor}; /// Object-safe interface for computing per-query MaxSim scores. pub trait MaxSimKernel: Send + Sync + std::fmt::Debug { @@ -18,7 +18,7 @@ pub trait MaxSimKernel: Send + Sync + std::fmt::Debug { /// [`MaxSimError::InvalidBufferLength`] if `scores.len() != self.nrows()`. fn compute_max_sim( &self, - doc: MatRef<'_, Standard>, + doc: MatRef<'_, RowMajor>, scores: &mut [f32], ) -> Result<(), MaxSimError>; } diff --git a/diskann-quantization/src/multi_vector/distance/kernels/f16.rs b/diskann-quantization/src/multi_vector/distance/kernels/f16.rs index a535c68dc..ccbe652c3 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/f16.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/f16.rs @@ -19,7 +19,7 @@ use super::Kernel; use super::TileBudget; use super::f32::{F32Kernel, max_ip_kernel}; use super::layouts; -use crate::multi_vector::{BlockTransposedRef, MatRef, Standard}; +use crate::multi_vector::{BlockTransposedRef, MatRef, RowMajor}; pub(crate) struct F16Entry; @@ -28,7 +28,7 @@ impl A, (), BlockTransposedRef<'_, half::f16, GROUP>, - MatRef<'_, Standard>, + MatRef<'_, RowMajor>, &mut [f32], > for F16Entry where @@ -44,7 +44,7 @@ where self, arch: A, lhs: BlockTransposedRef<'_, half::f16, GROUP>, - rhs: MatRef<'_, Standard>, + rhs: MatRef<'_, RowMajor>, scratch: &mut [f32], ) { max_ip_kernel(arch, lhs, rhs, scratch, TileBudget::default()); diff --git a/diskann-quantization/src/multi_vector/distance/kernels/f32/mod.rs b/diskann-quantization/src/multi_vector/distance/kernels/f32/mod.rs index a900ea356..07051e11a 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/f32/mod.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/f32/mod.rs @@ -23,7 +23,7 @@ use super::Kernel; use super::TileBudget; use super::layouts::{self, DescribeLayout}; use super::tiled_reduce::tiled_reduce; -use crate::multi_vector::{BlockTransposedRef, MatRef, Standard}; +use crate::multi_vector::{BlockTransposedRef, MatRef, RowMajor}; mod scalar; #[cfg(target_arch = "x86_64")] @@ -61,7 +61,7 @@ fn max_ip_kernel_panic(scratch_len: usize, padded_nrows: usize, a_ncols: usize, pub(super) fn max_ip_kernel( arch: A, a: BlockTransposedRef<'_, T, GROUP>, - b: MatRef<'_, Standard>, + b: MatRef<'_, RowMajor>, scratch: &mut [f32], budget: TileBudget, ) where @@ -86,7 +86,7 @@ pub(super) fn max_ip_kernel( // SAFETY: // - a.as_ptr() is valid for a.padded_nrows() * k elements of T. - // - MatRef> stores nrows * ncols contiguous T elements. + // - MatRef> stores nrows * ncols contiguous T elements. // - scratch.len() == a.padded_nrows() (checked above). // - a.padded_nrows() is always a multiple of GROUP, and the const assert above // verifies A_PANEL == GROUP at compile time. @@ -111,7 +111,7 @@ impl A, (), BlockTransposedRef<'_, f32, GROUP>, - MatRef<'_, Standard>, + MatRef<'_, RowMajor>, &mut [f32], > for F32Kernel where @@ -127,7 +127,7 @@ where self, arch: A, lhs: BlockTransposedRef<'_, f32, GROUP>, - rhs: MatRef<'_, Standard>, + rhs: MatRef<'_, RowMajor>, scratch: &mut [f32], ) { max_ip_kernel(arch, lhs, rhs, scratch, TileBudget::default()); diff --git a/diskann-quantization/src/multi_vector/distance/kernels/layouts.rs b/diskann-quantization/src/multi_vector/distance/kernels/layouts.rs index e1ec8dd36..e0306fcbb 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/layouts.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/layouts.rs @@ -45,7 +45,7 @@ impl Layout for BlockTransposed< type Element = T; } -/// Dense row-major tile layout. Matches [`MatRef>`](crate::multi_vector::MatRef). +/// Dense row-major tile layout. Matches [`MatRef>`](crate::multi_vector::MatRef). pub(super) struct RowMajor(PhantomData); impl RowMajor { @@ -86,7 +86,7 @@ impl DescribeLayout } } -impl DescribeLayout for crate::multi_vector::MatRef<'_, crate::multi_vector::Standard> { +impl DescribeLayout for crate::multi_vector::MatRef<'_, crate::multi_vector::RowMajor> { type Layout = RowMajor; fn layout(&self) -> Self::Layout { diff --git a/diskann-quantization/src/multi_vector/distance/kernels/tiled_reduce.rs b/diskann-quantization/src/multi_vector/distance/kernels/tiled_reduce.rs index 6e4b76e9c..5cd689d65 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/tiled_reduce.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/tiled_reduce.rs @@ -246,7 +246,7 @@ mod tests { use super::super::f32::{F32Kernel, max_ip_kernel}; use super::super::layouts; - use crate::multi_vector::{BlockTransposed, MatRef, Standard}; + use crate::multi_vector::{BlockTransposed, MatRef, RowMajor}; #[test] fn basic_panel_counts() { @@ -521,9 +521,9 @@ mod tests { layouts::RowMajor: ConvertTo as Kernel>::Right> + Layout, { - let a_mat = MatRef::new(Standard::new(a_nrows, dim).unwrap(), a_data).unwrap(); + let a_mat = MatRef::from_repr(RowMajor::new(a_nrows, dim).unwrap(), a_data).unwrap(); let a_bt = BlockTransposed::::from_matrix_view(a_mat.as_matrix_view()); - let b_mat = MatRef::new(Standard::new(b_nrows, dim).unwrap(), b_data).unwrap(); + let b_mat = MatRef::from_repr(RowMajor::new(b_nrows, dim).unwrap(), b_data).unwrap(); let mut scratch = vec![f32::MIN; a_bt.padded_nrows()]; max_ip_kernel::( @@ -741,9 +741,9 @@ mod tests { let b_data = gen_data(b_nrows * dim, ceil); let expected = naive(&a_data, a_nrows, &b_data, b_nrows, dim); - let a_mat = MatRef::new(Standard::new(a_nrows, dim).unwrap(), &a_data).unwrap(); + let a_mat = MatRef::from_repr(RowMajor::new(a_nrows, dim).unwrap(), &a_data).unwrap(); let a_bt = BlockTransposed::::from_matrix_view(a_mat.as_matrix_view()); - let b_mat = MatRef::new(Standard::new(b_nrows, dim).unwrap(), &b_data).unwrap(); + let b_mat = MatRef::from_repr(RowMajor::new(b_nrows, dim).unwrap(), &b_data).unwrap(); let mut scratch = vec![f32::MIN; a_bt.padded_nrows()]; max_ip_kernel::(arch, a_bt.as_view(), b_mat, &mut scratch, budget); diff --git a/diskann-quantization/src/multi_vector/distance/max_sim.rs b/diskann-quantization/src/multi_vector/distance/max_sim.rs index d9a4fb541..92715a4dd 100644 --- a/diskann-quantization/src/multi_vector/distance/max_sim.rs +++ b/diskann-quantization/src/multi_vector/distance/max_sim.rs @@ -28,7 +28,7 @@ pub enum MaxSimError { /// ``` /// /// Implements `DistanceFnMut` for various matrix types -/// (e.g., [`MatRef>`](crate::multi_vector::MatRef)). +/// (e.g., [`MatRef>`](crate::multi_vector::MatRef)). /// /// # Usage /// - Create with [`MaxSim::new`], providing a mutable scores buffer. diff --git a/diskann-quantization/src/multi_vector/distance/mod.rs b/diskann-quantization/src/multi_vector/distance/mod.rs index 178e8cad0..9e94c81bc 100644 --- a/diskann-quantization/src/multi_vector/distance/mod.rs +++ b/diskann-quantization/src/multi_vector/distance/mod.rs @@ -12,21 +12,21 @@ //! ``` //! use diskann_quantization::multi_vector::{ //! distance::{Chamfer, MaxSim, QueryMatRef}, -//! MatRef, Standard, +//! MatRef, RowMajor, //! }; //! use diskann_vector::{DistanceFunctionMut, PureDistanceFunction}; //! //! // Query: 2 vectors of dim 3 (wrapped as QueryMatRef) //! let query_data = [1.0f32, 0.0, 0.0, 0.0, 1.0, 0.0]; -//! let query: QueryMatRef<_> = MatRef::new( -//! Standard::new(2, 3).unwrap(), +//! let query: QueryMatRef<_> = MatRef::from_repr( +//! RowMajor::new(2, 3).unwrap(), //! &query_data, //! ).unwrap().into(); //! //! // Doc: 2 vectors of dim 3 //! let doc_data = [1.0f32, 0.0, 0.0, 0.0, 0.0, 1.0]; -//! let doc = MatRef::new( -//! Standard::new(2, 3).unwrap(), +//! let doc = MatRef::from_repr( +//! RowMajor::new(2, 3).unwrap(), //! &doc_data, //! ).unwrap(); //! diff --git a/diskann-quantization/src/multi_vector/matrix.rs b/diskann-quantization/src/multi_vector/matrix.rs deleted file mode 100644 index b0cf0954b..000000000 --- a/diskann-quantization/src/multi_vector/matrix.rs +++ /dev/null @@ -1,1970 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! Row-major matrix types for multi-vector representations. -//! -//! This module provides flexible matrix abstractions that support different underlying -//! storage formats through the [`Repr`] trait. The primary types are: -//! -//! - [`Mat`]: An owning matrix that manages its own memory. -//! - [`MatRef`]: An immutable borrowed view of matrix data. -//! - [`MatMut`]: A mutable borrowed view of matrix data. -//! -//! # Representations -//! -//! Representation types interact with the [`Mat`] family of types using the following traits: -//! -//! - [`Repr`]: Read-only matrix representation. -//! - [`ReprMut`]: Mutable matrix representation. -//! - [`ReprOwned`]: Owning matrix representation. -//! -//! Each trait refinement has a corresponding constructor: -//! -//! - [`NewRef`]: Construct a read-only [`MatRef`] view over a slice. -//! - [`NewMut`]: Construct a mutable [`MatMut`] matrix view over a slice. -//! - [`NewOwned`]: Construct a new owning [`Mat`]. -//! - -use std::{alloc::Layout, iter::FusedIterator, marker::PhantomData, ptr::NonNull}; - -use diskann_utils::{Reborrow, ReborrowMut, views::MatrixView}; -use thiserror::Error; - -use crate::utils; - -/// Representation trait describing the layout and access patterns for a matrix. -/// -/// Implementations define how raw bytes are interpreted as typed rows. This enables -/// matrices over different storage formats (dense, quantized, etc.) using a single -/// generic [`Mat`] type. -/// -/// # Associated Types -/// -/// - `Row<'a>`: The immutable row type (e.g., `&[f32]`, `&[f16]`). -/// -/// # Safety -/// -/// Implementations must ensure: -/// -/// - [`get_row`](Self::get_row) returns valid references for the given row index. -/// This call **must** be memory safe for `i < self.nrows()`, provided the caller upholds -/// the contract for the raw pointer. -/// -/// - The objects implicitly managed by this representation inherit the `Send` and `Sync` -/// attributes of `Repr`. That is, `Repr: Send` implies that the objects in backing memory -/// are [`Send`], and likewise with `Sync`. This is necessary to apply [`Send`] and [`Sync`] -/// bounds to [`Mat`], [`MatRef`], and [`MatMut`]. -pub unsafe trait Repr: Copy { - /// Immutable row reference type. - type Row<'a> - where - Self: 'a; - - /// Returns the number of rows in the matrix. - /// - /// # Safety Contract - /// - /// This function must be loosely pure in the sense that for any given instance of - /// `self`, `self.nrows()` must return the same value. - fn nrows(&self) -> usize; - - /// Returns the memory layout for a memory allocation containing [`Repr::nrows`] vectors - /// each with vector dimension [`Repr::ncols`]. - /// - /// # Safety Contract - /// - /// The [`Layout`] returned from this method must be consistent with the contract of - /// [`Repr::get_row`]. - fn layout(&self) -> Result; - - /// Returns an immutable reference to the `i`-th row. - /// - /// # Safety - /// - /// - `ptr` must point to a slice with a layout compatible with [`Repr::layout`]. - /// - The entire range for this slice must be within a single allocation. - /// - `i` must be less than [`Repr::nrows`]. - /// - The memory referenced by the returned [`Repr::Row`] must not be mutated for the - /// duration of lifetime `'a`. - /// - The lifetime for the returned [`Repr::Row`] is inferred from its usage. Correct - /// usage must properly tie the lifetime to a source. - unsafe fn get_row<'a>(self, ptr: NonNull, i: usize) -> Self::Row<'a>; -} - -/// Extension of [`Repr`] that supports mutable row access. -/// -/// # Associated Types -/// -/// - `RowMut<'a>`: The mutable row type (e.g., `&mut [f32]`). -/// -/// # Safety -/// -/// Implementors must ensure: -/// -/// - [`get_row_mut`](Self::get_row_mut) returns valid references for the given row index. -/// This call **must** be memory safe for `i < self.nrows()`, provided the caller upholds -/// the contract for the raw pointer. -/// -/// Additionally, since the implementation of the [`RowsMut`] iterator can give out rows -/// for all `i` in `0..self.nrows()`, the implementation of [`Self::get_row_mut`] must be -/// such that the result for disjoint `i` must not interfere with one another. -pub unsafe trait ReprMut: Repr { - /// Mutable row reference type. - type RowMut<'a> - where - Self: 'a; - - /// Returns a mutable reference to the i-th row. - /// - /// # Safety - /// - `ptr` must point to a slice with a layout compatible with [`Repr::layout`]. - /// - The entire range for this slice must be within a single allocation. - /// - `i` must be less than `self.nrows()`. - /// - The memory referenced by the returned [`ReprMut::RowMut`] must not be accessed - /// through any other reference for the duration of lifetime `'a`. - /// - The lifetime for the returned [`ReprMut::RowMut`] is inferred from its usage. - /// Correct usage must properly tie the lifetime to a source. - unsafe fn get_row_mut<'a>(self, ptr: NonNull, i: usize) -> Self::RowMut<'a>; -} - -/// Extension trait for [`Repr`] that supports deallocation of owned matrices. This is used -/// in conjunction with [`NewOwned`] to create matrices. -/// -/// Requires [`ReprMut`] since owned matrices should support mutation. -/// -/// # Safety -/// -/// Implementors must ensure that `drop` properly deallocates the memory in a way compatible -/// with all [`NewOwned`] implementations. -pub unsafe trait ReprOwned: ReprMut { - /// Deallocates memory at `ptr` and drops `self`. - /// - /// # Safety - /// - /// - `ptr` must have been obtained via [`NewOwned`] with the same value of `self`. - /// - This method may only be called once for such a pointer. - /// - After calling this method, the memory behind `ptr` may not be dereferenced at all. - unsafe fn drop(self, ptr: NonNull); -} - -/// A new-type version of `std::alloc::LayoutError` for cleaner error handling. -/// -/// This is basically the same as [`std::alloc::LayoutError`], but constructible in -/// use code to allow implementors of [`Repr::layout`] to return it for reasons other than -/// those derived from `std::alloc::Layout`'s methods. -#[derive(Debug, Clone, Copy)] -#[non_exhaustive] -pub struct LayoutError; - -impl LayoutError { - /// Construct a new opaque [`LayoutError`]. - pub fn new() -> Self { - Self - } -} - -impl Default for LayoutError { - fn default() -> Self { - Self::new() - } -} - -impl std::fmt::Display for LayoutError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "LayoutError") - } -} - -impl std::error::Error for LayoutError {} - -impl From for LayoutError { - fn from(_: std::alloc::LayoutError) -> Self { - LayoutError - } -} - -////////////////// -// Constructors // -////////////////// - -/// Create a new [`MatRef`] over a slice. -/// -/// # Safety -/// -/// Implementations must validate the length (and any other requirements) of the provided -/// slice to ensure it is compatible with the implementation of [`Repr`]. -pub unsafe trait NewRef: Repr { - /// Errors that can occur when initializing. - type Error; - - /// Create a new [`MatRef`] over `slice`. - fn new_ref(self, slice: &[T]) -> Result, Self::Error>; -} - -/// Create a new [`MatMut`] over a slice. -/// -/// # Safety -/// -/// Implementations must validate the length (and any other requirements) of the provided -/// slice to ensure it is compatible with the implementation of [`ReprMut`]. -pub unsafe trait NewMut: ReprMut { - /// Errors that can occur when initializing. - type Error; - - /// Create a new [`MatMut`] over `slice`. - fn new_mut(self, slice: &mut [T]) -> Result, Self::Error>; -} - -/// Create a new [`Mat`] from an initializer. -/// -/// # Safety -/// -/// Implementations must ensure that the returned [`Mat`] is compatible with -/// `Self`'s implementation of [`ReprOwned`]. -pub unsafe trait NewOwned: ReprOwned { - /// Errors that can occur when initializing. - type Error; - - /// Create a new [`Mat`] initialized with `init`. - fn new_owned(self, init: T) -> Result, Self::Error>; -} - -/// An initializer argument to [`NewOwned`] that uses a type's [`Default`] implementation -/// to initialize a matrix. -/// -/// ```rust -/// use diskann_quantization::multi_vector::{Mat, Standard, Defaulted}; -/// let mat = Mat::new(Standard::::new(4, 3).unwrap(), Defaulted).unwrap(); -/// for i in 0..4 { -/// assert!(mat.get_row(i).unwrap().iter().all(|&x| x == 0.0f32)); -/// } -/// ``` -#[derive(Debug, Clone, Copy)] -pub struct Defaulted; - -/// Create a new [`Mat`] cloned from a view. -pub trait NewCloned: ReprOwned { - /// Clone the contents behind `v`, returning a new owning [`Mat`]. - /// - /// Implementations should ensure the returned [`Mat`] is "semantically the same" as `v`. - fn new_cloned(v: MatRef<'_, Self>) -> Mat; -} - -////////////// -// Standard // -////////////// - -/// Metadata for dense row-major matrices. -/// -/// Rows are stored contiguously as `&[T]` slices. This is the default representation -/// type for standard floating-point multi-vectors. -/// -/// # Row Types -/// -/// - `Row<'a>`: `&'a [T]` -/// - `RowMut<'a>`: `&'a mut [T]` -#[derive(Debug)] -pub struct Standard { - nrows: usize, - ncols: usize, - _elem: PhantomData, -} - -// Hand-written so `Standard` is `Copy`/`Clone`/`PartialEq`/`Eq` for every `T`: it only -// stores two `usize` and a `PhantomData`, so derives would spuriously require the same -// bound on `T` (and the `Repr: Copy` supertrait must hold regardless of the element type). -impl Copy for Standard {} - -impl Clone for Standard { - fn clone(&self) -> Self { - *self - } -} - -impl PartialEq for Standard { - fn eq(&self, other: &Self) -> bool { - self.nrows == other.nrows && self.ncols == other.ncols - } -} - -impl Eq for Standard {} - -impl Standard { - /// Create a new `Standard` for data of type `T`. - /// - /// Successful construction requires: - /// - /// * The total number of elements determined by `nrows * ncols` does not exceed - /// `usize::MAX`. - /// * The total memory footprint defined by `ncols * nrows * size_of::()` does not - /// exceed `isize::MAX`. - pub fn new(nrows: usize, ncols: usize) -> Result { - Overflow::check::(nrows, ncols)?; - Ok(Self { - nrows, - ncols, - _elem: PhantomData, - }) - } - - /// Returns the number of total elements (`rows x cols`) in this matrix. - pub fn num_elements(&self) -> usize { - // Since we've constructed `self` - we know we cannot overflow. - self.nrows() * self.ncols() - } - - /// Returns `rows`, the number of rows in this matrix. - fn nrows(&self) -> usize { - self.nrows - } - - /// Returns `ncols`, the number of elements in a row of this matrix. - fn ncols(&self) -> usize { - self.ncols - } - - /// Checks the following: - /// - /// 1. Computation of the number of elements in `self` does not overflow. - /// 2. Argument `slice` has the expected number of elements. - fn check_slice(&self, slice: &[T]) -> Result<(), SliceError> { - let len = self.num_elements(); - - if slice.len() != len { - Err(SliceError::LengthMismatch { - expected: len, - found: slice.len(), - }) - } else { - Ok(()) - } - } - - /// Create a new [`Mat`] around the contents of `b` **without** any checks. - /// - /// # Safety - /// - /// The length of `b` must be exactly [`Standard::num_elements`]. - unsafe fn box_to_mat(self, b: Box<[T]>) -> Mat { - debug_assert_eq!(b.len(), self.num_elements(), "safety contract violated"); - - let ptr = utils::box_into_nonnull(b).cast::(); - - // SAFETY: `ptr` is properly aligned and points to a slice of the required length. - // Additionally, it is dropped via `Box::from_raw`, which is compatible with obtaining - // it from `Box::into_raw`. - unsafe { Mat::from_raw_parts(self, ptr) } - } -} - -/// Error for [`Standard::new`]. -#[derive(Debug, Clone, Copy)] -pub struct Overflow { - nrows: usize, - ncols: usize, - elsize: usize, -} - -impl Overflow { - /// Construct an `Overflow` error for the given dimensions and element type. - pub(crate) fn for_type(nrows: usize, ncols: usize) -> Self { - Self { - nrows, - ncols, - elsize: std::mem::size_of::(), - } - } - - /// Verify that `capacity` elements of type `T` fit within the `isize::MAX` byte - /// budget required by Rust's allocation APIs. - /// - /// On failure the error reports the original `(nrows, ncols)` dimensions rather - /// than the padded capacity. - pub(crate) fn check_byte_budget( - capacity: usize, - nrows: usize, - ncols: usize, - ) -> Result<(), Self> { - let bytes = std::mem::size_of::().saturating_mul(capacity); - if bytes <= isize::MAX as usize { - Ok(()) - } else { - Err(Self::for_type::(nrows, ncols)) - } - } - - pub(crate) fn check(nrows: usize, ncols: usize) -> Result<(), Self> { - // Guard the element count itself so that `num_elements()` can never overflow. - let capacity = nrows - .checked_mul(ncols) - .ok_or_else(|| Self::for_type::(nrows, ncols))?; - - Self::check_byte_budget::(capacity, nrows, ncols) - } -} - -impl std::fmt::Display for Overflow { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - if self.elsize == 0 { - write!( - f, - "ZST matrix with dimensions {} x {} has more than `usize::MAX` elements", - self.nrows, self.ncols, - ) - } else { - write!( - f, - "a matrix of size {} x {} with element size {} would exceed isize::MAX bytes", - self.nrows, self.ncols, self.elsize, - ) - } - } -} - -impl std::error::Error for Overflow {} - -/// Error types for [`Standard`]. -#[derive(Debug, Clone, Copy, Error)] -#[non_exhaustive] -pub enum SliceError { - #[error("Length mismatch: expected {expected}, found {found}")] - LengthMismatch { expected: usize, found: usize }, -} - -// SAFETY: The implementation correctly computes row offsets as `i * ncols` and -// constructs valid slices of the appropriate length. The `layout` method correctly -// reports the memory layout requirements. -unsafe impl Repr for Standard { - type Row<'a> - = &'a [T] - where - T: 'a; - - fn nrows(&self) -> usize { - self.nrows - } - - fn layout(&self) -> Result { - Ok(Layout::array::(self.num_elements())?) - } - - unsafe fn get_row<'a>(self, ptr: NonNull, i: usize) -> Self::Row<'a> { - debug_assert!(ptr.cast::().is_aligned()); - debug_assert!(i < self.nrows); - - // SAFETY: The caller asserts that `i` is less than `self.nrows()`. Since this type - // audits the constructors for `Mat` and friends, we know that there is room for at - // least `self.num_elements()` elements from the base pointer, so this access is safe. - let row_ptr = unsafe { ptr.as_ptr().cast::().add(i * self.ncols) }; - - // SAFETY: The logic is the same as the previous `unsafe` block. - unsafe { std::slice::from_raw_parts(row_ptr, self.ncols) } - } -} - -// SAFETY: The implementation correctly computes row offsets and constructs valid mutable -// slices. -unsafe impl ReprMut for Standard { - type RowMut<'a> - = &'a mut [T] - where - T: 'a; - - unsafe fn get_row_mut<'a>(self, ptr: NonNull, i: usize) -> Self::RowMut<'a> { - debug_assert!(ptr.cast::().is_aligned()); - debug_assert!(i < self.nrows); - - // SAFETY: The caller asserts that `i` is less than `self.nrows()`. Since this type - // audits the constructors for `Mat` and friends, we know that there is room for at - // least `self.num_elements()` elements from the base pointer, so this access is safe. - let row_ptr = unsafe { ptr.as_ptr().cast::().add(i * self.ncols) }; - - // SAFETY: The logic is the same as the previous `unsafe` block. Further, the caller - // attests that creating a mutable reference is safe. - unsafe { std::slice::from_raw_parts_mut(row_ptr, self.ncols) } - } -} - -// SAFETY: The drop implementation correctly reconstructs a Box from the raw pointer -// using the same length (nrows * ncols) that was used for allocation, allowing Box -// to properly deallocate the memory. -unsafe impl ReprOwned for Standard { - unsafe fn drop(self, ptr: NonNull) { - // SAFETY: The caller guarantees that `ptr` was obtained from an implementation of - // `NewOwned` for an equivalent instance of `self`. - // - // We ensure that `NewOwned` goes through boxes, so here we reconstruct a Box to - // let it handle deallocation. - unsafe { - let slice_ptr = std::ptr::slice_from_raw_parts_mut( - ptr.cast::().as_ptr(), - self.nrows * self.ncols, - ); - let _ = Box::from_raw(slice_ptr); - } - } -} - -// SAFETY: The implementation uses guarantees from `Box` to ensure that the pointer -// initialized by it is non-null and properly aligned to the underlying type. -unsafe impl NewOwned for Standard -where - T: Clone, -{ - type Error = crate::error::Infallible; - fn new_owned(self, value: T) -> Result, Self::Error> { - let b: Box<[T]> = std::iter::repeat_n(value, self.num_elements()).collect(); - - // SAFETY: By construction, `b` has length `self.num_elements()`. - Ok(unsafe { self.box_to_mat(b) }) - } -} - -// SAFETY: The implementation uses guarantees from `Box` to ensure that the pointer -// initialized by it is non-null and properly aligned to the underlying type. -unsafe impl NewOwned for Standard -where - T: Default, -{ - type Error = crate::error::Infallible; - fn new_owned(self, _: Defaulted) -> Result, Self::Error> { - let b: Box<[T]> = std::iter::repeat_with(T::default) - .take(self.num_elements()) - .collect(); - - // SAFETY: By construction, `b` has length `self.num_elements()`. - Ok(unsafe { self.box_to_mat(b) }) - } -} - -// SAFETY: This checks that the slice has the correct length, which is all that is -// required for [`Repr`]. -unsafe impl NewRef for Standard { - type Error = SliceError; - fn new_ref(self, data: &[T]) -> Result, Self::Error> { - self.check_slice(data)?; - - // SAFETY: The function `check_slice` verifies that `data` is compatible with - // the layout requirement of `Standard`. - // - // We've properly checked that the underlying pointer is okay. - Ok(unsafe { MatRef::from_raw_parts(self, utils::as_nonnull(data).cast::()) }) - } -} - -// SAFETY: This checks that the slice has the correct length, which is all that is -// required for [`ReprMut`]. -unsafe impl NewMut for Standard { - type Error = SliceError; - fn new_mut(self, data: &mut [T]) -> Result, Self::Error> { - self.check_slice(data)?; - - // SAFETY: The function `check_slice` verifies that `data` is compatible with - // the layout requirement of `Standard`. - // - // We've properly checked that the underlying pointer is okay. - Ok(unsafe { MatMut::from_raw_parts(self, utils::as_nonnull_mut(data).cast::()) }) - } -} - -impl NewCloned for Standard -where - T: Clone, -{ - fn new_cloned(v: MatRef<'_, Self>) -> Mat { - let b: Box<[T]> = v.as_slice().iter().cloned().collect(); - - // SAFETY: By construction, `b` has length `v.repr().num_elements()`. - unsafe { v.repr().box_to_mat(b) } - } -} - -///////// -// Mat // -///////// - -/// An owning matrix that manages its own memory. -/// -/// The matrix stores raw bytes interpreted according to representation type `T`. -/// Memory is automatically deallocated when the matrix is dropped. -#[derive(Debug)] -pub struct Mat { - ptr: NonNull, - repr: T, - _invariant: PhantomData T>, -} - -// SAFETY: [`Repr`] is required to propagate its `Send` bound. -unsafe impl Send for Mat where T: ReprOwned + Send {} - -// SAFETY: [`Repr`] is required to propagate its `Sync` bound. -unsafe impl Sync for Mat where T: ReprOwned + Sync {} - -impl Mat { - /// Create a new matrix using `init` as the initializer. - pub fn new(repr: T, init: U) -> Result>::Error> - where - T: NewOwned, - { - repr.new_owned(init) - } - - /// Returns the number of rows (vectors) in the matrix. - #[inline] - pub fn num_vectors(&self) -> usize { - self.repr.nrows() - } - - /// Returns a reference to the underlying representation. - pub fn repr(&self) -> &T { - &self.repr - } - - /// Returns the `i`th row if `i < self.num_vectors()`. - #[must_use] - pub fn get_row(&self, i: usize) -> Option> { - if i < self.num_vectors() { - // SAFETY: Bounds check passed, and the Mat was constructed - // with valid representation and pointer. - let row = unsafe { self.get_row_unchecked(i) }; - Some(row) - } else { - None - } - } - - pub(crate) unsafe fn get_row_unchecked(&self, i: usize) -> T::Row<'_> { - // SAFETY: Caller must ensure i < self.num_vectors(). The constructors for this type - // ensure that `ptr` is compatible with `T`. - unsafe { self.repr.get_row(self.ptr, i) } - } - - /// Returns the `i`th mutable row if `i < self.num_vectors()`. - #[must_use] - pub fn get_row_mut(&mut self, i: usize) -> Option> { - if i < self.num_vectors() { - // SAFETY: Bounds check passed, and we have exclusive access via &mut self. - Some(unsafe { self.get_row_mut_unchecked(i) }) - } else { - None - } - } - - pub(crate) unsafe fn get_row_mut_unchecked(&mut self, i: usize) -> T::RowMut<'_> { - // SAFETY: Caller asserts that `i < self.num_vectors()`. The constructors for this - // type ensure that `ptr` is compatible with `T`. - unsafe { self.repr.get_row_mut(self.ptr, i) } - } - - /// Returns an immutable view of the matrix. - #[inline] - pub fn as_view(&self) -> MatRef<'_, T> { - MatRef { - ptr: self.ptr, - repr: self.repr, - _lifetime: PhantomData, - } - } - - /// Returns a mutable view of the matrix. - #[inline] - pub fn as_view_mut(&mut self) -> MatMut<'_, T> { - MatMut { - ptr: self.ptr, - repr: self.repr, - _lifetime: PhantomData, - } - } - - /// Returns an iterator over immutable row references. - pub fn rows(&self) -> Rows<'_, T> { - Rows::new(self.reborrow()) - } - - /// Returns an iterator over mutable row references. - pub fn rows_mut(&mut self) -> RowsMut<'_, T> { - RowsMut::new(self.reborrow_mut()) - } - - /// Construct a new [`Mat`] over the raw pointer and representation without performing - /// any validity checks. - /// - /// # Safety - /// - /// Argument `ptr` must be: - /// - /// 1. Point to memory compatible with [`Repr::layout`]. - /// 2. Be compatible with the drop logic in [`ReprOwned`]. - pub(crate) unsafe fn from_raw_parts(repr: T, ptr: NonNull) -> Self { - Self { - ptr, - repr, - _invariant: PhantomData, - } - } - - /// Return the base pointer for the [`Mat`]. - pub fn as_raw_ptr(&self) -> *const u8 { - self.ptr.as_ptr() - } - - /// Return a mutable base pointer for the [`Mat`]. - pub(crate) fn as_raw_mut_ptr(&mut self) -> *mut u8 { - self.ptr.as_ptr() - } -} - -impl Drop for Mat { - fn drop(&mut self) { - // SAFETY: `ptr` was correctly initialized according to `layout` - // and we are guaranteed exclusive access to the data due to Rust borrow rules. - unsafe { self.repr.drop(self.ptr) }; - } -} - -impl Clone for Mat { - fn clone(&self) -> Self { - T::new_cloned(self.as_view()) - } -} - -impl Mat> { - /// Construct a [`Mat`] by calling `f` once per element in row-major order. - pub fn from_fn T>(repr: Standard, mut f: F) -> Self { - let b: Box<[T]> = (0..repr.num_elements()).map(|_| f()).collect(); - // SAFETY: `b` has length `repr.num_elements()` by construction. - unsafe { repr.box_to_mat(b) } - } - - /// Returns the raw dimension (columns) of the vectors in the matrix. - #[inline] - pub fn vector_dim(&self) -> usize { - self.repr.ncols() - } - - /// Return the backing data as a contiguous slice of `T`. - /// - /// The returned slice has `num_vectors() * vector_dim()` elements in row-major order. - #[inline] - pub fn as_slice(&self) -> &[T] { - self.as_view().as_slice() - } - - /// Return a [`MatrixView`] over the backing data. - #[inline] - pub fn as_matrix_view(&self) -> MatrixView<'_, T> { - self.as_view().as_matrix_view() - } -} - -//////////// -// MatRef // -//////////// - -/// An immutable borrowed view of a matrix. -/// -/// Provides read-only access to matrix data without ownership. Implements [`Copy`] -/// and can be freely cloned. -/// -/// # Type Parameter -/// - `T`: A [`Repr`] implementation defining the row layout. -/// -/// # Access -/// - [`get_row`](Self::get_row): Get an immutable row by index. -/// - [`rows`](Self::rows): Iterate over all rows. -#[derive(Debug, Clone, Copy)] -pub struct MatRef<'a, T: Repr> { - ptr: NonNull, - repr: T, - /// Marker to tie the lifetime to the borrowed data. - _lifetime: PhantomData<&'a T>, -} - -// SAFETY: [`Repr`] is required to propagate its `Send` bound. -unsafe impl Send for MatRef<'_, T> where T: Repr + Send {} - -// SAFETY: [`Repr`] is required to propagate its `Sync` bound. -unsafe impl Sync for MatRef<'_, T> where T: Repr + Sync {} - -impl<'a, T: Repr> MatRef<'a, T> { - /// Construct a new [`MatRef`] over `data`. - pub fn new(repr: T, data: &'a [U]) -> Result - where - T: NewRef, - { - repr.new_ref(data) - } - - /// Returns the number of rows (vectors) in the matrix. - #[inline] - pub fn num_vectors(&self) -> usize { - self.repr.nrows() - } - - /// Returns a reference to the underlying representation. - pub fn repr(&self) -> &T { - &self.repr - } - - /// Returns an immutable reference to the i-th row, or `None` if out of bounds. - #[must_use] - pub fn get_row(&self, i: usize) -> Option> { - if i < self.num_vectors() { - // SAFETY: Bounds check passed, and the MatRef was constructed - // with valid representation and pointer. - let row = unsafe { self.get_row_unchecked(i) }; - Some(row) - } else { - None - } - } - - /// Returns the i-th row without bounds checking. - /// - /// # Safety - /// - /// `i` must be less than `self.num_vectors()`. - #[inline] - pub(crate) unsafe fn get_row_unchecked(&self, i: usize) -> T::Row<'_> { - // SAFETY: Caller must ensure i < self.num_vectors(). - unsafe { self.repr.get_row(self.ptr, i) } - } - - /// Returns an iterator over immutable row references. - pub fn rows(&self) -> Rows<'_, T> { - Rows::new(*self) - } - - /// Return a [`Mat`] with the same contents as `self`. - pub fn to_owned(&self) -> Mat - where - T: NewCloned, - { - T::new_cloned(*self) - } - - /// Construct a new [`MatRef`] over the raw pointer and representation without performing - /// any validity checks. - /// - /// # Safety - /// - /// Argument `ptr` must point to memory compatible with [`Repr::layout`] and pass any - /// validity checks required by `T`. - pub unsafe fn from_raw_parts(repr: T, ptr: NonNull) -> Self { - Self { - ptr, - repr, - _lifetime: PhantomData, - } - } - - /// Return the base pointer for the [`MatRef`]. - pub fn as_raw_ptr(&self) -> *const u8 { - self.ptr.as_ptr() - } -} - -impl<'a, T> MatRef<'a, Standard> { - /// Returns the raw dimension (columns) of the vectors in the matrix. - #[inline] - pub fn vector_dim(&self) -> usize { - self.repr.ncols() - } - - /// Return the backing data as a contiguous slice of `T`. - /// - /// The returned slice has `num_vectors() * vector_dim()` elements in row-major order. - #[inline] - pub fn as_slice(&self) -> &'a [T] { - let len = self.repr.num_elements(); - // SAFETY: `Standard` guarantees `nrows * ncols` contiguous `T` elements - // starting at `self.ptr`. The lifetime `'a` is tied to the original data. - unsafe { std::slice::from_raw_parts(self.ptr.as_ptr().cast::(), len) } - } - - /// Return a [`MatrixView`] over the backing data. - #[allow(clippy::expect_used)] - #[inline] - pub fn as_matrix_view(&self) -> MatrixView<'a, T> { - // `Standard::new` validates that `nrows * ncols` does not overflow, - // so `try_from` is infallible here. - MatrixView::try_from(self.as_slice(), self.num_vectors(), self.vector_dim()) - .expect("Standard has valid dimensions") - } -} - -// Reborrow: Mat -> MatRef -impl<'this, T: ReprOwned> Reborrow<'this> for Mat { - type Target = MatRef<'this, T>; - - fn reborrow(&'this self) -> Self::Target { - self.as_view() - } -} - -// ReborrowMut: Mat -> MatMut -impl<'this, T: ReprOwned> ReborrowMut<'this> for Mat { - type Target = MatMut<'this, T>; - - fn reborrow_mut(&'this mut self) -> Self::Target { - self.as_view_mut() - } -} - -// Reborrow: MatRef -> MatRef (with shorter lifetime) -impl<'this, 'a, T: Repr> Reborrow<'this> for MatRef<'a, T> { - type Target = MatRef<'this, T>; - - fn reborrow(&'this self) -> Self::Target { - MatRef { - ptr: self.ptr, - repr: self.repr, - _lifetime: PhantomData, - } - } -} - -//////////// -// MatMut // -//////////// - -/// A mutable borrowed view of a matrix. -/// -/// Provides read-write access to matrix data without ownership. -/// -/// # Type Parameter -/// - `T`: A [`ReprMut`] implementation defining the row layout. -/// -/// # Access -/// - [`get_row`](Self::get_row): Get an immutable row by index. -/// - [`get_row_mut`](Self::get_row_mut): Get a mutable row by index. -/// - [`as_view`](Self::as_view): Reborrow as immutable [`MatRef`]. -/// - [`rows`](Self::rows), [`rows_mut`](Self::rows_mut): Iterate over rows. -#[derive(Debug)] -pub struct MatMut<'a, T: ReprMut> { - ptr: NonNull, - repr: T, - /// Marker to tie the lifetime to the mutably borrowed data. - _lifetime: PhantomData<&'a mut T>, -} - -// SAFETY: [`ReprMut`] is required to propagate its `Send` bound. -unsafe impl Send for MatMut<'_, T> where T: ReprMut + Send {} - -// SAFETY: [`ReprMut`] is required to propagate its `Sync` bound. -unsafe impl Sync for MatMut<'_, T> where T: ReprMut + Sync {} - -impl<'a, T: ReprMut> MatMut<'a, T> { - /// Construct a new [`MatMut`] over `data`. - pub fn new(repr: T, data: &'a mut [U]) -> Result - where - T: NewMut, - { - repr.new_mut(data) - } - - /// Returns the number of rows (vectors) in the matrix. - #[inline] - pub fn num_vectors(&self) -> usize { - self.repr.nrows() - } - - /// Returns a reference to the underlying representation. - pub fn repr(&self) -> &T { - &self.repr - } - - /// Returns an immutable reference to the i-th row, or `None` if out of bounds. - #[inline] - #[must_use] - pub fn get_row(&self, i: usize) -> Option> { - if i < self.num_vectors() { - // SAFETY: Bounds check passed. - Some(unsafe { self.get_row_unchecked(i) }) - } else { - None - } - } - - /// Returns the i-th row without bounds checking. - /// - /// # Safety - /// - /// `i` must be less than `self.num_vectors()`. - #[inline] - pub(crate) unsafe fn get_row_unchecked(&self, i: usize) -> T::Row<'_> { - // SAFETY: Caller must ensure i < self.num_vectors(). - unsafe { self.repr.get_row(self.ptr, i) } - } - - /// Returns a mutable reference to the `i`-th row, or `None` if out of bounds. - #[inline] - #[must_use] - pub fn get_row_mut(&mut self, i: usize) -> Option> { - if i < self.num_vectors() { - // SAFETY: Bounds check passed. - Some(unsafe { self.get_row_mut_unchecked(i) }) - } else { - None - } - } - - /// Returns a mutable reference to the i-th row without bounds checking. - /// - /// # Safety - /// - /// `i` must be less than [`num_vectors()`](Self::num_vectors). - #[inline] - pub(crate) unsafe fn get_row_mut_unchecked(&mut self, i: usize) -> T::RowMut<'_> { - // SAFETY: Caller asserts that `i < self.num_vectors()`. The constructors for this - // type ensure that `ptr` is compatible with `T`. - unsafe { self.repr.get_row_mut(self.ptr, i) } - } - - /// Reborrows as an immutable [`MatRef`]. - pub fn as_view(&self) -> MatRef<'_, T> { - MatRef { - ptr: self.ptr, - repr: self.repr, - _lifetime: PhantomData, - } - } - - /// Returns an iterator over immutable row references. - pub fn rows(&self) -> Rows<'_, T> { - Rows::new(self.reborrow()) - } - - /// Returns an iterator over mutable row references. - pub fn rows_mut(&mut self) -> RowsMut<'_, T> { - RowsMut::new(self.reborrow_mut()) - } - - /// Return a [`Mat`] with the same contents as `self`. - pub fn to_owned(&self) -> Mat - where - T: NewCloned, - { - T::new_cloned(self.as_view()) - } - - /// Construct a new [`MatMut`] over the raw pointer and representation without performing - /// any validity checks. - /// - /// # Safety - /// - /// Argument `ptr` must point to memory compatible with [`Repr::layout`]. - pub unsafe fn from_raw_parts(repr: T, ptr: NonNull) -> Self { - Self { - ptr, - repr, - _lifetime: PhantomData, - } - } - - /// Return the base pointer for the [`MatMut`]. - pub fn as_raw_ptr(&self) -> *const u8 { - self.ptr.as_ptr() - } - - /// Return a mutable base pointer for the [`MatMut`]. - pub(crate) fn as_raw_mut_ptr(&mut self) -> *mut u8 { - self.ptr.as_ptr() - } -} - -// Reborrow: MatMut -> MatRef -impl<'this, 'a, T: ReprMut> Reborrow<'this> for MatMut<'a, T> { - type Target = MatRef<'this, T>; - - fn reborrow(&'this self) -> Self::Target { - self.as_view() - } -} - -// ReborrowMut: MatMut -> MatMut (with shorter lifetime) -impl<'this, 'a, T: ReprMut> ReborrowMut<'this> for MatMut<'a, T> { - type Target = MatMut<'this, T>; - - fn reborrow_mut(&'this mut self) -> Self::Target { - MatMut { - ptr: self.ptr, - repr: self.repr, - _lifetime: PhantomData, - } - } -} - -impl<'a, T> MatMut<'a, Standard> { - /// Returns the raw dimension (columns) of the vectors in the matrix. - #[inline] - pub fn vector_dim(&self) -> usize { - self.repr.ncols() - } - - /// Return the backing data as a contiguous slice of `T`. - /// - /// The returned slice has `num_vectors() * vector_dim()` elements in row-major order. - #[inline] - pub fn as_slice(&self) -> &[T] { - self.as_view().as_slice() - } - - /// Return a [`MatrixView`] over the backing data. - #[inline] - pub fn as_matrix_view(&self) -> MatrixView<'_, T> { - self.as_view().as_matrix_view() - } -} - -////////// -// Rows // -////////// - -/// Iterator over immutable row references of a matrix. -/// -/// Created by [`Mat::rows`], [`MatRef::rows`], or [`MatMut::rows`]. -#[derive(Debug)] -pub struct Rows<'a, T: Repr> { - matrix: MatRef<'a, T>, - current: usize, -} - -impl<'a, T> Rows<'a, T> -where - T: Repr, -{ - fn new(matrix: MatRef<'a, T>) -> Self { - Self { matrix, current: 0 } - } -} - -impl<'a, T> Iterator for Rows<'a, T> -where - T: Repr + 'a, -{ - type Item = T::Row<'a>; - - fn next(&mut self) -> Option { - let current = self.current; - if current >= self.matrix.num_vectors() { - None - } else { - self.current += 1; - // SAFETY: We make sure through the above check that - // the access is within bounds. - // - // Extending the lifetime to `'a` is safe because the underlying - // MatRef has lifetime `'a`. - Some(unsafe { self.matrix.repr.get_row(self.matrix.ptr, current) }) - } - } - - fn size_hint(&self) -> (usize, Option) { - let remaining = self.matrix.num_vectors() - self.current; - (remaining, Some(remaining)) - } -} - -impl<'a, T> ExactSizeIterator for Rows<'a, T> where T: Repr + 'a {} -impl<'a, T> FusedIterator for Rows<'a, T> where T: Repr + 'a {} - -///////////// -// RowsMut // -///////////// - -/// Iterator over mutable row references of a matrix. -/// -/// Created by [`Mat::rows_mut`] or [`MatMut::rows_mut`]. -#[derive(Debug)] -pub struct RowsMut<'a, T: ReprMut> { - matrix: MatMut<'a, T>, - current: usize, -} - -impl<'a, T> RowsMut<'a, T> -where - T: ReprMut, -{ - fn new(matrix: MatMut<'a, T>) -> Self { - Self { matrix, current: 0 } - } -} - -impl<'a, T> Iterator for RowsMut<'a, T> -where - T: ReprMut + 'a, -{ - type Item = T::RowMut<'a>; - - fn next(&mut self) -> Option { - let current = self.current; - if current >= self.matrix.num_vectors() { - None - } else { - self.current += 1; - // SAFETY: We make sure through the above check that - // the access is within bounds. - // - // Extending the lifetime to `'a` is safe because: - // 1. The underlying MatMut has lifetime `'a`. - // 2. The iterator ensures that the mutable row indices are disjoint, so - // there is no aliasing as long as the implementation of `ReprMut` ensures - // there is not mutable sharing of the `RowMut` types. - Some(unsafe { self.matrix.repr.get_row_mut(self.matrix.ptr, current) }) - } - } - - fn size_hint(&self) -> (usize, Option) { - let remaining = self.matrix.num_vectors() - self.current; - (remaining, Some(remaining)) - } -} - -impl<'a, T> ExactSizeIterator for RowsMut<'a, T> where T: ReprMut + 'a {} -impl<'a, T> FusedIterator for RowsMut<'a, T> where T: ReprMut + 'a {} - -/////////// -// Tests // -/////////// - -#[cfg(test)] -mod tests { - use super::*; - - use std::fmt::Display; - - use diskann_utils::lazy_format; - - /// Helper to assert a type is Copy. - fn assert_copy(_: &T) {} - - // ── Variance assertions ────────────────────────────────────── - // - // These functions are never called. The test is that they compile: - // covariant positions must accept subtype coercions. - // - // The negative (invariance) counterparts live in - // `tests/compile-fail/multi/{mat,matmut}_invariant.rs`. - - /// `MatRef` is covariant in `'a`: a longer borrow can shorten. - fn _assert_matref_covariant_lifetime<'long: 'short, 'short, T: Repr>( - v: MatRef<'long, T>, - ) -> MatRef<'short, T> { - v - } - - /// `MatRef` is covariant in `T`: `Standard<&'long u8>` → `Standard<&'short u8>`. - fn _assert_matref_covariant_repr<'long: 'short, 'short, 'a>( - v: MatRef<'a, Standard<&'long u8>>, - ) -> MatRef<'a, Standard<&'short u8>> { - v - } - - /// `MatMut` is covariant in `'a`: a longer borrow can shorten. - fn _assert_matmut_covariant_lifetime<'long: 'short, 'short, T: ReprMut>( - v: MatMut<'long, T>, - ) -> MatMut<'short, T> { - v - } - - fn edge_cases(nrows: usize) -> Vec { - let max = usize::MAX; - - vec![ - nrows, - nrows + 1, - nrows + 11, - nrows + 20, - max / 2, - max.div_ceil(2), - max - 1, - max, - ] - } - - fn fill_mat(x: &mut Mat>, repr: Standard) { - assert_eq!(x.repr(), &repr); - assert_eq!(x.num_vectors(), repr.nrows()); - assert_eq!(x.vector_dim(), repr.ncols()); - - for i in 0..x.num_vectors() { - let row = x.get_row_mut(i).unwrap(); - assert_eq!(row.len(), repr.ncols()); - row.iter_mut() - .enumerate() - .for_each(|(j, r)| *r = 10 * i + j); - } - - for i in edge_cases(repr.nrows()).into_iter() { - assert!(x.get_row_mut(i).is_none()); - } - } - - fn fill_mat_mut(mut x: MatMut<'_, Standard>, repr: Standard) { - assert_eq!(x.repr(), &repr); - assert_eq!(x.num_vectors(), repr.nrows()); - assert_eq!(x.vector_dim(), repr.ncols()); - - for i in 0..x.num_vectors() { - let row = x.get_row_mut(i).unwrap(); - assert_eq!(row.len(), repr.ncols()); - - row.iter_mut() - .enumerate() - .for_each(|(j, r)| *r = 10 * i + j); - } - - for i in edge_cases(repr.nrows()).into_iter() { - assert!(x.get_row_mut(i).is_none()); - } - } - - fn fill_rows_mut(x: RowsMut<'_, Standard>, repr: Standard) { - assert_eq!(x.len(), repr.nrows()); - // Materialize all rows at once. - let mut all_rows: Vec<_> = x.collect(); - assert_eq!(all_rows.len(), repr.nrows()); - for (i, row) in all_rows.iter_mut().enumerate() { - assert_eq!(row.len(), repr.ncols()); - row.iter_mut() - .enumerate() - .for_each(|(j, r)| *r = 10 * i + j); - } - } - - fn check_mat(x: &Mat>, repr: Standard, ctx: &dyn Display) { - assert_eq!(x.repr(), &repr); - assert_eq!(x.num_vectors(), repr.nrows()); - assert_eq!(x.vector_dim(), repr.ncols()); - - for i in 0..x.num_vectors() { - let row = x.get_row(i).unwrap(); - - assert_eq!(row.len(), repr.ncols(), "ctx: {ctx}"); - row.iter().enumerate().for_each(|(j, r)| { - assert_eq!( - *r, - 10 * i + j, - "mismatched entry at row {}, col {} -- ctx: {}", - i, - j, - ctx - ) - }); - } - - for i in edge_cases(repr.nrows()).into_iter() { - assert!(x.get_row(i).is_none(), "ctx: {ctx}"); - } - } - - fn check_mat_ref(x: MatRef<'_, Standard>, repr: Standard, ctx: &dyn Display) { - assert_eq!(x.repr(), &repr); - assert_eq!(x.num_vectors(), repr.nrows()); - assert_eq!(x.vector_dim(), repr.ncols()); - - assert_copy(&x); - for i in 0..x.num_vectors() { - let row = x.get_row(i).unwrap(); - assert_eq!(row.len(), repr.ncols(), "ctx: {ctx}"); - - row.iter().enumerate().for_each(|(j, r)| { - assert_eq!( - *r, - 10 * i + j, - "mismatched entry at row {}, col {} -- ctx: {}", - i, - j, - ctx - ) - }); - } - - for i in edge_cases(repr.nrows()).into_iter() { - assert!(x.get_row(i).is_none(), "ctx: {ctx}"); - } - } - - fn check_mat_mut(x: MatMut<'_, Standard>, repr: Standard, ctx: &dyn Display) { - assert_eq!(x.repr(), &repr); - assert_eq!(x.num_vectors(), repr.nrows()); - assert_eq!(x.vector_dim(), repr.ncols()); - - for i in 0..x.num_vectors() { - let row = x.get_row(i).unwrap(); - assert_eq!(row.len(), repr.ncols(), "ctx: {ctx}"); - - row.iter().enumerate().for_each(|(j, r)| { - assert_eq!( - *r, - 10 * i + j, - "mismatched entry at row {}, col {} -- ctx: {}", - i, - j, - ctx - ) - }); - } - - for i in edge_cases(repr.nrows()).into_iter() { - assert!(x.get_row(i).is_none(), "ctx: {ctx}"); - } - } - - fn check_rows(x: Rows<'_, Standard>, repr: Standard, ctx: &dyn Display) { - assert_eq!(x.len(), repr.nrows(), "ctx: {ctx}"); - let all_rows: Vec<_> = x.collect(); - assert_eq!(all_rows.len(), repr.nrows(), "ctx: {ctx}"); - for (i, row) in all_rows.iter().enumerate() { - assert_eq!(row.len(), repr.ncols(), "ctx: {ctx}"); - row.iter().enumerate().for_each(|(j, r)| { - assert_eq!( - *r, - 10 * i + j, - "mismatched entry at row {}, col {} -- ctx: {}", - i, - j, - ctx - ) - }); - } - } - - ////////////// - // Standard // - ////////////// - - #[test] - fn standard_representation() { - let repr = Standard::::new(4, 3).unwrap(); - assert_eq!(repr.nrows(), 4); - assert_eq!(repr.ncols(), 3); - - let layout = repr.layout().unwrap(); - assert_eq!(layout.size(), 4 * 3 * std::mem::size_of::()); - assert_eq!(layout.align(), std::mem::align_of::()); - } - - #[test] - fn standard_zero_dimensions() { - for (nrows, ncols) in [(0, 0), (0, 5), (5, 0)] { - let repr = Standard::::new(nrows, ncols).unwrap(); - assert_eq!(repr.nrows(), nrows); - assert_eq!(repr.ncols(), ncols); - let layout = repr.layout().unwrap(); - assert_eq!(layout.size(), 0); - } - } - - #[test] - fn standard_check_slice() { - let repr = Standard::::new(3, 4).unwrap(); - - // Correct length succeeds - let data = vec![0u32; 12]; - assert!(repr.check_slice(&data).is_ok()); - - // Too short fails - let short = vec![0u32; 11]; - assert!(matches!( - repr.check_slice(&short), - Err(SliceError::LengthMismatch { - expected: 12, - found: 11 - }) - )); - - // Too long fails - let long = vec![0u32; 13]; - assert!(matches!( - repr.check_slice(&long), - Err(SliceError::LengthMismatch { - expected: 12, - found: 13 - }) - )); - - // Overflow case - let overflow_repr = Standard::::new(usize::MAX, 2).unwrap_err(); - assert!(matches!(overflow_repr, Overflow { .. })); - } - - #[test] - fn standard_new_rejects_element_count_overflow() { - // nrows * ncols overflows usize even though per-element size is small. - assert!(Standard::::new(usize::MAX, 2).is_err()); - assert!(Standard::::new(2, usize::MAX).is_err()); - assert!(Standard::::new(usize::MAX, usize::MAX).is_err()); - } - - #[test] - fn standard_new_rejects_byte_count_exceeding_isize_max() { - // Element count fits in usize, but total bytes exceed isize::MAX. - let half = (isize::MAX as usize / std::mem::size_of::()) + 1; - assert!(Standard::::new(half, 1).is_err()); - assert!(Standard::::new(1, half).is_err()); - } - - #[test] - fn standard_new_accepts_boundary_below_isize_max() { - // Largest allocation that still fits in isize::MAX bytes. - let max_elems = isize::MAX as usize / std::mem::size_of::(); - let repr = Standard::::new(max_elems, 1).unwrap(); - assert_eq!(repr.num_elements(), max_elems); - } - - #[test] - fn standard_new_zst_rejects_element_count_overflow() { - // For ZSTs the byte count is always 0, but element-count overflow - // must still be caught so that `num_elements()` never wraps. - assert!(Standard::<()>::new(usize::MAX, 2).is_err()); - assert!(Standard::<()>::new(usize::MAX / 2 + 1, 3).is_err()); - } - - #[test] - fn standard_new_zst_accepts_large_non_overflowing() { - // Large-but-valid ZST matrix: element count fits in usize. - let repr = Standard::<()>::new(usize::MAX, 1).unwrap(); - assert_eq!(repr.num_elements(), usize::MAX); - assert_eq!(repr.layout().unwrap().size(), 0); - } - - #[test] - fn standard_new_overflow_error_display() { - let err = Standard::::new(usize::MAX, 2).unwrap_err(); - let msg = err.to_string(); - assert!(msg.contains("would exceed isize::MAX bytes"), "{msg}"); - - let zst_err = Standard::<()>::new(usize::MAX, 2).unwrap_err(); - let zst_msg = zst_err.to_string(); - assert!(zst_msg.contains("ZST matrix"), "{zst_msg}"); - assert!(zst_msg.contains("usize::MAX"), "{zst_msg}"); - } - - ///////// - // Mat // - ///////// - - #[test] - fn mat_new_and_basic_accessors() { - let mat = Mat::new(Standard::::new(3, 4).unwrap(), 42usize).unwrap(); - let base: *const u8 = mat.as_raw_ptr(); - - assert_eq!(mat.num_vectors(), 3); - assert_eq!(mat.vector_dim(), 4); - - let repr = mat.repr(); - assert_eq!(repr.nrows(), 3); - assert_eq!(repr.ncols(), 4); - - for (i, r) in mat.rows().enumerate() { - assert_eq!(r, &[42, 42, 42, 42]); - let ptr = r.as_ptr().cast::(); - assert_eq!( - ptr, - base.wrapping_add(std::mem::size_of::() * mat.repr().ncols() * i), - ); - } - } - - #[test] - fn mat_new_with_default() { - let mat = Mat::new(Standard::::new(2, 3).unwrap(), Defaulted).unwrap(); - let base: *const u8 = mat.as_raw_ptr(); - - assert_eq!(mat.num_vectors(), 2); - for (i, row) in mat.rows().enumerate() { - assert!(row.iter().all(|&v| v == 0)); - - let ptr = row.as_ptr().cast::(); - assert_eq!( - ptr, - base.wrapping_add(std::mem::size_of::() * mat.repr().ncols() * i), - ); - } - } - - const ROWS: &[usize] = &[0, 1, 2, 3, 5, 10]; - const COLS: &[usize] = &[0, 1, 2, 3, 5, 10]; - - #[test] - fn test_mat() { - for nrows in ROWS { - for ncols in COLS { - let repr = Standard::::new(*nrows, *ncols).unwrap(); - let ctx = &lazy_format!("nrows = {}, ncols = {}", nrows, ncols); - - // Populate the matrix using `&mut Mat` - { - let ctx = &lazy_format!("{ctx} - direct"); - let mut mat = Mat::new(repr, Defaulted).unwrap(); - - assert_eq!(mat.num_vectors(), *nrows); - assert_eq!(mat.vector_dim(), *ncols); - - fill_mat(&mut mat, repr); - - check_mat(&mat, repr, ctx); - check_mat_ref(mat.reborrow(), repr, ctx); - check_mat_mut(mat.reborrow_mut(), repr, ctx); - check_rows(mat.rows(), repr, ctx); - - // Check reborrow preserves pointers. - assert_eq!(mat.as_raw_ptr(), mat.reborrow().as_raw_ptr()); - assert_eq!(mat.as_raw_ptr(), mat.reborrow_mut().as_raw_ptr()); - } - - // Populate the matrix using `MatMut` - { - let ctx = &lazy_format!("{ctx} - matmut"); - let mut mat = Mat::new(repr, Defaulted).unwrap(); - let matmut = mat.reborrow_mut(); - - assert_eq!(matmut.num_vectors(), *nrows); - assert_eq!(matmut.vector_dim(), *ncols); - - fill_mat_mut(matmut, repr); - - check_mat(&mat, repr, ctx); - check_mat_ref(mat.reborrow(), repr, ctx); - check_mat_mut(mat.reborrow_mut(), repr, ctx); - check_rows(mat.rows(), repr, ctx); - } - - // Populate the matrix using `RowsMut` - { - let ctx = &lazy_format!("{ctx} - rows_mut"); - let mut mat = Mat::new(repr, Defaulted).unwrap(); - fill_rows_mut(mat.rows_mut(), repr); - - check_mat(&mat, repr, ctx); - check_mat_ref(mat.reborrow(), repr, ctx); - check_mat_mut(mat.reborrow_mut(), repr, ctx); - check_rows(mat.rows(), repr, ctx); - } - } - } - } - - #[test] - fn test_mat_clone() { - for nrows in ROWS { - for ncols in COLS { - let repr = Standard::::new(*nrows, *ncols).unwrap(); - let ctx = &lazy_format!("nrows = {}, ncols = {}", nrows, ncols); - - let mut mat = Mat::new(repr, Defaulted).unwrap(); - fill_mat(&mut mat, repr); - - // Clone via Mat::clone - { - let ctx = &lazy_format!("{ctx} - Mat::clone"); - let cloned = mat.clone(); - - assert_eq!(cloned.num_vectors(), *nrows); - assert_eq!(cloned.vector_dim(), *ncols); - - check_mat(&cloned, repr, ctx); - check_mat_ref(cloned.reborrow(), repr, ctx); - check_rows(cloned.rows(), repr, ctx); - - // Cloned allocation is independent. - if repr.num_elements() > 0 { - assert_ne!(mat.as_raw_ptr(), cloned.as_raw_ptr()); - } - } - - // Clone via MatRef::to_owned - { - let ctx = &lazy_format!("{ctx} - MatRef::to_owned"); - let owned = mat.as_view().to_owned(); - - check_mat(&owned, repr, ctx); - check_mat_ref(owned.reborrow(), repr, ctx); - check_rows(owned.rows(), repr, ctx); - - if repr.num_elements() > 0 { - assert_ne!(mat.as_raw_ptr(), owned.as_raw_ptr()); - } - } - - // Clone via MatMut::to_owned - { - let ctx = &lazy_format!("{ctx} - MatMut::to_owned"); - let owned = mat.as_view_mut().to_owned(); - - check_mat(&owned, repr, ctx); - check_mat_ref(owned.reborrow(), repr, ctx); - check_rows(owned.rows(), repr, ctx); - - if repr.num_elements() > 0 { - assert_ne!(mat.as_raw_ptr(), owned.as_raw_ptr()); - } - } - } - } - } - - #[test] - fn test_mat_refmut() { - for nrows in ROWS { - for ncols in COLS { - let repr = Standard::::new(*nrows, *ncols).unwrap(); - let ctx = &lazy_format!("nrows = {}, ncols = {}", nrows, ncols); - - // Populate the matrix using `&mut Mat` - { - let ctx = &lazy_format!("{ctx} - by matmut"); - let mut b: Box<[_]> = (0..repr.num_elements()).map(|_| 0usize).collect(); - let ptr = b.as_ptr().cast::(); - let mut matmut = MatMut::new(repr, &mut b).unwrap(); - - assert_eq!( - ptr, - matmut.as_raw_ptr(), - "underlying memory should be preserved", - ); - - fill_mat_mut(matmut.reborrow_mut(), repr); - - check_mat_mut(matmut.reborrow_mut(), repr, ctx); - check_mat_ref(matmut.reborrow(), repr, ctx); - check_rows(matmut.rows(), repr, ctx); - check_rows(matmut.reborrow().rows(), repr, ctx); - - let matref = MatRef::new(repr, &b).unwrap(); - check_mat_ref(matref, repr, ctx); - check_mat_ref(matref.reborrow(), repr, ctx); - check_rows(matref.rows(), repr, ctx); - } - - // Populate the matrix using `RowsMut` - { - let ctx = &lazy_format!("{ctx} - by rows"); - let mut b: Box<[_]> = (0..repr.num_elements()).map(|_| 0usize).collect(); - let ptr = b.as_ptr().cast::(); - let mut matmut = MatMut::new(repr, &mut b).unwrap(); - - assert_eq!( - ptr, - matmut.as_raw_ptr(), - "underlying memory should be preserved", - ); - - fill_rows_mut(matmut.rows_mut(), repr); - - check_mat_mut(matmut.reborrow_mut(), repr, ctx); - check_mat_ref(matmut.reborrow(), repr, ctx); - check_rows(matmut.rows(), repr, ctx); - check_rows(matmut.reborrow().rows(), repr, ctx); - - let matref = MatRef::new(repr, &b).unwrap(); - check_mat_ref(matref, repr, ctx); - check_mat_ref(matref.reborrow(), repr, ctx); - check_rows(matref.rows(), repr, ctx); - } - } - } - } - - ////////////////// - // Constructors // - ////////////////// - - #[test] - fn test_standard_new_owned() { - let rows = [0, 1, 2, 3, 5, 10]; - let cols = [0, 1, 2, 3, 5, 10]; - - for nrows in rows { - for ncols in cols { - let m = Mat::new(Standard::new(nrows, ncols).unwrap(), 1usize).unwrap(); - let rows_iter = m.rows(); - let len = <_ as ExactSizeIterator>::len(&rows_iter); - assert_eq!(len, nrows); - for r in rows_iter { - assert_eq!(r.len(), ncols); - assert!(r.iter().all(|i| *i == 1usize)); - } - } - } - } - - #[test] - fn test_mat_from_fn() { - let rows = [0, 1, 2, 5]; - let cols = [0, 1, 3, 7]; - - for nrows in rows { - for ncols in cols { - let mut counter = 0u32; - let m = Mat::from_fn(Standard::new(nrows, ncols).unwrap(), || { - let v = counter; - counter += 1; - v - }); - - assert_eq!(counter as usize, nrows * ncols); - for (i, row) in m.rows().enumerate() { - assert_eq!(row.len(), ncols); - for (j, &v) in row.iter().enumerate() { - assert_eq!(v, (i * ncols + j) as u32); - } - } - } - } - } - - #[test] - fn matref_new_slice_length_error() { - let repr = Standard::::new(3, 4).unwrap(); - - // Correct length succeeds - let data = vec![0u32; 12]; - assert!(MatRef::new(repr, &data).is_ok()); - - // Too short fails - let short = vec![0u32; 11]; - assert!(matches!( - MatRef::new(repr, &short), - Err(SliceError::LengthMismatch { - expected: 12, - found: 11 - }) - )); - - // Too long fails - let long = vec![0u32; 13]; - assert!(matches!( - MatRef::new(repr, &long), - Err(SliceError::LengthMismatch { - expected: 12, - found: 13 - }) - )); - } - - #[test] - fn matmut_new_slice_length_error() { - let repr = Standard::::new(3, 4).unwrap(); - - // Correct length succeeds - let mut data = vec![0u32; 12]; - assert!(MatMut::new(repr, &mut data).is_ok()); - - // Too short fails - let mut short = vec![0u32; 11]; - assert!(matches!( - MatMut::new(repr, &mut short), - Err(SliceError::LengthMismatch { - expected: 12, - found: 11 - }) - )); - - // Too long fails - let mut long = vec![0u32; 13]; - assert!(matches!( - MatMut::new(repr, &mut long), - Err(SliceError::LengthMismatch { - expected: 12, - found: 13 - }) - )); - } - - #[test] - fn as_matrix_view_roundtrip() { - let data = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0]; - - // MatRef - let matref = MatRef::new(Standard::new(2, 3).unwrap(), &data).unwrap(); - let view = matref.as_matrix_view(); - assert_eq!(view.nrows(), 2); - assert_eq!(view.ncols(), 3); - for row in 0..2 { - for col in 0..3 { - assert_eq!(view[(row, col)], data[row * 3 + col]); - } - } - assert_eq!(matref.as_slice(), &data); - - // Mat - let mut mat = Mat::new(Standard::::new(2, 3).unwrap(), 0.0f32).unwrap(); - for i in 0..2 { - let r = mat.get_row_mut(i).unwrap(); - for j in 0..3 { - r[j] = data[i * 3 + j]; - } - } - let view = mat.as_matrix_view(); - assert_eq!(view.nrows(), 2); - assert_eq!(view.ncols(), 3); - for row in 0..2 { - for col in 0..3 { - assert_eq!(view[(row, col)], data[row * 3 + col]); - } - } - assert_eq!(mat.as_slice(), &data); - - // MatMut - let mut buf = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0]; - let matmut = MatMut::new(Standard::new(2, 3).unwrap(), &mut buf).unwrap(); - let view = matmut.as_matrix_view(); - assert_eq!(view.nrows(), 2); - assert_eq!(view.ncols(), 3); - for row in 0..2 { - for col in 0..3 { - assert_eq!(view[(row, col)], data[row * 3 + col]); - } - } - assert_eq!(matmut.as_slice(), &data); - } - - #[test] - fn test_standard_non_copy_element() { - let repr = Standard::::new(2, 3).unwrap(); - - // Owned fill via NewOwned (Clone). - let filled = Mat::new(repr, String::from("x")).unwrap(); - assert_eq!(filled.num_vectors(), 2); - assert!(filled.rows().flatten().all(|s| s == "x")); - - // NewOwned (Clone + Default). - let defaulted = Mat::new(repr, Defaulted).unwrap(); - assert!(defaulted.rows().flatten().all(String::is_empty)); - - // from_fn. - let mut counter = 0usize; - let mut mat = Mat::from_fn(repr, || { - let s = counter.to_string(); - counter += 1; - s - }); - assert_eq!(counter, 6); - assert_eq!(mat.get_row(1).unwrap()[0], "3"); - - // Mutation via get_row_mut. - mat.get_row_mut(0).unwrap()[0] = String::from("mutated"); - assert_eq!(mat.get_row(0).unwrap()[0], "mutated"); - - // Clone via NewCloned (Clone): independent allocation, equal contents. - let cloned = mat.clone(); - assert_ne!(mat.as_raw_ptr(), cloned.as_raw_ptr()); - assert_eq!(cloned.get_row(0).unwrap()[0], "mutated"); - - // Immutable view over a non-Copy slice (NewRef). - let data = [String::from("a"), String::from("b")]; - let view = MatRef::new(Standard::new(2, 1).unwrap(), &data).unwrap(); - assert_eq!(view.get_row(1).unwrap()[0], "b"); - - // Mutable view over a non-Copy slice (NewMut). - let mut data_mut = [String::from("a"), String::from("b")]; - let mut view_mut = MatMut::new(Standard::new(1, 2).unwrap(), &mut data_mut).unwrap(); - view_mut.get_row_mut(0).unwrap()[1] = String::from("z"); - assert_eq!(data_mut[1], "z"); - } -} diff --git a/diskann-quantization/src/multi_vector/mod.rs b/diskann-quantization/src/multi_vector/mod.rs index 0afe5180e..7c0b8989f 100644 --- a/diskann-quantization/src/multi_vector/mod.rs +++ b/diskann-quantization/src/multi_vector/mod.rs @@ -13,13 +13,13 @@ //! ``` //! use diskann_quantization::multi_vector::{ //! distance::QueryMatRef, -//! Chamfer, Mat, MatMut, MatRef, MaxSim, Standard, +//! Chamfer, Mat, MatMut, MatRef, MaxSim, RowMajor, //! }; //! use diskann_utils::ReborrowMut; //! use diskann_vector::{DistanceFunctionMut, PureDistanceFunction}; //! //! // Create an owned matrix (2 vectors, dim 3, initialized to 0.0) -//! let mut owned = Mat::new(Standard::new(2, 3).unwrap(), 0.0f32).unwrap(); +//! let mut owned = Mat::from_repr(RowMajor::new(2, 3).unwrap(), 0.0f32).unwrap(); //! assert_eq!(owned.num_vectors(), 2); //! //! // Modify via mutable view @@ -33,11 +33,11 @@ //! let doc_data = [1.0f32, 0.0, 0.0, 1.0]; //! //! // Wrap query as QueryMatRef for type-safe asymmetric distance -//! let query: QueryMatRef<_> = MatRef::new( -//! Standard::new(2, 2).unwrap(), +//! let query: QueryMatRef<_> = MatRef::from_repr( +//! RowMajor::new(2, 2).unwrap(), //! &query_data, //! ).unwrap().into(); -//! let doc = MatRef::new(Standard::new(2, 2).unwrap(), &doc_data).unwrap(); +//! let doc = MatRef::from_repr(RowMajor::new(2, 2).unwrap(), &doc_data).unwrap(); //! //! // Chamfer distance (sum of max similarities) //! let distance = Chamfer::evaluate(query, doc); @@ -53,14 +53,14 @@ pub mod block_transposed; pub mod distance; -pub(crate) mod matrix; +pub(crate) use diskann_utils::views; pub use block_transposed::{BlockTransposed, BlockTransposedMut, BlockTransposedRef}; +pub use diskann_utils::views::{ + LayoutError, Mat, MatMut, MatRef, NewCloned, NewDefault, NewMut, NewOwned, NewRef, Overflow, + Repr, ReprMut, ReprOwned, RowMajor, SliceError, +}; pub use distance::{ BoxErase, Chamfer, Erase, MaxSim, MaxSimElement, MaxSimError, MaxSimIsa, MaxSimKernel, NotSupported, ProjectedEigen, QueryMatRef, build_max_sim, }; -pub use matrix::{ - Defaulted, LayoutError, Mat, MatMut, MatRef, NewCloned, NewMut, NewOwned, NewRef, Overflow, - Repr, ReprMut, ReprOwned, SliceError, Standard, -}; diff --git a/diskann-quantization/src/product/tables/basic.rs b/diskann-quantization/src/product/tables/basic.rs index 9469fa58d..69d53a81b 100644 --- a/diskann-quantization/src/product/tables/basic.rs +++ b/diskann-quantization/src/product/tables/basic.rs @@ -5,7 +5,7 @@ use crate::traits::CompressInto; use crate::views::{ChunkOffsetsBase, ChunkOffsetsView}; -use diskann_utils::views::{DenseData, MatrixBase, MatrixView}; +use diskann_utils::views::{DenseData, Matrix, MatrixView}; use diskann_vector::{PureDistanceFunction, distance::SquaredL2}; use thiserror::Error; @@ -22,23 +22,40 @@ use thiserror::Error; /// /// # Invariants /// -/// * `offsets.dim() == pivots.nrows()`: The dimensionality of the two must agree. +/// * `offsets.dim() == pivots.ncols()`: The dimensionality of the two must agree. #[derive(Debug, Clone)] -pub struct BasicTableBase +pub struct BasicTableBase where - T: DenseData, U: DenseData, { - pivots: MatrixBase, + pivots: M, offsets: ChunkOffsetsBase, } /// A `BasicTableBase` that owns its contents. -pub type BasicTable = BasicTableBase, Box<[usize]>>; +pub type BasicTable = BasicTableBase, Box<[usize]>>; /// A `BasicTableBase` that references its contents. Construction of such a table will /// not result in a memory allocation. -pub type BasicTableView<'a> = BasicTableBase<&'a [f32], &'a [usize]>; +pub type BasicTableView<'a> = BasicTableBase, &'a [usize]>; + +/// Bridges owned ([`Matrix`]) and borrowed ([`MatrixView`]) pivot storage, exposing a +/// [`MatrixView`] so a [`BasicTableBase`] reads its dimensions and rows from one place. +pub trait PivotTable { + fn view(&self) -> MatrixView<'_, f32>; +} + +impl PivotTable for Matrix { + fn view(&self) -> MatrixView<'_, f32> { + self.as_matrix_view() + } +} + +impl PivotTable for MatrixView<'_, f32> { + fn view(&self) -> MatrixView<'_, f32> { + self.as_matrix_view() + } +} #[derive(Error, Debug)] #[non_exhaustive] @@ -52,9 +69,9 @@ pub enum BasicTableError { PivotsEmpty, } -impl BasicTableBase +impl BasicTableBase where - T: DenseData, + M: PivotTable, U: DenseData, { /// Construct a new `BasicTableBase` over the pivot table and offsets. @@ -62,11 +79,11 @@ where /// # Error /// /// Returns an error if `pivots.ncols() != offsets.dim()` or if `pivots.nrows() == 0`. - pub fn new( - pivots: MatrixBase, - offsets: ChunkOffsetsBase, - ) -> Result { - let pivot_dim = pivots.ncols(); + pub fn new(pivots: M, offsets: ChunkOffsetsBase) -> Result { + let (num_pivots, pivot_dim) = { + let view = pivots.view(); + (view.nrows(), view.ncols()) + }; let offsets_dim = offsets.dim(); if pivot_dim != offsets_dim { @@ -74,7 +91,7 @@ where pivot_dim, offsets_dim, }) - } else if pivots.nrows() == 0 { + } else if num_pivots == 0 { Err(BasicTableError::PivotsEmpty) } else { Ok(Self { pivots, offsets }) @@ -83,7 +100,7 @@ where /// Return a view over the pivot table. pub fn view_pivots(&self) -> MatrixView<'_, f32> { - self.pivots.as_view() + self.pivots.view() } /// Return a view over the schema offsets. @@ -93,7 +110,7 @@ where /// Return the number of pivots in each PQ chunk. pub fn ncenters(&self) -> usize { - self.pivots.nrows() + self.view_pivots().nrows() } /// Return the number of PQ chunks. @@ -103,7 +120,7 @@ where /// Return the dimensionality of the full-precision vectors associated with this table. pub fn dim(&self) -> usize { - self.pivots.ncols() + self.view_pivots().ncols() } } @@ -120,9 +137,9 @@ pub enum TableCompressionError { InfinityOrNaN(usize), } -impl CompressInto<&[f32], &mut [u8]> for BasicTableBase +impl CompressInto<&[f32], &mut [u8]> for BasicTableBase where - T: DenseData, + M: PivotTable, U: DenseData, { type Error = TableCompressionError; @@ -169,13 +186,14 @@ where return Err(Self::Error::InvalidOutputDim(self.nchunks(), to.len())); } + let pivots = self.view_pivots(); to.iter_mut().enumerate().try_for_each(|(chunk, to)| { let mut min_distance = f32::INFINITY; let mut min_index = usize::MAX; let range = self.offsets.at(chunk); let slice = &from[range.clone()]; - self.pivots.row_iter().enumerate().for_each(|(index, row)| { + pivots.row_iter().enumerate().for_each(|(index, row)| { let distance: f32 = SquaredL2::evaluate(slice, &row[range.clone()]); if distance < min_distance { min_distance = distance; diff --git a/diskann-quantization/src/product/tables/test.rs b/diskann-quantization/src/product/tables/test.rs index c8faec00c..32261c2dc 100644 --- a/diskann-quantization/src/product/tables/test.rs +++ b/diskann-quantization/src/product/tables/test.rs @@ -6,7 +6,7 @@ // A collection of test helpers to ensure uniformity across tables. use diskann_utils::views::Matrix; #[cfg(not(miri))] -use diskann_utils::views::{MatrixView, MutMatrixView}; +use diskann_utils::views::{MatrixView, MatrixViewMut}; #[cfg(not(miri))] use rand::seq::IndexedRandom; use rand::{ @@ -298,7 +298,7 @@ pub(super) fn check_pqtable_batch_compression_errors( build: &dyn Fn(Matrix, ChunkOffsets) -> T, context: &dyn std::fmt::Display, ) where - T: for<'a> CompressInto, MutMatrixView<'a, u8>>, + T: for<'a> CompressInto, MatrixViewMut<'a, u8>>, { let dim = 10; let num_chunks = 3; @@ -419,7 +419,7 @@ pub(super) fn check_pqtable_batch_compression_errors( let mut buf = Matrix::::new(0.0, num_points, offsets.dim()); let mut output = Matrix::::new(0, num_points, offsets.len()); - fn clear(mut x: MutMatrixView) { + fn clear(mut x: MatrixViewMut) { x.as_mut_slice().iter_mut().for_each(|i| *i = T::default()); } diff --git a/diskann-quantization/src/product/tables/transposed/table.rs b/diskann-quantization/src/product/tables/transposed/table.rs index bc527fd8a..ea5567b82 100644 --- a/diskann-quantization/src/product/tables/transposed/table.rs +++ b/diskann-quantization/src/product/tables/transposed/table.rs @@ -13,7 +13,7 @@ use crate::{ }; use diskann_utils::{ strided, - views::{self, MatrixView, MutMatrixView}, + views::{self, MatrixView, MatrixViewMut}, }; use thiserror::Error; @@ -187,7 +187,7 @@ impl TransposedTable { let range = self.offsets.at(i); if let Some(chunk_dim) = NonZeroUsize::new(range.len()) { // Construct a view for the packing buffer for this chunk. - let mut packing_view = views::MutMatrixView::try_from( + let mut packing_view = views::MatrixViewMut::try_from( &mut packing_buffer[..SUB_BATCH_SIZE * chunk_dim.get()], SUB_BATCH_SIZE, chunk_dim.get(), @@ -282,7 +282,7 @@ impl TransposedTable { /// * `query.len() != self.dim()`. /// * `partisl.nrows() != self.nchunks()`. /// * `partisl.ncols() != self.ncenters()`. - pub fn process_into(&self, query: &[f32], mut partials: MutMatrixView<'_, f32>) + pub fn process_into(&self, query: &[f32], mut partials: MatrixViewMut<'_, f32>) where T: pivots::ProcessInto, { @@ -422,7 +422,7 @@ pub enum TableBatchCompressionError { InfinityOrNaN(usize, usize), } -impl CompressInto, MutMatrixView<'_, u8>> for TransposedTable +impl CompressInto, MatrixViewMut<'_, u8>> for TransposedTable where T: Copy + Into, { @@ -465,7 +465,7 @@ where fn compress_into( &self, from: MatrixView<'_, T>, - mut to: MutMatrixView<'_, u8>, + mut to: MatrixViewMut<'_, u8>, ) -> Result<(), Self::Error> { if self.ncenters() > 256 { return Err(Self::Error::CannotCompressToByte(self.ncenters())); diff --git a/diskann-quantization/src/utils.rs b/diskann-quantization/src/utils.rs index 2e30ea3ae..baf5286a3 100644 --- a/diskann-quantization/src/utils.rs +++ b/diskann-quantization/src/utils.rs @@ -41,7 +41,7 @@ pub(crate) fn as_nonnull_mut(slice: &mut [T]) -> NonNull { /// /// This is the owned-allocation counterpart of [`as_nonnull`]: it consumes the /// box, preventing its destructor from running, and hands back a non-null -/// pointer suitable for storing inside a [`Mat`](super::multi_vector::matrix::Mat). +/// pointer suitable for storing inside a [`Mat`](super::multi_vector::views::Mat). /// /// To reclaim the memory later, reconstruct the `Box` via /// `Box::from_raw(std::ptr::slice_from_raw_parts_mut(ptr.as_ptr(), len))`. diff --git a/diskann-utils/Cargo.toml b/diskann-utils/Cargo.toml index ea358601e..283849c2a 100644 --- a/diskann-utils/Cargo.toml +++ b/diskann-utils/Cargo.toml @@ -32,6 +32,7 @@ workspace = true cfg-if.workspace = true rand.workspace = true rstest.workspace = true +trybuild = "1.0.101" [features] diff --git a/diskann-utils/src/lib.rs b/diskann-utils/src/lib.rs index 089e0dece..1f94e7adb 100644 --- a/diskann-utils/src/lib.rs +++ b/diskann-utils/src/lib.rs @@ -18,8 +18,8 @@ pub mod io; pub mod object_pool; pub mod sampling; -// Views pub mod strided; + pub mod views; mod lazystring; diff --git a/diskann-utils/src/strided.rs b/diskann-utils/src/strided.rs index efbfd621e..30a9f89ab 100644 --- a/diskann-utils/src/strided.rs +++ b/diskann-utils/src/strided.rs @@ -9,11 +9,11 @@ use std::{ }; use thiserror::Error; -use crate::views::{self, DenseData, MutDenseData}; +use crate::views::{DenseData, MatMut, MatRef, MutDenseData, RowMajor}; /// A row-major strided matrix. /// -/// This is a generalization of the `MatrixBase` class as it does not mandate a dense +/// This is a generalization of the dense [`Matrix`](crate::views::Matrix) type as it does not mandate a dense /// layout in memory. /// /// ```text @@ -75,7 +75,7 @@ pub struct TryFromErrorLight { data.as_slice().len(), linear_length(self.nrows, self.ncols, self.cstride) )] -pub struct TryFromError { +pub struct TryFromError { data: T, nrows: usize, ncols: usize, @@ -94,7 +94,7 @@ impl fmt::Debug for TryFromError { } } -impl TryFromError { +impl TryFromError { /// Consume the error and return the base data. pub fn into_inner(self) -> T { self.data @@ -497,17 +497,23 @@ where } } -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> From>> for StridedBase<&'a [T]> { + fn from(m: MatRef<'a, RowMajor>) -> Self { + let (nrows, ncols) = (m.nrows(), m.ncols()); + Self { + data: m.into(), + nrows, + ncols, + cstride: ncols, + } + } +} + +impl<'a, T> From>> for StridedBase<&'a mut [T]> { + fn from(m: MatMut<'a, RowMajor>) -> Self { + let (nrows, ncols) = (m.nrows(), m.ncols()); Self { - data: matrix.into_inner().into(), + data: m.into(), nrows, ncols, cstride: ncols, @@ -518,6 +524,7 @@ where #[cfg(test)] mod tests { use super::*; + use crate::views; #[test] fn test_linear_length() { diff --git a/diskann-utils/src/views.rs b/diskann-utils/src/views.rs index a9352918c..3ab1b47aa 100644 --- a/diskann-utils/src/views.rs +++ b/diskann-utils/src/views.rs @@ -3,139 +3,1226 @@ * Licensed under the MIT license. */ -use std::{ - fmt, - ops::{Index, IndexMut}, -}; +//! Row-major matrix types for multi-vector representations. +//! +//! This module provides flexible matrix abstractions that support different underlying +//! storage formats through the [`Repr`] trait. The primary types are: +//! +//! - [`Mat`]: An owning matrix that manages its own memory. +//! - [`MatRef`]: An immutable borrowed view of matrix data. +//! - [`MatMut`]: A mutable borrowed view of matrix data. +//! +//! # Representations +//! +//! Representation types interact with the [`Mat`] family of types using the following traits: +//! +//! - [`Repr`]: Read-only matrix representation. +//! - [`ReprMut`]: Mutable matrix representation. +//! - [`ReprOwned`]: Owning matrix representation. +//! +//! Each trait refinement has a corresponding constructor: +//! +//! - [`NewRef`]: Construct a read-only [`MatRef`] view over a slice. +//! - [`NewMut`]: Construct a mutable [`MatMut`] matrix view over a slice. +//! - [`NewOwned`]: Construct a new owning [`Mat`]. +//! + +use std::{alloc::Layout, iter::FusedIterator, marker::PhantomData, ptr::NonNull}; + +use crate::{Reborrow, ReborrowMut}; +use thiserror::Error; + +#[cfg(feature = "rayon")] +use rayon::iter::ParallelIterator; + +// Pointer helpers, kept private so the matrix framework is self-contained in diskann-utils. +fn as_nonnull(slice: &[T]) -> NonNull { + // SAFETY: Slices guarantee non-null pointers. + unsafe { NonNull::new_unchecked(slice.as_ptr().cast_mut()) } +} + +fn as_nonnull_mut(slice: &mut [T]) -> NonNull { + // SAFETY: Slices guarantee non-null pointers. + unsafe { NonNull::new_unchecked(slice.as_mut_ptr()) } +} + +fn box_into_nonnull(b: Box<[T]>) -> NonNull { + // SAFETY: `Box::into_raw` guarantees the returned pointer is non-null. + unsafe { NonNull::new_unchecked(Box::into_raw(b).cast::()) } +} + +/// Representation trait describing the layout and access patterns for a matrix. +/// +/// Implementations define how raw bytes are interpreted as typed rows. This enables +/// matrices over different storage formats (dense, quantized, etc.) using a single +/// generic [`Mat`] type. +/// +/// # Associated Types +/// +/// - `Row<'a>`: The immutable row type (e.g., `&[f32]`, `&[f16]`). +/// +/// # Safety +/// +/// Implementations must ensure: +/// +/// - [`get_row`](Self::get_row) returns valid references for the given row index. +/// This call **must** be memory safe for `i < self.nrows()`, provided the caller upholds +/// the contract for the raw pointer. +/// +/// - The objects implicitly managed by this representation inherit the `Send` and `Sync` +/// attributes of `Repr`. That is, `Repr: Send` implies that the objects in backing memory +/// are [`Send`], and likewise with `Sync`. This is necessary to apply [`Send`] and [`Sync`] +/// bounds to [`Mat`], [`MatRef`], and [`MatMut`]. +pub unsafe trait Repr: Copy { + /// Immutable row reference type. + type Row<'a> + where + Self: 'a; + + /// Returns the number of rows in the matrix. + /// + /// # Safety Contract + /// + /// This function must be loosely pure in the sense that for any given instance of + /// `self`, `self.nrows()` must return the same value. + fn nrows(&self) -> usize; + + /// Returns the memory layout for a memory allocation containing [`Repr::nrows`] vectors + /// each with vector dimension [`Repr::ncols`]. + /// + /// # Safety Contract + /// + /// The [`Layout`] returned from this method must be consistent with the contract of + /// [`Repr::get_row`]. + fn layout(&self) -> Result; + + /// Returns an immutable reference to the `i`-th row. + /// + /// # Safety + /// + /// - `ptr` must point to a slice with a layout compatible with [`Repr::layout`]. + /// - The entire range for this slice must be within a single allocation. + /// - `i` must be less than [`Repr::nrows`]. + /// - The memory referenced by the returned [`Repr::Row`] must not be mutated for the + /// duration of lifetime `'a`. + /// - The lifetime for the returned [`Repr::Row`] is inferred from its usage. Correct + /// usage must properly tie the lifetime to a source. + unsafe fn get_row<'a>(self, ptr: NonNull, i: usize) -> Self::Row<'a>; +} + +/// Extension of [`Repr`] that supports mutable row access. +/// +/// # Associated Types +/// +/// - `RowMut<'a>`: The mutable row type (e.g., `&mut [f32]`). +/// +/// # Safety +/// +/// Implementors must ensure: +/// +/// - [`get_row_mut`](Self::get_row_mut) returns valid references for the given row index. +/// This call **must** be memory safe for `i < self.nrows()`, provided the caller upholds +/// the contract for the raw pointer. +/// +/// Additionally, since the implementation of the [`RowsMut`] iterator can give out rows +/// for all `i` in `0..self.nrows()`, the implementation of [`Self::get_row_mut`] must be +/// such that the result for disjoint `i` must not interfere with one another. +pub unsafe trait ReprMut: Repr { + /// Mutable row reference type. + type RowMut<'a> + where + Self: 'a; + + /// Returns a mutable reference to the i-th row. + /// + /// # Safety + /// - `ptr` must point to a slice with a layout compatible with [`Repr::layout`]. + /// - The entire range for this slice must be within a single allocation. + /// - `i` must be less than `self.nrows()`. + /// - The memory referenced by the returned [`ReprMut::RowMut`] must not be accessed + /// through any other reference for the duration of lifetime `'a`. + /// - The lifetime for the returned [`ReprMut::RowMut`] is inferred from its usage. + /// Correct usage must properly tie the lifetime to a source. + unsafe fn get_row_mut<'a>(self, ptr: NonNull, i: usize) -> Self::RowMut<'a>; +} + +/// Extension trait for [`Repr`] that supports deallocation of owned matrices. This is used +/// in conjunction with [`NewOwned`] to create matrices. +/// +/// Requires [`ReprMut`] since owned matrices should support mutation. +/// +/// # Safety +/// +/// Implementors must ensure that `drop` properly deallocates the memory in a way compatible +/// with all [`NewOwned`] implementations. +pub unsafe trait ReprOwned: ReprMut { + /// Deallocates memory at `ptr` and drops `self`. + /// + /// # Safety + /// + /// - `ptr` must have been obtained via [`NewOwned`] with the same value of `self`. + /// - This method may only be called once for such a pointer. + /// - After calling this method, the memory behind `ptr` may not be dereferenced at all. + unsafe fn drop(self, ptr: NonNull); +} + +/// A new-type version of `std::alloc::LayoutError` for cleaner error handling. +/// +/// This is basically the same as [`std::alloc::LayoutError`], but constructible in +/// use code to allow implementors of [`Repr::layout`] to return it for reasons other than +/// those derived from `std::alloc::Layout`'s methods. +#[derive(Debug, Clone, Copy)] +#[non_exhaustive] +pub struct LayoutError; + +impl LayoutError { + /// Construct a new opaque [`LayoutError`]. + pub fn new() -> Self { + Self + } +} + +impl Default for LayoutError { + fn default() -> Self { + Self::new() + } +} + +impl std::fmt::Display for LayoutError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "LayoutError") + } +} + +impl std::error::Error for LayoutError {} + +impl From for LayoutError { + fn from(_: std::alloc::LayoutError) -> Self { + LayoutError + } +} + +////////////////// +// Constructors // +////////////////// + +/// Create a new [`MatRef`] over a slice. +/// +/// # Safety +/// +/// Implementations must validate the length (and any other requirements) of the provided +/// slice to ensure it is compatible with the implementation of [`Repr`]. +pub unsafe trait NewRef: Repr { + /// Errors that can occur when initializing. + type Error; + + /// Create a new [`MatRef`] over `slice`. + fn new_ref(self, slice: &[T]) -> Result, Self::Error>; +} + +/// Create a new [`MatMut`] over a slice. +/// +/// # Safety +/// +/// Implementations must validate the length (and any other requirements) of the provided +/// slice to ensure it is compatible with the implementation of [`ReprMut`]. +pub unsafe trait NewMut: ReprMut { + /// Errors that can occur when initializing. + type Error; + + /// Create a new [`MatMut`] over `slice`. + fn new_mut(self, slice: &mut [T]) -> Result, Self::Error>; +} + +/// Create a new [`Mat`] from an initializer. +/// +/// # Safety +/// +/// Implementations must ensure that the returned [`Mat`] is compatible with +/// `Self`'s implementation of [`ReprOwned`]. +pub unsafe trait NewOwned: ReprOwned { + /// Errors that can occur when initializing. + type Error; + + /// Create a new [`Mat`] initialized with `init`. + fn new_owned(self, init: T) -> Result, Self::Error>; +} + +/// Create a new [`Mat`] with every element default-initialized. +/// +/// This is a distinct trait from [`NewOwned`] (rather than a `NewOwned`) so that a repr +/// can offer both a value-fill and a default-fill constructor without the two impls overlapping +/// under coherence. +/// +/// # Safety +/// +/// Implementations must ensure that the returned [`Mat`] is compatible with +/// `Self`'s implementation of [`ReprOwned`]. +pub unsafe trait NewDefault: ReprOwned { + /// Errors that can occur when initializing. + type Error; + + /// Create a new [`Mat`] with each element set to its [`Default`]. + fn new_default(self) -> Result, Self::Error>; +} + +/// Create a new [`Mat`] cloned from a view. +pub trait NewCloned: ReprOwned { + /// Clone the contents behind `v`, returning a new owning [`Mat`]. + /// + /// Implementations should ensure the returned [`Mat`] is "semantically the same" as `v`. + fn new_cloned(v: MatRef<'_, Self>) -> Mat; +} + +////////////// +// RowMajor // +////////////// + +/// Metadata for dense row-major matrices. +/// +/// Rows are stored contiguously as `&[T]` slices. This is the default representation +/// type for standard floating-point multi-vectors. +/// +/// # Row Types +/// +/// - `Row<'a>`: `&'a [T]` +/// - `RowMut<'a>`: `&'a mut [T]` +#[derive(Debug)] +pub struct RowMajor { + nrows: usize, + ncols: usize, + _elem: PhantomData, +} + +// Hand-written so `RowMajor` is `Copy`/`Clone`/`PartialEq`/`Eq` for every `T`: it only +// stores two `usize` and a `PhantomData`, so derives would spuriously require the same +// bound on `T` (and the `Repr: Copy` supertrait must hold regardless of the element type). +impl Copy for RowMajor {} + +impl Clone for RowMajor { + fn clone(&self) -> Self { + *self + } +} + +impl PartialEq for RowMajor { + fn eq(&self, other: &Self) -> bool { + self.nrows == other.nrows && self.ncols == other.ncols + } +} + +impl Eq for RowMajor {} + +impl RowMajor { + /// Create a new `RowMajor` for data of type `T`. + /// + /// Successful construction requires: + /// + /// * The total number of elements determined by `nrows * ncols` does not exceed + /// `usize::MAX`. + /// * The total memory footprint defined by `ncols * nrows * size_of::()` does not + /// exceed `isize::MAX`. + pub fn new(nrows: usize, ncols: usize) -> Result { + Overflow::check::(nrows, ncols)?; + Ok(Self { + nrows, + ncols, + _elem: PhantomData, + }) + } + + /// Returns the number of total elements (`rows x cols`) in this matrix. + pub fn num_elements(&self) -> usize { + // Since we've constructed `self` - we know we cannot overflow. + self.nrows() * self.ncols() + } + + /// Returns `rows`, the number of rows in this matrix. + fn nrows(&self) -> usize { + self.nrows + } + + /// Returns `ncols`, the number of elements in a row of this matrix. + fn ncols(&self) -> usize { + self.ncols + } + + /// Checks the following: + /// + /// 1. Computation of the number of elements in `self` does not overflow. + /// 2. Argument `slice` has the expected number of elements. + fn check_slice(&self, slice: &[T]) -> Result<(), SliceError> { + let len = self.num_elements(); + + if slice.len() != len { + Err(SliceError::LengthMismatch { + expected: len, + found: slice.len(), + }) + } else { + Ok(()) + } + } + + /// Create a new [`Mat`] around the contents of `b` **without** any checks. + /// + /// # Safety + /// + /// The length of `b` must be exactly [`RowMajor::num_elements`]. + unsafe fn box_to_mat(self, b: Box<[T]>) -> Mat { + debug_assert_eq!(b.len(), self.num_elements(), "safety contract violated"); + + let ptr = box_into_nonnull(b).cast::(); + + // SAFETY: `ptr` is properly aligned and points to a slice of the required length. + // Additionally, it is dropped via `Box::from_raw`, which is compatible with obtaining + // it from `Box::into_raw`. + unsafe { Mat::from_raw_parts(self, ptr) } + } +} + +/// Error for [`RowMajor::new`]. +#[derive(Debug, Clone, Copy)] +pub struct Overflow { + nrows: usize, + ncols: usize, + elsize: usize, +} + +impl Overflow { + /// Construct an `Overflow` error for the given dimensions and element type. + pub fn for_type(nrows: usize, ncols: usize) -> Self { + Self { + nrows, + ncols, + elsize: std::mem::size_of::(), + } + } + + /// Verify that `capacity` elements of type `T` fit within the `isize::MAX` byte + /// budget required by Rust's allocation APIs. + /// + /// On failure the error reports the original `(nrows, ncols)` dimensions rather + /// than the padded capacity. + pub fn check_byte_budget(capacity: usize, nrows: usize, ncols: usize) -> Result<(), Self> { + let bytes = std::mem::size_of::().saturating_mul(capacity); + if bytes <= isize::MAX as usize { + Ok(()) + } else { + Err(Self::for_type::(nrows, ncols)) + } + } + + pub(crate) fn check(nrows: usize, ncols: usize) -> Result<(), Self> { + // Guard the element count itself so that `num_elements()` can never overflow. + let capacity = nrows + .checked_mul(ncols) + .ok_or_else(|| Self::for_type::(nrows, ncols))?; + + Self::check_byte_budget::(capacity, nrows, ncols) + } +} + +impl std::fmt::Display for Overflow { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + if self.elsize == 0 { + write!( + f, + "ZST matrix with dimensions {} x {} has more than `usize::MAX` elements", + self.nrows, self.ncols, + ) + } else { + write!( + f, + "a matrix of size {} x {} with element size {} would exceed isize::MAX bytes", + self.nrows, self.ncols, self.elsize, + ) + } + } +} + +impl std::error::Error for Overflow {} + +/// Error types for [`RowMajor`]. +#[derive(Debug, Clone, Copy, Error)] +#[non_exhaustive] +pub enum SliceError { + #[error("Length mismatch: expected {expected}, found {found}")] + LengthMismatch { expected: usize, found: usize }, +} + +// SAFETY: The implementation correctly computes row offsets as `i * ncols` and +// constructs valid slices of the appropriate length. The `layout` method correctly +// reports the memory layout requirements. +unsafe impl Repr for RowMajor { + type Row<'a> + = &'a [T] + where + T: 'a; + + fn nrows(&self) -> usize { + self.nrows + } + + fn layout(&self) -> Result { + Ok(Layout::array::(self.num_elements())?) + } + + unsafe fn get_row<'a>(self, ptr: NonNull, i: usize) -> Self::Row<'a> { + debug_assert!(ptr.cast::().is_aligned()); + debug_assert!(i < self.nrows); + + // SAFETY: The caller asserts that `i` is less than `self.nrows()`. Since this type + // audits the constructors for `Mat` and friends, we know that there is room for at + // least `self.num_elements()` elements from the base pointer, so this access is safe. + let row_ptr = unsafe { ptr.as_ptr().cast::().add(i * self.ncols) }; + + // SAFETY: The logic is the same as the previous `unsafe` block. + unsafe { std::slice::from_raw_parts(row_ptr, self.ncols) } + } +} + +// SAFETY: The implementation correctly computes row offsets and constructs valid mutable +// slices. +unsafe impl ReprMut for RowMajor { + type RowMut<'a> + = &'a mut [T] + where + T: 'a; + + unsafe fn get_row_mut<'a>(self, ptr: NonNull, i: usize) -> Self::RowMut<'a> { + debug_assert!(ptr.cast::().is_aligned()); + debug_assert!(i < self.nrows); + + // SAFETY: The caller asserts that `i` is less than `self.nrows()`. Since this type + // audits the constructors for `Mat` and friends, we know that there is room for at + // least `self.num_elements()` elements from the base pointer, so this access is safe. + let row_ptr = unsafe { ptr.as_ptr().cast::().add(i * self.ncols) }; + + // SAFETY: The logic is the same as the previous `unsafe` block. Further, the caller + // attests that creating a mutable reference is safe. + unsafe { std::slice::from_raw_parts_mut(row_ptr, self.ncols) } + } +} + +// SAFETY: The drop implementation correctly reconstructs a Box from the raw pointer +// using the same length (nrows * ncols) that was used for allocation, allowing Box +// to properly deallocate the memory. +unsafe impl ReprOwned for RowMajor { + unsafe fn drop(self, ptr: NonNull) { + // SAFETY: The caller guarantees that `ptr` was obtained from an implementation of + // `NewOwned` for an equivalent instance of `self`. + // + // We ensure that `NewOwned` goes through boxes, so here we reconstruct a Box to + // let it handle deallocation. + unsafe { + let slice_ptr = std::ptr::slice_from_raw_parts_mut( + ptr.cast::().as_ptr(), + self.nrows * self.ncols, + ); + let _ = Box::from_raw(slice_ptr); + } + } +} + +// SAFETY: The implementation uses guarantees from `Box` to ensure that the pointer +// initialized by it is non-null and properly aligned to the underlying type. +unsafe impl NewOwned for RowMajor +where + T: Clone, +{ + type Error = std::convert::Infallible; + fn new_owned(self, value: T) -> Result, Self::Error> { + let b: Box<[T]> = std::iter::repeat_n(value, self.num_elements()).collect(); + + // SAFETY: By construction, `b` has length `self.num_elements()`. + Ok(unsafe { self.box_to_mat(b) }) + } +} + +// SAFETY: The implementation uses guarantees from `Box` to ensure that the pointer +// initialized by it is non-null and properly aligned to the underlying type. +unsafe impl NewDefault for RowMajor +where + T: Default, +{ + type Error = std::convert::Infallible; + fn new_default(self) -> Result, Self::Error> { + let b: Box<[T]> = std::iter::repeat_with(T::default) + .take(self.num_elements()) + .collect(); + + // SAFETY: By construction, `b` has length `self.num_elements()`. + Ok(unsafe { self.box_to_mat(b) }) + } +} + +// SAFETY: This checks that the slice has the correct length, which is all that is +// required for [`Repr`]. +unsafe impl NewRef for RowMajor { + type Error = SliceError; + fn new_ref(self, data: &[T]) -> Result, Self::Error> { + self.check_slice(data)?; + + // SAFETY: The function `check_slice` verifies that `data` is compatible with + // the layout requirement of `RowMajor`. + // + // We've properly checked that the underlying pointer is okay. + Ok(unsafe { MatRef::from_raw_parts(self, as_nonnull(data).cast::()) }) + } +} + +// SAFETY: This checks that the slice has the correct length, which is all that is +// required for [`ReprMut`]. +unsafe impl NewMut for RowMajor { + type Error = SliceError; + fn new_mut(self, data: &mut [T]) -> Result, Self::Error> { + self.check_slice(data)?; + + // SAFETY: The function `check_slice` verifies that `data` is compatible with + // the layout requirement of `RowMajor`. + // + // We've properly checked that the underlying pointer is okay. + Ok(unsafe { MatMut::from_raw_parts(self, as_nonnull_mut(data).cast::()) }) + } +} + +impl NewCloned for RowMajor +where + T: Clone, +{ + fn new_cloned(v: MatRef<'_, Self>) -> Mat { + let b: Box<[T]> = v.as_slice().iter().cloned().collect(); + + // SAFETY: By construction, `b` has length `v.repr().num_elements()`. + unsafe { v.repr().box_to_mat(b) } + } +} + +///////// +// Mat // +///////// + +/// An owning matrix that manages its own memory. +/// +/// The matrix stores raw bytes interpreted according to representation type `T`. +/// Memory is automatically deallocated when the matrix is dropped. +#[derive(Debug)] +pub struct Mat { + ptr: NonNull, + repr: T, + _invariant: PhantomData T>, +} + +// SAFETY: [`Repr`] is required to propagate its `Send` bound. +unsafe impl Send for Mat where T: ReprOwned + Send {} + +// SAFETY: [`Repr`] is required to propagate its `Sync` bound. +unsafe impl Sync for Mat where T: ReprOwned + Sync {} + +impl Mat { + /// Create a new matrix using `init` as the initializer. + pub fn from_repr(repr: T, init: U) -> Result>::Error> + where + T: NewOwned, + { + repr.new_owned(init) + } + + /// Create a new matrix with every element default-initialized. + /// + /// ```rust + /// use diskann_utils::views::{Mat, RowMajor}; + /// let mat = Mat::from_default(RowMajor::::new(4, 3).unwrap()).unwrap(); + /// for i in 0..4 { + /// assert!(mat.get_row(i).unwrap().iter().all(|&x| x == 0.0f32)); + /// } + /// ``` + pub fn from_default(repr: T) -> Result::Error> + where + T: NewDefault, + { + repr.new_default() + } + + /// Returns the number of rows (vectors) in the matrix. + #[inline] + pub fn num_vectors(&self) -> usize { + self.repr.nrows() + } + + /// Returns a reference to the underlying representation. + pub fn repr(&self) -> &T { + &self.repr + } + + /// Returns the `i`th row if `i < self.num_vectors()`. + #[must_use] + pub fn get_row(&self, i: usize) -> Option> { + if i < self.num_vectors() { + // SAFETY: Bounds check passed, and the Mat was constructed + // with valid representation and pointer. + let row = unsafe { self.get_row_unchecked(i) }; + Some(row) + } else { + None + } + } + + /// Returns the `i`th row without bounds checking. + /// + /// # Safety + /// + /// `i` must be less than `self.num_vectors()`. + pub unsafe fn get_row_unchecked(&self, i: usize) -> T::Row<'_> { + // SAFETY: Caller must ensure i < self.num_vectors(). The constructors for this type + // ensure that `ptr` is compatible with `T`. + unsafe { self.repr.get_row(self.ptr, i) } + } + + /// Returns the `i`th mutable row if `i < self.num_vectors()`. + #[must_use] + pub fn get_row_mut(&mut self, i: usize) -> Option> { + if i < self.num_vectors() { + // SAFETY: Bounds check passed, and we have exclusive access via &mut self. + Some(unsafe { self.get_row_mut_unchecked(i) }) + } else { + None + } + } + + pub(crate) unsafe fn get_row_mut_unchecked(&mut self, i: usize) -> T::RowMut<'_> { + // SAFETY: Caller asserts that `i < self.num_vectors()`. The constructors for this + // type ensure that `ptr` is compatible with `T`. + unsafe { self.repr.get_row_mut(self.ptr, i) } + } + + /// Returns an immutable view of the matrix. + #[inline] + pub fn as_view(&self) -> MatRef<'_, T> { + MatRef { + ptr: self.ptr, + repr: self.repr, + _lifetime: PhantomData, + } + } + + /// Returns a mutable view of the matrix. + #[inline] + pub fn as_view_mut(&mut self) -> MatMut<'_, T> { + MatMut { + ptr: self.ptr, + repr: self.repr, + _lifetime: PhantomData, + } + } + + /// Returns an iterator over immutable row references. + pub fn rows(&self) -> Rows<'_, T> { + Rows::new(self.reborrow()) + } + + /// Returns an iterator over mutable row references. + pub fn rows_mut(&mut self) -> RowsMut<'_, T> { + RowsMut::new(self.reborrow_mut()) + } + + /// Construct a new [`Mat`] over the raw pointer and representation without performing + /// any validity checks. + /// + /// # Safety + /// + /// Argument `ptr` must be: + /// + /// 1. Point to memory compatible with [`Repr::layout`]. + /// 2. Be compatible with the drop logic in [`ReprOwned`]. + pub unsafe fn from_raw_parts(repr: T, ptr: NonNull) -> Self { + Self { + ptr, + repr, + _invariant: PhantomData, + } + } + + /// Return the base pointer for the [`Mat`]. + pub fn as_raw_ptr(&self) -> *const u8 { + self.ptr.as_ptr() + } + + /// Return a mutable base pointer for the [`Mat`]. + pub fn as_raw_mut_ptr(&mut self) -> *mut u8 { + self.ptr.as_ptr() + } +} + +impl Drop for Mat { + fn drop(&mut self) { + // SAFETY: `ptr` was correctly initialized according to `layout` + // and we are guaranteed exclusive access to the data due to Rust borrow rules. + unsafe { self.repr.drop(self.ptr) }; + } +} + +impl Clone for Mat { + fn clone(&self) -> Self { + T::new_cloned(self.as_view()) + } +} + +// Delegation macro for the dense (`RowMajor`) read API. The canonical read +// implementations live once on `MatRef<'_, RowMajor>` (`MatrixView`); the +// owning `Mat` and the mutable view forward to them through `self.as_view()`, so +// each method has a single body. Mirrors the `delegate_to_ref!` pattern used by +// `diskann-quantization`'s block-transposed layout. +macro_rules! delegate_read { + ($(#[$m:meta])* $vis:vis fn $name:ident(&self $(, $a:ident: $t:ty)*) $(-> $r:ty)?) => { + #[doc = "Delegates to the canonical immutable-view implementation."] + $(#[$m])* + #[inline] + $vis fn $name(&self $(, $a: $t)*) $(-> $r)? { + self.as_view().$name($($a),*) + } + }; + ($(#[$m:meta])* unsafe $vis:vis fn $name:ident(&self $(, $a:ident: $t:ty)*) $(-> $r:ty)?) => { + #[doc = "Delegates to the canonical immutable-view implementation."] + $(#[$m])* + #[inline] + $vis unsafe fn $name(&self $(, $a: $t)*) $(-> $r)? { + // SAFETY: the caller upholds the delegated method's safety contract. + unsafe { self.as_view().$name($($a),*) } + } + }; +} + +impl Mat> { + /// Construct a [`Mat`] by filling each element in row-major order from `gen`. + pub fn new>(mut gen: U, nrows: usize, ncols: usize) -> Self { + let repr = RowMajor::new(nrows, ncols).unwrap_or_else(|e| panic!("{e}")); + let b: Box<[T]> = (0..repr.num_elements()).map(|_| gen.generate()).collect(); + // SAFETY: `b` has length `repr.num_elements()` by construction. + unsafe { repr.box_to_mat(b) } + } + + /// Returns the raw dimension (columns) of the vectors in the matrix. + #[inline] + pub fn vector_dim(&self) -> usize { + self.repr.ncols() + } + + delegate_read!(pub fn as_slice(&self) -> &[T]); + delegate_read!(pub fn as_matrix_view(&self) -> MatrixView<'_, T>); +} + +//////////// +// MatRef // +//////////// + +/// An immutable borrowed view of a matrix. +/// +/// Provides read-only access to matrix data without ownership. Implements [`Copy`] +/// and can be freely cloned. +/// +/// # Type Parameter +/// - `T`: A [`Repr`] implementation defining the row layout. +/// +/// # Access +/// - [`get_row`](Self::get_row): Get an immutable row by index. +/// - [`rows`](Self::rows): Iterate over all rows. +#[derive(Debug, Clone, Copy)] +pub struct MatRef<'a, T: Repr> { + ptr: NonNull, + repr: T, + /// Marker to tie the lifetime to the borrowed data. + _lifetime: PhantomData<&'a T>, +} + +// SAFETY: [`Repr`] is required to propagate its `Send` bound. +unsafe impl Send for MatRef<'_, T> where T: Repr + Send {} + +// SAFETY: [`Repr`] is required to propagate its `Sync` bound. +unsafe impl Sync for MatRef<'_, T> where T: Repr + Sync {} + +impl<'a, T: Repr> MatRef<'a, T> { + /// Construct a new [`MatRef`] over `data`. + pub fn from_repr(repr: T, data: &'a [U]) -> Result + where + T: NewRef, + { + repr.new_ref(data) + } + + /// Returns the number of rows (vectors) in the matrix. + #[inline] + pub fn num_vectors(&self) -> usize { + self.repr.nrows() + } + + /// Returns a reference to the underlying representation. + pub fn repr(&self) -> &T { + &self.repr + } + + /// Returns an immutable reference to the i-th row, or `None` if out of bounds. + #[must_use] + pub fn get_row(&self, i: usize) -> Option> { + if i < self.num_vectors() { + // SAFETY: Bounds check passed, and the MatRef was constructed + // with valid representation and pointer. + let row = unsafe { self.get_row_unchecked(i) }; + Some(row) + } else { + None + } + } + + /// Returns the `i`th row without bounds checking. + /// + /// # Safety + /// + /// `i` must be less than `self.num_vectors()`. + #[inline] + pub unsafe fn get_row_unchecked(&self, i: usize) -> T::Row<'_> { + // SAFETY: Caller must ensure i < self.num_vectors(). + unsafe { self.repr.get_row(self.ptr, i) } + } + + /// Returns an iterator over immutable row references. + pub fn rows(&self) -> Rows<'_, T> { + Rows::new(*self) + } + + /// Return a [`Mat`] with the same contents as `self`. + pub fn to_owned(&self) -> Mat + where + T: NewCloned, + { + T::new_cloned(*self) + } + + /// Construct a new [`MatRef`] over the raw pointer and representation without performing + /// any validity checks. + /// + /// # Safety + /// + /// Argument `ptr` must point to memory compatible with [`Repr::layout`] and pass any + /// validity checks required by `T`. + pub unsafe fn from_raw_parts(repr: T, ptr: NonNull) -> Self { + Self { + ptr, + repr, + _lifetime: PhantomData, + } + } + + /// Return the base pointer for the [`MatRef`]. + pub fn as_raw_ptr(&self) -> *const u8 { + self.ptr.as_ptr() + } +} + +impl<'a, T> MatRef<'a, RowMajor> { + /// Returns the raw dimension (columns) of the vectors in the matrix. + #[inline] + pub fn vector_dim(&self) -> usize { + self.repr.ncols() + } + + /// Return the backing data as a contiguous slice of `T`. + /// + /// The returned slice has `num_vectors() * vector_dim()` elements in row-major order. + #[inline] + pub fn as_slice(&self) -> &'a [T] { + let len = self.repr.num_elements(); + // SAFETY: `RowMajor` guarantees `nrows * ncols` contiguous `T` elements starting + // at `self.ptr`, valid for the view's lifetime `'a`. + unsafe { std::slice::from_raw_parts(self.ptr.as_ptr().cast::(), len) } + } -#[cfg(feature = "rayon")] -use rayon::prelude::{IndexedParallelIterator, ParallelIterator, ParallelSlice, ParallelSliceMut}; -use thiserror::Error; + /// Return a [`MatrixView`] over the backing data. + #[inline] + pub fn as_matrix_view(&self) -> MatrixView<'a, T> { + *self + } +} -/// Various view types (types such as [`MatrixView`] that add semantic meaning to blobs -/// of data) need both immutable and mutable variants. -/// -/// This trait can be implemented by wrappers for immutable and mutable slice references, -/// allowing for a common code path for immutable and mutable view types. -/// -/// The main goal is to provide a way of retrieving an underlying dense slice, which can -/// then be used as the building block for higher level abstractions. -/// -/// # Safety -/// -/// This trait is unsafe because it requires `as_slice` to be idempotent (and unsafe code -/// relies on this). -/// -/// In other words: `as_slice` must **always** return the same slice with the same length. -pub unsafe trait DenseData { - type Elem; +// Reborrow: Mat -> MatRef +impl<'this, T: ReprOwned> Reborrow<'this> for Mat { + type Target = MatRef<'this, T>; - /// Return the underlying data as a slice. - fn as_slice(&self) -> &[Self::Elem]; + fn reborrow(&'this self) -> Self::Target { + self.as_view() + } } -/// A mutable companion to `DenseData`. -/// -/// This trait allows mutable methods on view types to be selectively enabled when data -/// underlying the type is mutable. -/// -/// # Safety +// ReborrowMut: Mat -> MatMut +impl<'this, T: ReprOwned> ReborrowMut<'this> for Mat { + type Target = MatMut<'this, T>; + + fn reborrow_mut(&'this mut self) -> Self::Target { + self.as_view_mut() + } +} + +// Reborrow: MatRef -> MatRef (with shorter lifetime) +impl<'this, 'a, T: Repr> Reborrow<'this> for MatRef<'a, T> { + type Target = MatRef<'this, T>; + + fn reborrow(&'this self) -> Self::Target { + MatRef { + ptr: self.ptr, + repr: self.repr, + _lifetime: PhantomData, + } + } +} + +//////////// +// MatMut // +//////////// + +/// A mutable borrowed view of a matrix. /// -/// This trait is unsafe because it requires `as_slice` to be idempotent (and unsafe code -/// relies on this). +/// Provides read-write access to matrix data without ownership. /// -/// In other words: `as_slice` must **always** return the same slice with the same length. +/// # Type Parameter +/// - `T`: A [`ReprMut`] implementation defining the row layout. /// -/// Additionally, the returned slice must span the exact same memory as `as_slice`. -pub unsafe trait MutDenseData: DenseData { - fn as_mut_slice(&mut self) -> &mut [Self::Elem]; +/// # Access +/// - [`get_row`](Self::get_row): Get an immutable row by index. +/// - [`get_row_mut`](Self::get_row_mut): Get a mutable row by index. +/// - [`as_view`](Self::as_view): Reborrow as immutable [`MatRef`]. +/// - [`rows`](Self::rows), [`rows_mut`](Self::rows_mut): Iterate over rows. +#[derive(Debug)] +pub struct MatMut<'a, T: ReprMut> { + ptr: NonNull, + repr: T, + /// Marker to tie the lifetime to the mutably borrowed data. + _lifetime: PhantomData<&'a mut T>, } -// SAFETY: This fulfills the idempotency requirement. -unsafe impl DenseData for &[T] { - type Elem = T; - fn as_slice(&self) -> &[Self::Elem] { - self +// SAFETY: [`ReprMut`] is required to propagate its `Send` bound. +unsafe impl Send for MatMut<'_, T> where T: ReprMut + Send {} + +// SAFETY: [`ReprMut`] is required to propagate its `Sync` bound. +unsafe impl Sync for MatMut<'_, T> where T: ReprMut + Sync {} + +impl<'a, T: ReprMut> MatMut<'a, T> { + /// Construct a new [`MatMut`] over `data`. + pub fn from_repr(repr: T, data: &'a mut [U]) -> Result + where + T: NewMut, + { + repr.new_mut(data) + } + + /// Returns the number of rows (vectors) in the matrix. + #[inline] + pub fn num_vectors(&self) -> usize { + self.repr.nrows() + } + + /// Returns a reference to the underlying representation. + pub fn repr(&self) -> &T { + &self.repr + } + + /// Returns an immutable reference to the i-th row, or `None` if out of bounds. + #[inline] + #[must_use] + pub fn get_row(&self, i: usize) -> Option> { + if i < self.num_vectors() { + // SAFETY: Bounds check passed. + Some(unsafe { self.get_row_unchecked(i) }) + } else { + None + } + } + + /// Returns the i-th row without bounds checking. + /// + /// # Safety + /// + /// `i` must be less than `self.num_vectors()`. + #[inline] + pub unsafe fn get_row_unchecked(&self, i: usize) -> T::Row<'_> { + // SAFETY: Caller must ensure i < self.num_vectors(). + unsafe { self.repr.get_row(self.ptr, i) } + } + + /// Returns a mutable reference to the `i`-th row, or `None` if out of bounds. + #[inline] + #[must_use] + pub fn get_row_mut(&mut self, i: usize) -> Option> { + if i < self.num_vectors() { + // SAFETY: Bounds check passed. + Some(unsafe { self.get_row_mut_unchecked(i) }) + } else { + None + } + } + + /// Returns a mutable reference to the i-th row without bounds checking. + /// + /// # Safety + /// + /// `i` must be less than [`num_vectors()`](Self::num_vectors). + #[inline] + pub(crate) unsafe fn get_row_mut_unchecked(&mut self, i: usize) -> T::RowMut<'_> { + // SAFETY: Caller asserts that `i < self.num_vectors()`. The constructors for this + // type ensure that `ptr` is compatible with `T`. + unsafe { self.repr.get_row_mut(self.ptr, i) } + } + + /// Reborrows as an immutable [`MatRef`]. + pub fn as_view(&self) -> MatRef<'_, T> { + MatRef { + ptr: self.ptr, + repr: self.repr, + _lifetime: PhantomData, + } + } + + /// Returns an iterator over immutable row references. + pub fn rows(&self) -> Rows<'_, T> { + Rows::new(self.reborrow()) + } + + /// Returns an iterator over mutable row references. + pub fn rows_mut(&mut self) -> RowsMut<'_, T> { + RowsMut::new(self.reborrow_mut()) + } + + /// Return a [`Mat`] with the same contents as `self`. + pub fn to_owned(&self) -> Mat + where + T: NewCloned, + { + T::new_cloned(self.as_view()) + } + + /// Construct a new [`MatMut`] over the raw pointer and representation without performing + /// any validity checks. + /// + /// # Safety + /// + /// Argument `ptr` must point to memory compatible with [`Repr::layout`]. + pub unsafe fn from_raw_parts(repr: T, ptr: NonNull) -> Self { + Self { + ptr, + repr, + _lifetime: PhantomData, + } + } + + /// Return the base pointer for the [`MatMut`]. + pub fn as_raw_ptr(&self) -> *const u8 { + self.ptr.as_ptr() + } + + /// Return a mutable base pointer for the [`MatMut`]. + pub fn as_raw_mut_ptr(&mut self) -> *mut u8 { + self.ptr.as_ptr() } } -// SAFETY: This fulfills the idempotency requirement. -unsafe impl DenseData for &mut [T] { - type Elem = T; - fn as_slice(&self) -> &[Self::Elem] { - self +// Reborrow: MatMut -> MatRef +impl<'this, 'a, T: ReprMut> Reborrow<'this> for MatMut<'a, T> { + type Target = MatRef<'this, T>; + + fn reborrow(&'this self) -> Self::Target { + self.as_view() } } -// SAFETY: This fulfills the idempotency requirement and returns a slice spanning the same -// range as `as_slice`. -unsafe impl MutDenseData for &mut [T] { - fn as_mut_slice(&mut self) -> &mut [Self::Elem] { - self +// ReborrowMut: MatMut -> MatMut (with shorter lifetime) +impl<'this, 'a, T: ReprMut> ReborrowMut<'this> for MatMut<'a, T> { + type Target = MatMut<'this, T>; + + fn reborrow_mut(&'this mut self) -> Self::Target { + MatMut { + ptr: self.ptr, + repr: self.repr, + _lifetime: PhantomData, + } } } -// SAFETY: This fulfills the idempotency requirement. -unsafe impl DenseData for Box<[T]> { - type Elem = T; - fn as_slice(&self) -> &[Self::Elem] { - self +impl<'a, T> MatMut<'a, RowMajor> { + /// Returns the raw dimension (columns) of the vectors in the matrix. + #[inline] + pub fn vector_dim(&self) -> usize { + self.repr.ncols() } + + delegate_read!(pub fn as_slice(&self) -> &[T]); + delegate_read!(pub fn as_matrix_view(&self) -> MatrixView<'_, T>); } -// SAFETY: This fulfills the idempotency requirement and returns a slice spanning the same -// memory as `as_slice`. -unsafe impl MutDenseData for Box<[T]> { - fn as_mut_slice(&mut self) -> &mut [Self::Elem] { - self +////////////////////////////// +// Dense (RowMajor) API + aliases +////////////////////////////// + +/// A generator for initializing matrix entries via [`Matrix::new`]. +pub trait Generator { + fn generate(&mut self) -> T; +} + +impl Generator for T +where + T: Clone, +{ + fn generate(&mut self) -> T { + self.clone() } } -//////////// -// Matrix // -//////////// +/// A [`Generator`] that invokes a closure to initialize each element. +pub struct Init(pub F); -/// A view over dense chunk of memory, interpreting that memory as a 2-dimensional matrix -/// laid out in row-major order. -/// -/// When this class view immutable memory, it is `Copy`. -#[derive(Debug, Clone, Copy, PartialEq)] -pub struct MatrixBase +impl Generator for Init where - T: DenseData, + F: FnMut() -> T, { - data: T, - nrows: usize, - ncols: usize, + fn generate(&mut self) -> T { + (self.0)() + } } +/// Owned dense row-major matrix. +pub type Matrix = Mat>; + +/// Shared dense row-major view. +pub type MatrixView<'a, T> = MatRef<'a, RowMajor>; + +/// Mutable dense row-major view. +pub type MatrixViewMut<'a, T> = MatMut<'a, RowMajor>; + +/// A lightweight, `'static` version of [`TryFromError`] that drops the offending data. #[derive(Debug, Error)] #[non_exhaustive] -#[error( - "tried to construct a matrix view with {nrows} rows and {ncols} columns over a slice \ - of length {len}" -)] +#[error("cannot view a slice of length {len} as a {nrows}x{ncols} matrix")] pub struct TryFromErrorLight { - len: usize, - nrows: usize, - ncols: usize, + pub len: usize, + pub nrows: usize, + pub ncols: usize, } +/// Error returned when a slice length is incompatible with the requested dimensions. +/// +/// Carries the original `data` so a failed construction does not deallocate it; recover it +/// with [`TryFromError::into_inner`]. #[derive(Error)] #[non_exhaustive] -#[error( - "tried to construct a matrix view with {nrows} rows and {ncols} columns over a slice \ - of length {}", data.as_slice().len() -)] +#[error("cannot view a slice of length {} as a {nrows}x{ncols} matrix", data.as_slice().len())] pub struct TryFromError { data: T, nrows: usize, ncols: usize, } -// Manually implement `fmt::Debug` so we don't require `T::Debug`. -impl fmt::Debug for TryFromError { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { +// Hand-written so `TryFromError` does not require `T: Debug`. +impl std::fmt::Debug for TryFromError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("TryFromError") .field("data_len", &self.data.as_slice().len()) .field("nrows", &self.nrows) @@ -145,13 +1232,12 @@ impl fmt::Debug for TryFromError { } impl TryFromError { - /// Consume the error and return the base data. + /// Consume the error and return the original data, without deallocating it. pub fn into_inner(self) -> T { self.data } - /// Return a variation of `Self` that is guaranteed to be `'static` by removing the - /// data that was passed to the original constructor. + /// Return a `'static` copy of this error that drops the offending data. pub fn as_static(&self) -> TryFromErrorLight { TryFromErrorLight { len: self.data.as_slice().len(), @@ -161,1711 +1247,1392 @@ impl TryFromError { } } -/// A generator for initializing the entries in a matrix via `Matrix::new`. -pub trait Generator { - fn generate(&mut self) -> T; -} +impl<'a, T> MatRef<'a, RowMajor> { + /// View `data` as an `nrows x ncols` matrix. + pub fn try_from( + data: &'a [T], + nrows: usize, + ncols: usize, + ) -> Result> { + let repr = match RowMajor::new(nrows, ncols) { + Ok(repr) if repr.num_elements() == data.len() => repr, + _ => return Err(TryFromError { data, nrows, ncols }), + }; + // SAFETY: `data.len()` was checked to equal `repr.num_elements()`. + Ok(unsafe { MatRef::from_raw_parts(repr, as_nonnull(data).cast::()) }) + } -impl Generator for T -where - T: Clone, -{ - fn generate(&mut self) -> T { - self.clone() + /// View `data` as a single-row matrix. + pub fn row_vector(data: &'a [T]) -> Self { + let repr = RowMajor { + nrows: 1, + ncols: data.len(), + _elem: PhantomData, + }; + // SAFETY: `data` is exactly `1 * data.len()` contiguous `T`, matching `repr`. + unsafe { MatRef::from_raw_parts(repr, as_nonnull(data).cast::()) } } -} -/// A matrix initializer that invokes the provided lambda to initialize each element. -pub struct Init(pub F); + /// View `data` as a single-column matrix. + pub fn column_vector(data: &'a [T]) -> Self { + let repr = RowMajor { + nrows: data.len(), + ncols: 1, + _elem: PhantomData, + }; + // SAFETY: `data` is exactly `data.len() * 1` contiguous `T`, matching `repr`. + unsafe { MatRef::from_raw_parts(repr, as_nonnull(data).cast::()) } + } -impl Generator for Init -where - F: FnMut() -> T, -{ - fn generate(&mut self) -> T { - (self.0)() + /// Number of rows. + #[inline] + pub fn nrows(&self) -> usize { + self.num_vectors() + } + + /// Number of columns. + #[inline] + pub fn ncols(&self) -> usize { + self.vector_dim() + } + + /// Row `i`. Panics if `i >= self.nrows()`. + pub fn row(&self, i: usize) -> &'a [T] { + assert!( + i < self.nrows(), + "tried to access row {i} of a matrix with {} rows", + self.nrows() + ); + let ncols = self.ncols(); + let start = i * ncols; + // SAFETY: `i < self.nrows()` was asserted, and the backing slice holds + // `nrows * ncols` elements, so `start..start + ncols` is in bounds. + unsafe { self.as_slice().get_unchecked(start..start + ncols) } + } + + /// Iterator over rows. + pub fn row_iter(&self) -> impl ExactSizeIterator { + self.as_slice().chunks_exact(self.ncols()) } -} -impl MatrixBase> { - /// Construct a new Matrix initialized with the contents of the generator. + /// Iterator over sub-matrices of up to `batchsize` rows. /// - /// Elements are initialized in memory order. - pub fn new(mut generator: U, nrows: usize, ncols: usize) -> Self + /// # Panics + /// Panics if `batchsize == 0`. + pub fn window_iter(&self, batchsize: usize) -> impl Iterator> { + assert!(batchsize != 0, "window_iter batchsize cannot be zero"); + let ncols = self.ncols(); + self.as_slice() + .chunks(ncols * batchsize) + .map(move |d| MatrixView::try_from(d, d.len() / ncols, ncols).expect("exact chunk")) + } + + /// Parallel iterator over rows. + #[cfg(feature = "rayon")] + pub fn par_row_iter(&self) -> impl rayon::iter::IndexedParallelIterator where - U: Generator, + T: Sync, { - let data: Box<[T]> = (0..nrows * ncols).map(|_| generator.generate()).collect(); - debug_assert_eq!(data.len(), nrows * ncols); - Self { data, nrows, ncols } + use rayon::slice::ParallelSlice; + self.as_slice().par_chunks_exact(self.ncols()) } -} -impl MatrixBase -where - T: DenseData, -{ - /// Try to construct a `MatrixBase` over the provided base. If the size of the base - /// is incorrect, return a `TryFromError` containing the base. + /// Parallel iterator over sub-matrices of up to `batchsize` rows. /// - /// The length of the base must be equal to `nrows * ncols`. - pub fn try_from(data: T, nrows: usize, ncols: usize) -> Result> { - let len = data.as_slice().len(); - if len != nrows * ncols { - Err(TryFromError { data, nrows, ncols }) - } else { - Ok(Self { data, nrows, ncols }) - } + /// # Panics + /// Panics if `batchsize == 0`. + #[cfg(feature = "rayon")] + pub fn par_window_iter( + &self, + batchsize: usize, + ) -> impl rayon::iter::IndexedParallelIterator> + where + T: Send + Sync, + { + use rayon::slice::ParallelSlice; + assert!(batchsize != 0, "par_window_iter batchsize cannot be zero"); + let ncols = self.ncols(); + self.as_slice() + .par_chunks(ncols * batchsize) + .map(move |d| MatrixView::try_from(d, d.len() / ncols, ncols).expect("exact chunk")) } - /// Return the number of columns in the matrix. - pub fn ncols(&self) -> usize { - self.ncols + /// View of the rows in `rows`, or `None` if out of bounds. + pub fn subview(&self, rows: std::ops::Range) -> Option> { + let ncols = self.ncols(); + let lo = rows.start.checked_mul(ncols)?; + let hi = rows.end.checked_mul(ncols)?; + let data = self.as_slice().get(lo..hi)?; + let repr = RowMajor { + nrows: rows.len(), + ncols, + _elem: PhantomData, + }; + // SAFETY: `data` is exactly `rows.len() * ncols` contiguous `T`, matching `repr`. + Some(unsafe { MatRef::from_raw_parts(repr, as_nonnull(data).cast::()) }) } - /// Return the number of rows in the matrix. - pub fn nrows(&self) -> usize { - self.nrows + /// Element at `(row, col)`, or `None` if out of bounds. + pub fn try_get(&self, row: usize, col: usize) -> Option<&'a T> { + if row >= self.nrows() || col >= self.ncols() { + None + } else { + // SAFETY: `row` and `col` were just verified in-bounds. + Some(unsafe { self.get_unchecked(row, col) }) + } } - /// Create a new [`Matrix`] by applying the closure `f` to each element. - /// - /// The returned matrix has the same shape as `self`. + /// Base pointer of the backing data. + pub fn as_ptr(&self) -> *const T { + self.as_slice().as_ptr() + } + + /// Apply `f` to every element, producing a new owned matrix of the same shape. pub fn map(&self, f: F) -> Matrix where - F: FnMut(&T::Elem) -> R, + F: FnMut(&T) -> R, { - let data: Box<[_]> = self.as_slice().iter().map(f).collect(); - Matrix { - data, + let data: Box<[R]> = self.as_slice().iter().map(f).collect(); + let repr = RowMajor { nrows: self.nrows(), ncols: self.ncols(), - } + _elem: PhantomData, + }; + // SAFETY: `data` has exactly `nrows * ncols` elements (one per source element). + unsafe { repr.box_to_mat(data) } } +} - /// Return the underlying data as a slice. - pub fn as_slice(&self) -> &[T::Elem] { - self.data.as_slice() +impl<'a, T> MatMut<'a, RowMajor> { + /// View `data` as an `nrows x ncols` mutable matrix. + pub fn try_from( + data: &'a mut [T], + nrows: usize, + ncols: usize, + ) -> Result> { + let repr = match RowMajor::new(nrows, ncols) { + Ok(repr) if repr.num_elements() == data.len() => repr, + _ => return Err(TryFromError { data, nrows, ncols }), + }; + // SAFETY: `data.len()` was checked to equal `repr.num_elements()`. + Ok(unsafe { MatMut::from_raw_parts(repr, as_nonnull_mut(data).cast::()) }) + } + + /// Number of rows. + #[inline] + pub fn nrows(&self) -> usize { + self.num_vectors() + } + + /// Number of columns. + #[inline] + pub fn ncols(&self) -> usize { + self.vector_dim() } - /// Return the underlying data as a mutable slice. - pub fn as_mut_slice(&mut self) -> &mut [T::Elem] + // Reads delegate to the canonical `MatrixView` implementation. + delegate_read!(pub fn row(&self, i: usize) -> &[T]); + delegate_read!(pub fn row_iter(&self) -> impl ExactSizeIterator); + delegate_read!(pub fn window_iter(&self, batchsize: usize) -> impl Iterator>); + delegate_read!(pub fn subview(&self, rows: std::ops::Range) -> Option>); + delegate_read!(pub fn try_get(&self, row: usize, col: usize) -> Option<&T>); + delegate_read!(pub fn as_ptr(&self) -> *const T); + + /// Parallel iterator over rows. + #[cfg(feature = "rayon")] + pub fn par_row_iter(&self) -> impl rayon::iter::IndexedParallelIterator where - T: MutDenseData, + T: Sync, { - self.data.as_mut_slice() + self.as_view().par_row_iter() } - /// Return row `row` as a slice. - /// - /// # Panic - /// - /// Panics if `row >= self.nrows()`. - pub fn row(&self, row: usize) -> &[T::Elem] { + /// Parallel iterator over sub-matrices of up to `batchsize` rows. + #[cfg(feature = "rayon")] + pub fn par_window_iter( + &self, + batchsize: usize, + ) -> impl rayon::iter::IndexedParallelIterator> + where + T: Send + Sync, + { + self.as_view().par_window_iter(batchsize) + } + + /// Apply `f` to every element, producing a new owned matrix of the same shape. + pub fn map(&self, f: F) -> Matrix + where + F: FnMut(&T) -> R, + { + self.as_view().map(f) + } + + /// Backing data as a mutable slice. + pub fn as_mut_slice(&mut self) -> &mut [T] { + let len = self.repr.num_elements(); + // SAFETY: the backing allocation holds `len` contiguous `T`, and `&mut self` grants + // exclusive access for the returned slice's lifetime. + unsafe { std::slice::from_raw_parts_mut(self.as_raw_mut_ptr().cast::(), len) } + } + + /// Mutable row `i`. Panics if `i >= self.nrows()`. + pub fn row_mut(&mut self, i: usize) -> &mut [T] { assert!( - row < self.nrows(), - "tried to access row {row} of a matrix with {} rows", + i < self.nrows(), + "tried to access row {i} of a matrix with {} rows", self.nrows() ); + let ncols = self.ncols(); + let start = i * ncols; + // SAFETY: `i < self.nrows()` was asserted, and the backing slice holds + // `nrows * ncols` elements, so `start..start + ncols` is in bounds. + unsafe { self.as_mut_slice().get_unchecked_mut(start..start + ncols) } + } + + /// Iterator over mutable rows. + pub fn row_iter_mut(&mut self) -> impl ExactSizeIterator { + let ncols = self.ncols(); + self.as_mut_slice().chunks_exact_mut(ncols) + } - // SAFETY: `row` is in-bounds. - unsafe { self.get_row_unchecked(row) } + /// Parallel iterator over mutable rows. + #[cfg(feature = "rayon")] + pub fn par_row_iter_mut(&mut self) -> impl rayon::iter::IndexedParallelIterator + where + T: Send, + { + use rayon::slice::ParallelSliceMut; + let ncols = self.ncols(); + self.as_mut_slice().par_chunks_exact_mut(ncols) } - /// Construct a new `MatrixBase` over the raw data. - /// - /// The returned `MatrixBase` will only have a single row with contents equal to `data`. - pub fn row_vector(data: T) -> Self { - let ncols = data.as_slice().len(); - Self { - data, + /// Base pointer of the backing data. + pub fn as_mut_ptr(&mut self) -> *mut T { + self.as_mut_slice().as_mut_ptr() + } +} + +impl Mat> { + /// Take ownership of `data`, interpreting it as an `nrows x ncols` matrix. + pub fn try_from( + data: Box<[T]>, + nrows: usize, + ncols: usize, + ) -> Result>> { + let repr = match RowMajor::new(nrows, ncols) { + Ok(repr) if repr.num_elements() == data.len() => repr, + _ => return Err(TryFromError { data, nrows, ncols }), + }; + // SAFETY: `data.len()` was checked to equal `repr.num_elements()`. + Ok(unsafe { repr.box_to_mat(data) }) + } + + /// Take ownership of `data` as a single-row matrix. + pub fn row_vector(data: Box<[T]>) -> Self { + let repr = RowMajor { nrows: 1, - ncols, - } + ncols: data.len(), + _elem: PhantomData, + }; + // SAFETY: `data` has exactly `1 * data.len()` elements matching `repr`. + unsafe { repr.box_to_mat(data) } } - /// Construct a new `MatrixBase` over the raw data. - /// - /// The returned `MatrixBase` will only have a single column with contents equal to `data`. - pub fn column_vector(data: T) -> Self { - let nrows = data.as_slice().len(); - Self { - data, - nrows, + /// Take ownership of `data` as a single-column matrix. + pub fn column_vector(data: Box<[T]>) -> Self { + let repr = RowMajor { + nrows: data.len(), ncols: 1, - } + _elem: PhantomData, + }; + // SAFETY: `data` has exactly `data.len() * 1` elements matching `repr`. + unsafe { repr.box_to_mat(data) } } - /// Return row `row` if `row < self.nrows()`. Otherwise, return `None`. - pub fn get_row(&self, row: usize) -> Option<&[T::Elem]> { - if row < self.nrows() { - // SAFETY: `row` is in-bounds. - Some(unsafe { self.get_row_unchecked(row) }) - } else { - None - } + /// Number of rows. + #[inline] + pub fn nrows(&self) -> usize { + self.num_vectors() } - /// 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(&self, row: usize) -> &[T::Elem] { - debug_assert!(row < self.nrows); - let ncols = self.ncols; - let start = row * ncols; - - 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) } + /// Number of columns. + #[inline] + pub fn ncols(&self) -> usize { + self.vector_dim() } - /// Return row `row` as a mutable slice. - /// - /// # Panics - /// - /// Panics if `row >= self.nrows()`. - pub fn row_mut(&mut self, row: usize) -> &mut [T::Elem] + // Reads delegate to the canonical `MatrixView` implementation. + delegate_read!(pub fn row(&self, i: usize) -> &[T]); + delegate_read!(pub fn row_iter(&self) -> impl ExactSizeIterator); + delegate_read!(pub fn window_iter(&self, batchsize: usize) -> impl Iterator>); + delegate_read!(pub fn subview(&self, rows: std::ops::Range) -> Option>); + delegate_read!(pub fn try_get(&self, row: usize, col: usize) -> Option<&T>); + delegate_read!(pub fn as_ptr(&self) -> *const T); + + /// Parallel iterator over rows. + #[cfg(feature = "rayon")] + pub fn par_row_iter(&self) -> impl rayon::iter::IndexedParallelIterator where - T: MutDenseData, + T: Sync, { - assert!( - row < self.nrows(), - "tried to access row {row} of a matrix with {} rows", - self.nrows() - ); + self.as_view().par_row_iter() + } - // SAFETY: `row` is in-bounds. - unsafe { self.get_row_unchecked_mut(row) } + /// Parallel iterator over sub-matrices of up to `batchsize` rows. + #[cfg(feature = "rayon")] + pub fn par_window_iter( + &self, + batchsize: usize, + ) -> impl rayon::iter::IndexedParallelIterator> + where + T: Send + Sync, + { + self.as_view().par_window_iter(batchsize) } - /// 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] + /// Apply `f` to every element, producing a new owned matrix of the same shape. + pub fn map(&self, f: F) -> Matrix where - T: MutDenseData, + F: FnMut(&T) -> R, { - debug_assert!(row < self.nrows); - let ncols = self.ncols; - let start = row * ncols; + self.as_view().map(f) + } - 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. - unsafe { - self.data - .as_mut_slice() - .get_unchecked_mut(start..start + ncols) - } + /// Backing data as a mutable slice. + pub fn as_mut_slice(&mut self) -> &mut [T] { + let len = self.repr.num_elements(); + // SAFETY: the backing allocation holds `len` contiguous `T`, and `&mut self` grants + // exclusive access for the returned slice's lifetime. + unsafe { std::slice::from_raw_parts_mut(self.as_raw_mut_ptr().cast::(), len) } } - /// Return a iterator over all rows in the matrix. - /// - /// Rows are yielded sequentially beginning with row 0. - pub fn row_iter(&self) -> impl ExactSizeIterator { - self.data.as_slice().chunks_exact(self.ncols()) + /// Mutable view over the whole matrix. + pub fn as_mut_view(&mut self) -> MatrixViewMut<'_, T> { + self.as_view_mut() } - /// Return a mutable iterator over all rows in the matrix. - /// - /// Rows are yielded sequentially beginning with row 0. - pub fn row_iter_mut(&mut self) -> impl ExactSizeIterator - where - T: MutDenseData, - { + /// Mutable row `i`. Panics if `i >= self.nrows()`. + pub fn row_mut(&mut self, i: usize) -> &mut [T] { + assert!( + i < self.nrows(), + "tried to access row {i} of a matrix with {} rows", + self.nrows() + ); let ncols = self.ncols(); - self.data.as_mut_slice().chunks_exact_mut(ncols) + let start = i * ncols; + // SAFETY: `i < self.nrows()` was asserted, and the backing slice holds + // `nrows * ncols` elements, so `start..start + ncols` is in bounds. + unsafe { self.as_mut_slice().get_unchecked_mut(start..start + ncols) } } - /// Return an iterator that divides the matrix into sub-matrices with (up to) - /// `batchsize` rows with `self.ncols()` columns. - /// - /// It is possible for yielded sub-matrices to have fewer than `batchsize` rows if the - /// number of rows in the parent matrix is not evenly divisible by `batchsize`. - /// - /// # Panics - /// - /// Panics if `batchsize = 0`. - pub fn window_iter(&self, batchsize: usize) -> impl Iterator> - where - T::Elem: Sync, - { - assert!(batchsize != 0, "window_iter batchsize cannot be zero"); + /// Iterator over mutable rows. + pub fn row_iter_mut(&mut self) -> impl ExactSizeIterator { let ncols = self.ncols(); - self.data - .as_slice() - .chunks(ncols * batchsize) - .map(move |data| { - let blobsize = data.len(); - let nrows = blobsize / ncols; - assert_eq!(blobsize % ncols, 0); - MatrixView { data, nrows, ncols } - }) + self.as_mut_slice().chunks_exact_mut(ncols) } - /// Return a parallel iterator that divides the matrix into sub-matrices with (up to) - /// `batchsize` rows with `self.ncols()` columns. - /// - /// This allows workers in parallel algorithms to work on dense subsets of the whole - /// matrix for better locality. - /// - /// It is possible for yielded sub-matrices to have fewer than `batchsize` rows if the - /// number of rows in the parent matrix is not evenly divisible by `batchsize`. - /// - /// # Panics - /// - /// Panics if `batchsize = 0`. + /// Parallel iterator over mutable rows. #[cfg(feature = "rayon")] - pub fn par_window_iter( - &self, - batchsize: usize, - ) -> impl IndexedParallelIterator> + pub fn par_row_iter_mut(&mut self) -> impl rayon::iter::IndexedParallelIterator where - T::Elem: Sync, + T: Send, { - assert!(batchsize != 0, "par_window_iter batchsize cannot be zero"); + use rayon::slice::ParallelSliceMut; let ncols = self.ncols(); - self.data - .as_slice() - .par_chunks(ncols * batchsize) - .map(move |data| { - let blobsize = data.len(); - let nrows = blobsize / ncols; - assert_eq!(blobsize % ncols, 0); - MatrixView { data, nrows, ncols } - }) + self.as_mut_slice().par_chunks_exact_mut(ncols) } - /// Return a parallel iterator that divides the matrix into mutable sub-matrices with - /// (up to) `batchsize` rows with `self.ncols()` columns. - /// - /// This allows workers in parallel algorithms to work on dense subsets of the whole - /// matrix for better locality. - /// - /// It is possible for yielded sub-matrices to have fewer than `batchsize` rows if the - /// number of rows in the parent matrix is not evenly divisible by `batchsize`. - /// - /// # Panics - /// - /// Panics if `batchsize = 0`. + /// Parallel iterator over mutable sub-matrices of up to `batchsize` rows. #[cfg(feature = "rayon")] pub fn par_window_iter_mut( &mut self, batchsize: usize, - ) -> impl IndexedParallelIterator> + ) -> impl rayon::iter::IndexedParallelIterator> where - T: MutDenseData, - T::Elem: Send, + T: Send, { + use rayon::slice::ParallelSliceMut; assert!( batchsize != 0, "par_window_iter_mut batchsize cannot be zero" ); let ncols = self.ncols(); - self.data - .as_mut_slice() + self.as_mut_slice() .par_chunks_mut(ncols * batchsize) - .map(move |data| { - let blobsize = data.len(); - let nrows = blobsize / ncols; - assert_eq!(blobsize % ncols, 0); - MutMatrixView { data, nrows, ncols } + .map(move |d| { + let nrows = d.len() / ncols; + MatrixViewMut::try_from(d, nrows, ncols).expect("exact chunk") }) } - /// Return a parallel iterator over the rows of the matrix. - #[cfg(feature = "rayon")] - pub fn par_row_iter(&self) -> impl IndexedParallelIterator - where - T::Elem: Sync, - { - self.as_slice().par_chunks_exact(self.ncols()) + /// Mutable base pointer of the backing data. + pub fn as_mut_ptr(&mut self) -> *mut T { + self.as_mut_slice().as_mut_ptr() } +} - /// Return a parallel iterator over the rows of the matrix. - #[cfg(feature = "rayon")] - pub fn par_row_iter_mut(&mut self) -> impl IndexedParallelIterator - where - T: MutDenseData, - T::Elem: Send, - { - let ncols = self.ncols(); - self.as_mut_slice().par_chunks_exact_mut(ncols) +// Indexing by `(row, col)`. +impl std::ops::Index<(usize, usize)> for Mat> { + type Output = T; + fn index(&self, (row, col): (usize, usize)) -> &T { + self.try_get(row, col).expect("index out of bounds") } +} - /// Consume the matrix, returning the inner representation. - /// - /// This loses the information about the number of rows and columnts. - pub fn into_inner(self) -> T { - self.data +impl std::ops::IndexMut<(usize, usize)> for Mat> { + fn index_mut(&mut self, (row, col): (usize, usize)) -> &mut T { + assert!( + row < self.nrows() && col < self.ncols(), + "index out of bounds" + ); + // SAFETY: `row` and `col` were checked in-bounds above. + unsafe { self.get_unchecked_mut(row, col) } } +} - /// Return a view over the matrix. - pub fn as_view(&self) -> MatrixView<'_, T::Elem> { - MatrixBase { - data: self.as_slice(), - nrows: self.nrows(), - ncols: self.ncols(), - } +impl std::ops::Index<(usize, usize)> for MatRef<'_, RowMajor> { + type Output = T; + fn index(&self, (row, col): (usize, usize)) -> &T { + self.try_get(row, col).expect("index out of bounds") } +} - /// Return a mutable view over the matrix. - pub fn as_mut_view(&mut self) -> MutMatrixView<'_, T::Elem> - where - T: MutDenseData, - { - let nrows = self.nrows(); - let ncols = self.ncols(); - MatrixBase { - data: self.as_mut_slice(), - nrows, - ncols, - } +impl std::ops::Index<(usize, usize)> for MatMut<'_, RowMajor> { + type Output = T; + fn index(&self, (row, col): (usize, usize)) -> &T { + self.try_get(row, col).expect("index out of bounds") } +} - /// Return a view over the specified rows of the matrix. - /// - /// If the specified range is out of bounds, return `None`. - /// - /// ```rust - /// use diskann_utils::views::Matrix; - /// - /// let mut mat = Matrix::new(0usize, 4, 3); - /// - /// // Fill the matrix with some data. - /// mat.row_iter_mut().enumerate().for_each(|(i, row)| row.fill(i)); - /// - /// // Creating a subview into an offset portion of the matrix. - /// let subview = mat.subview(1..3).unwrap(); - /// assert_eq!(subview.nrows(), 2); - /// assert_eq!(subview.row(0), &[1, 1, 1]); - /// assert_eq!(subview.row(1), &[2, 2, 2]); - /// - /// // A trying to access out-of-bounds returns `None` - /// assert!(mat.subview(3..5).is_none()); - /// ``` - pub fn subview(&self, rows: std::ops::Range) -> Option> { - let ncols = self.ncols(); - - let lower = rows.start.checked_mul(ncols)?; - let upper = rows.end.checked_mul(ncols)?; +impl std::ops::IndexMut<(usize, usize)> for MatMut<'_, RowMajor> { + fn index_mut(&mut self, (row, col): (usize, usize)) -> &mut T { + assert!( + row < self.nrows() && col < self.ncols(), + "index out of bounds" + ); + // SAFETY: `row` and `col` were checked in-bounds above. + unsafe { self.get_unchecked_mut(row, col) } + } +} - if let Some(data) = self.as_slice().get(lower..upper) { - Some(MatrixBase { - data, - nrows: rows.len(), - ncols: self.ncols(), - }) - } else { - None - } +// Equality compares shape and contents. +impl PartialEq for Mat> { + fn eq(&self, other: &Self) -> bool { + self.nrows() == other.nrows() + && self.ncols() == other.ncols() + && self.as_slice() == other.as_slice() } +} - /// Return a pointer to the base of the matrix. - pub fn as_ptr(&self) -> *const T::Elem { - self.as_slice().as_ptr() +impl PartialEq for MatRef<'_, RowMajor> { + fn eq(&self, other: &Self) -> bool { + self.nrows() == other.nrows() + && self.ncols() == other.ncols() + && self.as_slice() == other.as_slice() } +} - /// 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 PartialEq for MatMut<'_, RowMajor> { + fn eq(&self, other: &Self) -> bool { + self.nrows() == other.nrows() + && self.ncols() == other.ncols() + && self.as_slice() == other.as_slice() } +} - /// Return the value at the specified `row` and `col`. - /// - /// If either index is out-of-bounds, return `None`. - pub fn try_get(&self, row: usize, col: usize) -> Option<&T::Elem> { - if row >= self.nrows() || col >= self.ncols() { - None - } else { - // SAFETY: We just verified that `row` and `col` are in-bounds. - Some(unsafe { self.get_unchecked(row, col) }) - } +impl<'a, T> MatRef<'a, RowMajor> { + /// Reborrow as a shorter-lived view. + pub fn as_view(&self) -> MatrixView<'_, T> { + self.reborrow() } - /// Returns a reference to an element without boundschecking. + /// Element at `(row, col)` without bounds checking. /// /// # 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.ncols + col) + /// `row < self.nrows()` and `col < self.ncols()`. + pub unsafe fn get_unchecked(&self, row: usize, col: usize) -> &'a T { + debug_assert!(row < self.nrows()); + debug_assert!(col < self.ncols()); + let ncols = self.ncols(); + // SAFETY: guaranteed in-bounds by the caller. + unsafe { self.as_slice().get_unchecked(row * ncols + col) } + } +} + +impl<'a, T> MatMut<'a, RowMajor> { + /// Reborrow as a shorter-lived mutable view. + pub fn as_mut_view(&mut self) -> MatrixViewMut<'_, T> { + self.reborrow_mut() } - /// Returns a mutable reference to an element without boundschecking. + delegate_read!( + /// Element at `(row, col)` without bounds checking. + /// + /// # Safety + /// `row < self.nrows()` and `col < self.ncols()`. + unsafe pub fn get_unchecked(&self, row: usize, col: usize) -> &T + ); + + /// Mutable element at `(row, col)` without bounds checking. /// /// # 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 ncols = self.ncols; - debug_assert!(row < self.nrows); - debug_assert!(col < self.ncols); - self.as_mut_slice().get_unchecked_mut(row * ncols + col) + /// `row < self.nrows()` and `col < self.ncols()`. + pub unsafe fn get_unchecked_mut(&mut self, row: usize, col: usize) -> &mut T { + let ncols = self.ncols(); + // SAFETY: guaranteed in-bounds by the caller. + unsafe { self.as_mut_slice().get_unchecked_mut(row * ncols + col) } } - pub fn to_owned(&self) -> Matrix + /// Parallel iterator over mutable sub-matrices of up to `batchsize` rows. + #[cfg(feature = "rayon")] + pub fn par_window_iter_mut( + &mut self, + batchsize: usize, + ) -> impl rayon::iter::IndexedParallelIterator> where - T::Elem: Clone, + T: Send, { - Matrix { - data: self.data.as_slice().into(), - nrows: self.nrows, - ncols: self.ncols, - } + use rayon::slice::ParallelSliceMut; + assert!( + batchsize != 0, + "par_window_iter_mut batchsize cannot be zero" + ); + let ncols = self.ncols(); + self.as_mut_slice() + .par_chunks_mut(ncols * batchsize) + .map(move |d| { + let nrows = d.len() / ncols; + MatrixViewMut::try_from(d, nrows, ncols).expect("exact chunk") + }) } } -/// Represents an owning, 2-dimensional view of a contiguous block of memory, -/// interpreted as a matrix in row-major order. -pub type Matrix = MatrixBase>; +impl Mat> { + delegate_read!( + /// Element at `(row, col)` without bounds checking. + /// + /// # Safety + /// `row < self.nrows()` and `col < self.ncols()`. + unsafe pub fn get_unchecked(&self, row: usize, col: usize) -> &T + ); -/// Represents a non-owning, 2-dimensional view of a contiguous block of memory, -/// interpreted as a matrix in row-major order. -/// -/// This type is useful for functions that need to read matrix data without taking ownership. -/// By accepting a `MatrixView`, such functions can operate on both owned matrices (by converting them -/// to a `MatrixView`) and existing non-owning views. -pub type MatrixView<'a, T> = MatrixBase<&'a [T]>; + /// Mutable element at `(row, col)` without bounds checking. + /// + /// # Safety + /// `row < self.nrows()` and `col < self.ncols()`. + pub unsafe fn get_unchecked_mut(&mut self, row: usize, col: usize) -> &mut T { + let ncols = self.ncols(); + // SAFETY: guaranteed in-bounds by the caller. + unsafe { self.as_mut_slice().get_unchecked_mut(row * ncols + col) } + } -/// Represents a mutable non-owning, 2-dimensional view of a contiguous block of memory, -/// interpreted as a matrix in row-major order. -/// -/// This type is useful for functions that need to modify matrix data without taking ownership. -/// By accepting a `MutMatrixView`, such functions can operate on both owned matrices (by converting them -/// to a `MutMatrixView`) and existing non-owning mutable views. -pub type MutMatrixView<'a, T> = MatrixBase<&'a mut [T]>; + /// Consume the matrix, returning its backing storage as a `Box<[T]>`. + pub fn into_inner(self) -> Box<[T]> { + let this = core::mem::ManuallyDrop::new(self); + let slice = this.as_slice(); + let (ptr, len) = (slice.as_ptr().cast_mut(), slice.len()); + // SAFETY: the matrix is backed by a `Box<[T]>` of `len` elements (see + // `RowMajor::box_to_mat`); `ManuallyDrop` prevents a double free via `Mat`'s `Drop`. + unsafe { Box::from_raw(std::ptr::slice_from_raw_parts_mut(ptr, len)) } + } +} -/// Allow matrix views to be converted directly to slices. -impl<'a, T> From> for &'a [T] { - fn from(view: MatrixView<'a, T>) -> Self { - view.data +// A dense view converts directly to its backing slice. +impl<'a, T> From>> for &'a [T] { + fn from(m: MatRef<'a, RowMajor>) -> Self { + m.as_slice() } } -/// Allow mutable matrix views to be converted directly to slices. -impl<'a, T> From> for &'a [T] { - fn from(view: MutMatrixView<'a, T>) -> Self { - view.data +// A dense mutable view converts directly to its backing slice. +impl<'a, T> From>> for &'a mut [T] { + fn from(m: MatMut<'a, RowMajor>) -> Self { + let len = m.repr.num_elements(); + // SAFETY: `m: MatMut<'a, RowMajor>` exclusively views `len` contiguous `T` valid for `'a`. + unsafe { std::slice::from_raw_parts_mut(m.ptr.as_ptr().cast::(), len) } } } -/// Return a reference to the item at entry `(row, col)` in the matrix. -/// -/// # Panics +////////// +// Rows // +////////// + +/// Iterator over immutable row references of a matrix. /// -/// Panics if `row >= self.nrows()` or `col >= self.ncols()`. -impl Index<(usize, usize)> for MatrixBase +/// Created by [`Mat::rows`], [`MatRef::rows`], or [`MatMut::rows`]. +#[derive(Debug)] +pub struct Rows<'a, T: Repr> { + matrix: MatRef<'a, T>, + current: usize, +} + +impl<'a, T> Rows<'a, T> where - T: DenseData, + T: Repr, { - 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) } + fn new(matrix: MatRef<'a, T>) -> Self { + Self { matrix, current: 0 } } } -/// 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 MatrixBase +impl<'a, T> Iterator for Rows<'a, T> where - T: MutDenseData, + T: Repr + 'a, { - 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() - ); + type Item = T::Row<'a>; - // SAFETY: We have checked that `row` and `col` are in-bounds. - unsafe { self.get_unchecked_mut(row, col) } + fn next(&mut self) -> Option { + let current = self.current; + if current >= self.matrix.num_vectors() { + None + } else { + self.current += 1; + // SAFETY: We make sure through the above check that + // the access is within bounds. + // + // Extending the lifetime to `'a` is safe because the underlying + // MatRef has lifetime `'a`. + Some(unsafe { self.matrix.repr.get_row(self.matrix.ptr, current) }) + } + } + + fn size_hint(&self) -> (usize, Option) { + let remaining = self.matrix.num_vectors() - self.current; + (remaining, Some(remaining)) } } -/////////// -// Tests // -/////////// +impl<'a, T> ExactSizeIterator for Rows<'a, T> where T: Repr + 'a {} +impl<'a, T> FusedIterator for Rows<'a, T> where T: Repr + 'a {} -#[cfg(test)] -mod tests { - use super::*; - use crate::lazy_format; +///////////// +// RowsMut // +///////////// - /// This function is only callable with copyable types. - /// - /// This lets us test for types we expect to be `Copy`. - fn is_copyable(_x: T) -> bool { - true - } - - /// Test the that provided representation yields a slice with the expected base pointer - /// and length. - fn test_dense_data_repr( - ptr: *const T, - len: usize, - repr: Repr, - context: &dyn std::fmt::Display, - ) where - T: Copy, - Repr: DenseData, - { - let retrieved = repr.as_slice(); - assert_eq!(retrieved.len(), len, "{}", context); - assert_eq!(retrieved.as_ptr(), ptr, "{}", context); - } +/// Iterator over mutable row references of a matrix. +/// +/// Created by [`Mat::rows_mut`] or [`MatMut::rows_mut`]. +#[derive(Debug)] +pub struct RowsMut<'a, T: ReprMut> { + matrix: MatMut<'a, T>, + current: usize, +} - /// Set the underlying data for the provided representation to the following: - /// - /// [base, base + increment, base + increment + increment, ...] - fn set_mut_dense_data_repr(repr: &mut Repr, base: T, increment: T) - where - T: Copy + std::ops::Add, - Repr: DenseData + MutDenseData, - { - let slice = repr.as_mut_slice(); - for i in 0..slice.len() { - if i == 0 { - slice[i] = base; - } else { - slice[i] = slice[i - 1] + increment; - } - } +impl<'a, T> RowsMut<'a, T> +where + T: ReprMut, +{ + fn new(matrix: MatMut<'a, T>) -> Self { + Self { matrix, current: 0 } } +} - #[test] - fn slice_implements_dense_data_repr() { - for len in 0..10 { - let context = lazy_format!("len = {}", len); - let data: Vec = vec![0.0; len]; - let slice = data.as_slice(); - test_dense_data_repr(slice.as_ptr(), slice.len(), slice, &context); +impl<'a, T> Iterator for RowsMut<'a, T> +where + T: ReprMut + 'a, +{ + type Item = T::RowMut<'a>; + + fn next(&mut self) -> Option { + let current = self.current; + if current >= self.matrix.num_vectors() { + None + } else { + self.current += 1; + // SAFETY: We make sure through the above check that + // the access is within bounds. + // + // Extending the lifetime to `'a` is safe because: + // 1. The underlying MatMut has lifetime `'a`. + // 2. The iterator ensures that the mutable row indices are disjoint, so + // there is no aliasing as long as the implementation of `ReprMut` ensures + // there is not mutable sharing of the `RowMut` types. + Some(unsafe { self.matrix.repr.get_row_mut(self.matrix.ptr, current) }) } } - #[test] - fn mut_slice_mplements_dense_data_repr() { - for len in 0..10 { - let context = lazy_format!("len = {}", len); - let mut data: Vec = vec![0.0; len]; - let slice = data.as_mut_slice(); - - let ptr = slice.as_ptr(); - let len = slice.len(); - test_dense_data_repr(ptr, len, slice, &context); - } + fn size_hint(&self) -> (usize, Option) { + let remaining = self.matrix.num_vectors() - self.current; + (remaining, Some(remaining)) } +} - #[test] - fn mut_slice_implements_mut_dense_data_repr() { - for len in 0..10 { - let context = lazy_format!("len = {}", len); - let mut data: Vec = vec![0.0; len]; - let mut slice = data.as_mut_slice(); - - let base = 2.0; - let increment = 1.0; - set_mut_dense_data_repr(&mut slice, base, increment); - - for (i, &v) in slice.iter().enumerate() { - let context = lazy_format!("entry {}, {}", i, context); - assert_eq!(v, base + increment * (i as f32), "{}", context); - } - } +impl<'a, T> ExactSizeIterator for RowsMut<'a, T> where T: ReprMut + 'a {} +impl<'a, T> FusedIterator for RowsMut<'a, T> where T: ReprMut + 'a {} + +/////////////// +// DenseData // +/////////////// + +/// Abstraction over a type that can yield a dense slice of its contents. +/// +/// # Safety +/// +/// `as_slice` must be idempotent: it must **always** return the same slice with the same +/// length (unsafe code relies on this). +pub unsafe trait DenseData { + type Elem; + + /// Return the underlying data as a slice. + fn as_slice(&self) -> &[Self::Elem]; +} + +/// A mutable companion to [`DenseData`]. +/// +/// # Safety +/// +/// `as_mut_slice` must be idempotent and must span the exact same memory as +/// [`DenseData::as_slice`]. +pub unsafe trait MutDenseData: DenseData { + fn as_mut_slice(&mut self) -> &mut [Self::Elem]; +} + +// SAFETY: fulfills the idempotency requirement. +unsafe impl DenseData for &[T] { + type Elem = T; + fn as_slice(&self) -> &[Self::Elem] { + self } +} - ///////////////// - // Matrix View // - ///////////////// +// SAFETY: fulfills the idempotency requirement. +unsafe impl DenseData for &mut [T] { + type Elem = T; + fn as_slice(&self) -> &[Self::Elem] { + self + } +} - #[test] - fn try_from_error_misc() { - let x = TryFromError::<&[f32]> { - data: &[], - nrows: 1, - ncols: 2, - }; +// SAFETY: fulfills the idempotency requirement and spans the same memory as `as_slice`. +unsafe impl MutDenseData for &mut [T] { + fn as_mut_slice(&mut self) -> &mut [Self::Elem] { + self + } +} - 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")); +// SAFETY: fulfills the idempotency requirement. +unsafe impl DenseData for Box<[T]> { + type Elem = T; + fn as_slice(&self) -> &[Self::Elem] { + self } +} - fn make_test_matrix() -> Vec { - // Construct a matrix with 4 rows of length 3. - // The expected layout is as follows: - // - // 0, 1, 2, - // 1, 2, 3, - // 2, 3, 4, - // 3, 4, 5 - // - vec![0, 1, 2, 1, 2, 3, 2, 3, 4, 3, 4, 5] +// SAFETY: fulfills the idempotency requirement and spans the same memory as `as_slice`. +unsafe impl MutDenseData for Box<[T]> { + fn as_mut_slice(&mut self) -> &mut [Self::Elem] { + self } +} - #[cfg(feature = "rayon")] - fn test_basic_indexing_parallel(m: MatrixView<'_, usize>) { - // Par window iters. - let batchsize = 2; - m.par_window_iter(batchsize) - .enumerate() - .for_each(|(i, submatrix)| { - assert_eq!(submatrix.nrows(), batchsize); - assert_eq!(submatrix.ncols(), m.ncols()); - - // Make sure we are in the correct window of the original matrix. - let base = i * batchsize; - assert_eq!(submatrix[(0, 0)], base); - assert_eq!(submatrix[(0, 1)], base + 1); - assert_eq!(submatrix[(0, 2)], base + 2); - - assert_eq!(submatrix[(1, 0)], base + 1); - assert_eq!(submatrix[(1, 1)], base + 2); - assert_eq!(submatrix[(1, 2)], base + 3); - }); +/////////// +// Tests // +/////////// - // Try again, but with a batch size of 3 to ensure that we correctly handle cases - // where the last block is under-sized. - let batchsize = 3; - m.par_window_iter(batchsize) - .enumerate() - .for_each(|(i, submatrix)| { - if i == 0 { - assert_eq!(submatrix.nrows(), batchsize); - assert_eq!(submatrix.ncols(), m.ncols()); - - // Check indexing - assert_eq!(submatrix[(0, 0)], 0); - assert_eq!(submatrix[(0, 1)], 1); - assert_eq!(submatrix[(0, 2)], 2); - - assert_eq!(submatrix[(1, 0)], 1); - assert_eq!(submatrix[(1, 1)], 2); - assert_eq!(submatrix[(1, 2)], 3); - - assert_eq!(submatrix[(2, 0)], 2); - assert_eq!(submatrix[(2, 1)], 3); - assert_eq!(submatrix[(2, 2)], 4); - } else { - assert_eq!(submatrix.nrows(), 1); - assert_eq!(submatrix.ncols(), m.ncols()); - - // Check indexing - assert_eq!(submatrix[(0, 0)], 3); - assert_eq!(submatrix[(0, 1)], 4); - assert_eq!(submatrix[(0, 2)], 5); - } - }); +#[cfg(test)] +mod tests { + use super::*; - // par-row-iter - let seen_rows: Box<[usize]> = m - .par_row_iter() - .enumerate() - .map(|(i, row)| { - let expected: Box<[usize]> = (0..m.ncols()).map(|j| j + i).collect(); - assert_eq!(row, &*expected); - i - }) - .collect(); + use std::fmt::Display; - let expected: Box<[usize]> = (0..m.nrows()).collect(); - assert_eq!(seen_rows, expected); - } + use crate::lazy_format; - fn test_basic_indexing(m: &MatrixBase) - where - T: DenseData + Sync, - { - assert_eq!(m.nrows(), 4); - assert_eq!(m.ncols(), 3); - - // Basic indexing - assert_eq!(m[(0, 0)], 0); - assert_eq!(m[(0, 1)], 1); - assert_eq!(m[(0, 2)], 2); - - assert_eq!(m[(1, 0)], 1); - assert_eq!(m[(1, 1)], 2); - assert_eq!(m[(1, 2)], 3); - - assert_eq!(m[(2, 0)], 2); - assert_eq!(m[(2, 1)], 3); - assert_eq!(m[(2, 2)], 4); - - assert_eq!(m[(3, 0)], 3); - assert_eq!(m[(3, 1)], 4); - assert_eq!(m[(3, 2)], 5); - - // Row indexing. - assert_eq!(m.row(0), &[0, 1, 2]); - assert_eq!(m.row(1), &[1, 2, 3]); - assert_eq!(m.row(2), &[2, 3, 4]); - assert_eq!(m.row(3), &[3, 4, 5]); - - let rows: Vec> = m.row_iter().map(|x| x.to_vec()).collect(); - assert_eq!(m.row(0), &rows[0]); - assert_eq!(m.row(1), &rows[1]); - assert_eq!(m.row(2), &rows[2]); - assert_eq!(m.row(3), &rows[3]); - - // Window Iters. - let batchsize = 2; - m.window_iter(batchsize) - .enumerate() - .for_each(|(i, submatrix)| { - assert_eq!(submatrix.nrows(), batchsize); - assert_eq!(submatrix.ncols(), m.ncols()); - - // Make sure we are in the correct window of the original matrix. - let base = i * batchsize; - assert_eq!(submatrix[(0, 0)], base); - assert_eq!(submatrix[(0, 1)], base + 1); - assert_eq!(submatrix[(0, 2)], base + 2); - - assert_eq!(submatrix[(1, 0)], base + 1); - assert_eq!(submatrix[(1, 1)], base + 2); - assert_eq!(submatrix[(1, 2)], base + 3); - }); + /// Helper to assert a type is Copy. + fn assert_copy(_: &T) {} - // Try again, but with a batch size of 3 to ensure that we correctly handle cases - // where the last block is under-sized. - let batchsize = 3; - m.window_iter(batchsize) - .enumerate() - .for_each(|(i, submatrix)| { - if i == 0 { - assert_eq!(submatrix.nrows(), batchsize); - assert_eq!(submatrix.ncols(), m.ncols()); - - // Check indexing - assert_eq!(submatrix[(0, 0)], 0); - assert_eq!(submatrix[(0, 1)], 1); - assert_eq!(submatrix[(0, 2)], 2); - - assert_eq!(submatrix[(1, 0)], 1); - assert_eq!(submatrix[(1, 1)], 2); - assert_eq!(submatrix[(1, 2)], 3); - - assert_eq!(submatrix[(2, 0)], 2); - assert_eq!(submatrix[(2, 1)], 3); - assert_eq!(submatrix[(2, 2)], 4); - } else { - assert_eq!(submatrix.nrows(), 1); - assert_eq!(submatrix.ncols(), m.ncols()); - - // Check indexing - assert_eq!(submatrix[(0, 0)], 3); - assert_eq!(submatrix[(0, 1)], 4); - assert_eq!(submatrix[(0, 2)], 5); - } - }); + // ── Variance assertions ────────────────────────────────────── + // + // These functions are never called. The test is that they compile: + // covariant positions must accept subtype coercions. + // + // The negative (invariance) counterparts live in + // `tests/compile-fail/multi/{mat,matmut}_invariant.rs`. - #[cfg(all(not(miri), feature = "rayon"))] - test_basic_indexing_parallel(m.as_view()); + /// `MatRef` is covariant in `'a`: a longer borrow can shorten. + fn _assert_matref_covariant_lifetime<'long: 'short, 'short, T: Repr>( + v: MatRef<'long, T>, + ) -> MatRef<'short, T> { + v } - #[test] - fn matrix_happy_path() { - let data = make_test_matrix(); - let m = Matrix::try_from(data.into(), 4, 3).unwrap(); - test_basic_indexing(&m); - - // Get the base pointer of the matrix and make sure view-conversion preserves this - // value. - let ptr = m.as_ptr(); - let view = m.as_view(); - assert!(is_copyable(view)); - assert_eq!(view.as_ptr(), ptr); - assert_eq!(view.nrows(), m.nrows()); - assert_eq!(view.ncols(), m.ncols()); - test_basic_indexing(&view); + /// `MatRef` is covariant in `T`: `RowMajor<&'long u8>` → `RowMajor<&'short u8>`. + fn _assert_matref_covariant_repr<'long: 'short, 'short, 'a>( + v: MatRef<'a, RowMajor<&'long u8>>, + ) -> MatRef<'a, RowMajor<&'short u8>> { + v } - #[test] - fn matrix_try_from_construction_error() { - let data = make_test_matrix(); - let ptr = data.as_ptr(); - let len = data.len(); - - let m = Matrix::try_from(data.into(), 5, 4); - assert!(m.is_err()); - let err = m.unwrap_err(); - assert_eq!( - err.to_string(), - "tried to construct a matrix view with 5 rows and 4 columns over a slice of length 12" - ); - - // Make sure that we can retrieve the original allocation from the interior. - let data = err.into_inner(); - assert_eq!(data.as_ptr(), ptr); - assert_eq!(data.len(), len); - - let m = MatrixView::try_from(&data, 5, 4); - assert!(m.is_err()); - assert_eq!( - m.unwrap_err().to_string(), - "tried to construct a matrix view with 5 rows and 4 columns over a slice of length 12" - ); + /// `MatMut` is covariant in `'a`: a longer borrow can shorten. + fn _assert_matmut_covariant_lifetime<'long: 'short, 'short, T: ReprMut>( + v: MatMut<'long, T>, + ) -> MatMut<'short, T> { + v } - #[test] - fn matrix_mut_view() { - let mut m = Matrix::::new(0, 4, 3); - assert_eq!(m.nrows(), 4); - assert_eq!(m.ncols(), 3); - assert!(m.as_slice().iter().all(|&i| i == 0)); - let ptr = m.as_ptr(); - let mut_ptr = m.as_mut_ptr(); - assert_eq!(ptr, mut_ptr); - - let mut view = m.as_mut_view(); - assert_eq!(view.nrows(), 4); - assert_eq!(view.ncols(), 3); - assert_eq!(view.as_ptr(), ptr); - assert_eq!(view.as_mut_ptr(), mut_ptr); + fn edge_cases(nrows: usize) -> Vec { + let max = usize::MAX; - // Construct the test matrix manually. - for i in 0..view.nrows() { - for j in 0..view.ncols() { - view[(i, j)] = i + j; - } + vec![ + nrows, + nrows + 1, + nrows + 11, + nrows + 20, + max / 2, + max.div_ceil(2), + max - 1, + max, + ] + } + + fn fill_mat(x: &mut Mat>, repr: RowMajor) { + assert_eq!(x.repr(), &repr); + assert_eq!(x.num_vectors(), repr.nrows()); + assert_eq!(x.vector_dim(), repr.ncols()); + + for i in 0..x.num_vectors() { + let row = x.get_row_mut(i).unwrap(); + assert_eq!(row.len(), repr.ncols()); + row.iter_mut() + .enumerate() + .for_each(|(j, r)| *r = 10 * i + j); } - // Drop the view and test the original matrix. - test_basic_indexing(&m); - - let inner = m.into_inner(); - assert_eq!(inner.as_ptr(), ptr); - assert_eq!(inner.len(), 4 * 3); + for i in edge_cases(repr.nrows()).into_iter() { + assert!(x.get_row_mut(i).is_none()); + } } - #[test] - fn matrix_view_zero_sizes() { - let data: Vec = vec![]; - // Zero rows, but non-zero columns. - let m = MatrixView::try_from(data.as_slice(), 0, 10).unwrap(); - assert_eq!(m.nrows(), 0); - assert_eq!(m.ncols(), 10); - - // Non-zero rows, but zero columns. - let m = MatrixView::try_from(data.as_slice(), 3, 0).unwrap(); - assert_eq!(m.nrows(), 3); - assert_eq!(m.ncols(), 0); - let empty: &[usize] = &[]; - assert_eq!(m.row(0), empty); - assert_eq!(m.row(1), empty); - assert_eq!(m.row(2), empty); - - // Zero rows and columns. - let m = MatrixView::try_from(data.as_slice(), 0, 0).unwrap(); - assert_eq!(m.nrows(), 0); - assert_eq!(m.ncols(), 0); - } + fn fill_mat_mut(mut x: MatMut<'_, RowMajor>, repr: RowMajor) { + assert_eq!(x.repr(), &repr); + assert_eq!(x.num_vectors(), repr.nrows()); + assert_eq!(x.vector_dim(), repr.ncols()); - #[test] - fn matrix_view_construction_elementwise() { - let mut m = Matrix::::new(0, 4, 3); + for i in 0..x.num_vectors() { + let row = x.get_row_mut(i).unwrap(); + assert_eq!(row.len(), repr.ncols()); - // Construct the test matrix manually. - for i in 0..m.nrows() { - for j in 0..m.ncols() { - m[(i, j)] = i + j; - } + row.iter_mut() + .enumerate() + .for_each(|(j, r)| *r = 10 * i + j); } - test_basic_indexing(&m); - } - #[test] - fn matrix_construction_by_row() { - let mut m = Matrix::::new(0, 4, 3); - assert!(m.as_slice().iter().all(|i| *i == 0)); - - let ncols = m.ncols(); - for i in 0..m.nrows() { - let row = m.row_mut(i); - assert_eq!(row.len(), ncols); - row[0] = i; - row[1] = i + 1; - row[2] = i + 2; + for i in edge_cases(repr.nrows()).into_iter() { + assert!(x.get_row_mut(i).is_none()); } - test_basic_indexing(&m); - } - - #[test] - fn matrix_construction_by_rowiter() { - let mut m = Matrix::::new(0, 4, 3); - assert!(m.as_slice().iter().all(|i| *i == 0)); - - let ncols = m.ncols(); - m.row_iter_mut().enumerate().for_each(|(i, row)| { - assert_eq!(row.len(), ncols); - row[0] = i; - row[1] = i + 1; - row[2] = i + 2; - }); - test_basic_indexing(&m); } - #[cfg(all(not(miri), feature = "rayon"))] - #[test] - fn matrix_construction_by_par_windows() { - let mut m = Matrix::::new(0, 4, 3); - assert!(m.as_slice().iter().all(|i| *i == 0)); - - let ncols = m.ncols(); - for batchsize in 1..=4 { - m.par_window_iter_mut(batchsize) + fn fill_rows_mut(x: RowsMut<'_, RowMajor>, repr: RowMajor) { + assert_eq!(x.len(), repr.nrows()); + // Materialize all rows at once. + let mut all_rows: Vec<_> = x.collect(); + assert_eq!(all_rows.len(), repr.nrows()); + for (i, row) in all_rows.iter_mut().enumerate() { + assert_eq!(row.len(), repr.ncols()); + row.iter_mut() .enumerate() - .for_each(|(i, mut submatrix)| { - let base = i * batchsize; - submatrix.row_iter_mut().enumerate().for_each(|(j, row)| { - assert_eq!(row.len(), ncols); - row[0] = base + j; - row[1] = base + j + 1; - row[2] = base + j + 2; - }); - }); - test_basic_indexing(&m); + .for_each(|(j, r)| *r = 10 * i + j); } } - #[test] - fn matrix_construction_happens_in_memory_order() { - let mut i = 0; - let ncols = 3; - let initializer = Init(|| { - let value = (i % ncols) + (i / ncols); - i += 1; - value - }); + fn check_mat(x: &Mat>, repr: RowMajor, ctx: &dyn Display) { + assert_eq!(x.repr(), &repr); + assert_eq!(x.num_vectors(), repr.nrows()); + assert_eq!(x.vector_dim(), repr.ncols()); - let m = Matrix::new(initializer, 4, 3); - test_basic_indexing(&m); - } + for i in 0..x.num_vectors() { + let row = x.get_row(i).unwrap(); - // Panics - #[test] - #[should_panic(expected = "tried to access row 3 of a matrix with 3 rows")] - fn test_get_row_panics() { - let m = Matrix::::new(0, 3, 7); - m.row(3); - } + assert_eq!(row.len(), repr.ncols(), "ctx: {ctx}"); + row.iter().enumerate().for_each(|(j, r)| { + assert_eq!( + *r, + 10 * i + j, + "mismatched entry at row {}, col {} -- ctx: {}", + i, + j, + ctx + ) + }); + } - #[test] - #[should_panic(expected = "tried to access row 3 of a matrix with 3 rows")] - fn test_get_row_mut_panics() { - let mut m = Matrix::::new(0, 3, 7); - m.row_mut(3); + for i in edge_cases(repr.nrows()).into_iter() { + assert!(x.get_row(i).is_none(), "ctx: {ctx}"); + } } - #[test] - #[should_panic(expected = "row 3 is out of bounds (max: 3)")] - fn test_index_panics_row() { - let m = Matrix::::new(0, 3, 7); - assert!(m.try_get(3, 2).is_none()); - let _ = m[(3, 2)]; - } + fn check_mat_ref(x: MatRef<'_, RowMajor>, repr: RowMajor, ctx: &dyn Display) { + assert_eq!(x.repr(), &repr); + assert_eq!(x.num_vectors(), repr.nrows()); + assert_eq!(x.vector_dim(), repr.ncols()); - #[test] - #[should_panic(expected = "col 7 is out of bounds (max: 7)")] - fn test_index_panics_col() { - let m = Matrix::::new(0, 3, 7); - assert!(m.try_get(2, 7).is_none()); - let _ = m[(2, 7)]; - } + assert_copy(&x); + for i in 0..x.num_vectors() { + let row = x.get_row(i).unwrap(); + assert_eq!(row.len(), repr.ncols(), "ctx: {ctx}"); - #[test] - #[should_panic(expected = "row 3 is out of bounds (max: 3)")] - fn test_index_mut_panics_row() { - let mut m = Matrix::::new(0, 3, 7); - m[(3, 2)] = 1; - } + row.iter().enumerate().for_each(|(j, r)| { + assert_eq!( + *r, + 10 * i + j, + "mismatched entry at row {}, col {} -- ctx: {}", + i, + j, + ctx + ) + }); + } - #[test] - #[should_panic(expected = "col 7 is out of bounds (max: 7)")] - fn test_index_mut_panics_col() { - let mut m = Matrix::::new(0, 3, 7); - m[(2, 7)] = 1; + for i in edge_cases(repr.nrows()).into_iter() { + assert!(x.get_row(i).is_none(), "ctx: {ctx}"); + } } - #[test] - #[cfg(feature = "rayon")] - #[should_panic(expected = "par_window_iter batchsize cannot be zero")] - fn test_par_window_iter_panics() { - let m = Matrix::::new(0, 4, 4); - let _ = m.par_window_iter(0); - } + fn check_mat_mut(x: MatMut<'_, RowMajor>, repr: RowMajor, ctx: &dyn Display) { + assert_eq!(x.repr(), &repr); + assert_eq!(x.num_vectors(), repr.nrows()); + assert_eq!(x.vector_dim(), repr.ncols()); - #[test] - #[cfg(feature = "rayon")] - #[should_panic(expected = "par_window_iter_mut batchsize cannot be zero")] - fn test_par_window_iter_mut_panics() { - let mut m = Matrix::::new(0, 4, 4); - let _ = m.par_window_iter_mut(0); - } + for i in 0..x.num_vectors() { + let row = x.get_row(i).unwrap(); + assert_eq!(row.len(), repr.ncols(), "ctx: {ctx}"); - // Additional tests for better coverage + row.iter().enumerate().for_each(|(j, r)| { + assert_eq!( + *r, + 10 * i + j, + "mismatched entry at row {}, col {} -- ctx: {}", + i, + j, + ctx + ) + }); + } - #[test] - fn test_box_slice_dense_data_impls() { - // Test Box<[T]> implementations - let data: Box<[f32]> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0].into(); - let ptr = data.as_ptr(); - let len = data.len(); - - // Test DenseData impl for Box<[T]> - test_dense_data_repr(ptr, len, data, &lazy_format!("Box<[T]> DenseData")); - - // Test MutDenseData impl for Box<[T]> - let mut data: Box<[f32]> = vec![0.0; 6].into(); - set_mut_dense_data_repr(&mut data, 1.0, 2.0); - for (i, &v) in data.iter().enumerate() { - assert_eq!( - v, - 1.0 + 2.0 * (i as f32), - "Box<[T]> MutDenseData at index {}", - i - ); + for i in edge_cases(repr.nrows()).into_iter() { + assert!(x.get_row(i).is_none(), "ctx: {ctx}"); } } - #[test] - fn test_try_from_error_light() { - let data = vec![1, 2, 3]; - let err = MatrixView::try_from(data.as_slice(), 2, 3).unwrap_err(); - - // Test as_static method - let static_err = err.as_static(); - assert_eq!(static_err.len, 3); - assert_eq!(static_err.nrows, 2); - assert_eq!(static_err.ncols, 3); - - // Test Display for TryFromErrorLight - let display_msg = format!("{}", static_err); - assert!(display_msg.contains("tried to construct a matrix view with 2 rows and 3 columns")); - assert!(display_msg.contains("slice of length 3")); - - // Test into_inner method - let recovered_data = err.into_inner(); - assert_eq!(recovered_data, data.as_slice()); + fn check_rows(x: Rows<'_, RowMajor>, repr: RowMajor, ctx: &dyn Display) { + assert_eq!(x.len(), repr.nrows(), "ctx: {ctx}"); + let all_rows: Vec<_> = x.collect(); + assert_eq!(all_rows.len(), repr.nrows(), "ctx: {ctx}"); + for (i, row) in all_rows.iter().enumerate() { + assert_eq!(row.len(), repr.ncols(), "ctx: {ctx}"); + row.iter().enumerate().for_each(|(j, r)| { + assert_eq!( + *r, + 10 * i + j, + "mismatched entry at row {}, col {} -- ctx: {}", + i, + j, + ctx + ) + }); + } } - #[test] - fn test_get_row_optional() { - let data = make_test_matrix(); - let m = MatrixView::try_from(data.as_slice(), 4, 3).unwrap(); + ////////////// + // RowMajor // + ////////////// - // Test successful get_row - assert_eq!(m.get_row(0), Some(&[0, 1, 2][..])); - assert_eq!(m.get_row(1), Some(&[1, 2, 3][..])); - assert_eq!(m.get_row(3), Some(&[3, 4, 5][..])); + #[test] + fn standard_representation() { + let repr = RowMajor::::new(4, 3).unwrap(); + assert_eq!(repr.nrows(), 4); + assert_eq!(repr.ncols(), 3); - // Test out-of-bounds get_row - assert_eq!(m.get_row(4), None); - assert_eq!(m.get_row(100), None); + let layout = repr.layout().unwrap(); + assert_eq!(layout.size(), 4 * 3 * std::mem::size_of::()); + assert_eq!(layout.align(), std::mem::align_of::()); } #[test] - fn test_unsafe_get_unchecked_methods() { - let data = make_test_matrix(); - let mut m = Matrix::try_from(data.into(), 4, 3).unwrap(); - - // Safety: derives from known size of matrix and access element ids - unsafe { - assert_eq!(*m.get_unchecked(0, 0), 0); - assert_eq!(*m.get_unchecked(1, 2), 3); - assert_eq!(*m.get_unchecked(3, 1), 4); - } - - // Safety: derives from known size of matrix and access element ids - unsafe { - *m.get_unchecked_mut(0, 0) = 100; - *m.get_unchecked_mut(1, 2) = 200; - } - - assert_eq!(m[(0, 0)], 100); - assert_eq!(m[(1, 2)], 200); - - // Safety: derives from known size of matrix and access element ids - unsafe { - let row0 = m.get_row_unchecked(0); - assert_eq!(row0[0], 100); - assert_eq!(row0[1], 1); - assert_eq!(row0[2], 2); - } - - // Safety: derives from known size of matrix and access element ids - unsafe { - let row1 = m.get_row_unchecked_mut(1); - row1[0] = 300; + fn standard_zero_dimensions() { + for (nrows, ncols) in [(0, 0), (0, 5), (5, 0)] { + let repr = RowMajor::::new(nrows, ncols).unwrap(); + assert_eq!(repr.nrows(), nrows); + assert_eq!(repr.ncols(), ncols); + let layout = repr.layout().unwrap(); + assert_eq!(layout.size(), 0); } - - assert_eq!(m[(1, 0)], 300); } #[test] - fn test_to_owned() { - let data = make_test_matrix(); - let view = MatrixView::try_from(data.as_slice(), 4, 3).unwrap(); - - // Test to_owned creates a proper clone - let owned = view.to_owned(); - assert_eq!(owned.nrows(), view.nrows()); - assert_eq!(owned.ncols(), view.ncols()); - assert_eq!(owned.as_slice(), view.as_slice()); - - // Verify it's actually owned (different memory location) - assert_ne!(owned.as_ptr(), view.as_ptr()); + fn standard_check_slice() { + let repr = RowMajor::::new(3, 4).unwrap(); + + // Correct length succeeds + let data = vec![0u32; 12]; + assert!(repr.check_slice(&data).is_ok()); + + // Too short fails + let short = vec![0u32; 11]; + assert!(matches!( + repr.check_slice(&short), + Err(SliceError::LengthMismatch { + expected: 12, + found: 11 + }) + )); + + // Too long fails + let long = vec![0u32; 13]; + assert!(matches!( + repr.check_slice(&long), + Err(SliceError::LengthMismatch { + expected: 12, + found: 13 + }) + )); - // Test the owned matrix works properly - test_basic_indexing(&owned); + // Overflow case + let overflow_repr = RowMajor::::new(usize::MAX, 2).unwrap_err(); + assert!(matches!(overflow_repr, Overflow { .. })); } #[test] - fn test_generator_trait_impls() { - // Test Generator impl for T where T: Clone - let mut gen = 42i32; - assert_eq!(gen.generate(), 42); - assert_eq!(gen.generate(), 42); // Should be same value since it's cloned - - // Test Generator impl for Init - let mut counter = 0; - let mut gen = Init(|| { - counter += 1; - counter - }); - assert_eq!(gen.generate(), 1); - assert_eq!(gen.generate(), 2); - assert_eq!(gen.generate(), 3); + fn standard_new_rejects_element_count_overflow() { + // nrows * ncols overflows usize even though per-element size is small. + assert!(RowMajor::::new(usize::MAX, 2).is_err()); + assert!(RowMajor::::new(2, usize::MAX).is_err()); + assert!(RowMajor::::new(usize::MAX, usize::MAX).is_err()); } #[test] - fn test_matrix_from_conversions() { - let data = make_test_matrix(); - let m = Matrix::try_from(data.into(), 4, 3).unwrap(); - - // Test MatrixView to slice conversion - let view = m.as_view(); - let slice: &[usize] = view.into(); - assert_eq!(slice.len(), 12); - assert_eq!(slice[0], 0); - assert_eq!(slice[11], 5); - - // Test MutMatrixView to slice conversion - let data2 = make_test_matrix(); - let mut m2 = Matrix::try_from(data2.into(), 4, 3).unwrap(); - let mut_view = m2.as_mut_view(); - let slice2: &[usize] = mut_view.into(); - assert_eq!(slice2.len(), 12); - assert_eq!(slice2[0], 0); - assert_eq!(slice2[11], 5); + fn standard_new_rejects_byte_count_exceeding_isize_max() { + // Element count fits in usize, but total bytes exceed isize::MAX. + let half = (isize::MAX as usize / std::mem::size_of::()) + 1; + assert!(RowMajor::::new(half, 1).is_err()); + assert!(RowMajor::::new(1, half).is_err()); } #[test] - fn test_matrix_construction_edge_cases() { - // Test 1x1 matrix - let m = Matrix::new(42, 1, 1); - assert_eq!(m.nrows(), 1); - assert_eq!(m.ncols(), 1); - assert_eq!(m[(0, 0)], 42); - assert_eq!(*m.try_get(0, 0).unwrap(), 42); - - // Test single row matrix - let m = Matrix::new(7, 1, 5); - assert_eq!(m.nrows(), 1); - assert_eq!(m.ncols(), 5); - assert!(m.as_slice().iter().all(|&x| x == 7)); - - // Test single column matrix - let m = Matrix::new(9, 5, 1); - assert_eq!(m.nrows(), 5); - assert_eq!(m.ncols(), 1); - assert!(m.as_slice().iter().all(|&x| x == 9)); + fn standard_new_accepts_boundary_below_isize_max() { + // Largest allocation that still fits in isize::MAX bytes. + let max_elems = isize::MAX as usize / std::mem::size_of::(); + let repr = RowMajor::::new(max_elems, 1).unwrap(); + assert_eq!(repr.num_elements(), max_elems); } #[test] - fn test_matrix_view_edge_cases_with_data() { - // Test matrix with actual data for edge cases - let data = vec![10, 20]; - - // 2x1 matrix - let m = MatrixView::try_from(data.as_slice(), 2, 1).unwrap(); - assert_eq!(m.nrows(), 2); - assert_eq!(m.ncols(), 1); - assert_eq!(m[(0, 0)], 10); - assert_eq!(m[(1, 0)], 20); - assert_eq!(*m.try_get(0, 0).unwrap(), 10); - assert_eq!(*m.try_get(1, 0).unwrap(), 20); - assert_eq!(m.row(0), &[10]); - assert_eq!(m.row(1), &[20]); - - // 1x2 matrix - let m = MatrixView::try_from(data.as_slice(), 1, 2).unwrap(); - assert_eq!(m.nrows(), 1); - assert_eq!(m.ncols(), 2); - assert_eq!(m[(0, 0)], 10); - assert_eq!(m[(0, 1)], 20); - assert_eq!(*m.try_get(0, 0).unwrap(), 10); - assert_eq!(*m.try_get(0, 1).unwrap(), 20); - assert_eq!(m.row(0), &[10, 20]); + fn standard_new_zst_rejects_element_count_overflow() { + // For ZSTs the byte count is always 0, but element-count overflow + // must still be caught so that `num_elements()` never wraps. + assert!(RowMajor::<()>::new(usize::MAX, 2).is_err()); + assert!(RowMajor::<()>::new(usize::MAX / 2 + 1, 3).is_err()); } #[test] - fn test_row_vector() { - let data = vec![1, 2, 3]; - let m = MatrixView::row_vector(data.as_slice()); - assert_eq!(m.nrows(), 1); - assert_eq!(m.ncols(), 3); - assert_eq!(m.as_slice(), &[1, 2, 3]); - assert_eq!(m.row(0), &[1, 2, 3]); - - // Empty - let empty: &[i32] = &[]; - let m = MatrixView::row_vector(empty); - assert_eq!(m.nrows(), 1); - assert_eq!(m.ncols(), 0); - - // Owned - let m = Matrix::row_vector(vec![10u64, 20].into_boxed_slice()); - assert_eq!(m.nrows(), 1); - assert_eq!(m.ncols(), 2); - assert_eq!(m[(0, 0)], 10); - assert_eq!(m[(0, 1)], 20); + fn standard_new_zst_accepts_large_non_overflowing() { + // Large-but-valid ZST matrix: element count fits in usize. + let repr = RowMajor::<()>::new(usize::MAX, 1).unwrap(); + assert_eq!(repr.num_elements(), usize::MAX); + assert_eq!(repr.layout().unwrap().size(), 0); } #[test] - fn test_column_vector() { - let data = vec![1, 2, 3]; - let m = MatrixView::column_vector(data.as_slice()); - assert_eq!(m.nrows(), 3); - assert_eq!(m.ncols(), 1); - assert_eq!(m.as_slice(), &[1, 2, 3]); - assert_eq!(m[(0, 0)], 1); - assert_eq!(m[(1, 0)], 2); - assert_eq!(m[(2, 0)], 3); - assert_eq!(m.row(0), &[1]); - assert_eq!(m.row(1), &[2]); - assert_eq!(m.row(2), &[3]); - - // Empty - let empty: &[i32] = &[]; - let m = MatrixView::column_vector(empty); - assert_eq!(m.nrows(), 0); - assert_eq!(m.ncols(), 1); - - // Owned - let m = Matrix::column_vector(vec![10u64, 20].into_boxed_slice()); - assert_eq!(m.nrows(), 2); - assert_eq!(m.ncols(), 1); - assert_eq!(m[(0, 0)], 10); - assert_eq!(m[(1, 0)], 20); + fn standard_new_overflow_error_display() { + let err = RowMajor::::new(usize::MAX, 2).unwrap_err(); + let msg = err.to_string(); + assert!(msg.contains("would exceed isize::MAX bytes"), "{msg}"); + + let zst_err = RowMajor::<()>::new(usize::MAX, 2).unwrap_err(); + let zst_msg = zst_err.to_string(); + assert!(zst_msg.contains("ZST matrix"), "{zst_msg}"); + assert!(zst_msg.contains("usize::MAX"), "{zst_msg}"); } #[test] - fn test_map() { - let m = Matrix::try_from(vec![1u32, 2, 3, 4].into(), 2, 2).unwrap(); - let doubled = m.map(|&x| x * 2); - assert_eq!(doubled.as_slice(), &[2, 4, 6, 8]); - assert_eq!(doubled.nrows(), 2); - assert_eq!(doubled.ncols(), 2); - - // Type-changing map - let as_f64 = m.map(|&x| x as f64); - assert_eq!(as_f64.as_slice(), &[1.0, 2.0, 3.0, 4.0]); + fn try_from_error_recovers_owned_data() { + // A failed owned `try_from` must hand the `Box` back, not deallocate it. + let data: Box<[i32]> = vec![1, 2, 3, 4, 5, 6].into_boxed_slice(); + let err = Mat::try_from(data, 4, 4).unwrap_err(); + assert_eq!(&*err.into_inner(), &[1, 2, 3, 4, 5, 6]); } #[test] - fn test_try_get() { - let m = Matrix::try_from(vec![1, 2, 3, 4, 5, 6].into(), 2, 3).unwrap(); - assert_eq!(m.try_get(0, 0), Some(&1)); - assert_eq!(m.try_get(1, 2), Some(&6)); - assert_eq!(m.try_get(2, 0), None); - assert_eq!(m.try_get(0, 3), None); + fn eq_distinguishes_row_count_for_zero_columns() { + // With zero columns the backing slice is empty for any row count, so equality + // must compare the row count too: a 5x0 and a 3x0 matrix are not equal. + let a = Mat::from_repr(RowMajor::::new(5, 0).unwrap(), 0).unwrap(); + let b = Mat::from_repr(RowMajor::::new(3, 0).unwrap(), 0).unwrap(); + assert!(a != b); + assert!(a == Mat::from_repr(RowMajor::::new(5, 0).unwrap(), 0).unwrap()); } - #[test] - fn test_subview() { - let data = make_test_matrix(); - let m = Matrix::try_from(data.into(), 4, 3).unwrap(); - - // Create a subview of the first two rows - { - let subview = m.subview(0..4).unwrap(); - assert_eq!(subview.nrows(), 4); - assert_eq!(subview.ncols(), 3); - - assert_eq!(subview.row(0), &[0, 1, 2]); - assert_eq!(subview.row(1), &[1, 2, 3]); - assert_eq!(subview.row(2), &[2, 3, 4]); - assert_eq!(subview.row(3), &[3, 4, 5]); - assert!(subview.get_row(4).is_none()); - } + ///////// + // Mat // + ///////// - // Sub view over a subset that touches the end. - { - let subview = m.subview(1..4).unwrap(); - assert_eq!(subview.nrows(), 3); - assert_eq!(subview.ncols(), 3); + #[test] + fn mat_new_and_basic_accessors() { + let mat = Mat::from_repr(RowMajor::::new(3, 4).unwrap(), 42usize).unwrap(); + let base: *const u8 = mat.as_raw_ptr(); - assert_eq!(subview.row(0), &[1, 2, 3]); - assert_eq!(subview.row(1), &[2, 3, 4]); - assert_eq!(subview.row(2), &[3, 4, 5]); - assert!(subview.get_row(3).is_none()); - } + assert_eq!(mat.num_vectors(), 3); + assert_eq!(mat.vector_dim(), 4); - // Sub view over a subset that is in the middle - { - let subview = m.subview(1..3).unwrap(); - assert_eq!(subview.nrows(), 2); - assert_eq!(subview.ncols(), 3); + let repr = mat.repr(); + assert_eq!(repr.nrows(), 3); + assert_eq!(repr.ncols(), 4); - assert_eq!(subview.row(0), &[1, 2, 3]); - assert_eq!(subview.row(1), &[2, 3, 4]); - assert!(subview.get_row(2).is_none()); + for (i, r) in mat.rows().enumerate() { + assert_eq!(r, &[42, 42, 42, 42]); + let ptr = r.as_ptr().cast::(); + assert_eq!( + ptr, + base.wrapping_add(std::mem::size_of::() * mat.repr().ncols() * i), + ); } + } - // Empty sub-view. - { - let subview = m.subview(2..2).unwrap(); - assert_eq!(subview.nrows(), 0); - assert_eq!(subview.ncols(), 3); - } + #[test] + fn mat_new_with_default() { + let mat = Mat::from_default(RowMajor::::new(2, 3).unwrap()).unwrap(); + let base: *const u8 = mat.as_raw_ptr(); - // Empty subview in bounds - { - let subview = m.subview(0..0).unwrap(); - assert_eq!(subview.nrows(), 0); - assert_eq!(subview.ncols(), 3); + assert_eq!(mat.num_vectors(), 2); + for (i, row) in mat.rows().enumerate() { + assert!(row.iter().all(|&v| v == 0)); - let subview = m.subview(4..4).unwrap(); - assert_eq!(subview.nrows(), 0); - assert_eq!(subview.ncols(), 3); + let ptr = row.as_ptr().cast::(); + assert_eq!( + ptr, + base.wrapping_add(std::mem::size_of::() * mat.repr().ncols() * i), + ); } + } - // Empty out-of-bounds subview - assert!(m.subview(5..5).is_none()); + const ROWS: &[usize] = &[0, 1, 2, 3, 5, 10]; + const COLS: &[usize] = &[0, 1, 2, 3, 5, 10]; - // View too-large - assert!(m.subview(0..6).is_none()); - assert!(m.subview(2..10).is_none()); + #[test] + fn test_mat() { + for nrows in ROWS { + for ncols in COLS { + let repr = RowMajor::::new(*nrows, *ncols).unwrap(); + let ctx = &lazy_format!("nrows = {}, ncols = {}", nrows, ncols); + + // Populate the matrix using `&mut Mat` + { + let ctx = &lazy_format!("{ctx} - direct"); + let mut mat = Mat::from_default(repr).unwrap(); + + assert_eq!(mat.num_vectors(), *nrows); + assert_eq!(mat.vector_dim(), *ncols); + + fill_mat(&mut mat, repr); + + check_mat(&mat, repr, ctx); + check_mat_ref(mat.reborrow(), repr, ctx); + check_mat_mut(mat.reborrow_mut(), repr, ctx); + check_rows(mat.rows(), repr, ctx); + + // Check reborrow preserves pointers. + assert_eq!(mat.as_raw_ptr(), mat.reborrow().as_raw_ptr()); + assert_eq!(mat.as_raw_ptr(), mat.reborrow_mut().as_raw_ptr()); + } - // View disjoint. - assert!(m.subview(10..100).is_none()); + // Populate the matrix using `MatMut` + { + let ctx = &lazy_format!("{ctx} - matmut"); + let mut mat = Mat::from_default(repr).unwrap(); + let matmut = mat.reborrow_mut(); - // Negative bounds - #[expect( - clippy::reversed_empty_ranges, - reason = "we want to make sure it doesn't work" - )] - let empty = 3..2; - assert!(m.subview(empty).is_none()); + assert_eq!(matmut.num_vectors(), *nrows); + assert_eq!(matmut.vector_dim(), *ncols); - #[expect( - clippy::reversed_empty_ranges, - reason = "we want to make sure it doesn't work" - )] - let empty = 3..1; - assert!(m.subview(empty).is_none()); + fill_mat_mut(matmut, repr); - // Bounds that overflow. - assert!(m.subview(usize::MAX - 1..usize::MAX).is_none()); - assert!(m.subview(0..usize::MAX).is_none()); - } + check_mat(&mat, repr, ctx); + check_mat_ref(mat.reborrow(), repr, ctx); + check_mat_mut(mat.reborrow_mut(), repr, ctx); + check_rows(mat.rows(), repr, ctx); + } - #[test] - #[cfg(all(not(miri), feature = "rayon"))] - fn test_parallel_methods_edge_cases() { - let data = make_test_matrix(); - let m = Matrix::try_from(data.into(), 4, 3).unwrap(); - - // Test par_window_iter with batchsize larger than matrix - let windows: Vec<_> = m.par_window_iter(10).collect(); - assert_eq!(windows.len(), 1); - assert_eq!(windows[0].nrows(), 4); - assert_eq!(windows[0].ncols(), 3); - - // Test par_row_iter - let rows: Vec<_> = m.par_row_iter().collect(); - assert_eq!(rows.len(), 4); - assert_eq!(rows[0], &[0, 1, 2]); - assert_eq!(rows[3], &[3, 4, 5]); - - // Test par_window_iter_mut and par_row_iter_mut - let mut m2 = Matrix::new(0, 4, 3); - - // Use par_row_iter_mut to set values - m2.par_row_iter_mut().enumerate().for_each(|(i, row)| { - for (j, elem) in row.iter_mut().enumerate() { - *elem = i + j; + // Populate the matrix using `RowsMut` + { + let ctx = &lazy_format!("{ctx} - rows_mut"); + let mut mat = Mat::from_default(repr).unwrap(); + fill_rows_mut(mat.rows_mut(), repr); + + check_mat(&mat, repr, ctx); + check_mat_ref(mat.reborrow(), repr, ctx); + check_mat_mut(mat.reborrow_mut(), repr, ctx); + check_rows(mat.rows(), repr, ctx); + } } - }); - test_basic_indexing(&m2); - - // Test par_window_iter_mut with larger batchsize - let mut m3 = Matrix::new(0, 4, 3); - m3.par_window_iter_mut(10) - .enumerate() - .for_each(|(_, mut window)| { - window.row_iter_mut().enumerate().for_each(|(i, row)| { - for (j, elem) in row.iter_mut().enumerate() { - *elem = i + j; - } - }); - }); - test_basic_indexing(&m3); + } } #[test] - fn test_matrix_pointers() { - let mut m = Matrix::new(42, 3, 4); - - // Test as_ptr and as_mut_ptr return the same address - let const_ptr = m.as_ptr(); - let mut_ptr = m.as_mut_ptr(); - assert_eq!(const_ptr, mut_ptr as *const _); + fn test_mat_clone() { + for nrows in ROWS { + for ncols in COLS { + let repr = RowMajor::::new(*nrows, *ncols).unwrap(); + let ctx = &lazy_format!("nrows = {}, ncols = {}", nrows, ncols); + + let mut mat = Mat::from_default(repr).unwrap(); + fill_mat(&mut mat, repr); + + // Clone via Mat::clone + { + let ctx = &lazy_format!("{ctx} - Mat::clone"); + let cloned = mat.clone(); + + assert_eq!(cloned.num_vectors(), *nrows); + assert_eq!(cloned.vector_dim(), *ncols); + + check_mat(&cloned, repr, ctx); + check_mat_ref(cloned.reborrow(), repr, ctx); + check_rows(cloned.rows(), repr, ctx); + + // Cloned allocation is independent. + if repr.num_elements() > 0 { + assert_ne!(mat.as_raw_ptr(), cloned.as_raw_ptr()); + } + } - // Test that view pointers match original - let view = m.as_view(); - assert_eq!(view.as_ptr(), const_ptr); + // Clone via MatRef::to_owned + { + let ctx = &lazy_format!("{ctx} - MatRef::to_owned"); + let owned = mat.as_view().to_owned(); - let mut mut_view = m.as_mut_view(); - assert_eq!(mut_view.as_ptr(), const_ptr); - assert_eq!(mut_view.as_mut_ptr(), mut_ptr); - } + check_mat(&owned, repr, ctx); + check_mat_ref(owned.reborrow(), repr, ctx); + check_rows(owned.rows(), repr, ctx); - #[test] - fn test_matrix_iteration_empty_cases() { - // Test construction of empty matrices (we don't iterate over 0x0 matrices - // since chunks_exact requires non-zero chunk size) - let empty_data: Vec = vec![]; - - // Matrix with 0 rows but non-zero cols can be constructed - let _empty_matrix = MatrixView::try_from(empty_data.as_slice(), 0, 5).unwrap(); - - // Test with actual single row to verify iterator works normally - let data = vec![1, 2, 3]; - let single_row = MatrixView::try_from(data.as_slice(), 1, 3).unwrap(); - let rows: Vec<_> = single_row.row_iter().collect(); - assert_eq!(rows.len(), 1); - assert_eq!(rows[0], &[1, 2, 3]); - - // Test iteration over matrix with multiple rows but single column - let data = vec![1, 2, 3]; - let single_col = MatrixView::try_from(data.as_slice(), 3, 1).unwrap(); - let rows: Vec<_> = single_col.row_iter().collect(); - assert_eq!(rows.len(), 3); - assert_eq!(rows[0], &[1]); - assert_eq!(rows[1], &[2]); - assert_eq!(rows[2], &[3]); - } + if repr.num_elements() > 0 { + assert_ne!(mat.as_raw_ptr(), owned.as_raw_ptr()); + } + } - #[test] - fn test_matrix_init_generator_various_types() { - // Test with different types and generators - use std::sync::atomic::{AtomicUsize, Ordering}; + // Clone via MatMut::to_owned + { + let ctx = &lazy_format!("{ctx} - MatMut::to_owned"); + let owned = mat.as_view_mut().to_owned(); - let counter = AtomicUsize::new(0); - let m = Matrix::new(Init(|| counter.fetch_add(1, Ordering::SeqCst)), 2, 3); + check_mat(&owned, repr, ctx); + check_mat_ref(owned.reborrow(), repr, ctx); + check_rows(owned.rows(), repr, ctx); - // Should be filled in memory order - assert_eq!(m[(0, 0)], 0); - assert_eq!(m[(0, 1)], 1); - assert_eq!(m[(0, 2)], 2); - assert_eq!(m[(1, 0)], 3); - assert_eq!(m[(1, 1)], 4); - assert_eq!(m[(1, 2)], 5); + if repr.num_elements() > 0 { + assert_ne!(mat.as_raw_ptr(), owned.as_raw_ptr()); + } + } + } + } } #[test] - fn test_debug_error_formatting() { - // Test Debug implementation for TryFromError - let data = vec![1, 2, 3]; - let err = Matrix::try_from(data.into(), 2, 3).unwrap_err(); + fn test_mat_refmut() { + for nrows in ROWS { + for ncols in COLS { + let repr = RowMajor::::new(*nrows, *ncols).unwrap(); + let ctx = &lazy_format!("nrows = {}, ncols = {}", nrows, ncols); + + // Populate the matrix using `&mut Mat` + { + let ctx = &lazy_format!("{ctx} - by matmut"); + let mut b: Box<[_]> = (0..repr.num_elements()).map(|_| 0usize).collect(); + let ptr = b.as_ptr().cast::(); + let mut matmut = MatMut::from_repr(repr, &mut b).unwrap(); - let debug_str = format!("{:?}", err); - assert!(debug_str.contains("TryFromError")); - assert!(debug_str.contains("data_len: 3")); - assert!(debug_str.contains("nrows: 2")); - assert!(debug_str.contains("ncols: 3")); - - // Ensure Debug doesn't require T: Debug by using a non-Debug type - #[derive(Clone, Debug)] - struct NonDebug(#[allow(dead_code)] i32); + assert_eq!( + ptr, + matmut.as_raw_ptr(), + "underlying memory should be preserved", + ); - let non_debug_data: Box<[NonDebug]> = vec![NonDebug(1), NonDebug(2)].into(); - let non_debug_err = Matrix::try_from(non_debug_data, 1, 3).unwrap_err(); - let debug_str = format!("{:?}", non_debug_err); - assert!(debug_str.contains("TryFromError")); - } + fill_mat_mut(matmut.reborrow_mut(), repr); - // Comprehensive tests for rayon-specific functionality + check_mat_mut(matmut.reborrow_mut(), repr, ctx); + check_mat_ref(matmut.reborrow(), repr, ctx); + check_rows(matmut.rows(), repr, ctx); + check_rows(matmut.reborrow().rows(), repr, ctx); - #[test] - #[cfg(feature = "rayon")] - fn test_par_window_iter_comprehensive() { - use rayon::prelude::*; - - // Create a larger test matrix for more comprehensive testing - let data: Vec = (0..24).collect(); // 6x4 matrix - let m = MatrixView::try_from(data.as_slice(), 6, 4).unwrap(); - - // Test various batch sizes - for batchsize in 1..=8 { - let context = lazy_format!("batchsize = {}", batchsize); - let windows: Vec<_> = m.par_window_iter(batchsize).collect(); - - // Calculate expected number of windows - let expected_windows = (m.nrows()).div_ceil(batchsize); - assert_eq!(windows.len(), expected_windows, "{}", context); - - // Verify each window's properties - let mut total_rows_seen = 0; - for (window_idx, window) in windows.iter().enumerate() { - let expected_rows = if window_idx == windows.len() - 1 { - // Last window may have fewer rows - m.nrows() - (windows.len() - 1) * batchsize - } else { - batchsize - }; + let matref = MatRef::from_repr(repr, &b).unwrap(); + check_mat_ref(matref, repr, ctx); + check_mat_ref(matref.reborrow(), repr, ctx); + check_rows(matref.rows(), repr, ctx); + } - assert_eq!( - window.nrows(), - expected_rows, - "window {} - {}", - window_idx, - context - ); - assert_eq!( - window.ncols(), - m.ncols(), - "window {} - {}", - window_idx, - context - ); + // Populate the matrix using `RowsMut` + { + let ctx = &lazy_format!("{ctx} - by rows"); + let mut b: Box<[_]> = (0..repr.num_elements()).map(|_| 0usize).collect(); + let ptr = b.as_ptr().cast::(); + let mut matmut = MatMut::from_repr(repr, &mut b).unwrap(); - // Verify data integrity - for (row_idx, row) in window.row_iter().enumerate() { - let global_row = window_idx * batchsize + row_idx; - let expected: Vec = - (0..m.ncols()).map(|j| global_row * m.ncols() + j).collect(); assert_eq!( - row, - expected.as_slice(), - "window {}, row {} - {}", - window_idx, - row_idx, - context + ptr, + matmut.as_raw_ptr(), + "underlying memory should be preserved", ); - } - total_rows_seen += window.nrows(); - } + fill_rows_mut(matmut.rows_mut(), repr); + + check_mat_mut(matmut.reborrow_mut(), repr, ctx); + check_mat_ref(matmut.reborrow(), repr, ctx); + check_rows(matmut.rows(), repr, ctx); + check_rows(matmut.reborrow().rows(), repr, ctx); - assert_eq!(total_rows_seen, m.nrows(), "{}", context); + let matref = MatRef::from_repr(repr, &b).unwrap(); + check_mat_ref(matref, repr, ctx); + check_mat_ref(matref.reborrow(), repr, ctx); + check_rows(matref.rows(), repr, ctx); + } + } } + } - // Test with batchsize equal to matrix rows - let windows: Vec<_> = m.par_window_iter(m.nrows()).collect(); - assert_eq!(windows.len(), 1); - assert_eq!(windows[0].nrows(), m.nrows()); - assert_eq!(windows[0].ncols(), m.ncols()); + ////////////////// + // Constructors // + ////////////////// - // Test with batchsize larger than matrix rows - let windows: Vec<_> = m.par_window_iter(m.nrows() * 2).collect(); - assert_eq!(windows.len(), 1); - assert_eq!(windows[0].nrows(), m.nrows()); - assert_eq!(windows[0].ncols(), m.ncols()); + #[test] + fn test_standard_new_owned() { + let rows = [0, 1, 2, 3, 5, 10]; + let cols = [0, 1, 2, 3, 5, 10]; + + for nrows in rows { + for ncols in cols { + let m = Mat::from_repr(RowMajor::new(nrows, ncols).unwrap(), 1usize).unwrap(); + let rows_iter = m.rows(); + let len = <_ as ExactSizeIterator>::len(&rows_iter); + assert_eq!(len, nrows); + for r in rows_iter { + assert_eq!(r.len(), ncols); + assert!(r.iter().all(|i| *i == 1usize)); + } + } + } } #[test] - #[cfg(feature = "rayon")] - fn test_par_window_iter_mut_comprehensive() { - use rayon::prelude::*; - - // Test various matrix sizes and batch sizes - for nrows in [1, 2, 3, 5, 8, 10] { - for ncols in [1, 3, 4] { - for batchsize in [1, 2, 3, 7] { - let context = lazy_format!("{}x{}, batchsize={}", nrows, ncols, batchsize); - - let mut m = Matrix::new(0usize, nrows, ncols); - - // Use par_window_iter_mut to fill matrix - m.par_window_iter_mut(batchsize).enumerate().for_each( - |(window_idx, mut window)| { - let base_row = window_idx * batchsize; - window - .row_iter_mut() - .enumerate() - .for_each(|(row_offset, row)| { - let global_row = base_row + row_offset; - for (col, elem) in row.iter_mut().enumerate() { - *elem = global_row * ncols + col; - } - }); - }, - ); + fn test_mat_from_fn() { + let rows = [0, 1, 2, 5]; + let cols = [0, 1, 3, 7]; + + for nrows in rows { + for ncols in cols { + let mut counter = 0u32; + let m = Mat::new( + Init(|| { + let v = counter; + counter += 1; + v + }), + nrows, + ncols, + ); - // Verify the matrix was filled correctly - for row in 0..nrows { - for col in 0..ncols { - let expected = row * ncols + col; - assert_eq!( - m[(row, col)], - expected, - "pos ({}, {}) - {}", - row, - col, - context - ); - } + assert_eq!(counter as usize, nrows * ncols); + for (i, row) in m.rows().enumerate() { + assert_eq!(row.len(), ncols); + for (j, &v) in row.iter().enumerate() { + assert_eq!(v, (i * ncols + j) as u32); } } } @@ -1873,311 +2640,233 @@ mod tests { } #[test] - #[cfg(feature = "rayon")] - fn test_par_row_iter_comprehensive() { - use rayon::prelude::*; - - // Create test matrix with predictable pattern - let nrows = 7; - let ncols = 5; - let data: Vec = (0..(nrows * ncols) as i32).collect(); - let m = MatrixView::try_from(data.as_slice(), nrows, ncols).unwrap(); - - // Test that par_row_iter preserves order and data - let collected_rows: Vec> = m.par_row_iter().map(|row| row.to_vec()).collect(); - - assert_eq!(collected_rows.len(), nrows); - - for (row_idx, row) in collected_rows.iter().enumerate() { - assert_eq!(row.len(), ncols); - let expected: Vec = ((row_idx * ncols)..((row_idx + 1) * ncols)) - .map(|x| x as i32) - .collect(); - assert_eq!(row, &expected, "row {} mismatch", row_idx); - } - - // Test parallel enumeration - let enumerated_rows: Vec<(usize, Vec)> = m - .par_row_iter() - .enumerate() - .map(|(idx, row)| (idx, row.to_vec())) - .collect(); - - // Sort by index to ensure we got all indices - let mut sorted_rows = enumerated_rows; - sorted_rows.sort_by_key(|(idx, _)| *idx); - - assert_eq!(sorted_rows.len(), nrows); - for (expected_idx, (actual_idx, row)) in sorted_rows.iter().enumerate() { - assert_eq!(*actual_idx, expected_idx); - assert_eq!(row.len(), ncols); - } - - // Test parallel reduction operations - let sum: i32 = m.par_row_iter().map(|row| row.iter().sum::()).sum(); - - let expected_sum: i32 = data.iter().sum(); - assert_eq!(sum, expected_sum); - - // Test parallel find operations - let target_row = 3; - let found_row = m - .par_row_iter() - .enumerate() - .find_any(|(idx, _)| *idx == target_row) - .map(|(_, row)| row.to_vec()); + fn matref_new_slice_length_error() { + let repr = RowMajor::::new(3, 4).unwrap(); + + // Correct length succeeds + let data = vec![0u32; 12]; + assert!(MatRef::from_repr(repr, &data).is_ok()); + + // Too short fails + let short = vec![0u32; 11]; + assert!(matches!( + MatRef::from_repr(repr, &short), + Err(SliceError::LengthMismatch { + expected: 12, + found: 11 + }) + )); + + // Too long fails + let long = vec![0u32; 13]; + assert!(matches!( + MatRef::from_repr(repr, &long), + Err(SliceError::LengthMismatch { + expected: 12, + found: 13 + }) + )); + } - assert!(found_row.is_some()); - let expected_row: Vec = ((target_row * ncols)..((target_row + 1) * ncols)) - .map(|x| x as i32) - .collect(); - assert_eq!(found_row.unwrap(), expected_row); + #[test] + fn matmut_new_slice_length_error() { + let repr = RowMajor::::new(3, 4).unwrap(); + + // Correct length succeeds + let mut data = vec![0u32; 12]; + assert!(MatMut::from_repr(repr, &mut data).is_ok()); + + // Too short fails + let mut short = vec![0u32; 11]; + assert!(matches!( + MatMut::from_repr(repr, &mut short), + Err(SliceError::LengthMismatch { + expected: 12, + found: 11 + }) + )); + + // Too long fails + let mut long = vec![0u32; 13]; + assert!(matches!( + MatMut::from_repr(repr, &mut long), + Err(SliceError::LengthMismatch { + expected: 12, + found: 13 + }) + )); } #[test] - #[cfg(feature = "rayon")] - fn test_par_row_iter_mut_comprehensive() { - use rayon::prelude::*; - use std::sync::atomic::{AtomicUsize, Ordering}; - - let nrows = 6; - let ncols = 4; - let mut m = Matrix::new(0u32, nrows, ncols); - - // Test parallel modification - m.par_row_iter_mut().enumerate().for_each(|(row_idx, row)| { - for (col_idx, elem) in row.iter_mut().enumerate() { - *elem = (row_idx * ncols + col_idx) as u32; - } - }); + fn as_matrix_view_roundtrip() { + let data = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0]; - // Verify modifications were applied correctly - for row in 0..nrows { - for col in 0..ncols { - let expected = (row * ncols + col) as u32; - assert_eq!(m[(row, col)], expected, "pos ({}, {})", row, col); + // MatRef + let matref = MatRef::from_repr(RowMajor::new(2, 3).unwrap(), &data).unwrap(); + let view = matref.as_matrix_view(); + assert_eq!(view.nrows(), 2); + assert_eq!(view.ncols(), 3); + for row in 0..2 { + for col in 0..3 { + assert_eq!(view[(row, col)], data[row * 3 + col]); } } - - // Test parallel accumulation with atomic counter - let counter = AtomicUsize::new(0); - m.par_row_iter_mut().for_each(|row| { - counter.fetch_add(1, Ordering::Relaxed); - // Multiply each element by 2 - for elem in row { - *elem *= 2; + assert_eq!(matref.as_slice(), &data); + + // Mat + let mut mat = Mat::from_repr(RowMajor::::new(2, 3).unwrap(), 0.0f32).unwrap(); + for i in 0..2 { + let r = mat.get_row_mut(i).unwrap(); + for j in 0..3 { + r[j] = data[i * 3 + j]; } - }); - - assert_eq!(counter.load(Ordering::Relaxed), nrows); - - // Verify all elements were doubled - for row in 0..nrows { - for col in 0..ncols { - let expected = ((row * ncols + col) * 2) as u32; - assert_eq!(m[(row, col)], expected, "doubled pos ({}, {})", row, col); + } + let view = mat.as_matrix_view(); + assert_eq!(view.nrows(), 2); + assert_eq!(view.ncols(), 3); + for row in 0..2 { + for col in 0..3 { + assert_eq!(view[(row, col)], data[row * 3 + col]); } } - } + assert_eq!(mat.as_slice(), &data); - #[test] - #[cfg(feature = "rayon")] - fn test_parallel_iterators_with_single_dimensions() { - use rayon::prelude::*; - - // Test single row matrix - let data = vec![1, 2, 3, 4, 5]; - let single_row = MatrixView::try_from(data.as_slice(), 1, 5).unwrap(); - - let windows: Vec<_> = single_row.par_window_iter(1).collect(); - assert_eq!(windows.len(), 1); - assert_eq!(windows[0].nrows(), 1); - assert_eq!(windows[0].ncols(), 5); - - let rows: Vec<_> = single_row.par_row_iter().collect(); - assert_eq!(rows.len(), 1); - assert_eq!(rows[0], &[1, 2, 3, 4, 5]); - - // Test single column matrix - let data = vec![1, 2, 3, 4, 5]; - let single_col = MatrixView::try_from(data.as_slice(), 5, 1).unwrap(); - - let windows: Vec<_> = single_col.par_window_iter(2).collect(); - assert_eq!(windows.len(), 3); // ceil(5/2) = 3 - assert_eq!(windows[0].nrows(), 2); - assert_eq!(windows[1].nrows(), 2); - assert_eq!(windows[2].nrows(), 1); // Last window has remainder - - let rows: Vec<_> = single_col.par_row_iter().collect(); - assert_eq!(rows.len(), 5); - for (i, row) in rows.iter().enumerate() { - assert_eq!(row, &[i + 1]); + // MatMut + let mut buf = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0]; + let matmut = MatMut::from_repr(RowMajor::new(2, 3).unwrap(), &mut buf).unwrap(); + let view = matmut.as_matrix_view(); + assert_eq!(view.nrows(), 2); + assert_eq!(view.ncols(), 3); + for row in 0..2 { + for col in 0..3 { + assert_eq!(view[(row, col)], data[row * 3 + col]); + } } - - // Test 1x1 matrix - let data = vec![42]; - let tiny = MatrixView::try_from(data.as_slice(), 1, 1).unwrap(); - - let windows: Vec<_> = tiny.par_window_iter(1).collect(); - assert_eq!(windows.len(), 1); - assert_eq!(windows[0][(0, 0)], 42); - - let rows: Vec<_> = tiny.par_row_iter().collect(); - assert_eq!(rows.len(), 1); - assert_eq!(rows[0], &[42]); + assert_eq!(matmut.as_slice(), &data); } #[test] - #[cfg(feature = "rayon")] - fn test_parallel_window_properties() { - use rayon::prelude::*; - - // Test that windows maintain proper matrix properties - let data: Vec = (0..30).collect(); - let m = MatrixView::try_from(data.as_slice(), 6, 5).unwrap(); - - // Test window indexing works correctly - m.par_window_iter(2) - .enumerate() - .for_each(|(window_idx, window)| { - for row_idx in 0..window.nrows() { - for col_idx in 0..window.ncols() { - let global_row = window_idx * 2 + row_idx; - let expected = global_row * 5 + col_idx; - assert_eq!( - window[(row_idx, col_idx)], - expected, - "window {}, pos ({}, {})", - window_idx, - row_idx, - col_idx - ); - } - } - }); + fn test_standard_non_copy_element() { + let repr = RowMajor::::new(2, 3).unwrap(); + + // Owned fill via NewOwned (Clone). + let filled = Mat::from_repr(repr, String::from("x")).unwrap(); + assert_eq!(filled.num_vectors(), 2); + assert!(filled.rows().flatten().all(|s| s == "x")); + + // NewDefault (Default). + let defaulted = Mat::from_default(repr).unwrap(); + assert!(defaulted.rows().flatten().all(String::is_empty)); + + // from_fn. + let mut counter = 0usize; + let mut mat = Mat::new( + Init(|| { + let s = counter.to_string(); + counter += 1; + s + }), + 2, + 3, + ); + assert_eq!(counter, 6); + assert_eq!(mat.get_row(1).unwrap()[0], "3"); - // Test window as_slice consistency - m.par_window_iter(3) - .enumerate() - .for_each(|(window_idx, window)| { - let slice = window.as_slice(); - assert_eq!(slice.len(), window.nrows() * window.ncols()); + // Mutation via get_row_mut. + mat.get_row_mut(0).unwrap()[0] = String::from("mutated"); + assert_eq!(mat.get_row(0).unwrap()[0], "mutated"); - for (slice_idx, &value) in slice.iter().enumerate() { - let row = slice_idx / window.ncols(); - let col = slice_idx % window.ncols(); - assert_eq!( - value, - window[(row, col)], - "window {}, slice_idx {}", - window_idx, - slice_idx - ); - } - }); + // Clone via NewCloned (Clone): independent allocation, equal contents. + let cloned = mat.clone(); + assert_ne!(mat.as_raw_ptr(), cloned.as_raw_ptr()); + assert_eq!(cloned.get_row(0).unwrap()[0], "mutated"); - // Test window row iteration - m.par_window_iter(2).for_each(|window| { - let rows_via_iter: Vec<_> = window.row_iter().collect(); - assert_eq!(rows_via_iter.len(), window.nrows()); + // Immutable view over a non-Copy slice (NewRef). + let data = [String::from("a"), String::from("b")]; + let view = MatRef::from_repr(RowMajor::new(2, 1).unwrap(), &data).unwrap(); + assert_eq!(view.get_row(1).unwrap()[0], "b"); - for (row_idx, row) in rows_via_iter.iter().enumerate() { - assert_eq!(row.len(), window.ncols()); - for (col_idx, &value) in row.iter().enumerate() { - assert_eq!(value, window[(row_idx, col_idx)]); - } - } - }); + // Mutable view over a non-Copy slice (NewMut). + let mut data_mut = [String::from("a"), String::from("b")]; + let mut view_mut = MatMut::from_repr(RowMajor::new(1, 2).unwrap(), &mut data_mut).unwrap(); + view_mut.get_row_mut(0).unwrap()[1] = String::from("z"); + assert_eq!(data_mut[1], "z"); } - #[test] - #[cfg(feature = "rayon")] - fn test_parallel_performance_characteristics() { - use rayon::prelude::*; - use std::sync::atomic::{AtomicUsize, Ordering}; - - // Create a larger matrix to test parallelism benefits - let nrows = 100; - let ncols = 10; - let mut m = Matrix::new(0usize, nrows, ncols); - - // Test that parallel operations can be chained - let work_counter = AtomicUsize::new(0); - - m.par_window_iter_mut(10) - .enumerate() - .for_each(|(window_idx, mut window)| { - work_counter.fetch_add(1, Ordering::Relaxed); - - // Nested parallel operation within window - window - .row_iter_mut() - .enumerate() - .for_each(|(row_offset, row)| { - let global_row = window_idx * 10 + row_offset; - for (col, elem) in row.iter_mut().enumerate() { - *elem = global_row * ncols + col; - } - }); - }); - - // Should have processed 10 windows (100 rows / 10 batch size) - assert_eq!(work_counter.load(Ordering::Relaxed), 10); - - // Verify correctness - for row in 0..nrows { - for col in 0..ncols { - assert_eq!(m[(row, col)], row * ncols + col); - } - } - - // Test parallel reduction across windows - let total_sum: usize = m - .par_window_iter(15) - .map(|window| { - window - .row_iter() - .map(|row| row.iter().sum::()) - .sum::() - }) - .sum(); + // ── Dense read API (canonical on `MatrixView`, delegated by `Matrix`/`MatrixViewMut`) ── - let expected_sum: usize = (0..(nrows * ncols)).sum(); - assert_eq!(total_sum, expected_sum); + #[test] + fn dense_subview_and_try_get() { + let m: Matrix = + Matrix::try_from(vec![1, 2, 3, 4, 5, 6].into_boxed_slice(), 3, 2).unwrap(); + let sv = m.subview(1..3).unwrap(); + assert_eq!((sv.nrows(), sv.ncols()), (2, 2)); + assert_eq!(sv.row(0), &[3, 4]); + assert_eq!(sv.row(1), &[5, 6]); + assert!(m.subview(2..4).is_none()); + assert_eq!(m.try_get(2, 1), Some(&6)); + assert!(m.try_get(3, 0).is_none()); + assert!(m.try_get(0, 2).is_none()); } #[test] - #[cfg(feature = "rayon")] - fn test_rayon_trait_bounds_validation() { - use rayon::prelude::*; + fn dense_map_and_vectors() { + let m: Matrix = Matrix::try_from(vec![1, 2, 3, 4].into_boxed_slice(), 2, 2).unwrap(); + let doubled = m.map(|&x| x * 2); + assert_eq!(doubled.as_slice(), &[2, 4, 6, 8]); + assert_eq!((doubled.nrows(), doubled.ncols()), (2, 2)); - // Test that the Sync/Send bounds work correctly - let data: Vec = (0..20).collect(); - let m = MatrixView::try_from(data.as_slice(), 4, 5).unwrap(); + let rv = Matrix::row_vector(vec![7, 8, 9].into_boxed_slice()); + assert_eq!((rv.nrows(), rv.ncols()), (1, 3)); + let cv = Matrix::column_vector(vec![7, 8, 9].into_boxed_slice()); + assert_eq!((cv.nrows(), cv.ncols()), (3, 1)); + } - // This should compile because u64 is Sync - let _: Vec<_> = m.par_window_iter(2).collect(); - let _: Vec<_> = m.par_row_iter().collect(); + #[test] + fn dense_window_iter_batches_and_remainder() { + let m: Matrix = + Matrix::try_from((0..9).collect::>().into_boxed_slice(), 3, 3).unwrap(); + let rows: Vec<_> = m.window_iter(2).map(|w| w.nrows()).collect(); + assert_eq!(rows, vec![2, 1]); + } - // Test with mutable matrix - let mut m = Matrix::new(0u64, 4, 5); + #[test] + #[should_panic(expected = "window_iter batchsize cannot be zero")] + fn dense_window_iter_zero_panics() { + let m: Matrix = Matrix::try_from(vec![1, 2].into_boxed_slice(), 1, 2).unwrap(); + let _ = m.window_iter(0).count(); + } - // This should compile because u64 is Send - m.par_window_iter_mut(2).for_each(|mut window| { - window.row_iter_mut().for_each(|row| { - for elem in row { - *elem = 42; - } - }); - }); + #[test] + #[should_panic(expected = "tried to access row 5 of a matrix with 1 rows")] + fn dense_row_out_of_bounds_panics() { + let m: Matrix = Matrix::try_from(vec![1, 2].into_boxed_slice(), 1, 2).unwrap(); + let _ = m.row(5); + } - m.par_row_iter_mut().for_each(|row| { - for elem in row { - *elem += 1; - } - }); + #[test] + fn dense_get_unchecked_reads_and_writes() { + let mut m: Matrix = + Matrix::try_from(vec![1, 2, 3, 4].into_boxed_slice(), 2, 2).unwrap(); + // SAFETY: `(1, 1)` is in bounds for a 2x2 matrix. + assert_eq!(unsafe { *m.get_unchecked(1, 1) }, 4); + // SAFETY: `(0, 0)` is in bounds. + unsafe { *m.get_unchecked_mut(0, 0) = 9 }; + assert_eq!(m.as_slice(), &[9, 2, 3, 4]); + } - // Verify all elements are 43 - assert!(m.as_slice().iter().all(|&x| x == 43)); + #[test] + fn dense_reads_delegate_consistently() { + let mut m: Matrix = + Matrix::try_from(vec![1, 2, 3, 4, 5, 6].into_boxed_slice(), 3, 2).unwrap(); + let owned = m.row(1).to_vec(); + let via_view = m.as_view().row(1).to_vec(); + let via_mut_view = m.as_mut_view().row(1).to_vec(); + assert_eq!(owned, &[3, 4]); + assert_eq!(owned, via_view); + assert_eq!(owned, via_mut_view); + // The mutable view exposes the full read set (subview/window_iter/as_ptr/...). + assert!(m.as_mut_view().subview(0..1).is_some()); + assert_eq!(m.as_mut_view().window_iter(2).count(), 2); } } diff --git a/diskann-quantization/tests/compile-fail/multi/mat_as_view.rs b/diskann-utils/tests/compile-fail/mat_as_view.rs similarity index 72% rename from diskann-quantization/tests/compile-fail/multi/mat_as_view.rs rename to diskann-utils/tests/compile-fail/mat_as_view.rs index 02d26d026..58fa66291 100644 --- a/diskann-quantization/tests/compile-fail/multi/mat_as_view.rs +++ b/diskann-utils/tests/compile-fail/mat_as_view.rs @@ -3,12 +3,12 @@ * Licensed under the MIT license. */ -use diskann_quantization::multi_vector::{Mat, Standard}; +use diskann_utils::views::{Mat, RowMajor}; // Test that `as_view` on Mat correctly captures an immutable borrow, // preventing mutation of the Mat while the view is in scope. fn main() { - let mut mat: Mat> = Mat::new(Standard::new(4, 3).unwrap(), 0.0f32).unwrap(); + let mut mat: Mat> = Mat::from_repr(RowMajor::new(4, 3).unwrap(), 0.0f32).unwrap(); let view = mat.as_view(); // This should fail: we cannot mutably borrow `mat` while `view` exists let _ = mat.as_view_mut(); diff --git a/diskann-quantization/tests/compile-fail/multi/mat_as_view.stderr b/diskann-utils/tests/compile-fail/mat_as_view.stderr similarity index 89% rename from diskann-quantization/tests/compile-fail/multi/mat_as_view.stderr rename to diskann-utils/tests/compile-fail/mat_as_view.stderr index d63e03325..84deb9546 100644 --- a/diskann-quantization/tests/compile-fail/multi/mat_as_view.stderr +++ b/diskann-utils/tests/compile-fail/mat_as_view.stderr @@ -1,5 +1,5 @@ error[E0502]: cannot borrow `mat` as mutable because it is also borrowed as immutable - --> tests/compile-fail/multi/mat_as_view.rs:14:13 + --> tests/compile-fail/mat_as_view.rs:14:13 | 12 | let view = mat.as_view(); | --- immutable borrow occurs here diff --git a/diskann-quantization/tests/compile-fail/multi/mat_as_view_mut.rs b/diskann-utils/tests/compile-fail/mat_as_view_mut.rs similarity index 73% rename from diskann-quantization/tests/compile-fail/multi/mat_as_view_mut.rs rename to diskann-utils/tests/compile-fail/mat_as_view_mut.rs index 534c8b5ac..4858d17ea 100644 --- a/diskann-quantization/tests/compile-fail/multi/mat_as_view_mut.rs +++ b/diskann-utils/tests/compile-fail/mat_as_view_mut.rs @@ -3,12 +3,12 @@ * Licensed under the MIT license. */ -use diskann_quantization::multi_vector::{Mat, Standard}; +use diskann_utils::views::{Mat, RowMajor}; // Test that `as_view_mut` on Mat correctly captures a mutable lifetime, // preventing the Mat from being used while the mutable view is in scope. fn main() { - let mut mat: Mat> = Mat::new(Standard::new(4, 3).unwrap(), 0.0f32).unwrap(); + let mut mat: Mat> = Mat::from_repr(RowMajor::new(4, 3).unwrap(), 0.0f32).unwrap(); let view = mat.as_view_mut(); // This should fail: we cannot use `mat` while `view` is still alive let _ = mat.num_vectors(); diff --git a/diskann-quantization/tests/compile-fail/multi/mat_as_view_mut.stderr b/diskann-utils/tests/compile-fail/mat_as_view_mut.stderr similarity index 88% rename from diskann-quantization/tests/compile-fail/multi/mat_as_view_mut.stderr rename to diskann-utils/tests/compile-fail/mat_as_view_mut.stderr index cf5432eed..7d1a37d8a 100644 --- a/diskann-quantization/tests/compile-fail/multi/mat_as_view_mut.stderr +++ b/diskann-utils/tests/compile-fail/mat_as_view_mut.stderr @@ -1,5 +1,5 @@ error[E0502]: cannot borrow `mat` as immutable because it is also borrowed as mutable - --> tests/compile-fail/multi/mat_as_view_mut.rs:14:13 + --> tests/compile-fail/mat_as_view_mut.rs:14:13 | 12 | let view = mat.as_view_mut(); | --- mutable borrow occurs here diff --git a/diskann-quantization/tests/compile-fail/multi/mat_get_row.rs b/diskann-utils/tests/compile-fail/mat_get_row.rs similarity index 72% rename from diskann-quantization/tests/compile-fail/multi/mat_get_row.rs rename to diskann-utils/tests/compile-fail/mat_get_row.rs index d07a49cf3..e37a7d058 100644 --- a/diskann-quantization/tests/compile-fail/multi/mat_get_row.rs +++ b/diskann-utils/tests/compile-fail/mat_get_row.rs @@ -3,12 +3,12 @@ * Licensed under the MIT license. */ -use diskann_quantization::multi_vector::{Mat, Standard}; +use diskann_utils::views::{Mat, RowMajor}; // Test that `get_row` on Mat correctly captures an immutable borrow, // preventing mutation of the Mat while the row is in scope. fn main() { - let mut mat: Mat> = Mat::new(Standard::new(4, 3).unwrap(), 0.0f32).unwrap(); + let mut mat: Mat> = Mat::from_repr(RowMajor::new(4, 3).unwrap(), 0.0f32).unwrap(); let row = mat.get_row(0).unwrap(); // This should fail: we cannot mutably borrow `mat` while `row` exists let _ = mat.get_row_mut(1); diff --git a/diskann-quantization/tests/compile-fail/multi/mat_get_row.stderr b/diskann-utils/tests/compile-fail/mat_get_row.stderr similarity index 89% rename from diskann-quantization/tests/compile-fail/multi/mat_get_row.stderr rename to diskann-utils/tests/compile-fail/mat_get_row.stderr index 005e16421..d3e944b9b 100644 --- a/diskann-quantization/tests/compile-fail/multi/mat_get_row.stderr +++ b/diskann-utils/tests/compile-fail/mat_get_row.stderr @@ -1,5 +1,5 @@ error[E0502]: cannot borrow `mat` as mutable because it is also borrowed as immutable - --> tests/compile-fail/multi/mat_get_row.rs:14:13 + --> tests/compile-fail/mat_get_row.rs:14:13 | 12 | let row = mat.get_row(0).unwrap(); | --- immutable borrow occurs here diff --git a/diskann-quantization/tests/compile-fail/multi/mat_get_row_mut.rs b/diskann-utils/tests/compile-fail/mat_get_row_mut.rs similarity index 72% rename from diskann-quantization/tests/compile-fail/multi/mat_get_row_mut.rs rename to diskann-utils/tests/compile-fail/mat_get_row_mut.rs index 469d9349e..9f1a1f5aa 100644 --- a/diskann-quantization/tests/compile-fail/multi/mat_get_row_mut.rs +++ b/diskann-utils/tests/compile-fail/mat_get_row_mut.rs @@ -3,12 +3,12 @@ * Licensed under the MIT license. */ -use diskann_quantization::multi_vector::{Mat, Standard}; +use diskann_utils::views::{Mat, RowMajor}; // Test that `get_row_mut` on Mat correctly captures a mutable // lifetime, preventing the Mat from being used while the row is in scope. fn main() { - let mut mat: Mat> = Mat::new(Standard::new(4, 3).unwrap(), 0.0f32).unwrap(); + let mut mat: Mat> = Mat::from_repr(RowMajor::new(4, 3).unwrap(), 0.0f32).unwrap(); let row = mat.get_row_mut(0).unwrap(); // This should fail: we cannot use `mat` while `row` is still borrowed let _ = mat.num_vectors(); diff --git a/diskann-quantization/tests/compile-fail/multi/mat_get_row_mut.stderr b/diskann-utils/tests/compile-fail/mat_get_row_mut.stderr similarity index 88% rename from diskann-quantization/tests/compile-fail/multi/mat_get_row_mut.stderr rename to diskann-utils/tests/compile-fail/mat_get_row_mut.stderr index 48640670e..0b8aa65e9 100644 --- a/diskann-quantization/tests/compile-fail/multi/mat_get_row_mut.stderr +++ b/diskann-utils/tests/compile-fail/mat_get_row_mut.stderr @@ -1,5 +1,5 @@ error[E0502]: cannot borrow `mat` as immutable because it is also borrowed as mutable - --> tests/compile-fail/multi/mat_get_row_mut.rs:14:13 + --> tests/compile-fail/mat_get_row_mut.rs:14:13 | 12 | let row = mat.get_row_mut(0).unwrap(); | --- mutable borrow occurs here diff --git a/diskann-quantization/tests/compile-fail/multi/mat_invariant.rs b/diskann-utils/tests/compile-fail/mat_invariant.rs similarity index 62% rename from diskann-quantization/tests/compile-fail/multi/mat_invariant.rs rename to diskann-utils/tests/compile-fail/mat_invariant.rs index 813de1817..bdc7e1c4c 100644 --- a/diskann-quantization/tests/compile-fail/multi/mat_invariant.rs +++ b/diskann-utils/tests/compile-fail/mat_invariant.rs @@ -3,13 +3,13 @@ * Licensed under the MIT license. */ -use diskann_quantization::multi_vector::{Mat, Standard}; +use diskann_utils::views::{Mat, RowMajor}; // Verify that `Mat` is invariant in any generic parameters. // // This must not compile because it would allow assigning references with a shorter lifetime // into the matrix -fn bad<'long, 'short>(v: Mat>) -> Mat> +fn bad<'long, 'short>(v: Mat>) -> Mat> where 'long: 'short, { @@ -18,6 +18,6 @@ where fn main() { let b = 0u8; - let m = Mat::new(Standard::new(4, 3).unwrap(), &b).unwrap(); + let m = Mat::from_repr(RowMajor::new(4, 3).unwrap(), &b).unwrap(); bad(m); } diff --git a/diskann-quantization/tests/compile-fail/multi/mat_invariant.stderr b/diskann-utils/tests/compile-fail/mat_invariant.stderr similarity index 66% rename from diskann-quantization/tests/compile-fail/multi/mat_invariant.stderr rename to diskann-utils/tests/compile-fail/mat_invariant.stderr index c27439227..d3531694e 100644 --- a/diskann-quantization/tests/compile-fail/multi/mat_invariant.stderr +++ b/diskann-utils/tests/compile-fail/mat_invariant.stderr @@ -1,7 +1,7 @@ error: lifetime may not live long enough - --> tests/compile-fail/multi/mat_invariant.rs:16:5 + --> tests/compile-fail/mat_invariant.rs:16:5 | -12 | fn bad<'long, 'short>(v: Mat>) -> Mat> +12 | fn bad<'long, 'short>(v: Mat>) -> Mat> | ----- ------ lifetime `'short` defined here | | | lifetime `'long` defined here @@ -10,6 +10,6 @@ error: lifetime may not live long enough | ^ function was supposed to return data with lifetime `'long` but it is returning data with lifetime `'short` | = help: consider adding the following bound: `'short: 'long` - = note: requirement occurs because of the type `Mat>`, which makes the generic argument `Standard<&u8>` invariant + = note: requirement occurs because of the type `Mat>`, which makes the generic argument `RowMajor<&u8>` invariant = note: the struct `Mat` is invariant over the parameter `T` = help: see for more information about variance diff --git a/diskann-quantization/tests/compile-fail/multi/mat_reborrow.rs b/diskann-utils/tests/compile-fail/mat_reborrow.rs similarity index 74% rename from diskann-quantization/tests/compile-fail/multi/mat_reborrow.rs rename to diskann-utils/tests/compile-fail/mat_reborrow.rs index c1c56b0f2..55fe80a39 100644 --- a/diskann-quantization/tests/compile-fail/multi/mat_reborrow.rs +++ b/diskann-utils/tests/compile-fail/mat_reborrow.rs @@ -3,13 +3,13 @@ * Licensed under the MIT license. */ -use diskann_quantization::multi_vector::{Mat, Standard}; +use diskann_utils::views::{Mat, RowMajor}; use diskann_utils::Reborrow; // Test that `reborrow` on Mat correctly captures an immutable borrow, // preventing mutation of the Mat while the reborrow is in scope. fn main() { - let mut mat: Mat> = Mat::new(Standard::new(4, 3).unwrap(), 0.0f32).unwrap(); + let mut mat: Mat> = Mat::from_repr(RowMajor::new(4, 3).unwrap(), 0.0f32).unwrap(); let view = mat.reborrow(); // This should fail: we cannot mutably borrow `mat` while `view` exists let _ = mat.as_view_mut(); diff --git a/diskann-quantization/tests/compile-fail/multi/mat_reborrow.stderr b/diskann-utils/tests/compile-fail/mat_reborrow.stderr similarity index 89% rename from diskann-quantization/tests/compile-fail/multi/mat_reborrow.stderr rename to diskann-utils/tests/compile-fail/mat_reborrow.stderr index 602a8412f..3c7058031 100644 --- a/diskann-quantization/tests/compile-fail/multi/mat_reborrow.stderr +++ b/diskann-utils/tests/compile-fail/mat_reborrow.stderr @@ -1,5 +1,5 @@ error[E0502]: cannot borrow `mat` as mutable because it is also borrowed as immutable - --> tests/compile-fail/multi/mat_reborrow.rs:15:13 + --> tests/compile-fail/mat_reborrow.rs:15:13 | 13 | let view = mat.reborrow(); | --- immutable borrow occurs here diff --git a/diskann-quantization/tests/compile-fail/multi/mat_reborrow_mut.rs b/diskann-utils/tests/compile-fail/mat_reborrow_mut.rs similarity index 73% rename from diskann-quantization/tests/compile-fail/multi/mat_reborrow_mut.rs rename to diskann-utils/tests/compile-fail/mat_reborrow_mut.rs index 98c49fb28..c734dd32c 100644 --- a/diskann-quantization/tests/compile-fail/multi/mat_reborrow_mut.rs +++ b/diskann-utils/tests/compile-fail/mat_reborrow_mut.rs @@ -3,13 +3,13 @@ * Licensed under the MIT license. */ -use diskann_quantization::multi_vector::{Mat, Standard}; +use diskann_utils::views::{Mat, RowMajor}; use diskann_utils::ReborrowMut; // Test that `reborrow_mut` on Mat correctly captures a mutable borrow, // preventing use of the Mat while the reborrow is in scope. fn main() { - let mut mat: Mat> = Mat::new(Standard::new(4, 3).unwrap(), 0.0f32).unwrap(); + let mut mat: Mat> = Mat::from_repr(RowMajor::new(4, 3).unwrap(), 0.0f32).unwrap(); let view = mat.reborrow_mut(); // This should fail: we cannot use `mat` while `view` exists let _ = mat.num_vectors(); diff --git a/diskann-quantization/tests/compile-fail/multi/mat_reborrow_mut.stderr b/diskann-utils/tests/compile-fail/mat_reborrow_mut.stderr similarity index 88% rename from diskann-quantization/tests/compile-fail/multi/mat_reborrow_mut.stderr rename to diskann-utils/tests/compile-fail/mat_reborrow_mut.stderr index 14f5e1cf6..f1f3a95ec 100644 --- a/diskann-quantization/tests/compile-fail/multi/mat_reborrow_mut.stderr +++ b/diskann-utils/tests/compile-fail/mat_reborrow_mut.stderr @@ -1,5 +1,5 @@ error[E0502]: cannot borrow `mat` as immutable because it is also borrowed as mutable - --> tests/compile-fail/multi/mat_reborrow_mut.rs:15:13 + --> tests/compile-fail/mat_reborrow_mut.rs:15:13 | 13 | let view = mat.reborrow_mut(); | --- mutable borrow occurs here diff --git a/diskann-quantization/tests/compile-fail/multi/mat_rows.rs b/diskann-utils/tests/compile-fail/mat_rows.rs similarity index 73% rename from diskann-quantization/tests/compile-fail/multi/mat_rows.rs rename to diskann-utils/tests/compile-fail/mat_rows.rs index 95be7e6a6..c968bb37d 100644 --- a/diskann-quantization/tests/compile-fail/multi/mat_rows.rs +++ b/diskann-utils/tests/compile-fail/mat_rows.rs @@ -3,12 +3,12 @@ * Licensed under the MIT license. */ -use diskann_quantization::multi_vector::{Mat, Standard}; +use diskann_utils::views::{Mat, RowMajor}; // Test that `rows` on Mat correctly captures an immutable borrow, // preventing mutation of the Mat while the iterator is in scope. fn main() { - let mut mat: Mat> = Mat::new(Standard::new(4, 3).unwrap(), 0.0f32).unwrap(); + let mut mat: Mat> = Mat::from_repr(RowMajor::new(4, 3).unwrap(), 0.0f32).unwrap(); let iter = mat.rows(); // This should fail: we cannot mutably borrow `mat` while `iter` exists let _ = mat.as_view_mut(); diff --git a/diskann-quantization/tests/compile-fail/multi/mat_rows.stderr b/diskann-utils/tests/compile-fail/mat_rows.stderr similarity index 89% rename from diskann-quantization/tests/compile-fail/multi/mat_rows.stderr rename to diskann-utils/tests/compile-fail/mat_rows.stderr index 6f6e1420d..422a6683c 100644 --- a/diskann-quantization/tests/compile-fail/multi/mat_rows.stderr +++ b/diskann-utils/tests/compile-fail/mat_rows.stderr @@ -1,5 +1,5 @@ error[E0502]: cannot borrow `mat` as mutable because it is also borrowed as immutable - --> tests/compile-fail/multi/mat_rows.rs:14:13 + --> tests/compile-fail/mat_rows.rs:14:13 | 12 | let iter = mat.rows(); | --- immutable borrow occurs here diff --git a/diskann-quantization/tests/compile-fail/multi/mat_rows_mut.rs b/diskann-utils/tests/compile-fail/mat_rows_mut.rs similarity index 74% rename from diskann-quantization/tests/compile-fail/multi/mat_rows_mut.rs rename to diskann-utils/tests/compile-fail/mat_rows_mut.rs index 0ce872998..09ea4eb73 100644 --- a/diskann-quantization/tests/compile-fail/multi/mat_rows_mut.rs +++ b/diskann-utils/tests/compile-fail/mat_rows_mut.rs @@ -3,12 +3,12 @@ * Licensed under the MIT license. */ -use diskann_quantization::multi_vector::{Mat, Standard}; +use diskann_utils::views::{Mat, RowMajor}; // Test that the `rows_mut` iterator correctly captures a mutable lifetime, // preventing the Mat from being used while the iterator is in scope. fn main() { - let mut mat: Mat> = Mat::new(Standard::new(4, 3).unwrap(), 0.0f32).unwrap(); + let mut mat: Mat> = Mat::from_repr(RowMajor::new(4, 3).unwrap(), 0.0f32).unwrap(); let iter = mat.rows_mut(); // This should fail: we cannot use `mat` while the mutable iterator is alive let _ = mat.num_vectors(); diff --git a/diskann-quantization/tests/compile-fail/multi/mat_rows_mut.stderr b/diskann-utils/tests/compile-fail/mat_rows_mut.stderr similarity index 89% rename from diskann-quantization/tests/compile-fail/multi/mat_rows_mut.stderr rename to diskann-utils/tests/compile-fail/mat_rows_mut.stderr index 30e04d6b5..bc2e12fc6 100644 --- a/diskann-quantization/tests/compile-fail/multi/mat_rows_mut.stderr +++ b/diskann-utils/tests/compile-fail/mat_rows_mut.stderr @@ -1,5 +1,5 @@ error[E0502]: cannot borrow `mat` as immutable because it is also borrowed as mutable - --> tests/compile-fail/multi/mat_rows_mut.rs:14:13 + --> tests/compile-fail/mat_rows_mut.rs:14:13 | 12 | let iter = mat.rows_mut(); | --- mutable borrow occurs here diff --git a/diskann-quantization/tests/compile-fail/multi/matmut_as_view_borrows.rs b/diskann-utils/tests/compile-fail/matmut_as_view_borrows.rs similarity index 65% rename from diskann-quantization/tests/compile-fail/multi/matmut_as_view_borrows.rs rename to diskann-utils/tests/compile-fail/matmut_as_view_borrows.rs index 6faacbf71..f1cac56f5 100644 --- a/diskann-quantization/tests/compile-fail/multi/matmut_as_view_borrows.rs +++ b/diskann-utils/tests/compile-fail/matmut_as_view_borrows.rs @@ -3,13 +3,13 @@ * Licensed under the MIT license. */ -use diskann_quantization::multi_vector::{Mat, MatMut, Standard}; +use diskann_utils::views::{Mat, MatMut, RowMajor}; // Test that `as_view` on MatMut correctly captures an immutable lifetime, // preventing mutating the MatMut while the immutable view is in scope. fn main() { - let mut mat: Mat> = Mat::new(Standard::new(4, 3).unwrap(), 0.0f32).unwrap(); - let mut view: MatMut<'_, Standard> = mat.as_view_mut(); + let mut mat: Mat> = Mat::from_repr(RowMajor::new(4, 3).unwrap(), 0.0f32).unwrap(); + let mut view: MatMut<'_, RowMajor> = mat.as_view_mut(); let immut_view = view.as_view(); // This should fail: we cannot mutate `view` while `immut_view` exists let _ = view.get_row_mut(0); diff --git a/diskann-quantization/tests/compile-fail/multi/matmut_as_view_borrows.stderr b/diskann-utils/tests/compile-fail/matmut_as_view_borrows.stderr similarity index 88% rename from diskann-quantization/tests/compile-fail/multi/matmut_as_view_borrows.stderr rename to diskann-utils/tests/compile-fail/matmut_as_view_borrows.stderr index 3a6a2933b..3c49c90b5 100644 --- a/diskann-quantization/tests/compile-fail/multi/matmut_as_view_borrows.stderr +++ b/diskann-utils/tests/compile-fail/matmut_as_view_borrows.stderr @@ -1,5 +1,5 @@ error[E0502]: cannot borrow `view` as mutable because it is also borrowed as immutable - --> tests/compile-fail/multi/matmut_as_view_borrows.rs:15:13 + --> tests/compile-fail/matmut_as_view_borrows.rs:15:13 | 13 | let immut_view = view.as_view(); | ---- immutable borrow occurs here diff --git a/diskann-quantization/tests/compile-fail/multi/matmut_get_row.rs b/diskann-utils/tests/compile-fail/matmut_get_row.rs similarity index 64% rename from diskann-quantization/tests/compile-fail/multi/matmut_get_row.rs rename to diskann-utils/tests/compile-fail/matmut_get_row.rs index e792b7ad5..c62e43391 100644 --- a/diskann-quantization/tests/compile-fail/multi/matmut_get_row.rs +++ b/diskann-utils/tests/compile-fail/matmut_get_row.rs @@ -3,13 +3,13 @@ * Licensed under the MIT license. */ -use diskann_quantization::multi_vector::{Mat, MatMut, Standard}; +use diskann_utils::views::{Mat, MatMut, RowMajor}; // Test that `get_row` on MatMut correctly captures an immutable borrow, // preventing mutation of the MatMut while the row is in scope. fn main() { - let mut mat: Mat> = Mat::new(Standard::new(4, 3).unwrap(), 0.0f32).unwrap(); - let mut view: MatMut<'_, Standard> = mat.as_view_mut(); + let mut mat: Mat> = Mat::from_repr(RowMajor::new(4, 3).unwrap(), 0.0f32).unwrap(); + let mut view: MatMut<'_, RowMajor> = mat.as_view_mut(); let row = view.get_row(0).unwrap(); // This should fail: we cannot mutably borrow `view` while `row` exists let _ = view.get_row_mut(1); diff --git a/diskann-quantization/tests/compile-fail/multi/matmut_get_row.stderr b/diskann-utils/tests/compile-fail/matmut_get_row.stderr similarity index 89% rename from diskann-quantization/tests/compile-fail/multi/matmut_get_row.stderr rename to diskann-utils/tests/compile-fail/matmut_get_row.stderr index 618b5e97f..bba5f4c7f 100644 --- a/diskann-quantization/tests/compile-fail/multi/matmut_get_row.stderr +++ b/diskann-utils/tests/compile-fail/matmut_get_row.stderr @@ -1,5 +1,5 @@ error[E0502]: cannot borrow `view` as mutable because it is also borrowed as immutable - --> tests/compile-fail/multi/matmut_get_row.rs:15:13 + --> tests/compile-fail/matmut_get_row.rs:15:13 | 13 | let row = view.get_row(0).unwrap(); | ---- immutable borrow occurs here diff --git a/diskann-quantization/tests/compile-fail/multi/matmut_get_row_mut.rs b/diskann-utils/tests/compile-fail/matmut_get_row_mut.rs similarity index 65% rename from diskann-quantization/tests/compile-fail/multi/matmut_get_row_mut.rs rename to diskann-utils/tests/compile-fail/matmut_get_row_mut.rs index c24afb6e2..c4b157c82 100644 --- a/diskann-quantization/tests/compile-fail/multi/matmut_get_row_mut.rs +++ b/diskann-utils/tests/compile-fail/matmut_get_row_mut.rs @@ -3,13 +3,13 @@ * Licensed under the MIT license. */ -use diskann_quantization::multi_vector::{Mat, MatMut, Standard}; +use diskann_utils::views::{Mat, MatMut, RowMajor}; // Test that `get_row_mut` on MatMut correctly captures a mutable lifetime, // preventing the MatMut from being used while the row is in scope. fn main() { - let mut mat: Mat> = Mat::new(Standard::new(4, 3).unwrap(), 0.0f32).unwrap(); - let mut view: MatMut<'_, Standard> = mat.as_view_mut(); + let mut mat: Mat> = Mat::from_repr(RowMajor::new(4, 3).unwrap(), 0.0f32).unwrap(); + let mut view: MatMut<'_, RowMajor> = mat.as_view_mut(); let row = view.get_row_mut(0).unwrap(); // This should fail: we cannot use `view` while `row` is still borrowed let _ = view.get_row_mut(1).unwrap(); diff --git a/diskann-quantization/tests/compile-fail/multi/matmut_get_row_mut.stderr b/diskann-utils/tests/compile-fail/matmut_get_row_mut.stderr similarity index 88% rename from diskann-quantization/tests/compile-fail/multi/matmut_get_row_mut.stderr rename to diskann-utils/tests/compile-fail/matmut_get_row_mut.stderr index 2d6974b26..831e433ad 100644 --- a/diskann-quantization/tests/compile-fail/multi/matmut_get_row_mut.stderr +++ b/diskann-utils/tests/compile-fail/matmut_get_row_mut.stderr @@ -1,5 +1,5 @@ error[E0499]: cannot borrow `view` as mutable more than once at a time - --> tests/compile-fail/multi/matmut_get_row_mut.rs:15:13 + --> tests/compile-fail/matmut_get_row_mut.rs:15:13 | 13 | let row = view.get_row_mut(0).unwrap(); | ---- first mutable borrow occurs here diff --git a/diskann-quantization/tests/compile-fail/multi/matmut_invariant.rs b/diskann-utils/tests/compile-fail/matmut_invariant.rs similarity index 60% rename from diskann-quantization/tests/compile-fail/multi/matmut_invariant.rs rename to diskann-utils/tests/compile-fail/matmut_invariant.rs index b5a0a5577..8caf3bde6 100644 --- a/diskann-quantization/tests/compile-fail/multi/matmut_invariant.rs +++ b/diskann-utils/tests/compile-fail/matmut_invariant.rs @@ -3,13 +3,13 @@ * Licensed under the MIT license. */ -use diskann_quantization::multi_vector::{Mat, MatMut, Standard}; +use diskann_utils::views::{Mat, MatMut, RowMajor}; // Verify that `MatMut` is invariant in any generic parameters. // // This must not compile because it would allow assigning references with a shorter lifetime // into the matrix -fn bad<'long, 'short, 'a>(v: MatMut<'a, Standard<&'long u8>>) -> MatMut<'a, Standard<&'short u8>> +fn bad<'long, 'short, 'a>(v: MatMut<'a, RowMajor<&'long u8>>) -> MatMut<'a, RowMajor<&'short u8>> where 'long: 'short, { @@ -18,6 +18,6 @@ where fn main() { let b = 0u8; - let mut m = Mat::new(Standard::new(4, 3).unwrap(), &b).unwrap(); + let mut m = Mat::from_repr(RowMajor::new(4, 3).unwrap(), &b).unwrap(); bad(m.as_view_mut()); } diff --git a/diskann-quantization/tests/compile-fail/multi/matmut_invariant.stderr b/diskann-utils/tests/compile-fail/matmut_invariant.stderr similarity index 65% rename from diskann-quantization/tests/compile-fail/multi/matmut_invariant.stderr rename to diskann-utils/tests/compile-fail/matmut_invariant.stderr index 762febdab..831826fd2 100644 --- a/diskann-quantization/tests/compile-fail/multi/matmut_invariant.stderr +++ b/diskann-utils/tests/compile-fail/matmut_invariant.stderr @@ -1,7 +1,7 @@ error: lifetime may not live long enough - --> tests/compile-fail/multi/matmut_invariant.rs:16:5 + --> tests/compile-fail/matmut_invariant.rs:16:5 | -12 | fn bad<'long, 'short, 'a>(v: MatMut<'a, Standard<&'long u8>>) -> MatMut<'a, Standard<&'short u8>> +12 | fn bad<'long, 'short, 'a>(v: MatMut<'a, RowMajor<&'long u8>>) -> MatMut<'a, RowMajor<&'short u8>> | ----- ------ lifetime `'short` defined here | | | lifetime `'long` defined here @@ -10,6 +10,6 @@ error: lifetime may not live long enough | ^ function was supposed to return data with lifetime `'long` but it is returning data with lifetime `'short` | = help: consider adding the following bound: `'short: 'long` - = note: requirement occurs because of the type `MatMut<'_, Standard<&u8>>`, which makes the generic argument `Standard<&u8>` invariant + = note: requirement occurs because of the type `MatMut<'_, RowMajor<&u8>>`, which makes the generic argument `RowMajor<&u8>` invariant = note: the struct `MatMut<'a, T>` is invariant over the parameter `T` = help: see for more information about variance diff --git a/diskann-quantization/tests/compile-fail/multi/matmut_reborrow.rs b/diskann-utils/tests/compile-fail/matmut_reborrow.rs similarity index 67% rename from diskann-quantization/tests/compile-fail/multi/matmut_reborrow.rs rename to diskann-utils/tests/compile-fail/matmut_reborrow.rs index f4bbd58b3..4bec9e756 100644 --- a/diskann-quantization/tests/compile-fail/multi/matmut_reborrow.rs +++ b/diskann-utils/tests/compile-fail/matmut_reborrow.rs @@ -3,14 +3,14 @@ * Licensed under the MIT license. */ -use diskann_quantization::multi_vector::{Mat, MatMut, Standard}; +use diskann_utils::views::{Mat, MatMut, RowMajor}; use diskann_utils::Reborrow; // Test that `reborrow` on MatMut correctly captures an immutable borrow, // preventing mutation of the MatMut while the reborrow is in scope. fn main() { - let mut mat: Mat> = Mat::new(Standard::new(4, 3).unwrap(), 0.0f32).unwrap(); - let mut view: MatMut<'_, Standard> = mat.as_view_mut(); + let mut mat: Mat> = Mat::from_repr(RowMajor::new(4, 3).unwrap(), 0.0f32).unwrap(); + let mut view: MatMut<'_, RowMajor> = mat.as_view_mut(); let immut_view = view.reborrow(); // This should fail: we cannot mutably borrow `view` while `immut_view` exists let _ = view.get_row_mut(0); diff --git a/diskann-quantization/tests/compile-fail/multi/matmut_reborrow.stderr b/diskann-utils/tests/compile-fail/matmut_reborrow.stderr similarity index 89% rename from diskann-quantization/tests/compile-fail/multi/matmut_reborrow.stderr rename to diskann-utils/tests/compile-fail/matmut_reborrow.stderr index e95118902..f7b5c850c 100644 --- a/diskann-quantization/tests/compile-fail/multi/matmut_reborrow.stderr +++ b/diskann-utils/tests/compile-fail/matmut_reborrow.stderr @@ -1,5 +1,5 @@ error[E0502]: cannot borrow `view` as mutable because it is also borrowed as immutable - --> tests/compile-fail/multi/matmut_reborrow.rs:16:13 + --> tests/compile-fail/matmut_reborrow.rs:16:13 | 14 | let immut_view = view.reborrow(); | ---- immutable borrow occurs here diff --git a/diskann-quantization/tests/compile-fail/multi/matmut_reborrow_mut.rs b/diskann-utils/tests/compile-fail/matmut_reborrow_mut.rs similarity index 68% rename from diskann-quantization/tests/compile-fail/multi/matmut_reborrow_mut.rs rename to diskann-utils/tests/compile-fail/matmut_reborrow_mut.rs index 7b26f0a5c..d47829abc 100644 --- a/diskann-quantization/tests/compile-fail/multi/matmut_reborrow_mut.rs +++ b/diskann-utils/tests/compile-fail/matmut_reborrow_mut.rs @@ -3,14 +3,14 @@ * Licensed under the MIT license. */ -use diskann_quantization::multi_vector::{Mat, MatMut, Standard}; +use diskann_utils::views::{Mat, MatMut, RowMajor}; use diskann_utils::ReborrowMut; // Test that `reborrow_mut` on MatMut correctly captures a mutable lifetime, // preventing the original MatMut from being used while the reborrow is in scope. fn main() { - let mut mat: Mat> = Mat::new(Standard::new(4, 3).unwrap(), 0.0f32).unwrap(); - let mut view: MatMut<'_, Standard> = mat.as_view_mut(); + let mut mat: Mat> = Mat::from_repr(RowMajor::new(4, 3).unwrap(), 0.0f32).unwrap(); + let mut view: MatMut<'_, RowMajor> = mat.as_view_mut(); let reborrowed = view.reborrow_mut(); // This should fail: we cannot use `view` while `reborrowed` is still alive let _ = view.num_vectors(); diff --git a/diskann-quantization/tests/compile-fail/multi/matmut_reborrow_mut.stderr b/diskann-utils/tests/compile-fail/matmut_reborrow_mut.stderr similarity index 88% rename from diskann-quantization/tests/compile-fail/multi/matmut_reborrow_mut.stderr rename to diskann-utils/tests/compile-fail/matmut_reborrow_mut.stderr index 4429ae199..c6bd91776 100644 --- a/diskann-quantization/tests/compile-fail/multi/matmut_reborrow_mut.stderr +++ b/diskann-utils/tests/compile-fail/matmut_reborrow_mut.stderr @@ -1,5 +1,5 @@ error[E0502]: cannot borrow `view` as immutable because it is also borrowed as mutable - --> tests/compile-fail/multi/matmut_reborrow_mut.rs:16:13 + --> tests/compile-fail/matmut_reborrow_mut.rs:16:13 | 14 | let reborrowed = view.reborrow_mut(); | ---- mutable borrow occurs here diff --git a/diskann-quantization/tests/compile-fail/multi/matmut_rows.rs b/diskann-utils/tests/compile-fail/matmut_rows.rs similarity index 66% rename from diskann-quantization/tests/compile-fail/multi/matmut_rows.rs rename to diskann-utils/tests/compile-fail/matmut_rows.rs index bbb37ef44..6a097f5e1 100644 --- a/diskann-quantization/tests/compile-fail/multi/matmut_rows.rs +++ b/diskann-utils/tests/compile-fail/matmut_rows.rs @@ -3,13 +3,13 @@ * Licensed under the MIT license. */ -use diskann_quantization::multi_vector::{Mat, MatMut, Standard}; +use diskann_utils::views::{Mat, MatMut, RowMajor}; // Test that `rows` on MatMut correctly captures an immutable borrow, // preventing mutation of the MatMut while the iterator is in scope. fn main() { - let mut mat: Mat> = Mat::new(Standard::new(4, 3).unwrap(), 0.0f32).unwrap(); - let mut view: MatMut<'_, Standard> = mat.as_view_mut(); + let mut mat: Mat> = Mat::from_repr(RowMajor::new(4, 3).unwrap(), 0.0f32).unwrap(); + let mut view: MatMut<'_, RowMajor> = mat.as_view_mut(); let iter = view.rows(); // This should fail: we cannot mutably borrow `view` while `iter` exists let _ = view.get_row_mut(0); diff --git a/diskann-quantization/tests/compile-fail/multi/matmut_rows.stderr b/diskann-utils/tests/compile-fail/matmut_rows.stderr similarity index 89% rename from diskann-quantization/tests/compile-fail/multi/matmut_rows.stderr rename to diskann-utils/tests/compile-fail/matmut_rows.stderr index 4b19eea9e..c27b14164 100644 --- a/diskann-quantization/tests/compile-fail/multi/matmut_rows.stderr +++ b/diskann-utils/tests/compile-fail/matmut_rows.stderr @@ -1,5 +1,5 @@ error[E0502]: cannot borrow `view` as mutable because it is also borrowed as immutable - --> tests/compile-fail/multi/matmut_rows.rs:15:13 + --> tests/compile-fail/matmut_rows.rs:15:13 | 13 | let iter = view.rows(); | ---- immutable borrow occurs here diff --git a/diskann-quantization/tests/compile-fail/multi/matmut_rows_mut.rs b/diskann-utils/tests/compile-fail/matmut_rows_mut.rs similarity index 66% rename from diskann-quantization/tests/compile-fail/multi/matmut_rows_mut.rs rename to diskann-utils/tests/compile-fail/matmut_rows_mut.rs index b4668e5fc..b5261b62f 100644 --- a/diskann-quantization/tests/compile-fail/multi/matmut_rows_mut.rs +++ b/diskann-utils/tests/compile-fail/matmut_rows_mut.rs @@ -3,13 +3,13 @@ * Licensed under the MIT license. */ -use diskann_quantization::multi_vector::{Mat, MatMut, Standard}; +use diskann_utils::views::{Mat, MatMut, RowMajor}; // Test that the `rows_mut` iterator on MatMut correctly captures a mutable lifetime, // preventing the MatMut from being used while the iterator is in scope. fn main() { - let mut mat: Mat> = Mat::new(Standard::new(4, 3).unwrap(), 0.0f32).unwrap(); - let mut view: MatMut<'_, Standard> = mat.as_view_mut(); + let mut mat: Mat> = Mat::from_repr(RowMajor::new(4, 3).unwrap(), 0.0f32).unwrap(); + let mut view: MatMut<'_, RowMajor> = mat.as_view_mut(); let iter = view.rows_mut(); // This should fail: we cannot use `view` while the mutable iterator is alive let _ = view.num_vectors(); diff --git a/diskann-quantization/tests/compile-fail/multi/matmut_rows_mut.stderr b/diskann-utils/tests/compile-fail/matmut_rows_mut.stderr similarity index 88% rename from diskann-quantization/tests/compile-fail/multi/matmut_rows_mut.stderr rename to diskann-utils/tests/compile-fail/matmut_rows_mut.stderr index 2fa4d882c..d7e6771e1 100644 --- a/diskann-quantization/tests/compile-fail/multi/matmut_rows_mut.stderr +++ b/diskann-utils/tests/compile-fail/matmut_rows_mut.stderr @@ -1,5 +1,5 @@ error[E0502]: cannot borrow `view` as immutable because it is also borrowed as mutable - --> tests/compile-fail/multi/matmut_rows_mut.rs:15:13 + --> tests/compile-fail/matmut_rows_mut.rs:15:13 | 13 | let iter = view.rows_mut(); | ---- mutable borrow occurs here diff --git a/diskann-quantization/tests/compile-fail/multi/matref_get_row.rs b/diskann-utils/tests/compile-fail/matref_get_row.rs similarity index 68% rename from diskann-quantization/tests/compile-fail/multi/matref_get_row.rs rename to diskann-utils/tests/compile-fail/matref_get_row.rs index 5478adee8..430f0daae 100644 --- a/diskann-quantization/tests/compile-fail/multi/matref_get_row.rs +++ b/diskann-utils/tests/compile-fail/matref_get_row.rs @@ -3,13 +3,13 @@ * Licensed under the MIT license. */ -use diskann_quantization::multi_vector::{Mat, MatRef, Standard}; +use diskann_utils::views::{Mat, MatRef, RowMajor}; // Test that `get_row` on MatRef returns a row with the correct lifetime, // and that an immutable borrow is held while the row is in scope. fn main() { - let mut mat: Mat> = Mat::new(Standard::new(4, 3).unwrap(), 0.0f32).unwrap(); - let view: MatRef<'_, Standard> = mat.as_view(); + let mut mat: Mat> = Mat::from_repr(RowMajor::new(4, 3).unwrap(), 0.0f32).unwrap(); + let view: MatRef<'_, RowMajor> = mat.as_view(); let row = view.get_row(0).unwrap(); // This should fail: we cannot mutably borrow `mat` while `row` exists // (since `row` holds a reference derived from `mat`) diff --git a/diskann-quantization/tests/compile-fail/multi/matref_get_row.stderr b/diskann-utils/tests/compile-fail/matref_get_row.stderr similarity index 75% rename from diskann-quantization/tests/compile-fail/multi/matref_get_row.stderr rename to diskann-utils/tests/compile-fail/matref_get_row.stderr index 1d83c312b..997928827 100644 --- a/diskann-quantization/tests/compile-fail/multi/matref_get_row.stderr +++ b/diskann-utils/tests/compile-fail/matref_get_row.stderr @@ -1,7 +1,7 @@ error[E0502]: cannot borrow `mat` as mutable because it is also borrowed as immutable - --> tests/compile-fail/multi/matref_get_row.rs:16:13 + --> tests/compile-fail/matref_get_row.rs:16:13 | -12 | let view: MatRef<'_, Standard> = mat.as_view(); +12 | let view: MatRef<'_, RowMajor> = mat.as_view(); | --- immutable borrow occurs here ... 16 | let _ = mat.as_view_mut(); diff --git a/diskann-quantization/tests/compile-fail/multi/matref_rows.rs b/diskann-utils/tests/compile-fail/matref_rows.rs similarity index 66% rename from diskann-quantization/tests/compile-fail/multi/matref_rows.rs rename to diskann-utils/tests/compile-fail/matref_rows.rs index 82f2dc02c..0a9fa150e 100644 --- a/diskann-quantization/tests/compile-fail/multi/matref_rows.rs +++ b/diskann-utils/tests/compile-fail/matref_rows.rs @@ -3,13 +3,13 @@ * Licensed under the MIT license. */ -use diskann_quantization::multi_vector::{Mat, MatRef, Standard}; +use diskann_utils::views::{Mat, MatRef, RowMajor}; // Test that `rows` on MatRef returns an iterator with the correct lifetime, // preventing mutation of the underlying Mat while iterating. fn main() { - let mut mat: Mat> = Mat::new(Standard::new(4, 3).unwrap(), 0.0f32).unwrap(); - let view: MatRef<'_, Standard> = mat.as_view(); + let mut mat: Mat> = Mat::from_repr(RowMajor::new(4, 3).unwrap(), 0.0f32).unwrap(); + let view: MatRef<'_, RowMajor> = mat.as_view(); let iter = view.rows(); // This should fail: we cannot mutably borrow `mat` while `iter` exists let _ = mat.as_view_mut(); diff --git a/diskann-quantization/tests/compile-fail/multi/matref_rows.stderr b/diskann-utils/tests/compile-fail/matref_rows.stderr similarity index 75% rename from diskann-quantization/tests/compile-fail/multi/matref_rows.stderr rename to diskann-utils/tests/compile-fail/matref_rows.stderr index d2a19baaa..0c6980b9e 100644 --- a/diskann-quantization/tests/compile-fail/multi/matref_rows.stderr +++ b/diskann-utils/tests/compile-fail/matref_rows.stderr @@ -1,7 +1,7 @@ error[E0502]: cannot borrow `mat` as mutable because it is also borrowed as immutable - --> tests/compile-fail/multi/matref_rows.rs:15:13 + --> tests/compile-fail/matref_rows.rs:15:13 | -12 | let view: MatRef<'_, Standard> = mat.as_view(); +12 | let view: MatRef<'_, RowMajor> = mat.as_view(); | --- immutable borrow occurs here ... 15 | let _ = mat.as_view_mut(); diff --git a/diskann-utils/tests/compile_fail.rs b/diskann-utils/tests/compile_fail.rs new file mode 100644 index 000000000..7f7885a34 --- /dev/null +++ b/diskann-utils/tests/compile_fail.rs @@ -0,0 +1,14 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +// Borrow/lifetime/variance guarantees of the `Mat`/`MatRef`/`MatMut` matrix framework. +// These are check-time errors, so no `pass` bootstrap is needed to force monomorphization. +#![cfg(not(miri))] + +#[test] +fn compile_fail() { + let t = trybuild::TestCases::new(); + t.compile_fail("tests/compile-fail/*.rs"); +} diff --git a/diskann/src/error/ann_error.rs b/diskann/src/error/ann_error.rs index c27e622c1..401c5d77d 100644 --- a/diskann/src/error/ann_error.rs +++ b/diskann/src/error/ann_error.rs @@ -595,10 +595,7 @@ impl From for ANNError { } } -impl From> for ANNError -where - T: diskann_utils::views::DenseData, -{ +impl From> for ANNError { #[track_caller] fn from(err: diskann_utils::views::TryFromError) -> Self { Self::from(err.as_static()) @@ -1541,16 +1538,6 @@ Caused by: assert_eq!(ann_err.kind(), ANNErrorKind::IOError); } - #[test] - fn from_try_from_error_light() { - let data: &[f32] = &[1.0, 2.0, 3.0]; - let light = diskann_utils::views::MatrixView::try_from(data, 2, 2) - .unwrap_err() - .as_static(); - let ann_err = ANNError::from(light); - assert_eq!(ann_err.kind(), ANNErrorKind::DimensionMismatchError); - } - #[test] fn from_try_from_error() { let data: &[f32] = &[1.0, 2.0, 3.0];