Skip to main content

nodedb_array/codec/
coord_rle.rs

1// SPDX-License-Identifier: Apache-2.0
2
3// Leading-axis run-length compression for DimDict index streams.
4//
5// When the leading axis has long runs of the same dict index (common in
6// genomic tiles where chrom is constant across thousands of positions),
7// RLE saves ~75 % vs raw indices. For tiles with no runs longer than 2,
8// we fall back to the delta codec so overhead is bounded.
9//
10// Wire format:
11//   RLE mode:   [u32 LE = 0xFFFF_FFFF marker][u32 LE run_count] ...
12//   Delta mode: delegates to coord_delta::encode_coord_axis / decode_coord_axis
13//
14// The RLE marker 0xFFFF_FFFF is safe because DimDict cardinality fits in
15// a u32 but realistic tile dicts never exceed ~10M distinct values; an
16// all-ones 4-byte value is reserved as the RLE sentinel.
17
18use crate::codec::coord_delta::{decode_coord_axis, encode_coord_axis};
19use crate::codec::limits::{
20    MAX_CELLS_PER_TILE, MAX_DICT_CARDINALITY, MAX_RLE_RUN_LEN, MAX_RLE_RUNS, check_decoded_size,
21    checked_capacity,
22};
23use crate::error::{ArrayError, ArrayResult};
24use crate::tile::sparse_tile::DimDict;
25use crate::types::coord::value::CoordValue;
26
27/// Minimum ratio of run savings to total indices before RLE is preferred.
28/// Using integer arithmetic: savings_cells / total_cells >= 1/2.
29const RLE_BENEFIT_NUMERATOR: usize = 1;
30const RLE_BENEFIT_DENOMINATOR: usize = 2;
31
32fn should_use_rle(indices: &[u32]) -> bool {
33    if indices.len() < 4 {
34        return false;
35    }
36    let mut run_count = 1usize;
37    for i in 1..indices.len() {
38        if indices[i] != indices[i - 1] {
39            run_count += 1;
40        }
41    }
42    // RLE cells needed: run_count * 8 bytes (value + length pairs)
43    // Delta cells approx: indices.len() * 1.5 bytes average
44    // Use simpler heuristic: worth it if run_count < indices.len() / 2
45    run_count * RLE_BENEFIT_DENOMINATOR < indices.len() * RLE_BENEFIT_NUMERATOR
46}
47
48/// RLE mode sentinel: u32::MAX. This value can never appear as a dict count
49/// because a DimDict with u32::MAX distinct values would require ~16 GB of
50/// coordinate storage in the tile before it reaches this codec.
51const RLE_MARKER: u32 = u32::MAX;
52
53/// Encode a DimDict using RLE if beneficial, falling back to delta otherwise.
54pub fn encode_coord_axis_rle(dict: &DimDict, out: &mut Vec<u8>) -> ArrayResult<()> {
55    if should_use_rle(&dict.indices) {
56        encode_rle(dict, out)
57    } else {
58        encode_coord_axis(dict, out)
59    }
60}
61
62fn encode_rle(dict: &DimDict, out: &mut Vec<u8>) -> ArrayResult<()> {
63    out.extend_from_slice(&RLE_MARKER.to_le_bytes());
64
65    // Compute runs over index stream.
66    let mut runs: Vec<(u32, u32)> = Vec::new(); // (value, length)
67    let mut i = 0;
68    while i < dict.indices.len() {
69        let val = dict.indices[i];
70        let mut len = 1u32;
71        while i + (len as usize) < dict.indices.len() && dict.indices[i + (len as usize)] == val {
72            len += 1;
73        }
74        runs.push((val, len));
75        i += len as usize;
76    }
77
78    out.extend_from_slice(&(runs.len() as u32).to_le_bytes());
79    for (val, len) in &runs {
80        out.extend_from_slice(&val.to_le_bytes());
81        out.extend_from_slice(&len.to_le_bytes());
82    }
83
84    // Dict values.
85    out.extend_from_slice(&(dict.values.len() as u32).to_le_bytes());
86    for cv in &dict.values {
87        let bytes = zerompk::to_msgpack_vec(cv).map_err(|e| ArrayError::SegmentCorruption {
88            detail: format!("rle coord dict encode: {e}"),
89        })?;
90        out.extend_from_slice(&(bytes.len() as u32).to_le_bytes());
91        out.extend_from_slice(&bytes);
92    }
93
94    Ok(())
95}
96
97/// Decode a DimDict from bytes previously encoded by `encode_coord_axis_rle`.
98pub fn decode_coord_axis_rle(data: &[u8], pos: &mut usize) -> ArrayResult<DimDict> {
99    if *pos + 4 > data.len() {
100        return Err(ArrayError::SegmentCorruption {
101            detail: "rle coord: truncated mode marker".into(),
102        });
103    }
104    let first_u32 = u32::from_le_bytes(
105        data[*pos..*pos + 4]
106            .try_into()
107            .expect("invariant: bounds check at line 96 guarantees 4 bytes available"),
108    );
109
110    if first_u32 == RLE_MARKER {
111        *pos += 4;
112        decode_rle(data, pos)
113    } else {
114        // Delta mode — don't advance pos; coord_delta reads dict_count first.
115        decode_coord_axis(data, pos)
116    }
117}
118
119fn decode_rle(data: &[u8], pos: &mut usize) -> ArrayResult<DimDict> {
120    if *pos + 4 > data.len() {
121        return Err(ArrayError::SegmentCorruption {
122            detail: "rle coord: truncated run count".into(),
123        });
124    }
125    let run_count = u32::from_le_bytes(
126        data[*pos..*pos + 4]
127            .try_into()
128            .expect("invariant: bounds check at line 113 guarantees 4 bytes available"),
129    ) as usize;
130    *pos += 4;
131    checked_capacity(
132        run_count,
133        MAX_RLE_RUNS,
134        data.len() - *pos,
135        8,
136        "rle run_count",
137    )?;
138
139    let mut indices: Vec<u32> = Vec::new();
140    let mut total_len: usize = 0;
141    for _ in 0..run_count {
142        if *pos + 8 > data.len() {
143            return Err(ArrayError::SegmentCorruption {
144                detail: "rle coord: truncated run entry".into(),
145            });
146        }
147        let val =
148            u32::from_le_bytes(data[*pos..*pos + 4].try_into().expect(
149                "invariant: bounds check at line 125 guarantees 8 bytes available; first 4",
150            ));
151        *pos += 4;
152        let len =
153            u32::from_le_bytes(data[*pos..*pos + 4].try_into().expect(
154                "invariant: bounds check at line 125 guarantees 8 bytes available; second 4",
155            )) as usize;
156        *pos += 4;
157        check_decoded_size(len, MAX_RLE_RUN_LEN, "rle run len")?;
158        total_len = total_len.saturating_add(len);
159        check_decoded_size(total_len, MAX_CELLS_PER_TILE, "rle indices total_len")?;
160        for _ in 0..len {
161            indices.push(val);
162        }
163    }
164
165    // Dict values.
166    if *pos + 4 > data.len() {
167        return Err(ArrayError::SegmentCorruption {
168            detail: "rle coord: truncated dict count".into(),
169        });
170    }
171    let dict_count = u32::from_le_bytes(
172        data[*pos..*pos + 4]
173            .try_into()
174            .expect("invariant: bounds check at line 143 guarantees 4 bytes available"),
175    ) as usize;
176    *pos += 4;
177    let capacity = checked_capacity(
178        dict_count,
179        MAX_DICT_CARDINALITY,
180        data.len() - *pos,
181        4,
182        "rle dict_count",
183    )?;
184
185    let mut values: Vec<CoordValue> = Vec::with_capacity(capacity);
186    for _ in 0..dict_count {
187        if *pos + 4 > data.len() {
188            return Err(ArrayError::SegmentCorruption {
189                detail: "rle coord: truncated dict entry len".into(),
190            });
191        }
192        let len = u32::from_le_bytes(
193            data[*pos..*pos + 4]
194                .try_into()
195                .expect("invariant: bounds check at line 154 guarantees 4 bytes available"),
196        ) as usize;
197        *pos += 4;
198        if *pos + len > data.len() {
199            return Err(ArrayError::SegmentCorruption {
200                detail: "rle coord: truncated dict entry bytes".into(),
201            });
202        }
203        let cv: CoordValue = zerompk::from_msgpack(&data[*pos..*pos + len]).map_err(|e| {
204            ArrayError::SegmentCorruption {
205                detail: format!("rle coord dict decode: {e}"),
206            }
207        })?;
208        *pos += len;
209        values.push(cv);
210    }
211
212    Ok(DimDict { values, indices })
213}
214
215#[cfg(test)]
216mod tests {
217    use super::*;
218    use crate::types::coord::value::CoordValue;
219
220    fn make_dict(indices: Vec<u32>, values: Vec<CoordValue>) -> DimDict {
221        DimDict { values, indices }
222    }
223
224    fn roundtrip(dict: &DimDict) -> DimDict {
225        let mut buf = Vec::new();
226        encode_coord_axis_rle(dict, &mut buf).unwrap();
227        let mut pos = 0;
228        decode_coord_axis_rle(&buf, &mut pos).unwrap()
229    }
230
231    #[test]
232    fn huge_run_count_in_tiny_payload_is_rejected_before_allocation() {
233        let mut data = Vec::new();
234        data.extend_from_slice(&RLE_MARKER.to_le_bytes());
235        data.extend_from_slice(&(MAX_RLE_RUNS as u32).to_le_bytes());
236        let result = decode_coord_axis_rle(&data, &mut 0);
237        assert!(matches!(
238            result,
239            Err(crate::error::ArrayError::SegmentCorruption { .. })
240        ));
241    }
242
243    #[test]
244    fn empty_dict_roundtrip() {
245        let d = make_dict(vec![], vec![]);
246        let out = roundtrip(&d);
247        assert_eq!(out.indices, d.indices);
248        assert_eq!(out.values, d.values);
249    }
250
251    #[test]
252    fn single_entry_roundtrip() {
253        let d = make_dict(vec![0], vec![CoordValue::Int64(5)]);
254        let out = roundtrip(&d);
255        assert_eq!(out.indices, d.indices);
256        assert_eq!(out.values, d.values);
257    }
258
259    #[test]
260    fn long_run_uses_rle() {
261        // 1000 cells all pointing at index 0 — should trigger RLE.
262        let d = make_dict(vec![0u32; 1000], vec![CoordValue::Int64(42)]);
263        let mut buf = Vec::new();
264        encode_coord_axis_rle(&d, &mut buf).unwrap();
265        // RLE mode: first 4 bytes are the RLE_MARKER sentinel.
266        let first = u32::from_le_bytes(buf[0..4].try_into().unwrap());
267        assert_eq!(first, u32::MAX, "expected RLE mode marker");
268
269        let out = roundtrip(&d);
270        assert_eq!(out.indices, d.indices);
271    }
272
273    #[test]
274    fn no_runs_falls_back_to_delta() {
275        // All distinct coords — no benefit from RLE.
276        let values: Vec<CoordValue> = (0..10).map(CoordValue::Int64).collect();
277        let indices: Vec<u32> = (0..10).collect();
278        let d = make_dict(indices, values);
279        let mut buf = Vec::new();
280        encode_coord_axis_rle(&d, &mut buf).unwrap();
281        // Delta mode: first 4 bytes are dict_count (10), not the RLE marker.
282        let first = u32::from_le_bytes(buf[0..4].try_into().unwrap());
283        assert_ne!(first, u32::MAX, "should be delta mode");
284
285        let out = roundtrip(&d);
286        assert_eq!(out.indices, d.indices);
287    }
288
289    #[test]
290    fn mixed_runs_roundtrip() {
291        // Two chromosomes, each repeated 500 times.
292        let mut indices = vec![0u32; 500];
293        indices.extend(vec![1u32; 500]);
294        let d = make_dict(indices, vec![CoordValue::Int64(1), CoordValue::Int64(2)]);
295        let out = roundtrip(&d);
296        assert_eq!(out.indices, d.indices);
297        assert_eq!(out.values, d.values);
298    }
299
300    #[test]
301    fn large_rle_roundtrip() {
302        let d = make_dict(vec![0u32; 100_000], vec![CoordValue::Int64(99)]);
303        let out = roundtrip(&d);
304        assert_eq!(out.indices.len(), 100_000);
305        assert!(out.indices.iter().all(|&i| i == 0));
306    }
307
308    #[test]
309    fn timestamp_coord_roundtrip() {
310        let values = vec![
311            CoordValue::TimestampMs(1_000_000),
312            CoordValue::TimestampMs(2_000_000),
313        ];
314        let indices: Vec<u32> = (0..10).map(|i| i % 2).collect();
315        let d = make_dict(indices, values);
316        let out = roundtrip(&d);
317        assert_eq!(out.indices, d.indices);
318        assert_eq!(out.values, d.values);
319    }
320}