Skip to main content

hermes_core/structures/fast_field/
codec.rs

1//! Fast-field compression codecs with auto-selection.
2//!
3//! Four codecs are available, and the writer picks the smallest at build time:
4//!
5//! | Codec            | ID | Description                                      |
6//! |------------------|----|--------------------------------------------------|
7//! | Constant         |  0 | No data bytes — all values identical              |
8//! | Bitpacked        |  1 | min-subtract + global bitpack                     |
9//! | Linear           |  2 | Regression line, bitpack residuals                |
10//! | BlockwiseLinear  |  3 | Per-512-block linear, bitpack residuals per block |
11
12use std::io::{self, Write};
13
14use byteorder::{LittleEndian, WriteBytesExt};
15
16use super::{bitpack_read, bitpack_write, bits_needed_u64};
17
18// ── Codec type tag ───────────────────────────────────────────────────────
19
20/// Codec identifier stored in the column data region (first byte).
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22#[repr(u8)]
23pub enum CodecType {
24    Constant = 0,
25    Bitpacked = 1,
26    Linear = 2,
27    BlockwiseLinear = 3,
28}
29
30impl CodecType {
31    pub fn from_u8(v: u8) -> Option<Self> {
32        match v {
33            0 => Some(Self::Constant),
34            1 => Some(Self::Bitpacked),
35            2 => Some(Self::Linear),
36            3 => Some(Self::BlockwiseLinear),
37            _ => None,
38        }
39    }
40}
41
42/// Block size for BlockwiseLinear codec (matching Tantivy).
43pub const BLOCKWISE_LINEAR_BLOCK_SIZE: usize = 512;
44
45/// Validate an auto-codec payload before exposing it through the infallible
46/// hot-path readers below. This keeps every bounds check out of per-document
47/// access while ensuring corrupt segment metadata cannot trigger slice panics.
48pub fn validate_auto(data: &[u8], expected_values: usize) -> io::Result<()> {
49    let (&codec_id, rest) = data.split_first().ok_or_else(|| {
50        io::Error::new(io::ErrorKind::UnexpectedEof, "fast field codec is missing")
51    })?;
52
53    let packed_len = |count: usize, bpv: u8| -> io::Result<usize> {
54        if bpv > 64 {
55            return Err(io::Error::new(
56                io::ErrorKind::InvalidData,
57                "fast field bit width exceeds 64",
58            ));
59        }
60        count
61            .checked_mul(bpv as usize)
62            .and_then(|bits| bits.checked_add(7))
63            .map(|bits| bits / 8)
64            .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "fast field size overflow"))
65    };
66
67    match CodecType::from_u8(codec_id) {
68        Some(CodecType::Constant) => {
69            if rest.len() != 8 {
70                return Err(io::Error::new(
71                    io::ErrorKind::InvalidData,
72                    "invalid constant fast field length",
73                ));
74            }
75        }
76        Some(CodecType::Bitpacked) => {
77            if rest.len() < 9 {
78                return Err(io::Error::new(
79                    io::ErrorKind::UnexpectedEof,
80                    "bitpacked fast field header is truncated",
81                ));
82            }
83            let expected_len = 9usize
84                .checked_add(packed_len(expected_values, rest[8])?)
85                .ok_or_else(|| {
86                    io::Error::new(io::ErrorKind::InvalidData, "fast field size overflow")
87                })?;
88            if rest.len() != expected_len {
89                return Err(io::Error::new(
90                    io::ErrorKind::InvalidData,
91                    "bitpacked fast field length is inconsistent",
92                ));
93            }
94        }
95        Some(CodecType::Linear) => {
96            if rest.len() < 29 {
97                return Err(io::Error::new(
98                    io::ErrorKind::UnexpectedEof,
99                    "linear fast field header is truncated",
100                ));
101            }
102            let count = u32::from_le_bytes(rest[16..20].try_into().unwrap()) as usize;
103            if count != expected_values || count < 2 {
104                return Err(io::Error::new(
105                    io::ErrorKind::InvalidData,
106                    "linear fast field value count is inconsistent",
107                ));
108            }
109            let expected_len = 29usize
110                .checked_add(packed_len(count, rest[28])?)
111                .ok_or_else(|| {
112                    io::Error::new(io::ErrorKind::InvalidData, "fast field size overflow")
113                })?;
114            if rest.len() != expected_len {
115                return Err(io::Error::new(
116                    io::ErrorKind::InvalidData,
117                    "linear fast field length is inconsistent",
118                ));
119            }
120        }
121        Some(CodecType::BlockwiseLinear) => {
122            if rest.len() < 8 {
123                return Err(io::Error::new(
124                    io::ErrorKind::UnexpectedEof,
125                    "blockwise fast field header is truncated",
126                ));
127            }
128            let count = u32::from_le_bytes(rest[0..4].try_into().unwrap()) as usize;
129            let num_blocks = u32::from_le_bytes(rest[4..8].try_into().unwrap()) as usize;
130            if count != expected_values || num_blocks != count.div_ceil(BLOCKWISE_LINEAR_BLOCK_SIZE)
131            {
132                return Err(io::Error::new(
133                    io::ErrorKind::InvalidData,
134                    "blockwise fast field counts are inconsistent",
135                ));
136            }
137
138            let mut pos = 8usize;
139            for block_idx in 0..num_blocks {
140                let header_end = pos.checked_add(29).ok_or_else(|| {
141                    io::Error::new(io::ErrorKind::InvalidData, "fast field offset overflow")
142                })?;
143                if header_end > rest.len() {
144                    return Err(io::Error::new(
145                        io::ErrorKind::UnexpectedEof,
146                        "blockwise fast field block header is truncated",
147                    ));
148                }
149                let bpv = rest[pos + 24];
150                let declared =
151                    u32::from_le_bytes(rest[pos + 25..header_end].try_into().unwrap()) as usize;
152                let block_start = block_idx * BLOCKWISE_LINEAR_BLOCK_SIZE;
153                let block_count = (count - block_start).min(BLOCKWISE_LINEAR_BLOCK_SIZE);
154                let expected = packed_len(block_count, bpv)?;
155                if declared != expected {
156                    return Err(io::Error::new(
157                        io::ErrorKind::InvalidData,
158                        "blockwise fast field packed length is inconsistent",
159                    ));
160                }
161                pos = header_end.checked_add(declared).ok_or_else(|| {
162                    io::Error::new(io::ErrorKind::InvalidData, "fast field offset overflow")
163                })?;
164                if pos > rest.len() {
165                    return Err(io::Error::new(
166                        io::ErrorKind::UnexpectedEof,
167                        "blockwise fast field data is truncated",
168                    ));
169                }
170            }
171            if pos != rest.len() {
172                return Err(io::Error::new(
173                    io::ErrorKind::InvalidData,
174                    "blockwise fast field contains trailing data",
175                ));
176            }
177        }
178        None => {
179            return Err(io::Error::new(
180                io::ErrorKind::InvalidData,
181                "unknown fast field codec",
182            ));
183        }
184    }
185    Ok(())
186}
187
188// ── Estimator trait ──────────────────────────────────────────────────────
189
190/// Estimates serialized size for a given codec.
191///
192/// Usage: call `collect(val)` for every value, then `finalize()`,
193/// then `estimate()` returns the byte count.
194pub trait CodecEstimator {
195    fn collect(&mut self, value: u64);
196    fn finalize(&mut self) {}
197    fn estimate(&self) -> Option<u64>;
198    fn serialize(&self, values: &[u64], writer: &mut dyn Write) -> io::Result<u64>;
199}
200
201// ── Constant codec ───────────────────────────────────────────────────────
202
203/// All values are identical → zero data bytes. Value stored in the codec header.
204#[derive(Default)]
205pub struct ConstantEstimator {
206    first: Option<u64>,
207    all_same: bool,
208}
209
210impl CodecEstimator for ConstantEstimator {
211    fn collect(&mut self, value: u64) {
212        match self.first {
213            None => {
214                self.first = Some(value);
215                self.all_same = true;
216            }
217            Some(f) => {
218                if value != f {
219                    self.all_same = false;
220                }
221            }
222        }
223    }
224
225    fn estimate(&self) -> Option<u64> {
226        if self.all_same {
227            // codec_id(1) + value(8) = 9 bytes
228            Some(9)
229        } else {
230            None
231        }
232    }
233
234    fn serialize(&self, values: &[u64], writer: &mut dyn Write) -> io::Result<u64> {
235        let val = if values.is_empty() { 0 } else { values[0] };
236        writer.write_u8(CodecType::Constant as u8)?;
237        writer.write_u64::<LittleEndian>(val)?;
238        Ok(9)
239    }
240}
241
242// ── Bitpacked codec ──────────────────────────────────────────────────────
243
244/// Min-subtract + global bitpack. This is the existing codec, now behind a tag.
245#[derive(Default)]
246pub struct BitpackedEstimator {
247    min: u64,
248    max: u64,
249    count: usize,
250    initialized: bool,
251}
252
253impl CodecEstimator for BitpackedEstimator {
254    fn collect(&mut self, value: u64) {
255        if !self.initialized {
256            self.min = value;
257            self.max = value;
258            self.initialized = true;
259        } else {
260            self.min = self.min.min(value);
261            self.max = self.max.max(value);
262        }
263        self.count += 1;
264    }
265
266    fn estimate(&self) -> Option<u64> {
267        if self.count == 0 {
268            return Some(0);
269        }
270        let range = self.max - self.min;
271        let bpv = bits_needed_u64(range) as u64;
272        // codec_id(1) + min(8) + bpv(1) + packed data
273        let data_bits = self.count as u64 * bpv;
274        let data_bytes = data_bits.div_ceil(8);
275        Some(1 + 8 + 1 + data_bytes)
276    }
277
278    fn serialize(&self, values: &[u64], writer: &mut dyn Write) -> io::Result<u64> {
279        let (min_value, bpv) = if values.is_empty() {
280            (0u64, 0u8)
281        } else {
282            let min_val = values.iter().copied().min().unwrap();
283            let max_val = values.iter().copied().max().unwrap();
284            (min_val, bits_needed_u64(max_val - min_val))
285        };
286
287        writer.write_u8(CodecType::Bitpacked as u8)?;
288        writer.write_u64::<LittleEndian>(min_value)?;
289        writer.write_u8(bpv)?;
290        let mut bytes_written = 10u64; // 1 + 8 + 1
291
292        if bpv > 0 && !values.is_empty() {
293            let shifted: Vec<u64> = values.iter().map(|&v| v - min_value).collect();
294            let mut packed = Vec::new();
295            bitpack_write(&shifted, bpv, &mut packed);
296            writer.write_all(&packed)?;
297            bytes_written += packed.len() as u64;
298        }
299        Ok(bytes_written)
300    }
301}
302
303/// Read a single value from a bitpacked-codec column.
304///
305/// `data` starts right after the codec_id byte (i.e. at min_value).
306#[inline]
307pub fn bitpacked_read(data: &[u8], index: usize) -> u64 {
308    let min_value = u64::from_le_bytes(data[0..8].try_into().unwrap());
309    let bpv = data[8];
310    if bpv == 0 {
311        return min_value;
312    }
313    let packed = &data[9..];
314    bitpack_read(packed, bpv, index).wrapping_add(min_value)
315}
316
317// ── Linear codec ─────────────────────────────────────────────────────────
318
319/// Fits y = slope * x + intercept across all values, stores residuals bitpacked.
320///
321/// Header: codec_id(1) + intercept(8) + slope_num(8) + slope_den(8) + bpv(1) + offset(8) = 34
322///
323/// Estimation uses O(1) memory by tracking value extremes during collection and
324/// computing worst-case residual bounds in `finalize()`.
325///
326/// **Limitation**: the per-column offset is stored as i64 (8 bytes). When values
327/// span nearly the full u64 range (e.g. `FAST_FIELD_MISSING` mixed with small
328/// values), residuals can exceed i64 bounds.  The estimator returns `None` in
329/// that case so the auto-selector falls back to bitpacked.
330#[derive(Default)]
331pub struct LinearEstimator {
332    count: usize,
333    first: u64,
334    last: u64,
335    min_val: u64,
336    max_val: u64,
337    min_residual: i64,
338    max_residual: i64,
339    values_collected: bool,
340    /// Set by `finalize()` when residuals exceed i64 range.
341    overflow: bool,
342}
343
344impl CodecEstimator for LinearEstimator {
345    fn collect(&mut self, value: u64) {
346        if !self.values_collected {
347            self.first = value;
348            self.min_val = value;
349            self.max_val = value;
350            self.values_collected = true;
351        } else {
352            self.min_val = self.min_val.min(value);
353            self.max_val = self.max_val.max(value);
354        }
355        self.last = value;
356        self.count += 1;
357    }
358
359    fn finalize(&mut self) {
360        if self.count < 2 {
361            return;
362        }
363        // Compute worst-case residual bounds from value extremes vs predicted line.
364        // The predicted line spans [first, last]. The worst-case residuals occur when
365        // the most extreme value is farthest from the nearest predicted value.
366        // Predicted values range from min(first,last) to max(first,last), so:
367        //   max_residual ≥ max_val - min(predicted) = max_val - min(first, last)
368        //   min_residual ≤ min_val - max(predicted) = min_val - max(first, last)
369        // This is a conservative bound (may slightly overestimate bpv vs exact).
370        let pred_min = self.first.min(self.last) as i128;
371        let pred_max = self.first.max(self.last) as i128;
372        let min_res = self.min_val as i128 - pred_max;
373        let max_res = self.max_val as i128 - pred_min;
374        // The offset is stored as i64 on disk.  If residuals exceed i64 range,
375        // this codec cannot represent the data — mark as overflow.
376        if min_res < i64::MIN as i128 || max_res > i64::MAX as i128 {
377            self.overflow = true;
378            return;
379        }
380        self.min_residual = min_res as i64;
381        self.max_residual = max_res as i64;
382    }
383
384    fn estimate(&self) -> Option<u64> {
385        if self.count < 2 || self.overflow {
386            return None;
387        }
388        // Check for overflow: if the range doesn't fit u64, this codec is not viable
389        let range = (self.max_residual as i128 - self.min_residual as i128) as u64;
390        let bpv = bits_needed_u64(range) as u64;
391        let data_bits = self.count as u64 * bpv;
392        let data_bytes = data_bits.div_ceil(8);
393        // codec_id(1) + first(8) + last(8) + num_values(4) + offset(8) + bpv(1) + packed
394        Some(1 + 8 + 8 + 4 + 8 + 1 + data_bytes)
395    }
396
397    fn serialize(&self, values: &[u64], writer: &mut dyn Write) -> io::Result<u64> {
398        let n = values.len();
399        if n < 2 {
400            return Err(io::Error::new(
401                io::ErrorKind::InvalidInput,
402                "linear needs ≥ 2 values",
403            ));
404        }
405        let first = values[0];
406        let last = values[n - 1];
407
408        // Compute residuals using i128 to avoid overflow
409        let mut min_residual = i128::MAX;
410        for (i, &val) in values.iter().enumerate() {
411            let predicted = interpolate(first, last, n, i);
412            let residual = val as i128 - predicted as i128;
413            min_residual = min_residual.min(residual);
414        }
415
416        // The offset field is i64 on disk — reject data that doesn't fit.
417        if min_residual < i64::MIN as i128 || min_residual > i64::MAX as i128 {
418            return Err(io::Error::new(
419                io::ErrorKind::InvalidInput,
420                "linear codec: residual offset exceeds i64 range",
421            ));
422        }
423        let min_residual_i64 = min_residual as i64;
424
425        // Shift residuals to non-negative
426        let shifted: Vec<u64> = values
427            .iter()
428            .enumerate()
429            .map(|(i, &val)| {
430                let predicted = interpolate(first, last, n, i);
431                let residual = val as i128 - predicted as i128;
432                (residual - min_residual) as u64
433            })
434            .collect();
435        let max_shifted = shifted.iter().copied().max().unwrap_or(0);
436        let bpv = bits_needed_u64(max_shifted);
437        writer.write_u8(CodecType::Linear as u8)?;
438        writer.write_u64::<LittleEndian>(first)?;
439        writer.write_u64::<LittleEndian>(last)?;
440        writer.write_u32::<LittleEndian>(n as u32)?;
441        writer.write_i64::<LittleEndian>(min_residual_i64)?;
442        writer.write_u8(bpv)?;
443        let mut bytes_written = 30u64; // 1+8+8+4+8+1
444
445        if bpv > 0 {
446            let mut packed = Vec::new();
447            bitpack_write(&shifted, bpv, &mut packed);
448            writer.write_all(&packed)?;
449            bytes_written += packed.len() as u64;
450        }
451
452        Ok(bytes_written)
453    }
454}
455
456/// Interpolate value at index `i` on the line from first to last over `n` values.
457#[inline]
458fn interpolate(first: u64, last: u64, n: usize, i: usize) -> u64 {
459    if n <= 1 {
460        return first;
461    }
462    // Use i128 to avoid overflow
463    let first = first as i128;
464    let last = last as i128;
465    let n = n as i128;
466    let i = i as i128;
467    let result = first + (last - first) * i / (n - 1);
468    result as u64
469}
470
471/// Read a single value from a linear-codec column.
472///
473/// `data` starts right after the codec_id byte.
474#[inline]
475pub fn linear_read(data: &[u8], index: usize) -> u64 {
476    let first = u64::from_le_bytes(data[0..8].try_into().unwrap());
477    let last = u64::from_le_bytes(data[8..16].try_into().unwrap());
478    let n = u32::from_le_bytes(data[16..20].try_into().unwrap()) as usize;
479    let offset = i64::from_le_bytes(data[20..28].try_into().unwrap());
480    let bpv = data[28];
481    let predicted = interpolate(first, last, n, index);
482    let residual = if bpv == 0 {
483        0u64
484    } else {
485        bitpack_read(&data[29..], bpv, index)
486    };
487    // Use i128 to avoid overflow with large values
488    (predicted as i128 + offset as i128 + residual as i128) as u64
489}
490
491// ── BlockwiseLinear codec ────────────────────────────────────────────────
492
493/// Per-512-element-block linear interpolation with per-block bitpacked residuals.
494///
495/// Header: codec_id(1) + num_values(4) + num_blocks(4)
496/// Per block: first(8) + last(8) + offset(8) + bpv(1) + packed_len(4) + packed_data
497#[derive(Default)]
498pub struct BlockwiseLinearEstimator {
499    values: Vec<u64>,
500}
501
502impl CodecEstimator for BlockwiseLinearEstimator {
503    fn collect(&mut self, value: u64) {
504        self.values.push(value);
505    }
506
507    fn estimate(&self) -> Option<u64> {
508        let n = self.values.len();
509        if n < 2 * BLOCKWISE_LINEAR_BLOCK_SIZE {
510            // Only useful when there are enough values to amortize the per-block headers
511            return None;
512        }
513
514        let num_blocks = n.div_ceil(BLOCKWISE_LINEAR_BLOCK_SIZE);
515        // Global header
516        let mut total = 9u64; // codec_id(1) + num_values(4) + num_blocks(4)
517
518        for b in 0..num_blocks {
519            let start = b * BLOCKWISE_LINEAR_BLOCK_SIZE;
520            let end = (start + BLOCKWISE_LINEAR_BLOCK_SIZE).min(n);
521            let block = &self.values[start..end];
522            let block_len = block.len();
523
524            if block_len < 2 {
525                // Block header + 0 data
526                total += 29; // first(8)+last(8)+offset(8)+bpv(1)+packed_len(4)
527                continue;
528            }
529
530            let first = block[0];
531            let last = block[block_len - 1];
532            let mut min_res = i128::MAX;
533            let mut max_res = i128::MIN;
534            for (i, &val) in block.iter().enumerate() {
535                let pred = interpolate(first, last, block_len, i);
536                let res = val as i128 - pred as i128;
537                min_res = min_res.min(res);
538                max_res = max_res.max(res);
539            }
540            // Per-block offset is stored as i64 — if any block's residuals
541            // exceed i64 range, this codec cannot represent the data.
542            if min_res < i64::MIN as i128 || max_res > i64::MAX as i128 {
543                return None;
544            }
545            let range = (max_res - min_res) as u64;
546            let bpv = bits_needed_u64(range) as u64;
547            let data_bits = block_len as u64 * bpv;
548            let data_bytes = data_bits.div_ceil(8);
549            total += 29 + data_bytes;
550        }
551
552        Some(total)
553    }
554
555    fn serialize(&self, values: &[u64], writer: &mut dyn Write) -> io::Result<u64> {
556        let n = values.len();
557        let num_blocks = n.div_ceil(BLOCKWISE_LINEAR_BLOCK_SIZE);
558
559        writer.write_u8(CodecType::BlockwiseLinear as u8)?;
560        writer.write_u32::<LittleEndian>(n as u32)?;
561        writer.write_u32::<LittleEndian>(num_blocks as u32)?;
562        let mut bytes_written = 9u64;
563
564        for b in 0..num_blocks {
565            let start = b * BLOCKWISE_LINEAR_BLOCK_SIZE;
566            let end = (start + BLOCKWISE_LINEAR_BLOCK_SIZE).min(n);
567            let block = &values[start..end];
568            let block_len = block.len();
569
570            let first = block[0];
571            let last = if block_len > 1 {
572                block[block_len - 1]
573            } else {
574                first
575            };
576
577            // Compute residuals using i128 to avoid overflow
578            let mut min_residual = i128::MAX;
579            if block_len >= 2 {
580                for (i, &val) in block.iter().enumerate() {
581                    let pred = interpolate(first, last, block_len, i);
582                    let res = val as i128 - pred as i128;
583                    min_residual = min_residual.min(res);
584                }
585            } else {
586                min_residual = 0;
587            }
588
589            let shifted: Vec<u64> = block
590                .iter()
591                .enumerate()
592                .map(|(i, &val)| {
593                    if block_len < 2 {
594                        return 0;
595                    }
596                    let pred = interpolate(first, last, block_len, i);
597                    let res = val as i128 - pred as i128;
598                    (res - min_residual) as u64
599                })
600                .collect();
601            let max_shifted = shifted.iter().copied().max().unwrap_or(0);
602            let bpv = bits_needed_u64(max_shifted);
603
604            // Per-block offset is stored as i64 — reject data that doesn't fit.
605            if min_residual < i64::MIN as i128 || min_residual > i64::MAX as i128 {
606                return Err(io::Error::new(
607                    io::ErrorKind::InvalidInput,
608                    "blockwise linear codec: per-block residual offset exceeds i64 range",
609                ));
610            }
611            let min_res_i64 = min_residual as i64;
612            writer.write_u64::<LittleEndian>(first)?;
613            writer.write_u64::<LittleEndian>(last)?;
614            writer.write_i64::<LittleEndian>(min_res_i64)?;
615            writer.write_u8(bpv)?;
616
617            let mut packed = Vec::new();
618            if bpv > 0 {
619                bitpack_write(&shifted, bpv, &mut packed);
620            }
621            writer.write_u32::<LittleEndian>(packed.len() as u32)?;
622            writer.write_all(&packed)?;
623            bytes_written += 29 + packed.len() as u64;
624        }
625
626        Ok(bytes_written)
627    }
628}
629
630/// Read a single value from a blockwise-linear-codec column.
631///
632/// `data` starts right after the codec_id byte.
633pub fn blockwise_linear_read(data: &[u8], index: usize) -> u64 {
634    let _num_values = u32::from_le_bytes(data[0..4].try_into().unwrap()) as usize;
635    let num_blocks = u32::from_le_bytes(data[4..8].try_into().unwrap()) as usize;
636
637    let target_block = index / BLOCKWISE_LINEAR_BLOCK_SIZE;
638    let index_in_block = index % BLOCKWISE_LINEAR_BLOCK_SIZE;
639
640    // Scan block headers to find the right one
641    let mut pos = 8usize;
642    for b in 0..num_blocks {
643        let first = u64::from_le_bytes(data[pos..pos + 8].try_into().unwrap());
644        let last = u64::from_le_bytes(data[pos + 8..pos + 16].try_into().unwrap());
645        let offset = i64::from_le_bytes(data[pos + 16..pos + 24].try_into().unwrap());
646        let bpv = data[pos + 24];
647        let packed_len = u32::from_le_bytes(data[pos + 25..pos + 29].try_into().unwrap()) as usize;
648
649        if b == target_block {
650            let block_start = b * BLOCKWISE_LINEAR_BLOCK_SIZE;
651            let block_end = ((b + 1) * BLOCKWISE_LINEAR_BLOCK_SIZE).min(_num_values);
652            let block_len = block_end - block_start;
653
654            let predicted = interpolate(first, last, block_len, index_in_block);
655            let residual = if bpv == 0 {
656                0u64
657            } else {
658                bitpack_read(&data[pos + 29..], bpv, index_in_block)
659            };
660            return (predicted as i128 + offset as i128 + residual as i128) as u64;
661        }
662
663        pos += 29 + packed_len;
664    }
665
666    0 // Should not reach here
667}
668
669// ── Auto-selection ───────────────────────────────────────────────────────
670
671/// Serialize values using the codec that produces the smallest output.
672///
673/// Returns the number of bytes written.
674pub fn serialize_auto(values: &[u64], writer: &mut dyn Write) -> io::Result<u64> {
675    let mut constant = ConstantEstimator::default();
676    let mut bitpacked = BitpackedEstimator::default();
677    let mut linear = LinearEstimator::default();
678    let mut blockwise = BlockwiseLinearEstimator::default();
679
680    // Pass 1: collect
681    for &v in values {
682        constant.collect(v);
683        bitpacked.collect(v);
684        linear.collect(v);
685        blockwise.collect(v);
686    }
687
688    // Finalize
689    constant.finalize();
690    bitpacked.finalize();
691    linear.finalize();
692    blockwise.finalize();
693
694    // Pick smallest
695    let candidates: Vec<(&dyn CodecEstimator, &str)> = vec![
696        (&constant, "constant"),
697        (&bitpacked, "bitpacked"),
698        (&linear, "linear"),
699        (&blockwise, "blockwise_linear"),
700    ];
701
702    let (best, _name) = candidates
703        .into_iter()
704        .filter_map(|(est, name)| est.estimate().map(|size| (est, name, size)))
705        .min_by_key(|&(_, _, size)| size)
706        .map(|(est, name, _)| (est, name))
707        .unwrap_or((&bitpacked as &dyn CodecEstimator, "bitpacked"));
708
709    best.serialize(values, writer)
710}
711
712/// Batch-read `count` consecutive values starting at `start_index` from bitpacked data.
713///
714/// `data` starts right after the codec_id byte (at min_value).
715/// Uses direct array access for byte-aligned bpv (8, 16, 32, 64) which the
716/// compiler auto-vectorizes to SIMD on aarch64 (NEON) and x86_64 (SSE/AVX).
717/// For arbitrary bpv, uses a tight scalar loop with the u64 fast-path.
718pub fn bitpacked_read_batch(data: &[u8], start_index: usize, out: &mut [u64]) {
719    let min_value = u64::from_le_bytes(data[0..8].try_into().unwrap());
720    let bpv = data[8];
721
722    if bpv == 0 {
723        out.iter_mut().for_each(|v| *v = min_value);
724        return;
725    }
726
727    let packed = &data[9..];
728
729    match bpv {
730        // Byte-aligned fast paths — compiler auto-vectorizes these loops
731        8 => {
732            for (i, v) in out.iter_mut().enumerate() {
733                let idx = start_index + i;
734                *v = (packed[idx] as u64).wrapping_add(min_value);
735            }
736        }
737        16 => {
738            for (i, v) in out.iter_mut().enumerate() {
739                let idx = start_index + i;
740                let byte_off = idx * 2;
741                let raw = u16::from_le_bytes([packed[byte_off], packed[byte_off + 1]]);
742                *v = (raw as u64).wrapping_add(min_value);
743            }
744        }
745        32 => {
746            for (i, v) in out.iter_mut().enumerate() {
747                let idx = start_index + i;
748                let byte_off = idx * 4;
749                let raw = u32::from_le_bytes(packed[byte_off..byte_off + 4].try_into().unwrap());
750                *v = (raw as u64).wrapping_add(min_value);
751            }
752        }
753        64 => {
754            for (i, v) in out.iter_mut().enumerate() {
755                let idx = start_index + i;
756                let byte_off = idx * 8;
757                let raw = u64::from_le_bytes(packed[byte_off..byte_off + 8].try_into().unwrap());
758                *v = raw.wrapping_add(min_value);
759            }
760        }
761        // Arbitrary bpv — tight scalar loop using u64 fast-path read
762        _ => {
763            for (i, v) in out.iter_mut().enumerate() {
764                *v = super::bitpack_read(packed, bpv, start_index + i).wrapping_add(min_value);
765            }
766        }
767    }
768}
769
770/// Batch-read `out.len()` consecutive values starting at `start_index` from auto-codec data.
771///
772/// Dispatches codec type once (vs. per-value in `auto_read`), enabling tight inner
773/// loops that the compiler auto-vectorizes for byte-aligned bitpacked columns.
774pub fn auto_read_batch(data: &[u8], start_index: usize, out: &mut [u64]) {
775    if data.is_empty() || out.is_empty() {
776        out.iter_mut().for_each(|v| *v = 0);
777        return;
778    }
779    let codec_id = data[0];
780    let rest = &data[1..];
781    match CodecType::from_u8(codec_id) {
782        Some(CodecType::Constant) => {
783            let val = u64::from_le_bytes(rest[0..8].try_into().unwrap());
784            out.iter_mut().for_each(|v| *v = val);
785        }
786        Some(CodecType::Bitpacked) => bitpacked_read_batch(rest, start_index, out),
787        Some(CodecType::Linear) => {
788            for (i, v) in out.iter_mut().enumerate() {
789                *v = linear_read(rest, start_index + i);
790            }
791        }
792        Some(CodecType::BlockwiseLinear) => {
793            for (i, v) in out.iter_mut().enumerate() {
794                *v = blockwise_linear_read(rest, start_index + i);
795            }
796        }
797        None => out.iter_mut().for_each(|v| *v = 0),
798    }
799}
800
801/// Read a single value from auto-codec encoded data.
802///
803/// The first byte identifies the codec.
804#[inline]
805pub fn auto_read(data: &[u8], index: usize) -> u64 {
806    if data.is_empty() {
807        return 0;
808    }
809    let codec_id = data[0];
810    let rest = &data[1..];
811    match CodecType::from_u8(codec_id) {
812        Some(CodecType::Constant) => {
813            // rest = value(8)
814            u64::from_le_bytes(rest[0..8].try_into().unwrap())
815        }
816        Some(CodecType::Bitpacked) => bitpacked_read(rest, index),
817        Some(CodecType::Linear) => linear_read(rest, index),
818        Some(CodecType::BlockwiseLinear) => blockwise_linear_read(rest, index),
819        None => 0,
820    }
821}
822
823// ── Tests ────────────────────────────────────────────────────────────────
824
825#[cfg(test)]
826mod tests {
827    use super::*;
828
829    fn roundtrip(values: &[u64]) -> Vec<u64> {
830        let mut buf = Vec::new();
831        serialize_auto(values, &mut buf).unwrap();
832        (0..values.len()).map(|i| auto_read(&buf, i)).collect()
833    }
834
835    #[test]
836    fn test_constant_codec() {
837        let values: Vec<u64> = vec![42; 100];
838        let mut buf = Vec::new();
839        serialize_auto(&values, &mut buf).unwrap();
840        assert_eq!(buf[0], CodecType::Constant as u8);
841        assert_eq!(buf.len(), 9);
842        assert_eq!(roundtrip(&values), values);
843    }
844
845    #[test]
846    fn test_bitpacked_codec() {
847        let values: Vec<u64> = (0..50).map(|i| 1000 + (i % 7) * 13).collect();
848        let result = roundtrip(&values);
849        assert_eq!(result, values);
850    }
851
852    #[test]
853    fn test_linear_codec_sequential() {
854        // Perfectly linear → 0 bpv residuals
855        let values: Vec<u64> = (0..1000).map(|i| 100 + i * 3).collect();
856        let mut buf = Vec::new();
857        serialize_auto(&values, &mut buf).unwrap();
858        // Should pick linear (smaller than bitpacked for sequential data)
859        assert_eq!(roundtrip(&values), values);
860    }
861
862    #[test]
863    fn test_blockwise_linear_codec() {
864        // Two distinct linear segments
865        let mut values: Vec<u64> = Vec::new();
866        for i in 0..1500 {
867            if i < 750 {
868                values.push(100 + i * 2);
869            } else {
870                values.push(5000 + (i - 750) * 5);
871            }
872        }
873        let result = roundtrip(&values);
874        assert_eq!(result, values);
875    }
876
877    #[test]
878    fn test_empty() {
879        let values: Vec<u64> = vec![];
880        let mut buf = Vec::new();
881        serialize_auto(&values, &mut buf).unwrap();
882        assert!(buf.len() <= 10);
883    }
884
885    #[test]
886    fn test_validate_rejects_truncated_and_inconsistent_payloads() {
887        assert!(validate_auto(&[], 1).is_err());
888        assert!(validate_auto(&[CodecType::Constant as u8], 1).is_err());
889
890        let mut bitpacked = vec![CodecType::Bitpacked as u8];
891        bitpacked.extend_from_slice(&0u64.to_le_bytes());
892        bitpacked.push(65);
893        assert!(validate_auto(&bitpacked, 1).is_err());
894
895        let mut valid = Vec::new();
896        serialize_auto(&[1, 2, 3, 4], &mut valid).unwrap();
897        assert!(validate_auto(&valid, 4).is_ok());
898        valid.pop();
899        assert!(validate_auto(&valid, 4).is_err());
900    }
901
902    #[test]
903    fn test_single_value() {
904        let values = vec![999u64];
905        assert_eq!(roundtrip(&values), values);
906    }
907
908    #[test]
909    fn test_two_values() {
910        let values = vec![10u64, 20];
911        assert_eq!(roundtrip(&values), values);
912    }
913
914    #[test]
915    fn test_large_range() {
916        let values = vec![0u64, u64::MAX / 2, u64::MAX];
917        assert_eq!(roundtrip(&values), values);
918    }
919
920    #[test]
921    fn test_timestamps_pick_linear_or_blockwise() {
922        // Simulate timestamps (monotonically increasing with small jitter)
923        let mut values: Vec<u64> = Vec::new();
924        let mut ts = 1_700_000_000u64;
925        for _ in 0..2000 {
926            values.push(ts);
927            ts += 1000 + (ts % 7); // ~1000 with jitter
928        }
929        let result = roundtrip(&values);
930        assert_eq!(result, values);
931    }
932
933    /// Helper: roundtrip via auto_read_batch and compare with per-element auto_read.
934    fn roundtrip_batch(values: &[u64]) {
935        let mut buf = Vec::new();
936        serialize_auto(values, &mut buf).unwrap();
937
938        // Batch read all values
939        let mut batch_out = vec![0u64; values.len()];
940        auto_read_batch(&buf, 0, &mut batch_out);
941        assert_eq!(batch_out, values, "batch read mismatch");
942
943        // Batch read a sub-range
944        if values.len() >= 10 {
945            let start = 3;
946            let count = values.len() - 6;
947            let mut sub = vec![0u64; count];
948            auto_read_batch(&buf, start, &mut sub);
949            assert_eq!(
950                sub,
951                &values[start..start + count],
952                "sub-range batch mismatch"
953            );
954        }
955    }
956
957    #[test]
958    fn test_batch_read_constant() {
959        roundtrip_batch(&vec![42u64; 100]);
960    }
961
962    #[test]
963    fn test_batch_read_bitpacked_8bit() {
964        // Values with range < 256 → 8-bit bpv
965        let values: Vec<u64> = (0..200).map(|i| 1000 + (i % 200)).collect();
966        roundtrip_batch(&values);
967    }
968
969    #[test]
970    fn test_batch_read_bitpacked_16bit() {
971        // Values with range fitting 16 bits
972        let values: Vec<u64> = (0..200).map(|i| 50000 + i * 100).collect();
973        roundtrip_batch(&values);
974    }
975
976    #[test]
977    fn test_batch_read_bitpacked_arbitrary() {
978        // Arbitrary bpv (e.g. 13 bits)
979        let values: Vec<u64> = (0..100).map(|i| 999 + (i * 37) % 8000).collect();
980        roundtrip_batch(&values);
981    }
982
983    #[test]
984    fn test_batch_read_linear() {
985        let values: Vec<u64> = (0..500).map(|i| 100 + i * 3).collect();
986        roundtrip_batch(&values);
987    }
988
989    #[test]
990    fn test_batch_read_blockwise() {
991        let mut values = Vec::new();
992        for i in 0..1500u64 {
993            values.push(if i < 750 {
994                100 + i * 2
995            } else {
996                5000 + (i - 750) * 5
997            });
998        }
999        roundtrip_batch(&values);
1000    }
1001
1002    /// Regression: zigzag-encoded i64 timestamps mixed with FAST_FIELD_MISSING (u64::MAX).
1003    /// The linear codec's min_residual clamping to i64 corrupts data when values
1004    /// span nearly the full u64 range.
1005    #[test]
1006    fn test_zigzag_timestamps_with_missing() {
1007        use super::super::{FAST_FIELD_MISSING, zigzag_encode};
1008
1009        // Simulate issued_at column: most docs have timestamps, some are missing
1010        let timestamps: Vec<i64> = vec![
1011            1724630400, // 2024-08-26
1012            1724716800, // 2024-08-27
1013            1724803200, // 2024-08-28
1014            1700000000, // 2023-11-14
1015            1680000000, // 2023-03-28
1016            1724630400, // duplicate
1017        ];
1018
1019        // Build values array: zigzag-encoded timestamps + some FAST_FIELD_MISSING gaps
1020        let mut values = Vec::new();
1021        for (i, &ts) in timestamps.iter().enumerate() {
1022            values.push(zigzag_encode(ts));
1023            // Insert a missing value after every 2nd doc
1024            if i % 2 == 1 {
1025                values.push(FAST_FIELD_MISSING);
1026            }
1027        }
1028
1029        let result = roundtrip(&values);
1030        assert_eq!(
1031            result, values,
1032            "zigzag timestamps + missing roundtrip failed"
1033        );
1034    }
1035
1036    /// Test each codec individually with zigzag-encoded values + FAST_FIELD_MISSING
1037    #[test]
1038    fn test_codecs_individually_with_zigzag_and_missing() {
1039        use super::super::{FAST_FIELD_MISSING, zigzag_encode};
1040
1041        let values: Vec<u64> = vec![
1042            zigzag_encode(1724630400), // 3449260800
1043            zigzag_encode(1700000000), // 3400000000
1044            FAST_FIELD_MISSING,
1045            zigzag_encode(1680000000), // 3360000000
1046            zigzag_encode(1724716800), // 3449433600
1047            FAST_FIELD_MISSING,
1048            zigzag_encode(1724630400), // 3449260800
1049            zigzag_encode(0),          // 0
1050        ];
1051
1052        // Test bitpacked directly
1053        {
1054            let mut est = BitpackedEstimator::default();
1055            for &v in &values {
1056                est.collect(v);
1057            }
1058            est.finalize();
1059            if est.estimate().is_some() {
1060                let mut buf = Vec::new();
1061                est.serialize(&values, &mut buf).unwrap();
1062                for (i, &expected) in values.iter().enumerate() {
1063                    let got = auto_read(&buf, i);
1064                    assert_eq!(
1065                        got, expected,
1066                        "bitpacked: index {} expected {} got {}",
1067                        i, expected, got
1068                    );
1069                }
1070            }
1071        }
1072
1073        // Test linear directly (needs ≥ 2 values)
1074        {
1075            let mut est = LinearEstimator::default();
1076            for &v in &values {
1077                est.collect(v);
1078            }
1079            est.finalize();
1080            if est.estimate().is_some() {
1081                let mut buf = Vec::new();
1082                est.serialize(&values, &mut buf).unwrap();
1083                for (i, &expected) in values.iter().enumerate() {
1084                    let got = auto_read(&buf, i);
1085                    assert_eq!(
1086                        got, expected,
1087                        "linear: index {} expected {} got {}",
1088                        i, expected, got
1089                    );
1090                }
1091            }
1092        }
1093
1094        // Test auto (whichever is selected)
1095        let result = roundtrip(&values);
1096        assert_eq!(result, values, "auto codec roundtrip failed");
1097    }
1098
1099    /// Regression: value that the user observed corrupted in production
1100    #[test]
1101    fn test_specific_issued_at_roundtrip() {
1102        use super::super::{FAST_FIELD_MISSING, zigzag_encode};
1103
1104        // Reproduce exact scenario: 100 docs, mix of timestamps and missing
1105        let mut values = Vec::new();
1106        let base_ts = 1724630400i64; // 2024-08-26 epoch
1107        for i in 0..100u64 {
1108            if i % 5 == 0 {
1109                // Every 5th doc has no issued_at
1110                values.push(FAST_FIELD_MISSING);
1111            } else {
1112                // Varying timestamps
1113                let ts = base_ts - (i as i64 * 86400); // one day apart
1114                values.push(zigzag_encode(ts));
1115            }
1116        }
1117
1118        let result = roundtrip(&values);
1119        for (i, (&expected, &got)) in values.iter().zip(result.iter()).enumerate() {
1120            assert_eq!(
1121                got,
1122                expected,
1123                "doc {}: expected {} (zigzag of {}), got {}",
1124                i,
1125                expected,
1126                if expected == FAST_FIELD_MISSING {
1127                    -1 // placeholder
1128                } else {
1129                    super::super::zigzag_decode(expected)
1130                },
1131                got
1132            );
1133        }
1134    }
1135
1136    /// Large-scale test: exercise blockwise linear codec with realistic timestamp data.
1137    /// Tests 10K, 50K, 100K docs to catch codec edge cases.
1138    #[test]
1139    fn test_large_scale_timestamp_roundtrip() {
1140        use super::super::{FAST_FIELD_MISSING, zigzag_encode};
1141
1142        for num_docs in [10_000, 50_000, 100_000] {
1143            let mut values = Vec::with_capacity(num_docs);
1144            let base_ts = 1724630400i64;
1145
1146            for i in 0..num_docs {
1147                if i % 7 == 0 {
1148                    values.push(FAST_FIELD_MISSING);
1149                } else {
1150                    // Timestamps spanning ~5 years, with some jitter
1151                    let ts = base_ts - (i as i64 * 3600) + ((i as i64 * 37) % 1000);
1152                    values.push(zigzag_encode(ts));
1153                }
1154            }
1155
1156            // Check which codec is selected
1157            let mut buf = Vec::new();
1158            serialize_auto(&values, &mut buf).unwrap();
1159            let codec_id = buf[0];
1160            let codec_name = match CodecType::from_u8(codec_id) {
1161                Some(CodecType::Constant) => "constant",
1162                Some(CodecType::Bitpacked) => "bitpacked",
1163                Some(CodecType::Linear) => "linear",
1164                Some(CodecType::BlockwiseLinear) => "blockwise_linear",
1165                None => "unknown",
1166            };
1167
1168            // Verify roundtrip
1169            let mut failures = Vec::new();
1170            for (i, &expected) in values.iter().enumerate() {
1171                let got = auto_read(&buf, i);
1172                if got != expected {
1173                    failures.push((i, expected, got));
1174                    if failures.len() >= 5 {
1175                        break;
1176                    }
1177                }
1178            }
1179
1180            assert!(
1181                failures.is_empty(),
1182                "num_docs={}, codec={}: {} failures. First 5: {:?}",
1183                num_docs,
1184                codec_name,
1185                failures.len(),
1186                failures
1187            );
1188        }
1189    }
1190
1191    /// Regression: blockwise linear codec selected for column where most blocks
1192    /// are efficient (sorted timestamps only) but a few blocks contain
1193    /// FAST_FIELD_MISSING, causing min_residual clamping corruption.
1194    #[test]
1195    fn test_blockwise_linear_with_clustered_missing() {
1196        use super::super::{FAST_FIELD_MISSING, zigzag_encode};
1197
1198        // 3000 values: first 512 are all FAST_FIELD_MISSING,
1199        // remaining 2488 are sorted timestamps (efficient linear blocks).
1200        // This should trigger blockwise linear selection overall,
1201        // but the first block has a mix that triggers the bug.
1202        let mut values = Vec::new();
1203
1204        // Block 0 (indices 0-511): mix of MISSING and timestamps
1205        // — first 100 are MISSING, rest are timestamps
1206        for i in 0..512 {
1207            if i < 100 {
1208                values.push(FAST_FIELD_MISSING);
1209            } else {
1210                let ts = 1724630400i64 + (i as i64 * 100);
1211                values.push(zigzag_encode(ts));
1212            }
1213        }
1214
1215        // Blocks 1-5 (indices 512-3071): sorted timestamps only
1216        for i in 512..3072 {
1217            let ts = 1724630400i64 + (i as i64 * 100);
1218            values.push(zigzag_encode(ts));
1219        }
1220
1221        let result = roundtrip(&values);
1222        let mut failures = Vec::new();
1223        for (i, (&expected, &got)) in values.iter().zip(result.iter()).enumerate() {
1224            if got != expected {
1225                failures.push((i, expected, got));
1226            }
1227        }
1228        assert!(
1229            failures.is_empty(),
1230            "blockwise linear with clustered missing: {} failures. First 5: {:?}",
1231            failures.len(),
1232            &failures[..failures.len().min(5)]
1233        );
1234    }
1235
1236    /// Test each codec FORCED with zigzag timestamps + FAST_FIELD_MISSING.
1237    /// This catches bugs that only manifest when a specific codec is forced.
1238    #[test]
1239    fn test_forced_codecs_with_timestamps_and_missing() {
1240        use super::super::{FAST_FIELD_MISSING, zigzag_encode};
1241
1242        let mut values = Vec::new();
1243        let base_ts = 1724630400i64;
1244        for i in 0..200 {
1245            if i % 5 == 0 {
1246                values.push(FAST_FIELD_MISSING);
1247            } else {
1248                let ts = base_ts - (i as i64 * 86400);
1249                values.push(zigzag_encode(ts));
1250            }
1251        }
1252
1253        // Force bitpacked
1254        {
1255            let est = BitpackedEstimator::default();
1256            let mut buf = Vec::new();
1257            est.serialize(&values, &mut buf).unwrap();
1258            for (i, &expected) in values.iter().enumerate() {
1259                let got = bitpacked_read(&buf[1..], i); // skip codec_id byte
1260                assert_eq!(got, expected, "forced bitpacked: index {} failed", i);
1261            }
1262        }
1263
1264        // Force linear — should error because FAST_FIELD_MISSING + timestamps
1265        // produce residuals exceeding i64 range
1266        {
1267            let est = LinearEstimator::default();
1268            let mut buf = Vec::new();
1269            let result = est.serialize(&values, &mut buf);
1270            assert!(
1271                result.is_err(),
1272                "linear codec should reject data with residuals exceeding i64"
1273            );
1274        }
1275
1276        // Force linear with values that DO fit in i64 (no FAST_FIELD_MISSING)
1277        {
1278            let safe_values: Vec<u64> = values
1279                .iter()
1280                .filter(|&&v| v != FAST_FIELD_MISSING)
1281                .copied()
1282                .collect();
1283            let est = LinearEstimator::default();
1284            let mut buf = Vec::new();
1285            est.serialize(&safe_values, &mut buf).unwrap();
1286            for (i, &expected) in safe_values.iter().enumerate() {
1287                let got = linear_read(&buf[1..], i);
1288                assert_eq!(got, expected, "forced linear (safe): index {} failed", i);
1289            }
1290        }
1291
1292        // Blockwise linear estimator should return None for data with FAST_FIELD_MISSING
1293        {
1294            let mut large_values = Vec::new();
1295            for i in 0..2000 {
1296                if i % 5 == 0 {
1297                    large_values.push(FAST_FIELD_MISSING);
1298                } else {
1299                    let ts = base_ts - (i as i64 * 86400);
1300                    large_values.push(zigzag_encode(ts));
1301                }
1302            }
1303            let mut est = BlockwiseLinearEstimator::default();
1304            for &v in &large_values {
1305                est.collect(v);
1306            }
1307            assert!(
1308                est.estimate().is_none(),
1309                "blockwise linear should reject data with per-block residuals exceeding i64"
1310            );
1311        }
1312    }
1313}