Skip to main content

diskann_quantization/multi_vector/
block_transposed.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6//! Block-transposed matrix types with configurable packing.
7//!
8//! This module provides block-transposed matrix types — [`BlockTransposed`] (owned),
9//! [`BlockTransposedRef`] (shared view), and [`BlockTransposedMut`] (mutable view) —
10//! where groups of `GROUP` rows are stored in transposed form to enable efficient SIMD
11//! processing. An optional packing factor `PACK` interleaves adjacent columns within
12//! each group, which can be used to feed SIMD instructions that operate on packed pairs
13//! (e.g. `vpmaddwd` with `PACK = 2`).
14//!
15//! # Layout
16//!
17//! ## `PACK = 1` (standard block-transpose)
18//!
19//! Given a logical matrix with rows `a`, `b`, `c`, `d`, `e` (each with `K` columns)
20//! and `GROUP = 3`:
21//!
22//! ```text
23//!            Group Size (3)
24//!            <---------->
25//!
26//!            +----------+    ^
27//!            | a0 b0 c0 |    |
28//!            | a1 b1 c1 |    |
29//!            | a2 b2 c2 |    | Block Size (K)
30//!  Block 0   | ...      |    |
31//!  (Full)    | aK bK cK |    |
32//!            +----------+    v
33//!            +----------+
34//!            | d0 e0 XX |
35//!  Block 1   | d1 e1 XX |
36//!  (Partial) | ...      |
37//!            | dK eK XX |
38//!            +----------+
39//! ```
40//!
41//! ## `PACK = 2` (super-packed)
42//!
43//! With `GROUP = 4`, `PACK = 2`, and a logical matrix with rows `a`, `b`, `c`, `d`,
44//! `e`, `f` (each with **5** columns — odd, to show padding), adjacent column-pairs
45//! are interleaved per row within each group panel:
46//!
47//! ```text
48//!              GROUP × PACK (4 × 2 = 8)
49//!              <----------------------------->
50//!
51//!              +-----------------------------+    ^
52//!              | a0 a1  b0 b1  c0 c1  d0 d1  |    |  col-pair (0, 1)
53//!              | a2 a3  b2 b3  c2 c3  d2 d3  |    |  col-pair (2, 3)
54//!    Block 0   | a4 __  b4 __  c4 __  d4 __  |    |  col-pair (4, pad)
55//!    (Full)    +-----------------------------+    v
56//!              +-----------------------------+
57//!              | e0 e1  f0 f1  XX XX  XX XX  |       col-pair (0, 1)
58//!    Block 1   | e2 e3  f2 f3  XX XX  XX XX  |       col-pair (2, 3)
59//!    (Partial) | e4 __  f4 __  XX XX  XX XX  |       col-pair (4, pad)
60//!              +-----------------------------+
61//!
62//!    __ = zero (column padding)    XX = zero (row padding)
63//!    padded_ncols = 6  (5 rounded up to next multiple of PACK)
64//!    Block Size  = padded_ncols / PACK = 3 physical rows per block
65//! ```
66//!
67//! Each physical row of a block holds one column-pair across all `GROUP` rows.
68//! For example, the first physical row stores columns `(0, 1)` for rows
69//! `a, b, c, d` interleaved as `[a0, a1, b0, b1, c0, c1, d0, d1]`.
70//!
71//! Because `ncols = 5` is odd (not a multiple of `PACK = 2`), the last
72//! column-pair `(4, pad)` is zero-padded: `[a4, 0, b4, 0, c4, 0, d4, 0]`.
73//!
74//! # Constraints
75//!
76//! - `GROUP > 0`
77//! - `PACK > 0`
78//! - `GROUP % PACK == 0`
79
80use std::{alloc::Layout, marker::PhantomData, ptr::NonNull};
81
82use diskann_utils::{
83    Reborrow, ReborrowMut,
84    strided::StridedView,
85    views::{MatrixView, MutMatrixView},
86};
87
88use super::matrix::{
89    Defaulted, LayoutError, Mat, MatMut, MatRef, NewCloned, NewMut, NewOwned, NewRef, Overflow,
90    Repr, ReprMut, ReprOwned, SliceError,
91};
92use crate::bits::{AsMutPtr, AsPtr, MutSlicePtr, SlicePtr};
93use crate::utils;
94
95/// Round `ncols` up to the next multiple of `PACK`.
96#[inline]
97fn padded_ncols<const PACK: usize>(ncols: usize) -> usize {
98    ncols.next_multiple_of(PACK)
99}
100
101/// Compute the total number of `T` elements required to store a block-transposed matrix
102/// of `nrows x ncols` with group size `GROUP` and packing factor `PACK`.
103///
104/// This is the **unchecked** flavor — it assumes the caller has already validated that
105/// the dimensions do not overflow (e.g. after construction). For use in the constructor,
106/// prefer [`checked_compute_capacity`].
107///
108/// Compile-time constraints (`GROUP > 0`, `PACK > 0`, `GROUP % PACK == 0`) are enforced
109/// by [`BlockTransposedRepr::_ASSERTIONS`]; this function does **not** duplicate them.
110#[inline]
111fn compute_capacity<const GROUP: usize, const PACK: usize>(nrows: usize, ncols: usize) -> usize {
112    nrows.next_multiple_of(GROUP) * padded_ncols::<PACK>(ncols)
113}
114
115/// Checked variant of [`compute_capacity`] that returns `None` if any intermediate
116/// arithmetic overflows. Used by the constructor to reject impossibly large dimensions
117/// before committing to an allocation.
118#[inline]
119fn checked_compute_capacity<const GROUP: usize, const PACK: usize>(
120    nrows: usize,
121    ncols: usize,
122) -> Option<usize> {
123    nrows
124        .checked_next_multiple_of(GROUP)?
125        .checked_mul(ncols.checked_next_multiple_of(PACK)?)
126}
127
128/// Compute the linear index for the element at logical `(row, col)` in a block-transposed
129/// layout with group size `GROUP`, packing factor `PACK`, and `ncols` logical columns.
130#[inline]
131fn linear_index<const GROUP: usize, const PACK: usize>(
132    row: usize,
133    col: usize,
134    ncols: usize,
135) -> usize {
136    let pncols = padded_ncols::<PACK>(ncols);
137    let block = row / GROUP;
138    let row_in_block = row % GROUP;
139    block * GROUP * pncols + (col / PACK) * GROUP * PACK + row_in_block * PACK + (col % PACK)
140}
141
142/// Compute the offset from a row's base pointer (at col=0) to the element at `col`.
143///
144/// This is purely a function of the column index and the const layout parameters, not
145/// of any particular matrix's dimensions.
146#[inline]
147fn col_offset<const GROUP: usize, const PACK: usize>(col: usize) -> usize {
148    (col / PACK) * GROUP * PACK + (col % PACK)
149}
150
151/// Internal layout descriptor for block-transposed matrices.
152///
153/// This is not part of the public API — use [`BlockTransposed`], [`BlockTransposedRef`],
154/// or [`BlockTransposedMut`] instead.
155#[derive(Debug, Clone, Copy, PartialEq, Eq)]
156pub(crate) struct BlockTransposedRepr<T, const GROUP: usize, const PACK: usize = 1> {
157    nrows: usize,
158    ncols: usize,
159    _elem: PhantomData<T>,
160}
161
162impl<T: Copy, const GROUP: usize, const PACK: usize> BlockTransposedRepr<T, GROUP, PACK> {
163    // Compile-time assertions — evaluated whenever any method references this constant.
164    const _ASSERTIONS: () = {
165        assert!(GROUP > 0, "group size GROUP must be positive");
166        assert!(PACK > 0, "packing factor PACK must be positive");
167        assert!(
168            GROUP.is_multiple_of(PACK),
169            "GROUP must be divisible by PACK"
170        );
171    };
172
173    /// Create a new `BlockTransposedRepr` descriptor.
174    ///
175    /// Successful construction requires that the total memory for the backing allocation
176    /// does not exceed `isize::MAX`.
177    pub fn new(nrows: usize, ncols: usize) -> Result<Self, Overflow> {
178        let () = Self::_ASSERTIONS;
179        let capacity = checked_compute_capacity::<GROUP, PACK>(nrows, ncols)
180            .ok_or_else(|| Overflow::for_type::<T>(nrows, ncols))?;
181        Overflow::check_byte_budget::<T>(capacity, nrows, ncols)?;
182        Ok(Self {
183            nrows,
184            ncols,
185            _elem: PhantomData,
186        })
187    }
188
189    // ── Query helpers ────────────────────────────────────────────────
190
191    /// The total number of `T` elements in the backing allocation (including padding).
192    #[inline]
193    fn storage_len(&self) -> usize {
194        compute_capacity::<GROUP, PACK>(self.nrows, self.ncols)
195    }
196
197    /// Number of logical rows.
198    #[inline]
199    fn nrows(&self) -> usize {
200        self.nrows
201    }
202
203    /// Number of logical columns (dimensionality).
204    #[inline]
205    pub fn ncols(&self) -> usize {
206        self.ncols
207    }
208
209    /// Number of physical (padded) columns — logical columns rounded up to
210    /// the next multiple of `PACK`.
211    #[inline]
212    pub fn padded_ncols(&self) -> usize {
213        padded_ncols::<PACK>(self.ncols)
214    }
215
216    /// Number of completely full blocks.
217    #[inline]
218    pub fn full_blocks(&self) -> usize {
219        self.nrows / GROUP
220    }
221
222    /// Total number of blocks including a possible partially-filled tail.
223    #[inline]
224    pub fn num_blocks(&self) -> usize {
225        self.nrows.div_ceil(GROUP)
226    }
227
228    /// Number of valid elements in the last block, or 0 if all blocks are full.
229    #[inline]
230    pub fn remainder(&self) -> usize {
231        self.nrows % GROUP
232    }
233
234    /// Total number of logical rows rounded up to the next multiple of `GROUP`.
235    ///
236    /// This is the number of "available" row slots in the backing allocation,
237    /// including zero-padded rows in the last (possibly partial) block.
238    #[inline]
239    pub fn padded_nrows(&self) -> usize {
240        self.num_blocks() * GROUP
241    }
242
243    /// The stride (in elements) between the start of consecutive blocks.
244    #[inline]
245    fn block_stride(&self) -> usize {
246        GROUP * self.padded_ncols()
247    }
248
249    /// The linear offset of the start of `block`.
250    #[inline]
251    fn block_offset(&self, block: usize) -> usize {
252        block * self.block_stride()
253    }
254
255    /// Verify that `slice` has exactly `self.storage_len()` elements.
256    fn check_slice(&self, slice: &[T]) -> Result<(), SliceError> {
257        let cap = self.storage_len();
258        if slice.len() != cap {
259            Err(SliceError::LengthMismatch {
260                expected: cap,
261                found: slice.len(),
262            })
263        } else {
264            Ok(())
265        }
266    }
267
268    /// Helper: wrap a `Box<[T]>` into a [`Mat`] without any further checks.
269    ///
270    /// # Safety
271    ///
272    /// `b.len()` must equal `self.storage_len()`.
273    unsafe fn box_to_mat(self, b: Box<[T]>) -> Mat<Self> {
274        debug_assert_eq!(b.len(), self.storage_len(), "safety contract violated");
275
276        let ptr = utils::box_into_nonnull(b).cast::<u8>();
277
278        // SAFETY: `ptr` is properly aligned and compatible with our layout.
279        unsafe { Mat::from_raw_parts(self, ptr) }
280    }
281}
282
283// ════════════════════════════════════════════════════════════════════
284// Row view types
285// ════════════════════════════════════════════════════════════════════
286
287/// An immutable view of a single logical row in a block-transposed matrix.
288///
289/// Because the elements of a logical row are strided (not contiguous), this struct
290/// provides indexed access and iteration over the row's elements.
291#[derive(Debug, Clone, Copy)]
292pub struct Row<'a, T, const GROUP: usize, const PACK: usize = 1> {
293    /// Pointer to the element at `(row, col=0)` in the backing allocation.
294    base: SlicePtr<'a, T>,
295    ncols: usize,
296}
297
298impl<T: Copy, const GROUP: usize, const PACK: usize> Row<'_, T, GROUP, PACK> {
299    /// Number of elements (columns) in this row.
300    #[inline]
301    pub fn len(&self) -> usize {
302        self.ncols
303    }
304
305    /// Whether the row is empty.
306    #[inline]
307    pub fn is_empty(&self) -> bool {
308        self.ncols == 0
309    }
310
311    /// Get a reference to the element at column `col`, or `None` if out of bounds.
312    #[inline]
313    pub fn get(&self, col: usize) -> Option<&T> {
314        if col < self.ncols {
315            // SAFETY: bounds checked, offset computed from validated layout.
316            Some(unsafe { &*self.base.as_ptr().add(col_offset::<GROUP, PACK>(col)) })
317        } else {
318            None
319        }
320    }
321
322    /// Return an iterator over the elements of this row.
323    #[inline]
324    pub fn iter(&self) -> RowIter<'_, T, GROUP, PACK> {
325        RowIter {
326            base: self.base,
327            col: 0,
328            ncols: self.ncols,
329        }
330    }
331}
332
333impl<T: Copy, const GROUP: usize, const PACK: usize> std::ops::Index<usize>
334    for Row<'_, T, GROUP, PACK>
335{
336    type Output = T;
337
338    #[inline]
339    #[allow(clippy::panic)] // Index is expected to panic on OOB
340    fn index(&self, col: usize) -> &Self::Output {
341        self.get(col)
342            .unwrap_or_else(|| panic!("column index {col} out of bounds (ncols = {})", self.ncols))
343    }
344}
345
346/// Iterator over the elements of a [`Row`].
347#[derive(Debug, Clone)]
348pub struct RowIter<'a, T, const GROUP: usize, const PACK: usize = 1> {
349    base: SlicePtr<'a, T>,
350    col: usize,
351    ncols: usize,
352}
353
354impl<T: Copy, const GROUP: usize, const PACK: usize> Iterator for RowIter<'_, T, GROUP, PACK> {
355    type Item = T;
356
357    #[inline]
358    fn next(&mut self) -> Option<Self::Item> {
359        if self.col >= self.ncols {
360            return None;
361        }
362        // SAFETY: col < ncols means the offset is within the backing allocation.
363        let val = unsafe { *self.base.as_ptr().add(col_offset::<GROUP, PACK>(self.col)) };
364        self.col += 1;
365        Some(val)
366    }
367
368    #[inline]
369    fn size_hint(&self) -> (usize, Option<usize>) {
370        let remaining = self.ncols - self.col;
371        (remaining, Some(remaining))
372    }
373}
374
375impl<T: Copy, const GROUP: usize, const PACK: usize> ExactSizeIterator
376    for RowIter<'_, T, GROUP, PACK>
377{
378}
379impl<T: Copy, const GROUP: usize, const PACK: usize> std::iter::FusedIterator
380    for RowIter<'_, T, GROUP, PACK>
381{
382}
383
384/// A mutable view of a single logical row in a block-transposed matrix.
385#[derive(Debug)]
386pub struct RowMut<'a, T, const GROUP: usize, const PACK: usize = 1> {
387    base: MutSlicePtr<'a, T>,
388    ncols: usize,
389}
390
391impl<T: Copy, const GROUP: usize, const PACK: usize> RowMut<'_, T, GROUP, PACK> {
392    /// Number of elements (columns) in this row.
393    #[inline]
394    pub fn len(&self) -> usize {
395        self.ncols
396    }
397
398    /// Whether the row is empty.
399    #[inline]
400    pub fn is_empty(&self) -> bool {
401        self.ncols == 0
402    }
403
404    /// Get a reference to the element at column `col`, or `None` if out of bounds.
405    #[inline]
406    pub fn get(&self, col: usize) -> Option<&T> {
407        if col < self.ncols {
408            // SAFETY: bounds checked.
409            Some(unsafe { &*self.base.as_ptr().add(col_offset::<GROUP, PACK>(col)) })
410        } else {
411            None
412        }
413    }
414
415    /// Get a mutable reference to the element at column `col`, or `None` if out of bounds.
416    #[inline]
417    pub fn get_mut(&mut self, col: usize) -> Option<&mut T> {
418        if col < self.ncols {
419            // SAFETY: bounds checked.
420            Some(unsafe { &mut *self.base.as_mut_ptr().add(col_offset::<GROUP, PACK>(col)) })
421        } else {
422            None
423        }
424    }
425
426    /// Set the element at column `col`.
427    ///
428    /// # Panics
429    ///
430    /// Panics if `col >= self.len()`.
431    #[inline]
432    pub fn set(&mut self, col: usize, value: T) {
433        assert!(
434            col < self.ncols,
435            "column index {col} out of bounds (ncols = {})",
436            self.ncols
437        );
438        // SAFETY: bounds checked.
439        unsafe { *self.base.as_mut_ptr().add(col_offset::<GROUP, PACK>(col)) = value };
440    }
441}
442
443impl<T: Copy, const GROUP: usize, const PACK: usize> std::ops::Index<usize>
444    for RowMut<'_, T, GROUP, PACK>
445{
446    type Output = T;
447
448    #[inline]
449    #[allow(clippy::panic)] // Index is expected to panic on OOB
450    fn index(&self, col: usize) -> &Self::Output {
451        self.get(col)
452            .unwrap_or_else(|| panic!("column index {col} out of bounds (ncols = {})", self.ncols))
453    }
454}
455
456impl<T: Copy, const GROUP: usize, const PACK: usize> std::ops::IndexMut<usize>
457    for RowMut<'_, T, GROUP, PACK>
458{
459    #[inline]
460    #[allow(clippy::panic)] // IndexMut is expected to panic on OOB
461    fn index_mut(&mut self, col: usize) -> &mut Self::Output {
462        let ncols = self.ncols;
463        self.get_mut(col)
464            .unwrap_or_else(|| panic!("column index {col} out of bounds (ncols = {ncols})"))
465    }
466}
467
468// ════════════════════════════════════════════════════════════════════
469// Repr / ReprMut / ReprOwned
470// ════════════════════════════════════════════════════════════════════
471
472// SAFETY: `get_row` produces a valid `Row` for valid indices. The layout
473// reports the correct capacity for the block-transposed backing allocation.
474unsafe impl<T: Copy, const GROUP: usize, const PACK: usize> Repr
475    for BlockTransposedRepr<T, GROUP, PACK>
476{
477    type Row<'a>
478        = Row<'a, T, GROUP, PACK>
479    where
480        Self: 'a;
481
482    fn nrows(&self) -> usize {
483        self.nrows
484    }
485
486    fn layout(&self) -> Result<Layout, LayoutError> {
487        Ok(Layout::array::<T>(self.storage_len())?)
488    }
489
490    unsafe fn get_row<'a>(self, ptr: NonNull<u8>, i: usize) -> Self::Row<'a> {
491        debug_assert!(i < self.nrows);
492
493        // When ncols == 0 the backing allocation is zero-sized, so we must not
494        // compute any pointer offset.  Return a dangling base instead.
495        if self.ncols == 0 {
496            return Row {
497                // SAFETY: The row is empty (ncols == 0) so the pointer will never be
498                // dereferenced. A dangling `NonNull` satisfies the non-null invariant.
499                base: unsafe { SlicePtr::new_unchecked(NonNull::dangling()) },
500                ncols: 0,
501            };
502        }
503
504        let base_ptr = ptr.as_ptr().cast::<T>();
505        let offset = linear_index::<GROUP, PACK>(i, 0, self.ncols);
506
507        // SAFETY: The caller asserts `i < self.nrows()`. The backing allocation has at
508        // least `self.storage_len()` elements, so the computed offset is in bounds.
509        let row_base = unsafe { base_ptr.add(offset) };
510
511        Row {
512            // SAFETY: `row_base` is derived from a `NonNull<u8>` with a valid offset,
513            // so it is non-null. The lifetime is tied to the caller's `'a`.
514            base: unsafe { SlicePtr::new_unchecked(NonNull::new_unchecked(row_base)) },
515            ncols: self.ncols,
516        }
517    }
518}
519
520// SAFETY: `get_row_mut` produces a valid `RowMut`. Disjoint row indices
521// produce disjoint base pointers because each row within a block starts at a unique
522// offset modulo GROUP.
523unsafe impl<T: Copy, const GROUP: usize, const PACK: usize> ReprMut
524    for BlockTransposedRepr<T, GROUP, PACK>
525{
526    type RowMut<'a>
527        = RowMut<'a, T, GROUP, PACK>
528    where
529        Self: 'a;
530
531    unsafe fn get_row_mut<'a>(self, ptr: NonNull<u8>, i: usize) -> Self::RowMut<'a> {
532        debug_assert!(i < self.nrows);
533
534        // When ncols == 0 the backing allocation is zero-sized, so we must not
535        // compute any pointer offset.  Return a dangling base instead.
536        if self.ncols == 0 {
537            return RowMut {
538                // SAFETY: The row is empty (ncols == 0) so the pointer will never be
539                // dereferenced. A dangling `NonNull` satisfies the non-null invariant.
540                base: unsafe { MutSlicePtr::new_unchecked(NonNull::dangling()) },
541                ncols: 0,
542            };
543        }
544
545        let base_ptr = ptr.as_ptr().cast::<T>();
546        let offset = linear_index::<GROUP, PACK>(i, 0, self.ncols);
547
548        // SAFETY: `i < self.nrows` (debug-asserted) guarantees the offset is within
549        // the backing allocation. Same reasoning as `get_row`.
550        let row_base = unsafe { base_ptr.add(offset) };
551
552        RowMut {
553            // SAFETY: `row_base` is derived from a `NonNull<u8>` with a valid offset,
554            // so it is non-null. The lifetime is tied to the caller's `'a`.
555            base: unsafe { MutSlicePtr::new_unchecked(NonNull::new_unchecked(row_base)) },
556            ncols: self.ncols,
557        }
558    }
559}
560
561// SAFETY: Memory is deallocated by reconstructing the `Box<[T]>` that was created during
562// `NewOwned`.
563unsafe impl<T: Copy, const GROUP: usize, const PACK: usize> ReprOwned
564    for BlockTransposedRepr<T, GROUP, PACK>
565{
566    unsafe fn drop(self, ptr: NonNull<u8>) {
567        // SAFETY: `ptr` was obtained from `Box::into_raw` with length `self.storage_len()`.
568        unsafe {
569            let slice_ptr =
570                std::ptr::slice_from_raw_parts_mut(ptr.cast::<T>().as_ptr(), self.storage_len());
571            let _ = Box::from_raw(slice_ptr);
572        }
573    }
574}
575
576// ════════════════════════════════════════════════════════════════════
577// Constructors
578// ════════════════════════════════════════════════════════════════════
579
580// SAFETY: The returned `Mat` contains a `Box` with exactly `self.storage_len()` elements.
581unsafe impl<T: Copy, const GROUP: usize, const PACK: usize> NewOwned<T>
582    for BlockTransposedRepr<T, GROUP, PACK>
583{
584    type Error = crate::error::Infallible;
585
586    fn new_owned(self, value: T) -> Result<Mat<Self>, Self::Error> {
587        let b: Box<[T]> = vec![value; self.storage_len()].into_boxed_slice();
588
589        // SAFETY: By construction, `b.len() == self.storage_len()`.
590        Ok(unsafe { self.box_to_mat(b) })
591    }
592}
593
594// SAFETY: This safely re-uses `<Self as NewOwned<T>>`.
595unsafe impl<T: Copy + Default, const GROUP: usize, const PACK: usize> NewOwned<Defaulted>
596    for BlockTransposedRepr<T, GROUP, PACK>
597{
598    type Error = crate::error::Infallible;
599
600    fn new_owned(self, _: Defaulted) -> Result<Mat<Self>, Self::Error> {
601        self.new_owned(T::default())
602    }
603}
604
605impl<T: Copy, const GROUP: usize, const PACK: usize> NewCloned
606    for BlockTransposedRepr<T, GROUP, PACK>
607{
608    fn new_cloned(v: MatRef<'_, Self>) -> Mat<Self> {
609        let b: Box<[T]> = BlockTransposedRef::new(v).as_slice().into();
610
611        // SAFETY: `b` was copied from the complete backing allocation and therefore
612        // has exactly `v.repr().storage_len()` elements.
613        unsafe { v.repr().box_to_mat(b) }
614    }
615}
616
617// SAFETY: This checks slice length against storage_len.
618unsafe impl<T: Copy, const GROUP: usize, const PACK: usize> NewRef<T>
619    for BlockTransposedRepr<T, GROUP, PACK>
620{
621    type Error = SliceError;
622
623    fn new_ref(self, data: &[T]) -> Result<MatRef<'_, Self>, Self::Error> {
624        self.check_slice(data)?;
625
626        // SAFETY: `check_slice` verified the length.
627        Ok(unsafe { MatRef::from_raw_parts(self, utils::as_nonnull(data).cast::<u8>()) })
628    }
629}
630
631// SAFETY: This checks slice length against storage_len.
632unsafe impl<T: Copy, const GROUP: usize, const PACK: usize> NewMut<T>
633    for BlockTransposedRepr<T, GROUP, PACK>
634{
635    type Error = SliceError;
636
637    fn new_mut(self, data: &mut [T]) -> Result<MatMut<'_, Self>, Self::Error> {
638        self.check_slice(data)?;
639
640        // SAFETY: `check_slice` verified the length.
641        Ok(unsafe { MatMut::from_raw_parts(self, utils::as_nonnull_mut(data).cast::<u8>()) })
642    }
643}
644
645// ════════════════════════════════════════════════════════════════════
646// Delegation macro
647// ════════════════════════════════════════════════════════════════════
648
649/// Generates a forwarding method that delegates to `self.as_view().$name(...)`.
650///
651/// The generated doc-comment links back to the canonical implementation on
652/// [`BlockTransposedRef`], so documentation stays in sync automatically.
653macro_rules! delegate_to_ref {
654    // Safe function.
655    ($(#[$m:meta])* $vis:vis fn $name:ident(&self $(, $a:ident: $t:ty)*) $(-> $r:ty)?) => {
656        #[doc = concat!("See [`BlockTransposedRef::", stringify!($name), "`].")]
657        $(#[$m])*
658        #[inline]
659        $vis fn $name(&self $(, $a: $t)*) $(-> $r)? {
660            self.as_view().$name($($a),*)
661        }
662    };
663    // Unsafe function.
664    ($(#[$m:meta])* unsafe $vis:vis fn $name:ident(&self $(, $a:ident: $t:ty)*) $(-> $r:ty)?) => {
665        #[doc = concat!("See [`BlockTransposedRef::", stringify!($name), "`].")]
666        $(#[$m])*
667        #[inline]
668        $vis unsafe fn $name(&self $(, $a: $t)*) $(-> $r)? {
669            // SAFETY: Caller upholds the safety contract of the delegated method.
670            unsafe { self.as_view().$name($($a),*) }
671        }
672    };
673}
674
675// ════════════════════════════════════════════════════════════════════
676// Public wrapper types
677// ════════════════════════════════════════════════════════════════════
678
679/// An owning block-transposed matrix.
680///
681/// Wraps an owned allocation of `T` elements laid out in block-transposed order.
682/// See the [module-level documentation](self) for layout details.
683///
684/// For shared and mutable views, see [`BlockTransposedRef`] and [`BlockTransposedMut`].
685///
686/// # Row Types
687///
688/// Because rows are not contiguous in memory, the row types are view structs:
689///
690/// - [`Row`] — a `Copy` handle supporting `Index<usize>` and `.iter()`.
691/// - [`RowMut`] — a mutable handle supporting `IndexMut<usize>`.
692#[derive(Debug, Clone)]
693pub struct BlockTransposed<T: Copy, const GROUP: usize, const PACK: usize = 1> {
694    data: Mat<BlockTransposedRepr<T, GROUP, PACK>>,
695}
696
697/// A shared (immutable) view of a block-transposed matrix.
698///
699/// Created by [`BlockTransposed::as_view`].
700#[derive(Debug, Clone, Copy)]
701pub struct BlockTransposedRef<'a, T: Copy, const GROUP: usize, const PACK: usize = 1> {
702    data: MatRef<'a, BlockTransposedRepr<T, GROUP, PACK>>,
703}
704
705/// A mutable view of a block-transposed matrix.
706///
707/// Created by [`BlockTransposed::as_view_mut`].
708pub struct BlockTransposedMut<'a, T: Copy, const GROUP: usize, const PACK: usize = 1> {
709    data: MatMut<'a, BlockTransposedRepr<T, GROUP, PACK>>,
710}
711
712// ── BlockTransposedRef (core read implementations) ───────────────
713
714impl<'a, T: Copy, const GROUP: usize, const PACK: usize> BlockTransposedRef<'a, T, GROUP, PACK> {
715    fn new(data: MatRef<'a, BlockTransposedRepr<T, GROUP, PACK>>) -> Self {
716        Self { data }
717    }
718
719    /// Returns the number of logical rows.
720    #[inline]
721    pub fn nrows(&self) -> usize {
722        self.data.repr().nrows()
723    }
724
725    /// Returns the number of logical columns (dimensionality).
726    #[inline]
727    pub fn ncols(&self) -> usize {
728        self.data.repr().ncols()
729    }
730
731    /// Returns the number of physical (padded) columns.
732    #[inline]
733    pub fn padded_ncols(&self) -> usize {
734        self.data.repr().padded_ncols()
735    }
736
737    /// Group size (blocking factor `GROUP`).
738    pub const fn group_size(&self) -> usize {
739        GROUP
740    }
741
742    /// Group size (blocking factor `GROUP`) as a `const` function on the *type*.
743    pub const fn const_group_size() -> usize {
744        GROUP
745    }
746
747    /// Packing factor `PACK`.
748    pub const fn pack_size(&self) -> usize {
749        PACK
750    }
751
752    /// Number of completely full blocks.
753    #[inline]
754    pub fn full_blocks(&self) -> usize {
755        self.data.repr().full_blocks()
756    }
757
758    /// Total number of blocks including any partially-filled tail.
759    #[inline]
760    pub fn num_blocks(&self) -> usize {
761        self.data.repr().num_blocks()
762    }
763
764    /// Number of valid elements in the last partially-full block, or 0 if all
765    /// blocks are full.
766    #[inline]
767    pub fn remainder(&self) -> usize {
768        self.data.repr().remainder()
769    }
770
771    /// Total number of logical rows rounded up to the next multiple of `GROUP`.
772    ///
773    /// This is the number of "available" row slots in the backing allocation,
774    /// including zero-padded rows in the last (possibly partial) block.
775    #[inline]
776    pub fn padded_nrows(&self) -> usize {
777        self.data.repr().padded_nrows()
778    }
779
780    /// Return a raw typed pointer to the start of the backing data.
781    #[inline]
782    pub fn as_ptr(&self) -> *const T {
783        self.data.as_raw_ptr().cast::<T>()
784    }
785
786    /// Return the backing data as a shared slice.
787    ///
788    /// The returned slice has `storage_len()` elements — this includes all padding
789    /// for partial blocks and column-group alignment.
790    #[inline]
791    pub fn as_slice(&self) -> &'a [T] {
792        let len = self.data.repr().storage_len();
793        // SAFETY: The backing allocation has exactly `storage_len()` elements of type T.
794        unsafe { std::slice::from_raw_parts(self.as_ptr(), len) }
795    }
796
797    /// Return a pointer to the start of the given block.
798    ///
799    /// The caller may assume that for the returned pointer `ptr`,
800    /// `[ptr, ptr + GROUP * padded_ncols)` points to valid memory, even for the
801    /// remainder block.
802    ///
803    /// # Safety
804    ///
805    /// `block` must be less than `self.num_blocks()`. No bounds check is
806    /// performed in release builds; callers must verify the index themselves
807    /// (e.g. by iterating `0..self.num_blocks()`).
808    #[inline]
809    pub unsafe fn block_ptr_unchecked(&self, block: usize) -> *const T {
810        debug_assert!(block < self.num_blocks());
811        // SAFETY: Caller asserts `block < self.num_blocks()`.
812        unsafe { self.as_ptr().add(self.data.repr().block_offset(block)) }
813    }
814
815    /// Return a view over a full block as a [`MatrixView`].
816    ///
817    /// The returned view has `padded_ncols / PACK` rows and `GROUP * PACK`
818    /// columns. For `PACK == 1` this simplifies to `ncols` rows and `GROUP`
819    /// columns (the standard transposed interpretation).
820    ///
821    /// # Panics
822    ///
823    /// Panics if `block >= self.full_blocks()`.
824    #[allow(clippy::expect_used)]
825    pub fn block(&self, block: usize) -> MatrixView<'a, T> {
826        assert!(block < self.full_blocks());
827        let offset = self.data.repr().block_offset(block);
828        let stride = self.data.repr().block_stride();
829        // SAFETY: `block < full_blocks()` (asserted above) guarantees
830        // `offset + stride` is within the backing allocation.
831        let data: &[T] = unsafe { std::slice::from_raw_parts(self.as_ptr().add(offset), stride) };
832        MatrixView::try_from(data, self.padded_ncols() / PACK, GROUP * PACK)
833            .expect("base data should have been sized correctly")
834    }
835
836    /// Return a view over the remainder block, or `None` if there is no
837    /// remainder.
838    ///
839    /// The returned view has the same dimensions as [`block()`](Self::block):
840    /// `padded_ncols / PACK` rows and `GROUP * PACK` columns.
841    #[allow(clippy::expect_used)]
842    pub fn remainder_block(&self) -> Option<MatrixView<'a, T>> {
843        if self.remainder() == 0 {
844            None
845        } else {
846            let offset = self.data.repr().block_offset(self.full_blocks());
847            let stride = self.data.repr().block_stride();
848            // SAFETY: The remainder block exists (`remainder() != 0`),
849            // so `offset + stride` is within the backing allocation.
850            let data: &[T] =
851                unsafe { std::slice::from_raw_parts(self.as_ptr().add(offset), stride) };
852            Some(
853                MatrixView::try_from(data, self.padded_ncols() / PACK, GROUP * PACK)
854                    .expect("base data should have been sized correctly"),
855            )
856        }
857    }
858
859    /// Retrieve the value at the logical `(row, col)`.
860    ///
861    /// # Panics
862    ///
863    /// Panics if `row >= self.nrows()` or `col >= self.ncols()`.
864    #[inline]
865    pub fn get_element(&self, row: usize, col: usize) -> T {
866        assert!(
867            row < self.nrows(),
868            "row {row} out of bounds (nrows = {})",
869            self.nrows()
870        );
871        assert!(
872            col < self.ncols(),
873            "col {col} out of bounds (ncols = {})",
874            self.ncols()
875        );
876        let idx = linear_index::<GROUP, PACK>(row, col, self.ncols());
877        // SAFETY: bounds checked above.
878        unsafe { *self.as_ptr().add(idx) }
879    }
880
881    /// Get an immutable row view, or `None` if `i` is out of bounds.
882    #[inline]
883    pub fn get_row(&self, i: usize) -> Option<Row<'_, T, GROUP, PACK>> {
884        self.data.get_row(i)
885    }
886}
887
888// ── BlockTransposedMut ───────────────────────────────────────────
889
890impl<'a, T: Copy, const GROUP: usize, const PACK: usize> BlockTransposedMut<'a, T, GROUP, PACK> {
891    fn new(data: MatMut<'a, BlockTransposedRepr<T, GROUP, PACK>>) -> Self {
892        Self { data }
893    }
894
895    /// Borrow as an immutable [`BlockTransposedRef`].
896    #[inline]
897    pub fn as_view(&self) -> BlockTransposedRef<'_, T, GROUP, PACK> {
898        BlockTransposedRef::new(self.data.as_view())
899    }
900
901    // ── Delegated read methods ───────────────────────────────────
902
903    delegate_to_ref!(pub fn nrows(&self) -> usize);
904    delegate_to_ref!(pub fn ncols(&self) -> usize);
905    delegate_to_ref!(pub fn padded_ncols(&self) -> usize);
906    delegate_to_ref!(pub fn full_blocks(&self) -> usize);
907    delegate_to_ref!(pub fn num_blocks(&self) -> usize);
908    delegate_to_ref!(pub fn remainder(&self) -> usize);
909    delegate_to_ref!(pub fn padded_nrows(&self) -> usize);
910    delegate_to_ref!(pub fn as_ptr(&self) -> *const T);
911    delegate_to_ref!(pub fn as_slice(&self) -> &[T]);
912    delegate_to_ref!(#[allow(clippy::missing_safety_doc)] unsafe pub fn block_ptr_unchecked(&self, block: usize) -> *const T);
913    delegate_to_ref!(#[allow(clippy::expect_used)] pub fn block(&self, block: usize) -> MatrixView<'_, T>);
914    delegate_to_ref!(#[allow(clippy::expect_used)] pub fn remainder_block(&self) -> Option<MatrixView<'_, T>>);
915    delegate_to_ref!(pub fn get_element(&self, row: usize, col: usize) -> T);
916
917    /// Group size (blocking factor `GROUP`).
918    pub const fn group_size(&self) -> usize {
919        GROUP
920    }
921
922    /// Group size as `const` function on the *type*.
923    pub const fn const_group_size() -> usize {
924        GROUP
925    }
926
927    /// Packing factor `PACK`.
928    pub const fn pack_size(&self) -> usize {
929        PACK
930    }
931
932    /// Get an immutable row view, or `None` if `i` is out of bounds.
933    #[inline]
934    pub fn get_row(&self, i: usize) -> Option<Row<'_, T, GROUP, PACK>> {
935        self.data.get_row(i)
936    }
937
938    // ── Mutable methods ──────────────────────────────────────────
939    //
940    // The `_inner` variants consume `self` by value so that the lifetime of
941    // the returned view is tied to `'a` (the underlying allocation), not to
942    // a temporary reborrow. Public `&mut self` methods reborrow into a
943    // short-lived `BlockTransposedMut` and then call the inner variant.
944
945    /// Return the backing data as a mutable slice.
946    ///
947    /// The returned slice has `storage_len()` elements (including all padding).
948    #[inline]
949    pub fn as_mut_slice(&mut self) -> &mut [T] {
950        self.reborrow_mut().mut_slice_inner()
951    }
952
953    fn mut_slice_inner(mut self) -> &'a mut [T] {
954        let len = self.data.repr().storage_len();
955        // SAFETY: We own exclusive access through `self`.
956        unsafe { std::slice::from_raw_parts_mut(self.data.as_raw_mut_ptr().cast::<T>(), len) }
957    }
958
959    /// Return a mutable view over a full block.
960    ///
961    /// # Panics
962    ///
963    /// Panics if `block >= self.full_blocks()`.
964    #[allow(clippy::expect_used)]
965    pub fn block_mut(&mut self, block: usize) -> MutMatrixView<'_, T> {
966        self.reborrow_mut().block_mut_inner(block)
967    }
968
969    #[allow(clippy::expect_used)]
970    fn block_mut_inner(mut self, block: usize) -> MutMatrixView<'a, T> {
971        let repr = *self.data.repr();
972        assert!(block < repr.full_blocks());
973        let offset = repr.block_offset(block);
974        let stride = repr.block_stride();
975        let pncols = repr.padded_ncols();
976        // SAFETY: `block < full_blocks()`, so the range is within the allocation.
977        let data: &mut [T] = unsafe {
978            std::slice::from_raw_parts_mut(
979                self.data.as_raw_mut_ptr().cast::<T>().add(offset),
980                stride,
981            )
982        };
983        MutMatrixView::try_from(data, pncols / PACK, GROUP * PACK)
984            .expect("base data should have been sized correctly")
985    }
986
987    /// Return a mutable view over the remainder block, or `None` if there is no
988    /// remainder.
989    #[allow(clippy::expect_used)]
990    pub fn remainder_block_mut(&mut self) -> Option<MutMatrixView<'_, T>> {
991        self.reborrow_mut().remainder_block_mut_inner()
992    }
993
994    #[allow(clippy::expect_used)]
995    fn remainder_block_mut_inner(mut self) -> Option<MutMatrixView<'a, T>> {
996        let repr = *self.data.repr();
997        if repr.remainder() == 0 {
998            None
999        } else {
1000            let offset = repr.block_offset(repr.full_blocks());
1001            let stride = repr.block_stride();
1002            let pncols = repr.padded_ncols();
1003            // SAFETY: Remainder block exists, so the range is within the allocation.
1004            let data: &mut [T] = unsafe {
1005                std::slice::from_raw_parts_mut(
1006                    self.data.as_raw_mut_ptr().cast::<T>().add(offset),
1007                    stride,
1008                )
1009            };
1010            Some(
1011                MutMatrixView::try_from(data, pncols / PACK, GROUP * PACK)
1012                    .expect("base data should have been sized correctly"),
1013            )
1014        }
1015    }
1016
1017    /// Get a mutable row view, or `None` if `i` is out of bounds.
1018    #[inline]
1019    pub fn get_row_mut(&mut self, i: usize) -> Option<RowMut<'_, T, GROUP, PACK>> {
1020        self.data.get_row_mut(i)
1021    }
1022
1023    // ── Private helpers ──────────────────────────────────────────
1024
1025    fn reborrow_mut(&mut self) -> BlockTransposedMut<'_, T, GROUP, PACK> {
1026        BlockTransposedMut::new(self.data.reborrow_mut())
1027    }
1028}
1029
1030// ── BlockTransposed (owned) ──────────────────────────────────────
1031
1032impl<T: Copy, const GROUP: usize, const PACK: usize> BlockTransposed<T, GROUP, PACK> {
1033    /// Borrow as an immutable [`BlockTransposedRef`].
1034    pub fn as_view(&self) -> BlockTransposedRef<'_, T, GROUP, PACK> {
1035        BlockTransposedRef::new(self.data.as_view())
1036    }
1037
1038    /// Borrow as a mutable [`BlockTransposedMut`].
1039    pub fn as_view_mut(&mut self) -> BlockTransposedMut<'_, T, GROUP, PACK> {
1040        BlockTransposedMut::new(self.data.as_view_mut())
1041    }
1042
1043    // ── Delegated read methods ───────────────────────────────────
1044
1045    delegate_to_ref!(pub fn nrows(&self) -> usize);
1046    delegate_to_ref!(pub fn ncols(&self) -> usize);
1047    delegate_to_ref!(pub fn padded_ncols(&self) -> usize);
1048    delegate_to_ref!(pub fn full_blocks(&self) -> usize);
1049    delegate_to_ref!(pub fn num_blocks(&self) -> usize);
1050    delegate_to_ref!(pub fn remainder(&self) -> usize);
1051    delegate_to_ref!(pub fn padded_nrows(&self) -> usize);
1052    delegate_to_ref!(pub fn as_ptr(&self) -> *const T);
1053    delegate_to_ref!(pub fn as_slice(&self) -> &[T]);
1054    delegate_to_ref!(#[allow(clippy::missing_safety_doc)] unsafe pub fn block_ptr_unchecked(&self, block: usize) -> *const T);
1055    delegate_to_ref!(#[allow(clippy::expect_used)] pub fn block(&self, block: usize) -> MatrixView<'_, T>);
1056    delegate_to_ref!(#[allow(clippy::expect_used)] pub fn remainder_block(&self) -> Option<MatrixView<'_, T>>);
1057    delegate_to_ref!(pub fn get_element(&self, row: usize, col: usize) -> T);
1058
1059    /// Group size (blocking factor `GROUP`).
1060    pub const fn group_size(&self) -> usize {
1061        GROUP
1062    }
1063
1064    /// Group size (blocking factor `GROUP`) as a `const` function on the *type*.
1065    pub const fn const_group_size() -> usize {
1066        GROUP
1067    }
1068
1069    /// Packing factor `PACK`.
1070    pub const fn pack_size(&self) -> usize {
1071        PACK
1072    }
1073
1074    /// Get an immutable row view, or `None` if `i` is out of bounds.
1075    #[inline]
1076    pub fn get_row(&self, i: usize) -> Option<Row<'_, T, GROUP, PACK>> {
1077        self.data.get_row(i)
1078    }
1079
1080    // ── Mutable methods (delegated to BlockTransposedMut) ────────
1081
1082    /// See [`BlockTransposedMut::as_mut_slice`].
1083    #[inline]
1084    pub fn as_mut_slice(&mut self) -> &mut [T] {
1085        self.as_view_mut().mut_slice_inner()
1086    }
1087
1088    /// See [`BlockTransposedMut::block_mut`].
1089    #[allow(clippy::expect_used)]
1090    pub fn block_mut(&mut self, block: usize) -> MutMatrixView<'_, T> {
1091        self.as_view_mut().block_mut_inner(block)
1092    }
1093
1094    /// See [`BlockTransposedMut::remainder_block_mut`].
1095    #[allow(clippy::expect_used)]
1096    pub fn remainder_block_mut(&mut self) -> Option<MutMatrixView<'_, T>> {
1097        self.as_view_mut().remainder_block_mut_inner()
1098    }
1099
1100    /// Get a mutable row view, or `None` if `i` is out of bounds.
1101    #[inline]
1102    pub fn get_row_mut(&mut self, i: usize) -> Option<RowMut<'_, T, GROUP, PACK>> {
1103        self.data.get_row_mut(i)
1104    }
1105}
1106
1107// ── Reborrow ─────────────────────────────────────────────────────
1108
1109impl<'this, T: Copy, const GROUP: usize, const PACK: usize> Reborrow<'this>
1110    for BlockTransposed<T, GROUP, PACK>
1111{
1112    type Target = BlockTransposedRef<'this, T, GROUP, PACK>;
1113
1114    #[inline]
1115    fn reborrow(&'this self) -> Self::Target {
1116        self.as_view()
1117    }
1118}
1119
1120// ── Factory methods ──────────────────────────────────────────────
1121
1122impl<T: Copy + Default, const GROUP: usize, const PACK: usize> BlockTransposed<T, GROUP, PACK> {
1123    /// Construct a default-initialized block-transposed matrix from dimensions.
1124    ///
1125    /// # Panics
1126    ///
1127    /// Panics if the dimensions overflow the allocation budget.
1128    #[allow(clippy::expect_used)]
1129    pub fn new(nrows: usize, ncols: usize) -> Self {
1130        let repr = BlockTransposedRepr::<T, GROUP, PACK>::new(nrows, ncols)
1131            .expect("dimensions should not overflow");
1132        Self {
1133            data: Mat::new(repr, Defaulted).expect("infallible"),
1134        }
1135    }
1136
1137    /// Fallible variant of [`new`](Self::new).
1138    pub fn try_new(nrows: usize, ncols: usize) -> Result<Self, Overflow> {
1139        let repr = BlockTransposedRepr::<T, GROUP, PACK>::new(nrows, ncols)?;
1140        Ok(Self {
1141            data: Mat::new(repr, Defaulted).expect("infallible"),
1142        })
1143    }
1144
1145    /// Construct a block-transposed matrix by copying data from a [`StridedView`].
1146    ///
1147    /// Each source element at `(row, col)` is placed at the correct offset in the
1148    /// block-transposed layout. Padding positions (both partial-block rows and
1149    /// column-group padding when `ncols % PACK != 0`) are filled with
1150    /// `T::default()`.
1151    ///
1152    /// The loop iterates in physical (block-transposed) order — block, column-group,
1153    /// row-within-block, pack-lane — so that writes to the backing allocation are
1154    /// sequential. Source reads stride across rows of the [`StridedView`], which is
1155    /// acceptable because read-side prefetch is more effective than write-side.
1156    pub fn from_strided(v: StridedView<'_, T>) -> Self {
1157        let nrows = v.nrows();
1158        let ncols = v.ncols();
1159        let mut mat = Self::new(nrows, ncols);
1160
1161        let repr = *mat.data.repr();
1162        let num_blocks = repr.num_blocks();
1163        let pncols = repr.padded_ncols();
1164        let num_col_groups = pncols / PACK;
1165
1166        // Walk the backing allocation in physical order so that writes are
1167        // sequential. The allocation is default-initialized, so padding positions
1168        // already hold `T::default()` and can be skipped.
1169        let mut dst = mat.data.as_raw_mut_ptr().cast::<T>();
1170        for block in 0..num_blocks {
1171            let row_base = block * GROUP;
1172            for cg in 0..num_col_groups {
1173                let col_base = cg * PACK;
1174                for rib in 0..GROUP {
1175                    let row = row_base + rib;
1176                    if row < nrows {
1177                        // SAFETY: row < nrows is checked by the enclosing `if` condition.
1178                        let src_row = unsafe { v.get_row_unchecked(row) };
1179                        for p in 0..PACK {
1180                            let col = col_base + p;
1181                            if col < ncols {
1182                                // SAFETY: dst advances sequentially through the
1183                                // backing allocation which has exactly `storage_len`
1184                                // elements, and our loop visits each position once.
1185                                // `col < ncols` is checked above, and `src_row` has
1186                                // exactly `ncols` elements.
1187                                unsafe { *dst = *src_row.get_unchecked(col) };
1188                            }
1189                            // SAFETY: dst advances sequentially through the
1190                            // backing allocation which has exactly `storage_len`
1191                            // elements, and our loop visits each position once.
1192                            dst = unsafe { dst.add(1) };
1193                        }
1194                    } else {
1195                        // SAFETY: Entire row is padding — skip PACK positions.
1196                        // dst remains within the allocation.
1197                        dst = unsafe { dst.add(PACK) };
1198                    }
1199                }
1200            }
1201        }
1202
1203        mat
1204    }
1205
1206    /// Construct a block-transposed matrix by copying data from a [`MatrixView`].
1207    pub fn from_matrix_view(v: MatrixView<'_, T>) -> Self {
1208        Self::from_strided(v.into())
1209    }
1210}
1211
1212// ════════════════════════════════════════════════════════════════════
1213// Index<(usize, usize)> for BlockTransposed
1214// ════════════════════════════════════════════════════════════════════
1215
1216impl<T: Copy, const GROUP: usize, const PACK: usize> std::ops::Index<(usize, usize)>
1217    for BlockTransposed<T, GROUP, PACK>
1218{
1219    type Output = T;
1220
1221    #[inline]
1222    fn index(&self, (row, col): (usize, usize)) -> &Self::Output {
1223        assert!(row < self.nrows());
1224        assert!(col < self.ncols());
1225        let idx = linear_index::<GROUP, PACK>(row, col, self.ncols());
1226        // SAFETY: bounds checked above and the backing allocation has `storage_len()` elements.
1227        unsafe { &*self.as_ptr().add(idx) }
1228    }
1229}
1230
1231// ════════════════════════════════════════════════════════════════════
1232// Tests
1233// ════════════════════════════════════════════════════════════════════
1234
1235#[cfg(test)]
1236mod tests {
1237    //! Test organisation:
1238    //!
1239    //!  1. **Helper functions** — `gen_*` element generators.
1240    //!  2. [`test_full_api`] — single parameterized function that exhaustively
1241    //!     exercises the full read + write API on all three wrapper types
1242    //!     (`BlockTransposed`, `BlockTransposedRef`, `BlockTransposedMut`).
1243    //!  3. **Test runners** — `#[test]` functions that call `test_full_api`
1244    //!     with various `(T, GROUP, PACK, nrows, ncols)` combinations.
1245    //!  4. [`test_block_layout_pack1`] — verifies that `PACK=1` blocks are
1246    //!     the standard row-to-column transposition.
1247    //!  5. **Focused tests** — edge cases that cannot be expressed as
1248    //!     parameters to `test_full_api` (`Send`/`Sync`, panic paths,
1249    //!     non-unit strides, concurrent mutation, etc.).
1250
1251    use diskann_utils::{lazy_format, views::Matrix};
1252
1253    use super::*;
1254    use crate::utils::div_round_up;
1255
1256    // ── Per-type element generators ──────────────────────────────────
1257    //
1258    // Each generator maps a flat index to a non-zero `T` value so that
1259    // `T::default()` (zero) can be used unambiguously to verify padding.
1260
1261    fn gen_f32(i: usize) -> f32 {
1262        (i + 1) as f32
1263    }
1264    fn gen_i32(i: usize) -> i32 {
1265        (i + 1) as i32
1266    }
1267    fn gen_u8(i: usize) -> u8 {
1268        ((i % 255) + 1) as u8
1269    }
1270
1271    #[test]
1272    fn clone_has_independent_backing_allocation() {
1273        let mut data = Matrix::new(0, 5, 3);
1274        data.as_mut_slice()
1275            .iter_mut()
1276            .enumerate()
1277            .for_each(|(i, value)| *value = (i + 1) as i32);
1278        let mut original = BlockTransposed::<i32, 4, 2>::from_matrix_view(data.as_view());
1279        let column_padding = linear_index::<4, 2>(0, 3, original.ncols());
1280        let row_padding = linear_index::<4, 2>(5, 0, original.ncols());
1281        let row_and_column_padding = linear_index::<4, 2>(5, 3, original.ncols());
1282        original.as_mut_slice()[column_padding] = -10;
1283        original.as_mut_slice()[row_padding] = -11;
1284        original.as_mut_slice()[row_and_column_padding] = -12;
1285
1286        let mut cloned = original.clone();
1287
1288        assert_eq!(cloned.as_slice(), original.as_slice());
1289        assert_eq!(cloned.as_slice()[column_padding], -10);
1290        assert_eq!(cloned.as_slice()[row_padding], -11);
1291        assert_eq!(cloned.as_slice()[row_and_column_padding], -12);
1292        assert_ne!(cloned.as_ptr(), original.as_ptr());
1293
1294        cloned.get_row_mut(0).unwrap()[0] = -1;
1295        assert_eq!(original[(0, 0)], 1);
1296        assert_eq!(cloned[(0, 0)], -1);
1297    }
1298
1299    // ── Unified parameterized test ──────────────────────────────────
1300
1301    /// Exhaustive test for the full `BlockTransposed` / `BlockTransposedRef` /
1302    /// `BlockTransposedMut` API surface, parameterized over element type `T`,
1303    /// group size `GROUP`, and packing factor `PACK`.
1304    ///
1305    /// Exercises: construction, query helpers, `Index` / `get_element`,
1306    /// immutable row views (`Row`), mutable row views (`RowMut`),
1307    /// `as_slice` / `as_mut_slice`, block views (immutable and mutable),
1308    /// `remainder_block` / `remainder_block_mut`, `block_ptr_unchecked`,
1309    /// `from_matrix_view`, OOB `get_row` returns, and both column and row
1310    /// padding verification.
1311    fn test_full_api<
1312        T: Copy + Default + PartialEq + std::fmt::Debug + 'static,
1313        const GROUP: usize,
1314        const PACK: usize,
1315    >(
1316        nrows: usize,
1317        ncols: usize,
1318        gen_element: fn(usize) -> T,
1319    ) {
1320        let context = lazy_format!(
1321            "T={}, GROUP={}, PACK={}, nrows={}, ncols={}",
1322            std::any::type_name::<T>(),
1323            GROUP,
1324            PACK,
1325            nrows,
1326            ncols,
1327        );
1328
1329        // ── Construction ─────────────────────────────────────────
1330
1331        let mut data = Matrix::new(T::default(), nrows, ncols);
1332        data.as_mut_slice()
1333            .iter_mut()
1334            .enumerate()
1335            .for_each(|(i, d)| *d = gen_element(i));
1336
1337        let mut transpose = BlockTransposed::<T, GROUP, PACK>::from_strided(data.as_view().into());
1338
1339        let expected_padded = div_round_up(ncols, PACK) * PACK;
1340        let expected_remainder = nrows % GROUP;
1341        let storage_len = transpose.as_slice().len();
1342
1343        // ── Query methods on owned type ──────────────────────────
1344
1345        assert_eq!(transpose.nrows(), nrows, "{}", context);
1346        assert_eq!(transpose.ncols(), ncols, "{}", context);
1347        assert_eq!(transpose.group_size(), GROUP, "{}", context);
1348        assert_eq!(
1349            BlockTransposed::<T, GROUP, PACK>::const_group_size(),
1350            GROUP,
1351            "{}",
1352            context
1353        );
1354        assert_eq!(transpose.pack_size(), PACK, "{}", context);
1355        assert_eq!(transpose.full_blocks(), nrows / GROUP, "{}", context);
1356        assert_eq!(
1357            transpose.num_blocks(),
1358            div_round_up(nrows, GROUP),
1359            "{}",
1360            context,
1361        );
1362        assert_eq!(transpose.remainder(), expected_remainder, "{}", context);
1363        assert_eq!(transpose.padded_ncols(), expected_padded, "{}", context);
1364
1365        // ── Element access (owned) ───────────────────────────────
1366
1367        for row in 0..nrows {
1368            for col in 0..ncols {
1369                assert_eq!(
1370                    data[(row, col)],
1371                    transpose[(row, col)],
1372                    "Index at ({}, {}) -- {}",
1373                    row,
1374                    col,
1375                    context,
1376                );
1377                assert_eq!(
1378                    data[(row, col)],
1379                    transpose.get_element(row, col),
1380                    "get_element at ({}, {}) -- {}",
1381                    row,
1382                    col,
1383                    context,
1384                );
1385            }
1386        }
1387
1388        // ── Immutable row views (owned) ──────────────────────────
1389
1390        let view = transpose.as_view();
1391        for row in 0..nrows {
1392            let row_view = view.get_row(row).unwrap();
1393            assert_eq!(row_view.len(), ncols, "{}", context);
1394            assert_eq!(row_view.is_empty(), ncols == 0, "{}", context);
1395            for col in 0..ncols {
1396                assert_eq!(
1397                    data[(row, col)],
1398                    row_view[col],
1399                    "row view at ({}, {}) -- {}",
1400                    row,
1401                    col,
1402                    context,
1403                );
1404            }
1405            // Row::get — in-bounds + OOB.
1406            if ncols > 0 {
1407                assert_eq!(row_view.get(0), Some(&data[(row, 0)]), "{}", context);
1408            }
1409            assert_eq!(row_view.get(ncols), None, "{}", context);
1410
1411            // Iterator + ExactSizeIterator.
1412            let iter = row_view.iter();
1413            assert_eq!(iter.len(), ncols, "{}", context);
1414            let (lo, hi) = iter.size_hint();
1415            assert_eq!(lo, ncols, "{}", context);
1416            assert_eq!(hi, Some(ncols), "{}", context);
1417
1418            let collected: Vec<T> = row_view.iter().collect();
1419            assert_eq!(collected.len(), ncols, "{}", context);
1420            for col in 0..ncols {
1421                assert_eq!(data[(row, col)], collected[col], "{}", context);
1422            }
1423        }
1424        // OOB row returns None.
1425        assert!(view.get_row(nrows).is_none(), "{}", context);
1426        let _ = view;
1427
1428        // ── BlockTransposedRef API ───────────────────────────────
1429
1430        {
1431            let view = transpose.as_view();
1432            assert_eq!(view.nrows(), nrows, "{}", context);
1433            assert_eq!(view.ncols(), ncols, "{}", context);
1434            assert_eq!(view.padded_ncols(), expected_padded, "{}", context);
1435            assert_eq!(view.group_size(), GROUP, "{}", context);
1436            assert_eq!(
1437                BlockTransposedRef::<T, GROUP, PACK>::const_group_size(),
1438                GROUP,
1439            );
1440            assert_eq!(view.pack_size(), PACK, "{}", context);
1441            assert_eq!(view.full_blocks(), nrows / GROUP, "{}", context);
1442            assert_eq!(view.num_blocks(), div_round_up(nrows, GROUP), "{}", context,);
1443            assert_eq!(view.remainder(), expected_remainder, "{}", context);
1444            assert_eq!(view.as_ptr(), transpose.as_ptr(), "{}", context);
1445            assert_eq!(view.as_slice(), transpose.as_slice(), "{}", context);
1446
1447            for row in 0..nrows {
1448                for col in 0..ncols {
1449                    assert_eq!(
1450                        data[(row, col)],
1451                        view.get_element(row, col),
1452                        "Ref get_element at ({}, {}) -- {}",
1453                        row,
1454                        col,
1455                        context,
1456                    );
1457                }
1458                let row_view = view.get_row(row).unwrap();
1459                for col in 0..ncols {
1460                    assert_eq!(data[(row, col)], row_view[col], "{}", context);
1461                }
1462            }
1463            assert!(view.get_row(nrows).is_none(), "{}", context);
1464        }
1465
1466        // ── BlockTransposedMut read API ──────────────────────────
1467
1468        let expected_ptr = transpose.as_ptr();
1469        {
1470            let mut_view = transpose.as_view_mut();
1471            assert_eq!(mut_view.nrows(), nrows, "{}", context);
1472            assert_eq!(mut_view.ncols(), ncols, "{}", context);
1473            assert_eq!(mut_view.padded_ncols(), expected_padded, "{}", context);
1474            assert_eq!(mut_view.group_size(), GROUP, "{}", context);
1475            assert_eq!(
1476                BlockTransposedMut::<T, GROUP, PACK>::const_group_size(),
1477                GROUP,
1478            );
1479            assert_eq!(mut_view.pack_size(), PACK, "{}", context);
1480            assert_eq!(mut_view.full_blocks(), nrows / GROUP, "{}", context);
1481            assert_eq!(
1482                mut_view.num_blocks(),
1483                div_round_up(nrows, GROUP),
1484                "{}",
1485                context,
1486            );
1487            assert_eq!(mut_view.remainder(), expected_remainder, "{}", context);
1488            assert_eq!(mut_view.as_ptr(), expected_ptr, "{}", context);
1489            assert_eq!(mut_view.as_slice().len(), storage_len, "{}", context);
1490
1491            for row in 0..nrows {
1492                for col in 0..ncols {
1493                    assert_eq!(
1494                        data[(row, col)],
1495                        mut_view.get_element(row, col),
1496                        "Mut get_element at ({}, {}) -- {}",
1497                        row,
1498                        col,
1499                        context,
1500                    );
1501                }
1502                let row_view = mut_view.get_row(row).unwrap();
1503                for col in 0..ncols {
1504                    assert_eq!(data[(row, col)], row_view[col], "{}", context);
1505                }
1506            }
1507            assert!(mut_view.get_row(nrows).is_none(), "{}", context);
1508        }
1509
1510        // ── BlockTransposedMut::as_view() ────────────────────────
1511
1512        {
1513            let mut_view = transpose.as_view_mut();
1514            let ref_from_mut = mut_view.as_view();
1515            assert_eq!(ref_from_mut.nrows(), nrows, "{}", context);
1516            for row in 0..nrows {
1517                for col in 0..ncols {
1518                    assert_eq!(
1519                        data[(row, col)],
1520                        ref_from_mut.get_element(row, col),
1521                        "{}",
1522                        context,
1523                    );
1524                }
1525            }
1526        }
1527
1528        // ── as_mut_slice ─────────────────────────────────────────
1529
1530        // Through BlockTransposedMut.
1531        {
1532            let mut mut_view = transpose.as_view_mut();
1533            assert_eq!(mut_view.as_mut_slice().len(), storage_len, "{}", context);
1534        }
1535        // Through BlockTransposed (owned).
1536        assert_eq!(transpose.as_mut_slice().len(), storage_len, "{}", context);
1537
1538        // ── Immutable block views on all three types ─────────────
1539
1540        let expected_block_nrows = expected_padded / PACK;
1541        let expected_block_ncols = GROUP * PACK;
1542
1543        for b in 0..transpose.full_blocks() {
1544            let block_data: Vec<T>;
1545            let ptr: *const T;
1546            {
1547                let block = transpose.block(b);
1548                assert_eq!(block.nrows(), expected_block_nrows, "{}", context);
1549                assert_eq!(block.ncols(), expected_block_ncols, "{}", context);
1550
1551                // SAFETY: b < full_blocks <= num_blocks.
1552                ptr = unsafe { transpose.block_ptr_unchecked(b) };
1553                assert_eq!(ptr, block.as_slice().as_ptr(), "{}", context);
1554
1555                block_data = block.as_slice().to_vec();
1556            }
1557
1558            // Same block via Ref.
1559            {
1560                let view = transpose.as_view();
1561                assert_eq!(view.block(b).as_slice(), &block_data[..], "{}", context);
1562                // SAFETY: `b` is in range `0..num_blocks` by the loop bound.
1563                assert_eq!(unsafe { view.block_ptr_unchecked(b) }, ptr, "{}", context);
1564            }
1565
1566            // Same block via Mut (read path).
1567            {
1568                let mut_view = transpose.as_view_mut();
1569                assert_eq!(mut_view.block(b).as_slice(), &block_data[..], "{}", context);
1570                assert_eq!(
1571                    // SAFETY: `b` is in range `0..num_blocks` by the loop bound.
1572                    unsafe { mut_view.block_ptr_unchecked(b) },
1573                    ptr,
1574                    "{}",
1575                    context,
1576                );
1577            }
1578        }
1579
1580        // Remainder block (immutable, all three types).
1581        if expected_remainder != 0 {
1582            let remainder_data: Vec<T>;
1583            let ptr: *const T;
1584            let fb = transpose.full_blocks();
1585            {
1586                let block = transpose.remainder_block().unwrap();
1587                assert_eq!(block.nrows(), expected_block_nrows, "{}", context);
1588                assert_eq!(block.ncols(), expected_block_ncols, "{}", context);
1589
1590                // SAFETY: fb < num_blocks (remainder exists).
1591                ptr = unsafe { transpose.block_ptr_unchecked(fb) };
1592                assert_eq!(ptr, block.as_slice().as_ptr(), "{}", context);
1593
1594                remainder_data = block.as_slice().to_vec();
1595            }
1596
1597            // Via Ref.
1598            {
1599                let view = transpose.as_view();
1600                let ref_block = view.remainder_block().unwrap();
1601                assert_eq!(ref_block.as_slice(), &remainder_data[..], "{}", context);
1602            }
1603            // Via Mut (read path).
1604            {
1605                let mut_view = transpose.as_view_mut();
1606                let mut_block = mut_view.remainder_block().unwrap();
1607                assert_eq!(mut_block.as_slice(), &remainder_data[..], "{}", context);
1608            }
1609        } else {
1610            assert!(transpose.remainder_block().is_none(), "{}", context);
1611            {
1612                let view = transpose.as_view();
1613                assert!(view.remainder_block().is_none(), "{}", context);
1614            }
1615            {
1616                let mut_view = transpose.as_view_mut();
1617                assert!(mut_view.remainder_block().is_none(), "{}", context);
1618            }
1619        }
1620
1621        // ── Mutable block views via BlockTransposedMut ───────────
1622
1623        {
1624            let mut mut_view = transpose.as_view_mut();
1625            for b in 0..mut_view.full_blocks() {
1626                let block_mut = mut_view.block_mut(b);
1627                assert_eq!(block_mut.nrows(), expected_block_nrows, "{}", context);
1628                assert_eq!(block_mut.ncols(), expected_block_ncols, "{}", context);
1629            }
1630            if expected_remainder != 0 {
1631                let rem = mut_view.remainder_block_mut().unwrap();
1632                assert_eq!(rem.nrows(), expected_block_nrows, "{}", context);
1633                assert_eq!(rem.ncols(), expected_block_ncols, "{}", context);
1634            } else {
1635                assert!(mut_view.remainder_block_mut().is_none(), "{}", context);
1636            }
1637        }
1638
1639        // Mutable block views via owned BlockTransposed.
1640        for b in 0..transpose.full_blocks() {
1641            let block_mut = transpose.block_mut(b);
1642            assert_eq!(block_mut.nrows(), expected_block_nrows, "{}", context);
1643            assert_eq!(block_mut.ncols(), expected_block_ncols, "{}", context);
1644        }
1645        if expected_remainder != 0 {
1646            let rem = transpose.remainder_block_mut().unwrap();
1647            assert_eq!(rem.nrows(), expected_block_nrows, "{}", context);
1648            assert_eq!(rem.ncols(), expected_block_ncols, "{}", context);
1649        } else {
1650            assert!(transpose.remainder_block_mut().is_none(), "{}", context);
1651        }
1652
1653        // ── Mutable row views via BlockTransposedMut ─────────────
1654
1655        {
1656            let mut mut_view = transpose.as_view_mut();
1657            for row in 0..nrows {
1658                let row_view = mut_view.get_row_mut(row).unwrap();
1659                assert_eq!(row_view.len(), ncols, "{}", context);
1660                assert_eq!(row_view.is_empty(), ncols == 0, "{}", context);
1661                for col in 0..ncols {
1662                    assert_eq!(data[(row, col)], row_view[col], "{}", context);
1663                }
1664            }
1665            assert!(mut_view.get_row_mut(nrows).is_none(), "{}", context);
1666        }
1667
1668        // ── Row::get, RowMut::get, RowMut::get_mut ──────────────
1669
1670        if nrows > 0 && ncols > 0 {
1671            // Row::get OOB.
1672            {
1673                let view = transpose.as_view();
1674                let row = view.get_row(0).unwrap();
1675                assert_eq!(row.get(ncols), None, "{}", context);
1676                assert_eq!(row.get(usize::MAX), None, "{}", context);
1677            }
1678
1679            // RowMut::get OOB.
1680            let row = transpose.get_row_mut(0).unwrap();
1681            assert_eq!(row.get(ncols), None, "{}", context);
1682
1683            // RowMut::get_mut — mutate and verify.
1684            let mut row = transpose.get_row_mut(0).unwrap();
1685            let sentinel = gen_element(usize::MAX / 2);
1686            let original = row[0];
1687            if let Some(v) = row.get_mut(0) {
1688                *v = sentinel;
1689            }
1690            assert_eq!(row.get_mut(ncols), None, "{}", context);
1691            // Explicit scope end so the mutable borrow is released before the next access.
1692            let _ = row;
1693            assert_eq!(transpose.get_element(0, 0), sentinel, "{}", context);
1694            // Restore original.
1695            transpose.get_row_mut(0).unwrap().set(0, original);
1696        }
1697
1698        // ── Zero out via block_mut / remainder_block_mut ─────────
1699
1700        for b in 0..transpose.full_blocks() {
1701            transpose.block_mut(b).as_mut_slice().fill(T::default());
1702        }
1703        if transpose.remainder() != 0 {
1704            transpose
1705                .remainder_block_mut()
1706                .unwrap()
1707                .as_mut_slice()
1708                .fill(T::default());
1709        }
1710        assert!(
1711            transpose.as_slice().iter().all(|v| *v == T::default()),
1712            "not fully zeroed -- {}",
1713            context,
1714        );
1715
1716        // ── Padding verification (fresh construction) ────────────
1717
1718        let transpose = BlockTransposed::<T, GROUP, PACK>::from_strided(data.as_view().into());
1719        let raw = transpose.as_slice();
1720
1721        // Column padding.
1722        for row in 0..nrows {
1723            for col in ncols..expected_padded {
1724                let idx = linear_index::<GROUP, PACK>(row, col, ncols);
1725                assert_eq!(
1726                    raw[idx],
1727                    T::default(),
1728                    "col padding at ({}, {}) -- {}",
1729                    row,
1730                    col,
1731                    context,
1732                );
1733            }
1734        }
1735
1736        // Row padding (within partial blocks).
1737        let padded_nrows = nrows.next_multiple_of(GROUP);
1738        for row in nrows..padded_nrows {
1739            for col in 0..expected_padded {
1740                let idx = linear_index::<GROUP, PACK>(row, col, ncols);
1741                assert_eq!(
1742                    raw[idx],
1743                    T::default(),
1744                    "row padding at ({}, {}) -- {}",
1745                    row,
1746                    col,
1747                    context,
1748                );
1749            }
1750        }
1751
1752        // ── padded_nrows() returns padded row count ──────────────
1753
1754        assert_eq!(
1755            transpose.as_view().padded_nrows(),
1756            padded_nrows,
1757            "padded_nrows() mismatch -- {}",
1758            context,
1759        );
1760
1761        // ── from_matrix_view produces identical results ──────────
1762
1763        if nrows > 0 && ncols > 0 {
1764            let via_matrix = BlockTransposed::<T, GROUP, PACK>::from_matrix_view(data.as_view());
1765            assert_eq!(via_matrix.as_slice(), transpose.as_slice(), "{}", context);
1766        }
1767    }
1768
1769    // ════════════════════════════════════════════════════════════════
1770    // Test runners — each combination gets the full API surface.
1771    // ════════════════════════════════════════════════════════════════
1772
1773    #[test]
1774    fn test_api_pack1_group16() {
1775        // Miri: boundary rows around GROUP=16 block transitions;
1776        // full run: exhaustive sweep.
1777        let rows: Vec<usize> = if cfg!(miri) {
1778            vec![0, 1, 15, 16, 17, 33]
1779        } else {
1780            (0..128).collect()
1781        };
1782        let cols: Vec<usize> = if cfg!(miri) {
1783            vec![0, 1, 2]
1784        } else {
1785            (0..5).collect()
1786        };
1787        for &nrows in &rows {
1788            for &ncols in &cols {
1789                test_full_api::<f32, 16, 1>(nrows, ncols, gen_f32);
1790            }
1791        }
1792    }
1793
1794    #[test]
1795    fn test_api_pack1_group8() {
1796        // Miri: boundary rows around GROUP=8 block transitions;
1797        // full run: exhaustive sweep.
1798        let rows: Vec<usize> = if cfg!(miri) {
1799            vec![0, 1, 7, 8, 9, 17]
1800        } else {
1801            (0..128).collect()
1802        };
1803        let cols: Vec<usize> = if cfg!(miri) {
1804            vec![0, 1, 2]
1805        } else {
1806            (0..5).collect()
1807        };
1808        for &nrows in &rows {
1809            for &ncols in &cols {
1810                test_full_api::<f32, 8, 1>(nrows, ncols, gen_f32);
1811            }
1812        }
1813    }
1814
1815    #[test]
1816    fn test_api_pack2() {
1817        // Miri: boundary rows around GROUP=4/8/16 transitions;
1818        // cols hit PACK=2 boundary (even/odd). Full run: exhaustive.
1819        let rows: Vec<usize> = if cfg!(miri) {
1820            vec![0, 1, 3, 4, 5, 7, 8, 9, 15, 16, 17]
1821        } else {
1822            (0..48).collect()
1823        };
1824        let cols: Vec<usize> = if cfg!(miri) {
1825            vec![0, 1, 2, 3, 4, 5]
1826        } else {
1827            (0..9).collect()
1828        };
1829        for &nrows in &rows {
1830            for &ncols in &cols {
1831                test_full_api::<f32, 4, 2>(nrows, ncols, gen_f32);
1832                test_full_api::<f32, 8, 2>(nrows, ncols, gen_f32);
1833                test_full_api::<f32, 16, 2>(nrows, ncols, gen_f32);
1834            }
1835        }
1836    }
1837
1838    #[test]
1839    fn test_api_pack4() {
1840        // Miri: boundary rows around GROUP=4/8/16 transitions;
1841        // cols hit PACK=4 boundary (0,1,3,4,5,8). Full run: exhaustive.
1842        let rows: Vec<usize> = if cfg!(miri) {
1843            vec![0, 1, 3, 4, 5, 7, 8, 9, 15, 16, 17]
1844        } else {
1845            (0..48).collect()
1846        };
1847        let cols: Vec<usize> = if cfg!(miri) {
1848            vec![0, 1, 3, 4, 5, 8]
1849        } else {
1850            (0..9).collect()
1851        };
1852        for &nrows in &rows {
1853            for &ncols in &cols {
1854                test_full_api::<f32, 4, 4>(nrows, ncols, gen_f32);
1855                test_full_api::<f32, 8, 4>(nrows, ncols, gen_f32);
1856                test_full_api::<f32, 16, 4>(nrows, ncols, gen_f32);
1857            }
1858        }
1859    }
1860
1861    /// Exercise the unified test with non-`f32` element types.
1862    #[test]
1863    fn test_api_non_f32() {
1864        // i32:  PACK=1 and PACK=2
1865        test_full_api::<i32, 4, 1>(10, 7, gen_i32);
1866        test_full_api::<i32, 8, 2>(12, 5, gen_i32);
1867
1868        // u8:   PACK=1 and PACK=2
1869        test_full_api::<u8, 4, 2>(12, 5, gen_u8);
1870        test_full_api::<u8, 8, 1>(10, 7, gen_u8);
1871    }
1872
1873    // ════════════════════════════════════════════════════════════════
1874    // Block layout verification (PACK=1 only)
1875    // ════════════════════════════════════════════════════════════════
1876
1877    /// Verify that for PACK=1, each block is the standard row-to-column
1878    /// transposition of a GROUP-row slice of the source matrix.
1879    fn test_block_layout_pack1<
1880        T: Copy + Default + PartialEq + std::fmt::Debug + 'static,
1881        const GROUP: usize,
1882    >(
1883        nrows: usize,
1884        ncols: usize,
1885        gen_element: fn(usize) -> T,
1886    ) {
1887        let mut data = Matrix::new(T::default(), nrows, ncols);
1888        data.as_mut_slice()
1889            .iter_mut()
1890            .enumerate()
1891            .for_each(|(i, d)| *d = gen_element(i));
1892
1893        let transpose = BlockTransposed::<T, GROUP, 1>::from_strided(data.as_view().into());
1894
1895        // Full blocks.
1896        for b in 0..transpose.full_blocks() {
1897            let block = transpose.block(b);
1898            for i in 0..block.nrows() {
1899                for j in 0..block.ncols() {
1900                    assert_eq!(
1901                        block[(i, j)],
1902                        data[(GROUP * b + j, i)],
1903                        "block {} at ({}, {}) -- GROUP={}, nrows={}, ncols={}",
1904                        b,
1905                        i,
1906                        j,
1907                        GROUP,
1908                        nrows,
1909                        ncols,
1910                    );
1911                }
1912            }
1913        }
1914
1915        // Remainder block.
1916        if transpose.remainder() != 0 {
1917            let fb = transpose.full_blocks();
1918            let block = transpose.remainder_block().unwrap();
1919            for i in 0..block.nrows() {
1920                for j in 0..transpose.remainder() {
1921                    assert_eq!(
1922                        block[(i, j)],
1923                        data[(GROUP * fb + j, i)],
1924                        "remainder at ({}, {}) -- GROUP={}, nrows={}, ncols={}",
1925                        i,
1926                        j,
1927                        GROUP,
1928                        nrows,
1929                        ncols,
1930                    );
1931                }
1932            }
1933        }
1934    }
1935
1936    #[test]
1937    fn test_block_layout_pack1_group16() {
1938        let rows: Vec<usize> = if cfg!(miri) {
1939            vec![0, 1, 15, 16, 17, 33]
1940        } else {
1941            (0..128).collect()
1942        };
1943        let cols: Vec<usize> = if cfg!(miri) {
1944            vec![0, 1, 2]
1945        } else {
1946            (0..5).collect()
1947        };
1948        for &nrows in &rows {
1949            for &ncols in &cols {
1950                test_block_layout_pack1::<f32, 16>(nrows, ncols, gen_f32);
1951            }
1952        }
1953    }
1954
1955    #[test]
1956    fn test_block_layout_pack1_group8() {
1957        let rows: Vec<usize> = if cfg!(miri) {
1958            vec![0, 1, 7, 8, 9, 17]
1959        } else {
1960            (0..128).collect()
1961        };
1962        let cols: Vec<usize> = if cfg!(miri) {
1963            vec![0, 1, 2]
1964        } else {
1965            (0..5).collect()
1966        };
1967        for &nrows in &rows {
1968            for &ncols in &cols {
1969                test_block_layout_pack1::<f32, 8>(nrows, ncols, gen_f32);
1970            }
1971        }
1972    }
1973
1974    // ════════════════════════════════════════════════════════════════
1975    // Focused tests (not part of the unified parameterized test)
1976    // ════════════════════════════════════════════════════════════════
1977
1978    // ── Send / Sync static assertions ───────────────────────────────
1979
1980    #[test]
1981    fn test_row_view_send_sync() {
1982        fn assert_send<T: Send>() {}
1983        fn assert_sync<T: Sync>() {}
1984
1985        assert_send::<Row<'_, f32, 16>>();
1986        assert_sync::<Row<'_, f32, 16>>();
1987        assert_send::<Row<'_, u8, 8, 2>>();
1988        assert_sync::<Row<'_, u8, 8, 2>>();
1989
1990        assert_send::<RowMut<'_, f32, 16>>();
1991        assert_sync::<RowMut<'_, f32, 16>>();
1992        assert_send::<RowMut<'_, i32, 4, 4>>();
1993        assert_sync::<RowMut<'_, i32, 4, 4>>();
1994    }
1995
1996    // ── NewRef / NewMut from raw slices ─────────────────────────────
1997
1998    #[test]
1999    fn test_new_ref_and_new_mut() {
2000        let nrows = 5;
2001        let ncols = 3;
2002        let repr = BlockTransposedRepr::<f32, 4>::new(nrows, ncols).unwrap();
2003
2004        let mat = BlockTransposed::<f32, 4>::new(nrows, ncols);
2005        let raw: &[f32] = mat.as_slice();
2006
2007        let mat_ref = BlockTransposedRef::new(repr.new_ref(raw).unwrap());
2008        assert_eq!(mat_ref.nrows(), nrows);
2009        assert_eq!(mat_ref.ncols(), ncols);
2010        for row in 0..nrows {
2011            for col in 0..ncols {
2012                assert_eq!(mat_ref.get_element(row, col), mat.get_element(row, col));
2013            }
2014        }
2015
2016        let mut buf = raw.to_vec();
2017        let mat_mut = BlockTransposedMut::new(repr.new_mut(&mut buf).unwrap());
2018        assert_eq!(mat_mut.nrows(), nrows);
2019        assert_eq!(mat_mut.ncols(), ncols);
2020
2021        // Wrong-length slice should fail.
2022        let mut short = vec![0.0_f32; 2];
2023        assert!(repr.new_ref(&short).is_err());
2024        assert!(repr.new_mut(&mut short).is_err());
2025    }
2026
2027    // ── Row view edge cases ─────────────────────────────────────────
2028
2029    #[test]
2030    fn test_row_view_empty() {
2031        /// Verify that immutable and mutable empty-row views are sound for a
2032        /// given `GROUP`/`PACK` combination.
2033        fn check_empty<const GROUP: usize, const PACK: usize>() {
2034            let mut mat = BlockTransposed::<f32, GROUP, PACK>::new(4, 0);
2035
2036            // Immutable views.
2037            let view = mat.as_view();
2038            for i in 0..4 {
2039                let row = view.get_row(i).unwrap();
2040                assert!(row.is_empty());
2041                assert_eq!(row.len(), 0);
2042                assert_eq!(row.iter().count(), 0);
2043            }
2044
2045            // Mutable views.
2046            for i in 0..4 {
2047                let row = mat.get_row_mut(i).unwrap();
2048                assert!(row.is_empty());
2049                assert_eq!(row.len(), 0);
2050            }
2051        }
2052
2053        check_empty::<16, 1>(); // default PACK
2054        check_empty::<4, 2>(); // PACK > 1
2055        check_empty::<4, 4>(); // PACK == GROUP
2056    }
2057
2058    // ── Bounds-checking panic tests ─────────────────────────────────
2059
2060    #[test]
2061    #[should_panic(expected = "column index 3 out of bounds")]
2062    fn test_row_view_index_oob() {
2063        let mat = BlockTransposed::<f32, 4>::new(4, 3);
2064        let view = mat.as_view();
2065        let row = view.get_row(0).unwrap();
2066        let _ = row[3];
2067    }
2068
2069    #[test]
2070    #[should_panic(expected = "column index 3 out of bounds")]
2071    fn test_row_view_mut_index_oob() {
2072        let mut mat = BlockTransposed::<f32, 4>::new(4, 3);
2073        let row = mat.get_row_mut(0).unwrap();
2074        let _ = row[3];
2075    }
2076
2077    #[test]
2078    #[should_panic(expected = "column index 3 out of bounds")]
2079    fn test_row_view_mut_index_mut_oob() {
2080        let mut mat = BlockTransposed::<f32, 4>::new(4, 3);
2081        let mut row = mat.get_row_mut(0).unwrap();
2082        row[3] = 1.0;
2083    }
2084
2085    #[test]
2086    #[should_panic(expected = "column index 3 out of bounds")]
2087    fn test_row_view_set_oob() {
2088        let mut mat = BlockTransposed::<f32, 4>::new(4, 3);
2089        let mut row = mat.get_row_mut(0).unwrap();
2090        row.set(3, 1.0);
2091    }
2092
2093    #[test]
2094    #[should_panic(expected = "row 4 out of bounds")]
2095    fn test_get_element_row_oob() {
2096        let mat = BlockTransposed::<f32, 4>::new(4, 3);
2097        mat.get_element(4, 0);
2098    }
2099
2100    #[test]
2101    #[should_panic(expected = "col 3 out of bounds")]
2102    fn test_get_element_col_oob() {
2103        let mat = BlockTransposed::<f32, 4>::new(4, 3);
2104        mat.get_element(0, 3);
2105    }
2106
2107    #[test]
2108    #[should_panic(expected = "assertion failed")]
2109    fn test_index_tuple_row_oob() {
2110        let mat = BlockTransposed::<f32, 4>::new(4, 3);
2111        let _ = mat[(4, 0)];
2112    }
2113
2114    #[test]
2115    #[should_panic(expected = "assertion failed")]
2116    fn test_index_tuple_col_oob() {
2117        let mat = BlockTransposed::<f32, 4>::new(4, 3);
2118        let _ = mat[(0, 3)];
2119    }
2120
2121    #[test]
2122    #[should_panic]
2123    fn test_block_oob() {
2124        let mat = BlockTransposed::<f32, 4>::new(4, 3);
2125        let _ = mat.block(1);
2126    }
2127
2128    #[test]
2129    #[should_panic]
2130    fn test_block_mut_oob() {
2131        let mut mat = BlockTransposed::<f32, 4>::new(4, 3);
2132        let _ = mat.block_mut(1);
2133    }
2134
2135    // ── from_strided with non-unit stride ───────────────────────────
2136
2137    #[test]
2138    fn test_from_strided_nonunit_stride() {
2139        use diskann_utils::strided::StridedView;
2140
2141        const GROUP: usize = 4;
2142        const PACK: usize = 2;
2143        let nrows = 5;
2144        let ncols = 3;
2145        let cstride = 8;
2146
2147        let required_len = (nrows - 1) * cstride + ncols;
2148        let mut flat = vec![0.0_f32; required_len];
2149        for row in 0..nrows {
2150            for col in 0..ncols {
2151                flat[row * cstride + col] = (row * 100 + col + 1) as f32;
2152            }
2153        }
2154
2155        let strided = StridedView::try_shrink_from(&flat, nrows, ncols, cstride)
2156            .expect("should construct strided view");
2157        let transpose = BlockTransposed::<f32, GROUP, PACK>::from_strided(strided);
2158
2159        assert_eq!(transpose.nrows(), nrows);
2160        assert_eq!(transpose.ncols(), ncols);
2161
2162        for row in 0..nrows {
2163            for col in 0..ncols {
2164                let expected = (row * 100 + col + 1) as f32;
2165                assert_eq!(
2166                    transpose[(row, col)],
2167                    expected,
2168                    "mismatch at ({}, {})",
2169                    row,
2170                    col,
2171                );
2172            }
2173        }
2174
2175        let padded_ncols = ncols.next_multiple_of(PACK);
2176        let raw: &[f32] = transpose.as_slice();
2177        for row in 0..nrows {
2178            for col in ncols..padded_ncols {
2179                let idx = linear_index::<GROUP, PACK>(row, col, ncols);
2180                assert_eq!(
2181                    raw[idx], 0.0,
2182                    "column-padding at ({}, {}) should be zero",
2183                    row, col,
2184                );
2185            }
2186        }
2187    }
2188
2189    // ── Concurrent multi-row mutation ───────────────────────────────
2190
2191    #[test]
2192    fn test_concurrent_row_mutation() {
2193        const GROUP: usize = 8;
2194        const PACK: usize = 2;
2195
2196        let (nrows, ncols, num_threads) = if cfg!(miri) { (8, 4, 2) } else { (64, 16, 4) };
2197
2198        let mut mat = BlockTransposed::<f32, GROUP, PACK>::new(nrows, ncols);
2199        let rows: Vec<RowMut<'_, f32, GROUP, PACK>> = mat.data.rows_mut().collect();
2200        let rows_per_thread = nrows / num_threads;
2201        let mut rows = rows.into_boxed_slice();
2202
2203        std::thread::scope(|s| {
2204            let mut remaining = &mut rows[..];
2205            for thread_id in 0..num_threads {
2206                let chunk_len = if thread_id == num_threads - 1 {
2207                    remaining.len()
2208                } else {
2209                    rows_per_thread
2210                };
2211                let (chunk, rest) = remaining.split_at_mut(chunk_len);
2212                remaining = rest;
2213                let start_row = thread_id * rows_per_thread;
2214
2215                s.spawn(move || {
2216                    for (offset, row_view) in chunk.iter_mut().enumerate() {
2217                        let row = start_row + offset;
2218                        for col in 0..ncols {
2219                            let value = (thread_id * 10000 + row * 100 + col) as f32;
2220                            row_view.set(col, value);
2221                        }
2222                    }
2223                });
2224            }
2225        });
2226
2227        for row in 0..nrows {
2228            let thread_id = (row / rows_per_thread).min(num_threads - 1);
2229            for col in 0..ncols {
2230                let expected = (thread_id * 10000 + row * 100 + col) as f32;
2231                assert_eq!(
2232                    mat.get_element(row, col),
2233                    expected,
2234                    "mismatch at ({}, {})",
2235                    row,
2236                    col,
2237                );
2238            }
2239        }
2240    }
2241}