hurray_core/layout/addressing/
coo.rs

1//! COO (Coordinate) sparse layout element lookup.
2//!
3//! Implements element lookup from `docs/spec/layouts/coo.md § Storage Order`:
4//! locate the stored non-zero whose coordinate tuple equals the query, using a
5//! lexicographic binary search when the entries are sorted, or a linear scan
6//! otherwise. Mirrors the standalone-function API validated by
7//! [`super::csf::element_offset`].
8
9use std::cmp::Ordering;
10
11use crate::{Error, Result};
12
13/// Looks up the storage offset (`values` buffer index) of a logical index in a COO
14/// tensor, or returns `None` if the element is structurally absent (implicit zero).
15///
16/// # Arguments
17///
18/// - `query` — logical index `[i0, …, i_{rank-1}]`; its length defines the rank.
19/// - `is_sorted` — the COO descriptor's `is_sorted` flag. When `true`, the stored
20///   entries are in strictly increasing lexicographic order and a binary search is used;
21///   when `false`, a linear scan is used.
22/// - `indices` — the `indices` buffer (buffer 1) as `uint64` values: `nnz × rank`
23///   coordinates in row-major order, so entry `r`'s coordinates are
24///   `indices[r*rank .. r*rank + rank]`.
25///
26/// The returned offset indexes both the `values` buffer and the entry's row in `indices`.
27///
28/// # Errors
29///
30/// - [`Error::IndexRankMismatch`] — `query` is empty (rank 0).
31/// - [`Error::InvalidLayout`] — `indices.len()` is not a multiple of the rank.
32///
33/// The caller is responsible for validating each query coordinate against the tensor
34/// shape; an out-of-bounds coordinate simply reports a structural zero (`None`).
35///
36/// # Examples
37///
38/// ```
39/// use hurray_core::layout::addressing::coo::element_offset;
40///
41/// // 4×4 matrix with three sorted non-zeros: (0,1), (2,0), (2,3).
42/// let indices: &[u64] = &[0, 1, /* */ 2, 0, /* */ 2, 3];
43///
44/// assert_eq!(element_offset(&[2, 0], true, indices).unwrap(), Some(1));
45/// assert_eq!(element_offset(&[2, 3], true, indices).unwrap(), Some(2));
46/// assert_eq!(element_offset(&[1, 1], true, indices).unwrap(), None); // structural zero
47/// ```
48pub fn element_offset(query: &[u64], is_sorted: bool, indices: &[u64]) -> Result<Option<u64>> {
49    let rank = query.len();
50    if rank == 0 {
51        return Err(Error::IndexRankMismatch {
52            index_rank: 0,
53            shape_rank: 0,
54        });
55    }
56    if !indices.len().is_multiple_of(rank) {
57        return Err(Error::InvalidLayout(format!(
58            "coo: indices.len()={} is not a multiple of rank={rank}",
59            indices.len()
60        )));
61    }
62    let nnz = indices.len() / rank;
63    let coords = |r: usize| &indices[r * rank..r * rank + rank];
64
65    if is_sorted {
66        // Lexicographic binary search: slice `Ord` compares element-wise, which is exactly
67        // the spec's dimension-0-major lexicographic order.
68        let (mut lo, mut hi) = (0usize, nnz);
69        while lo < hi {
70            let mid = lo + (hi - lo) / 2;
71            match coords(mid).cmp(query) {
72                Ordering::Less => lo = mid + 1,
73                Ordering::Greater => hi = mid,
74                Ordering::Equal => return Ok(Some(mid as u64)),
75            }
76        }
77        Ok(None)
78    } else {
79        for r in 0..nnz {
80            if coords(r) == query {
81                return Ok(Some(r as u64));
82            }
83        }
84        Ok(None)
85    }
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91
92    // 4×4 matrix, sorted non-zeros at (0,1), (2,0), (2,3), (3,3).
93    const INDICES: &[u64] = &[0, 1, 2, 0, 2, 3, 3, 3];
94
95    #[test]
96    fn sorted_lookup_hits() {
97        assert_eq!(element_offset(&[0, 1], true, INDICES).unwrap(), Some(0));
98        assert_eq!(element_offset(&[2, 0], true, INDICES).unwrap(), Some(1));
99        assert_eq!(element_offset(&[2, 3], true, INDICES).unwrap(), Some(2));
100        assert_eq!(element_offset(&[3, 3], true, INDICES).unwrap(), Some(3));
101    }
102
103    #[test]
104    fn sorted_lookup_structural_zeros() {
105        assert_eq!(element_offset(&[0, 0], true, INDICES).unwrap(), None);
106        assert_eq!(element_offset(&[1, 1], true, INDICES).unwrap(), None);
107        assert_eq!(element_offset(&[2, 2], true, INDICES).unwrap(), None);
108        // Between stored coordinates of row 3.
109        assert_eq!(element_offset(&[3, 0], true, INDICES).unwrap(), None);
110    }
111
112    #[test]
113    fn unsorted_lookup_matches_sorted() {
114        // Same non-zeros, permuted; linear scan must still find them.
115        let unsorted: &[u64] = &[2, 3, 0, 1, 3, 3, 2, 0];
116        assert_eq!(element_offset(&[0, 1], false, unsorted).unwrap(), Some(1));
117        assert_eq!(element_offset(&[2, 3], false, unsorted).unwrap(), Some(0));
118        assert_eq!(element_offset(&[3, 3], false, unsorted).unwrap(), Some(2));
119        assert_eq!(element_offset(&[1, 1], false, unsorted).unwrap(), None);
120    }
121
122    #[test]
123    fn rank_3_lookup() {
124        // Shape [2,3,4]; non-zeros (0,0,1), (0,2,3), (1,1,0), sorted.
125        let indices: &[u64] = &[0, 0, 1, 0, 2, 3, 1, 1, 0];
126        assert_eq!(element_offset(&[0, 2, 3], true, indices).unwrap(), Some(1));
127        assert_eq!(element_offset(&[1, 1, 0], true, indices).unwrap(), Some(2));
128        assert_eq!(element_offset(&[1, 1, 2], true, indices).unwrap(), None);
129    }
130
131    #[test]
132    fn empty_tensor_is_all_zeros() {
133        assert_eq!(element_offset(&[0, 0], true, &[]).unwrap(), None);
134        assert_eq!(element_offset(&[0, 0], false, &[]).unwrap(), None);
135    }
136
137    #[test]
138    fn rank_zero_query_is_rejected() {
139        assert!(matches!(
140            element_offset(&[], true, &[]),
141            Err(Error::IndexRankMismatch { .. })
142        ));
143    }
144
145    #[test]
146    fn indices_not_multiple_of_rank_is_rejected() {
147        // rank 2, but 3 index values.
148        assert!(matches!(
149            element_offset(&[0, 0], true, &[0, 1, 2]),
150            Err(Error::InvalidLayout(_))
151        ));
152    }
153}
hurray_core/layout/addressing/coo.rs

hurray_core/layout/addressing/
coo.rs
