hurray_core/layout/addressing/
csr.rs

1//! CSR (Compressed Sparse Row) sparse layout element lookup.
2//!
3//! Implements element lookup from `docs/spec/layouts/csr.md`: within a row, the stored
4//! column indices are strictly increasing, so a non-zero is found by binary-searching the
5//! row's slice of `col_indices`. Mirrors the standalone-function API validated by
6//! [`super::csf::element_offset`].
7
8use crate::{Error, Result};
9
10/// Looks up the storage offset (`values` / `col_indices` buffer index) of logical index
11/// `(row, col)` in a CSR matrix, or returns `None` if the element is a structural zero.
12///
13/// # Arguments
14///
15/// - `query` — logical index `[row, col]`; MUST have length 2 (CSR is rank-2).
16/// - `col_indices` — the `col_indices` buffer (buffer 1) as `uint64`: the column of each
17///   non-zero, in row-major storage order (`nnz` entries).
18/// - `row_ptr` — the `row_ptr` buffer (buffer 2) as `uint64`: `nrows + 1` entries, where
19///   `row_ptr[i]` is the first storage index of row `i` and `row_ptr[nrows] = nnz`.
20///
21/// # Errors
22///
23/// - [`Error::IndexRankMismatch`] — `query.len() != 2`.
24/// - [`Error::IndexOutOfRange`] — `row` is not a valid row (`row >= nrows`).
25/// - [`Error::InvalidLayout`] — `row_ptr` is empty, or a `row_ptr` entry points outside
26///   `col_indices` / is non-monotone for the queried row.
27///
28/// The caller is responsible for validating `col` against `shape[1]`; an out-of-range
29/// column simply reports a structural zero (`None`).
30///
31/// # Examples
32///
33/// ```
34/// use hurray_core::layout::addressing::csr::element_offset;
35///
36/// // 3×3 matrix:
37/// //   row 0: (0,0), (0,2)
38/// //   row 1: —
39/// //   row 2: (2,1)
40/// let col_indices: &[u64] = &[0, 2, 1];
41/// let row_ptr: &[u64] = &[0, 2, 2, 3];
42///
43/// assert_eq!(element_offset(&[0, 2], col_indices, row_ptr).unwrap(), Some(1));
44/// assert_eq!(element_offset(&[2, 1], col_indices, row_ptr).unwrap(), Some(2));
45/// assert_eq!(element_offset(&[1, 0], col_indices, row_ptr).unwrap(), None); // empty row
46/// ```
47pub fn element_offset(query: &[u64], col_indices: &[u64], row_ptr: &[u64]) -> Result<Option<u64>> {
48    if query.len() != 2 {
49        return Err(Error::IndexRankMismatch {
50            index_rank: query.len(),
51            shape_rank: 2,
52        });
53    }
54    let (row, col) = (query[0], query[1]);
55
56    if row_ptr.is_empty() {
57        return Err(Error::InvalidLayout(
58            "csr: row_ptr must have at least nrows+1 = 1 entries".into(),
59        ));
60    }
61    let nrows = (row_ptr.len() - 1) as u64;
62    if row >= nrows {
63        return Err(Error::IndexOutOfRange {
64            dim: 0,
65            index: row,
66            size: nrows,
67        });
68    }
69
70    let start = row_ptr[row as usize] as usize;
71    let end = row_ptr[row as usize + 1] as usize;
72    if start > end || end > col_indices.len() {
73        return Err(Error::InvalidLayout(format!(
74            "csr: row_ptr[{row}..{}] = [{start}..{end}] out of bounds for col_indices.len()={}",
75            row + 1,
76            col_indices.len()
77        )));
78    }
79
80    match col_indices[start..end].binary_search(&col) {
81        Ok(k) => Ok(Some((start + k) as u64)),
82        Err(_) => Ok(None),
83    }
84}
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89
90    // 4×4 matrix:
91    //   row 0: (0,0)=a, (0,3)=b
92    //   row 1: (1,1)=c
93    //   row 2: —
94    //   row 3: (3,0)=d, (3,2)=e, (3,3)=f
95    // values order: a,b,c,d,e,f
96    const COL_INDICES: &[u64] = &[0, 3, 1, 0, 2, 3];
97    const ROW_PTR: &[u64] = &[0, 2, 3, 3, 6];
98
99    #[test]
100    fn lookup_hits() {
101        assert_eq!(
102            element_offset(&[0, 0], COL_INDICES, ROW_PTR).unwrap(),
103            Some(0)
104        );
105        assert_eq!(
106            element_offset(&[0, 3], COL_INDICES, ROW_PTR).unwrap(),
107            Some(1)
108        );
109        assert_eq!(
110            element_offset(&[1, 1], COL_INDICES, ROW_PTR).unwrap(),
111            Some(2)
112        );
113        assert_eq!(
114            element_offset(&[3, 0], COL_INDICES, ROW_PTR).unwrap(),
115            Some(3)
116        );
117        assert_eq!(
118            element_offset(&[3, 2], COL_INDICES, ROW_PTR).unwrap(),
119            Some(4)
120        );
121        assert_eq!(
122            element_offset(&[3, 3], COL_INDICES, ROW_PTR).unwrap(),
123            Some(5)
124        );
125    }
126
127    #[test]
128    fn structural_zeros() {
129        assert_eq!(element_offset(&[0, 1], COL_INDICES, ROW_PTR).unwrap(), None);
130        assert_eq!(element_offset(&[1, 0], COL_INDICES, ROW_PTR).unwrap(), None);
131        assert_eq!(element_offset(&[2, 2], COL_INDICES, ROW_PTR).unwrap(), None); // empty row
132        assert_eq!(element_offset(&[3, 1], COL_INDICES, ROW_PTR).unwrap(), None);
133    }
134
135    #[test]
136    fn wrong_rank_rejected() {
137        assert!(matches!(
138            element_offset(&[0], COL_INDICES, ROW_PTR),
139            Err(Error::IndexRankMismatch { .. })
140        ));
141        assert!(matches!(
142            element_offset(&[0, 0, 0], COL_INDICES, ROW_PTR),
143            Err(Error::IndexRankMismatch { .. })
144        ));
145    }
146
147    #[test]
148    fn row_out_of_range_rejected() {
149        assert!(matches!(
150            element_offset(&[4, 0], COL_INDICES, ROW_PTR),
151            Err(Error::IndexOutOfRange { dim: 0, .. })
152        ));
153    }
154
155    #[test]
156    fn empty_row_ptr_rejected() {
157        assert!(matches!(
158            element_offset(&[0, 0], COL_INDICES, &[]),
159            Err(Error::InvalidLayout(_))
160        ));
161    }
162
163    #[test]
164    fn empty_matrix() {
165        // 2×2 all-zero: row_ptr = [0,0,0], no stored columns.
166        assert_eq!(element_offset(&[0, 0], &[], &[0, 0, 0]).unwrap(), None);
167        assert_eq!(element_offset(&[1, 1], &[], &[0, 0, 0]).unwrap(), None);
168    }
169}
hurray_core/layout/addressing/csr.rs

hurray_core/layout/addressing/
csr.rs
