1use 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
24fn 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
49fn 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
66const DICT_TAG_INT64_FASTLANES: u8 = 1; const DICT_TAG_MSGPACK: u8 = 2; fn 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
88pub fn encode_coord_axis(dict: &DimDict, out: &mut Vec<u8>) -> ArrayResult<()> {
90 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 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
130pub fn decode_coord_axis(data: &[u8], pos: &mut usize) -> ArrayResult<DimDict> {
132 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 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 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}