metrale_grammar/fsm/
compact_array.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2//
3// Compact2DArray — a Compressed-Sparse-Row (CSR) 2D array.
4// Port of `cpp/support/compact_2d_array.h`, restricted to the FSM use
5// case (rows of `FsmEdge`). It stores a jagged 2D array as two flat
6// vectors so that all rows are contiguous in memory — this layout is
7// load-bearing for the compiled-grammar memory footprint.
8
9/// A CSR-packed jagged 2D array.
10///
11/// `data` holds every row's elements contiguously; `indptr` holds the
12/// start offset of each row, with a trailing entry equal to `data.len()`
13/// so row `i` is `data[indptr[i]..indptr[i+1]]`.
14#[derive(Debug, Clone, PartialEq, Eq)]
15pub struct Compact2DArray<T> {
16    data: Vec<T>,
17    indptr: Vec<i32>,
18}
19
20impl<T: Clone> Default for Compact2DArray<T> {
21    fn default() -> Self {
22        Self::new()
23    }
24}
25
26impl<T: Clone> Compact2DArray<T> {
27    /// Create an empty array (zero rows).
28    pub fn new() -> Self {
29        Self {
30            data: Vec::new(),
31            indptr: vec![0],
32        }
33    }
34
35    /// Number of rows.
36    pub fn len(&self) -> usize {
37        self.indptr.len() - 1
38    }
39
40    /// True when there are no rows.
41    pub fn is_empty(&self) -> bool {
42        self.len() == 0
43    }
44
45    /// Total number of stored elements across all rows.
46    pub fn total_elems(&self) -> usize {
47        self.data.len()
48    }
49
50    /// Borrow row `i` as a slice. Panics if `i` is out of bounds.
51    pub fn row(&self, i: usize) -> &[T] {
52        let start = self.indptr[i] as usize;
53        let end = self.indptr[i + 1] as usize;
54        &self.data[start..end]
55    }
56
57    /// Append a new row from a slice; returns the new row's index.
58    pub fn push_row(&mut self, row: &[T]) -> usize {
59        self.data.extend_from_slice(row);
60        self.indptr.push(self.data.len() as i32);
61        self.len() - 1
62    }
63
64    /// Iterate over every row as a slice.
65    pub fn iter_rows(&self) -> impl Iterator<Item = &[T]> {
66        (0..self.len()).map(move |i| self.row(i))
67    }
68
69    /// Approximate heap memory size in bytes (mirrors C++ `MemorySize`).
70    pub fn memory_size(&self) -> usize {
71        self.data.len() * std::mem::size_of::<T>() + self.indptr.len() * std::mem::size_of::<i32>()
72    }
73
74    /// Raw element backing store (CSR `data` array).
75    pub fn raw_data(&self) -> &[T] {
76        &self.data
77    }
78
79    /// Raw row-offset array (CSR `indptr` array).
80    pub fn raw_indptr(&self) -> &[i32] {
81        &self.indptr
82    }
83
84    /// Reconstruct from raw CSR parts. Returns `None` if the parts are
85    /// inconsistent (used by deserialization).
86    pub fn from_raw(data: Vec<T>, indptr: Vec<i32>) -> Option<Self> {
87        if indptr.is_empty() || indptr[0] != 0 {
88            return None;
89        }
90        if *indptr.last().unwrap() as usize != data.len() {
91            return None;
92        }
93        if indptr.windows(2).any(|w| w[0] > w[1]) {
94            return None;
95        }
96        Some(Self { data, indptr })
97    }
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103
104    #[test]
105    fn empty_array() {
106        let a: Compact2DArray<i32> = Compact2DArray::new();
107        assert_eq!(a.len(), 0);
108        assert!(a.is_empty());
109        assert_eq!(a.total_elems(), 0);
110    }
111
112    #[test]
113    fn push_and_read_rows() {
114        let mut a: Compact2DArray<i32> = Compact2DArray::new();
115        assert_eq!(a.push_row(&[1, 2, 3]), 0);
116        assert_eq!(a.push_row(&[]), 1);
117        assert_eq!(a.push_row(&[4]), 2);
118        assert_eq!(a.len(), 3);
119        assert_eq!(a.row(0), &[1, 2, 3]);
120        assert_eq!(a.row(1), &[] as &[i32]);
121        assert_eq!(a.row(2), &[4]);
122        assert_eq!(a.total_elems(), 4);
123    }
124
125    #[test]
126    fn iter_rows_yields_all() {
127        let mut a: Compact2DArray<i32> = Compact2DArray::new();
128        a.push_row(&[1]);
129        a.push_row(&[2, 3]);
130        let collected: Vec<Vec<i32>> = a.iter_rows().map(|r| r.to_vec()).collect();
131        assert_eq!(collected, vec![vec![1], vec![2, 3]]);
132    }
133
134    #[test]
135    fn raw_csr_roundtrip() {
136        let mut a: Compact2DArray<i32> = Compact2DArray::new();
137        a.push_row(&[1, 2]);
138        a.push_row(&[3]);
139        let data = a.raw_data().to_vec();
140        let indptr = a.raw_indptr().to_vec();
141        let b = Compact2DArray::from_raw(data, indptr).unwrap();
142        assert_eq!(a, b);
143    }
144
145    #[test]
146    fn from_raw_rejects_inconsistent() {
147        assert!(Compact2DArray::<i32>::from_raw(vec![1], vec![]).is_none());
148        assert!(Compact2DArray::<i32>::from_raw(vec![1], vec![1, 1]).is_none());
149        assert!(Compact2DArray::<i32>::from_raw(vec![1, 2], vec![0, 5]).is_none());
150        assert!(Compact2DArray::<i32>::from_raw(vec![1], vec![0, 1]).is_some());
151    }
152
153    #[test]
154    fn memory_size_nonzero() {
155        let mut a: Compact2DArray<i32> = Compact2DArray::new();
156        a.push_row(&[1, 2, 3]);
157        assert!(a.memory_size() >= 3 * 4);
158    }
159}
metrale_grammar/fsm/compact_array.rs

metrale_grammar/fsm/
compact_array.rs
