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