Skip to main content

nodedb_array/codec/
coord_delta.rs

1// SPDX-License-Identifier: Apache-2.0
2
3// Per-axis zigzag-varint delta encode/decode over a `DimDict`'s index stream.
4//
5// The DimDict stores (dict_values, indices). For each cell, we record the
6// dict index. Delta-encoding the index stream compresses monotonic coordinate
7// access patterns (common in range scans) from 4 bytes/cell to 1-2 bytes.
8// Non-monotonic and repeated-value patterns are handled correctly — the
9// zigzag mapping keeps backward steps small too.
10//
11// Wire format per axis:
12//   [u32 LE] dict entry count (= cardinality of distinct values)
13//   [N * msgpack bytes] serialized CoordValue dict entries (zerompk)
14//   [u32 LE] index count (= cell count)
15//   [zigzag-varint deltas ...] index deltas (first value as-is, then deltas)
16
17use crate::codec::limits::{
18    MAX_CELLS_PER_TILE, MAX_DICT_CARDINALITY, check_decoded_size, checked_capacity,
19};
20use crate::error::{ArrayError, ArrayResult};
21use crate::tile::sparse_tile::DimDict;
22use crate::types::coord::value::CoordValue;
23
24// ---------------------------------------------------------------------------
25// Zigzag varint helpers
26// ---------------------------------------------------------------------------
27
28fn zigzag_encode(v: i64) -> u64 {
29    ((v << 1) ^ (v >> 63)) as u64
30}
31
32fn zigzag_decode(v: u64) -> i64 {
33    ((v >> 1) as i64) ^ (-((v & 1) as i64))
34}
35
36fn write_varint(out: &mut Vec<u8>, mut v: u64) {
37    loop {
38        let b = (v & 0x7F) as u8;
39        v >>= 7;
40        if v == 0 {
41            out.push(b);
42            break;
43        } else {
44            out.push(b | 0x80);
45        }
46    }
47}
48
49/// Returns (value, bytes_consumed). Returns None on truncated input.
50fn read_varint(data: &[u8]) -> Option<(u64, usize)> {
51    let mut val: u64 = 0;
52    let mut shift = 0u32;
53    for (i, &b) in data.iter().enumerate() {
54        val |= ((b & 0x7F) as u64) << shift;
55        if b & 0x80 == 0 {
56            return Some((val, i + 1));
57        }
58        shift += 7;
59        if shift >= 70 {
60            return None;
61        }
62    }
63    None
64}
65
66// ---------------------------------------------------------------------------
67// Public API
68// ---------------------------------------------------------------------------
69
70/// Tag byte preceding the dict-values block. Selects the encoding format
71/// for the distinct coordinate values stored by the `DimDict`.
72///
73/// Tag 0 is reserved/forbidden — a stray zero byte is a parse error.
74const DICT_TAG_INT64_FASTLANES: u8 = 1; // homogeneous Int64 dict via fastlanes
75const DICT_TAG_MSGPACK: u8 = 2; // per-value zerompk msgpack (string/mixed CoordValue)
76
77fn try_encode_int64_dict(values: &[CoordValue]) -> Option<Vec<i64>> {
78    let mut out = Vec::with_capacity(values.len());
79    for v in values {
80        match v {
81            CoordValue::Int64(n) => out.push(*n),
82            _ => return None,
83        }
84    }
85    Some(out)
86}
87
88/// Encode one axis of a `DimDict` into a byte vector.
89pub fn encode_coord_axis(dict: &DimDict, out: &mut Vec<u8>) -> ArrayResult<()> {
90    // Dict entries: tag + count + values. Per-axis homogeneous-Int64 dicts
91    // are batch-encoded with fastlanes; mixed/string dicts fall back to
92    // per-value zerompk msgpack.
93    if let Some(ints) = try_encode_int64_dict(&dict.values) {
94        out.push(DICT_TAG_INT64_FASTLANES);
95        out.extend_from_slice(&(ints.len() as u32).to_le_bytes());
96        let encoded = nodedb_codec::fastlanes::encode(&ints).map_err(|error| {
97            ArrayError::SegmentCorruption {
98                detail: format!("coordinate FastLanes encode: {error}"),
99            }
100        })?;
101        out.extend_from_slice(&(encoded.len() as u32).to_le_bytes());
102        out.extend_from_slice(&encoded);
103    } else {
104        out.push(DICT_TAG_MSGPACK);
105        let dict_count = dict.values.len() as u32;
106        out.extend_from_slice(&dict_count.to_le_bytes());
107        for cv in &dict.values {
108            let bytes = zerompk::to_msgpack_vec(cv).map_err(|e| ArrayError::SegmentCorruption {
109                detail: format!("coord dict encode: {e}"),
110            })?;
111            let len = bytes.len() as u32;
112            out.extend_from_slice(&len.to_le_bytes());
113            out.extend_from_slice(&bytes);
114        }
115    }
116
117    // Index stream: count + zigzag-varint deltas.
118    let idx_count = dict.indices.len() as u32;
119    out.extend_from_slice(&idx_count.to_le_bytes());
120    let mut prev: i64 = 0;
121    for &idx in &dict.indices {
122        let cur = idx as i64;
123        let delta = cur - prev;
124        write_varint(out, zigzag_encode(delta));
125        prev = cur;
126    }
127    Ok(())
128}
129
130/// Decode one axis from `data[pos..]`, advancing `pos` past consumed bytes.
131pub fn decode_coord_axis(data: &[u8], pos: &mut usize) -> ArrayResult<DimDict> {
132    // Dict entries — tag-dispatched.
133    if *pos >= data.len() {
134        return Err(ArrayError::SegmentCorruption {
135            detail: "coord axis: truncated dict tag".into(),
136        });
137    }
138    let dict_tag = data[*pos];
139    *pos += 1;
140
141    let dict_values: Vec<CoordValue> =
142        match dict_tag {
143            DICT_TAG_MSGPACK => {
144                if *pos + 4 > data.len() {
145                    return Err(ArrayError::SegmentCorruption {
146                        detail: "coord axis: truncated dict count".into(),
147                    });
148                }
149                let dict_count = u32::from_le_bytes(
150                    data[*pos..*pos + 4]
151                        .try_into()
152                        .expect("invariant: preceding bounds check guarantees 4 bytes available"),
153                ) as usize;
154                *pos += 4;
155                let capacity = checked_capacity(
156                    dict_count,
157                    MAX_DICT_CARDINALITY,
158                    data.len() - *pos,
159                    4,
160                    "coord_delta dict_count",
161                )?;
162
163                let mut out: Vec<CoordValue> = Vec::with_capacity(capacity);
164                for _ in 0..dict_count {
165                    if *pos + 4 > data.len() {
166                        return Err(ArrayError::SegmentCorruption {
167                            detail: "coord axis: truncated dict entry len".into(),
168                        });
169                    }
170                    let len =
171                        u32::from_le_bytes(data[*pos..*pos + 4].try_into().expect(
172                            "invariant: preceding bounds check guarantees 4 bytes available",
173                        )) as usize;
174                    *pos += 4;
175                    if *pos + len > data.len() {
176                        return Err(ArrayError::SegmentCorruption {
177                            detail: "coord axis: truncated dict entry bytes".into(),
178                        });
179                    }
180                    let cv: CoordValue =
181                        zerompk::from_msgpack(&data[*pos..*pos + len]).map_err(|e| {
182                            ArrayError::SegmentCorruption {
183                                detail: format!("coord dict decode: {e}"),
184                            }
185                        })?;
186                    *pos += len;
187                    out.push(cv);
188                }
189                out
190            }
191            DICT_TAG_INT64_FASTLANES => {
192                if *pos + 8 > data.len() {
193                    return Err(ArrayError::SegmentCorruption {
194                        detail: "coord axis: truncated int64-dict header".into(),
195                    });
196                }
197                let dict_count = u32::from_le_bytes(data[*pos..*pos + 4].try_into().expect(
198                    "invariant: preceding bounds check guarantees 8 bytes available; first 4",
199                )) as usize;
200                *pos += 4;
201                check_decoded_size(
202                    dict_count,
203                    MAX_DICT_CARDINALITY,
204                    "coord_delta int64-dict count",
205                )?;
206                let payload_len = u32::from_le_bytes(data[*pos..*pos + 4].try_into().expect(
207                    "invariant: preceding bounds check guarantees 8 bytes available; second 4",
208                )) as usize;
209                *pos += 4;
210                if *pos + payload_len > data.len() {
211                    return Err(ArrayError::SegmentCorruption {
212                        detail: "coord axis: truncated int64-dict payload".into(),
213                    });
214                }
215                let ints = nodedb_codec::fastlanes::decode(&data[*pos..*pos + payload_len])
216                    .map_err(|e| ArrayError::SegmentCorruption {
217                        detail: format!("coord dict fastlanes decode: {e}"),
218                    })?;
219                *pos += payload_len;
220                if ints.len() != dict_count {
221                    return Err(ArrayError::SegmentCorruption {
222                        detail: format!(
223                            "coord dict count mismatch: declared {dict_count}, decoded {}",
224                            ints.len()
225                        ),
226                    });
227                }
228                ints.into_iter().map(CoordValue::Int64).collect()
229            }
230            other => {
231                return Err(ArrayError::SegmentCorruption {
232                    detail: format!("coord axis: unknown dict tag {other:#04x}"),
233                });
234            }
235        };
236
237    // Index stream.
238    if *pos + 4 > data.len() {
239        return Err(ArrayError::SegmentCorruption {
240            detail: "coord axis: truncated index count".into(),
241        });
242    }
243    let idx_count = u32::from_le_bytes(
244        data[*pos..*pos + 4]
245            .try_into()
246            .expect("invariant: preceding bounds check guarantees 4 bytes available"),
247    ) as usize;
248    *pos += 4;
249    let capacity = checked_capacity(
250        idx_count,
251        MAX_CELLS_PER_TILE,
252        data.len() - *pos,
253        1,
254        "coord_delta idx_count",
255    )?;
256
257    let mut indices: Vec<u32> = Vec::with_capacity(capacity);
258    let mut prev: i64 = 0;
259    for _ in 0..idx_count {
260        let (zz, consumed) =
261            read_varint(&data[*pos..]).ok_or_else(|| ArrayError::SegmentCorruption {
262                detail: "coord axis: truncated index varint".into(),
263            })?;
264        *pos += consumed;
265        let delta = zigzag_decode(zz);
266        let cur = prev + delta;
267        indices.push(cur as u32);
268        prev = cur;
269    }
270
271    Ok(DimDict {
272        values: dict_values,
273        indices,
274    })
275}
276
277#[cfg(test)]
278mod tests {
279    use super::*;
280    use crate::types::coord::value::CoordValue;
281
282    fn roundtrip(dict: &DimDict) -> DimDict {
283        let mut buf = Vec::new();
284        encode_coord_axis(dict, &mut buf).unwrap();
285        let mut pos = 0;
286        decode_coord_axis(&buf, &mut pos).unwrap()
287    }
288
289    fn make_dict(coords: &[i64]) -> DimDict {
290        let mut d = DimDict {
291            values: vec![],
292            indices: vec![],
293        };
294        for &c in coords {
295            let cv = CoordValue::Int64(c);
296            if let Some(idx) = d.values.iter().position(|x| x == &cv) {
297                d.indices.push(idx as u32);
298            } else {
299                d.indices.push(d.values.len() as u32);
300                d.values.push(cv);
301            }
302        }
303        d
304    }
305
306    #[test]
307    fn empty_dict_roundtrip() {
308        let d = DimDict {
309            values: vec![],
310            indices: vec![],
311        };
312        let out = roundtrip(&d);
313        assert_eq!(out.values, d.values);
314        assert_eq!(out.indices, d.indices);
315    }
316
317    #[test]
318    fn single_entry_roundtrip() {
319        let d = make_dict(&[42]);
320        let out = roundtrip(&d);
321        assert_eq!(out.values, d.values);
322        assert_eq!(out.indices, d.indices);
323    }
324
325    #[test]
326    fn monotonic_coords_roundtrip() {
327        let coords: Vec<i64> = (0..1000).collect();
328        let d = make_dict(&coords);
329        let out = roundtrip(&d);
330        assert_eq!(out.indices, d.indices);
331        assert_eq!(out.values.len(), d.values.len());
332    }
333
334    #[test]
335    fn non_monotonic_coords_roundtrip() {
336        let coords = vec![100i64, 5, 200, 5, 100, 300];
337        let d = make_dict(&coords);
338        let out = roundtrip(&d);
339        assert_eq!(out.indices, d.indices);
340        assert_eq!(out.values, d.values);
341    }
342
343    #[test]
344    fn repeated_coords_roundtrip() {
345        let coords = vec![7i64; 100];
346        let d = make_dict(&coords);
347        let out = roundtrip(&d);
348        assert_eq!(out.indices, d.indices);
349        assert_eq!(out.values.len(), 1);
350    }
351
352    #[test]
353    fn large_coord_roundtrip() {
354        let coords: Vec<i64> = (0..100_000).step_by(7).collect();
355        let d = make_dict(&coords);
356        let out = roundtrip(&d);
357        assert_eq!(out.indices, d.indices);
358    }
359
360    #[test]
361    fn string_coord_roundtrip() {
362        let mut d = DimDict {
363            values: vec![],
364            indices: vec![],
365        };
366        for s in &["ALT", "REF", "ALT", "DEL"] {
367            let cv = CoordValue::String(s.to_string());
368            if let Some(idx) = d.values.iter().position(|x| x == &cv) {
369                d.indices.push(idx as u32);
370            } else {
371                d.indices.push(d.values.len() as u32);
372                d.values.push(cv);
373            }
374        }
375        let out = roundtrip(&d);
376        assert_eq!(out.indices, d.indices);
377        assert_eq!(out.values, d.values);
378    }
379
380    #[test]
381    fn huge_msgpack_dictionary_count_in_tiny_payload_is_rejected_before_allocation() {
382        let mut data = Vec::new();
383        data.push(DICT_TAG_MSGPACK);
384        data.extend_from_slice(&(MAX_DICT_CARDINALITY as u32).to_le_bytes());
385        let result = decode_coord_axis(&data, &mut 0);
386        assert!(matches!(
387            result,
388            Err(crate::error::ArrayError::SegmentCorruption { .. })
389        ));
390    }
391
392    #[test]
393    fn zero_tag_byte_is_corruption() {
394        // Tag 0 is reserved/forbidden. A payload starting with 0x00 must
395        // be rejected as SegmentCorruption, not silently decoded.
396        let data = [0x00u8; 16];
397        let result = decode_coord_axis(&data, &mut 0);
398        assert!(
399            matches!(
400                result,
401                Err(crate::error::ArrayError::SegmentCorruption { .. })
402            ),
403            "expected SegmentCorruption for tag 0x00, got {result:?}"
404        );
405    }
406}