Skip to main content

ipfrs_semantic/
embedding_compression_codec.rs

1//! # Embedding Compression Codec
2//!
3//! Production-quality codec for compressing dense embedding vectors using multiple
4//! pure-Rust algorithms: scalar quantization, product quantization, delta coding,
5//! run-length encoding, and hybrid PQ+RLE.
6//!
7//! ## Supported Methods
8//!
9//! - `ScalarQuantization` — uniform min-max quantization to N bits
10//! - `ProductQuantization` — split into subvectors, quantize each independently
11//! - `DeltaCoding` — delta-encode sorted values, then scalar quantize
12//! - `RunLengthEncoding` — RLE on quantized values
13//! - `HybridPQ` — product quantization + RLE on residuals
14//!
15//! ## Example
16//!
17//! ```rust
18//! use ipfrs_semantic::{EmbeddingCompressionCodec, EccMethod, EccCodecConfig};
19//!
20//! let mut codec = EmbeddingCompressionCodec::new();
21//! let id = codec.register_codec("my-sq8", EccMethod::ScalarQuantization, 8, 64);
22//! let embedding = vec![0.1f64; 128];
23//! let compressed = codec.compress(id, &embedding).unwrap();
24//! let decompressed = codec.decompress(&compressed).unwrap();
25//! let mse = codec.reconstruction_error(&embedding, &decompressed);
26//! assert!(mse < 1e-3);
27//! ```
28
29use std::collections::{HashMap, VecDeque};
30use std::time::{SystemTime, UNIX_EPOCH};
31
32use thiserror::Error;
33
34// ---------------------------------------------------------------------------
35// PRNG — xorshift64 (no external rand dependency for internal use)
36// ---------------------------------------------------------------------------
37
38#[allow(dead_code)]
39#[inline]
40fn xorshift64(state: &mut u64) -> u64 {
41    let mut x = *state;
42    x ^= x << 13;
43    x ^= x >> 7;
44    x ^= x << 17;
45    *state = x;
46    x
47}
48
49// ---------------------------------------------------------------------------
50// Public type aliases (as required by spec)
51// ---------------------------------------------------------------------------
52
53/// Type alias for the codec itself.
54pub type EccEmbeddingCompressionCodec = EmbeddingCompressionCodec;
55
56/// Numeric identifier for a registered codec.
57pub type EccCodecId = u32;
58
59// ---------------------------------------------------------------------------
60// Error type
61// ---------------------------------------------------------------------------
62
63/// Errors produced by [`EmbeddingCompressionCodec`].
64#[derive(Debug, Error, Clone, PartialEq)]
65pub enum EccError {
66    /// Requested codec id was not found in the registry.
67    #[error("codec id {0} not found in registry")]
68    CodecNotFound(EccCodecId),
69
70    /// Embedding has zero length.
71    #[error("embedding must not be empty")]
72    EmptyEmbedding,
73
74    /// Compressed payload is corrupt or truncated.
75    #[error("compressed data is corrupt: {0}")]
76    CorruptData(String),
77
78    /// Dimension mismatch between original and decompressed.
79    #[error("dimension mismatch: expected {expected}, got {got}")]
80    DimensionMismatch { expected: usize, got: usize },
81
82    /// Block size is incompatible with the embedding dimension.
83    #[error("block size {block_size} does not divide embedding dim {dim} evenly")]
84    BlockSizeMismatch { block_size: usize, dim: usize },
85
86    /// Quantization bit width is not supported (only 4, 8, 16).
87    #[error("unsupported quantize_bits {0}: must be 4, 8, or 16")]
88    UnsupportedBitWidth(u8),
89
90    /// General internal error.
91    #[error("internal error: {0}")]
92    Internal(String),
93}
94
95// ---------------------------------------------------------------------------
96// EccMethod
97// ---------------------------------------------------------------------------
98
99/// Compression algorithm used by a codec.
100#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
101pub enum EccMethod {
102    /// Uniform min-max scalar quantization to N bits.
103    ScalarQuantization,
104    /// Split vector into sub-blocks, scalar-quantize each independently.
105    ProductQuantization,
106    /// Delta-encode values sorted by magnitude, then scalar-quantize.
107    DeltaCoding,
108    /// Run-length encoding applied on quantized symbols.
109    RunLengthEncoding,
110    /// Product quantization followed by RLE on residual symbols.
111    HybridPQ,
112}
113
114impl std::fmt::Display for EccMethod {
115    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
116        match self {
117            EccMethod::ScalarQuantization => write!(f, "ScalarQuantization"),
118            EccMethod::ProductQuantization => write!(f, "ProductQuantization"),
119            EccMethod::DeltaCoding => write!(f, "DeltaCoding"),
120            EccMethod::RunLengthEncoding => write!(f, "RunLengthEncoding"),
121            EccMethod::HybridPQ => write!(f, "HybridPQ"),
122        }
123    }
124}
125
126// ---------------------------------------------------------------------------
127// Configuration structs
128// ---------------------------------------------------------------------------
129
130/// Configuration for creating a new codec.
131#[derive(Debug, Clone)]
132pub struct EccCodecConfig {
133    /// Name for display / lookup.
134    pub name: String,
135    /// Compression method.
136    pub method: EccMethod,
137    /// Number of bits per quantized value (4, 8, or 16).
138    pub quantize_bits: u8,
139    /// Whether to apply delta coding as a pre-processing step.
140    pub use_delta_coding: bool,
141    /// Sub-block size for PQ / HybridPQ.
142    pub block_size: usize,
143}
144
145impl Default for EccCodecConfig {
146    fn default() -> Self {
147        Self {
148            name: "default-sq8".to_string(),
149            method: EccMethod::ScalarQuantization,
150            quantize_bits: 8,
151            use_delta_coding: false,
152            block_size: 8,
153        }
154    }
155}
156
157/// Registered codec specification stored in the codec registry.
158#[derive(Debug, Clone)]
159pub struct EccCodecSpec {
160    /// Unique numeric id.
161    pub id: EccCodecId,
162    /// Human-readable name.
163    pub name: String,
164    /// Compression method.
165    pub method: EccMethod,
166    /// Quantization bit-width.
167    pub bits: u8,
168    /// Sub-block size.
169    pub block_size: usize,
170}
171
172// ---------------------------------------------------------------------------
173// Compressed payload
174// ---------------------------------------------------------------------------
175
176/// Output of a compression operation.
177#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
178pub struct EccCompressed {
179    /// Id of the codec that produced this payload.
180    pub codec_id: EccCodecId,
181    /// Method used.
182    pub method: EccMethod,
183    /// Raw compressed bytes.
184    pub data: Vec<u8>,
185    /// Original vector length.
186    pub original_dim: usize,
187    /// Minimum value observed in the original vector (for de-quantization).
188    pub min_val: f64,
189    /// Maximum value observed in the original vector (for de-quantization).
190    pub max_val: f64,
191}
192
193// ---------------------------------------------------------------------------
194// Audit / statistics
195// ---------------------------------------------------------------------------
196
197/// One entry in the codec's compression audit log.
198#[derive(Debug, Clone)]
199pub struct EccCompressionRecord {
200    /// Unix timestamp (seconds) when the compression occurred.
201    pub ts: u64,
202    /// Size of the original vector in bytes (f64 × dim × 8).
203    pub original_bytes: usize,
204    /// Size of the compressed payload in bytes.
205    pub compressed_bytes: usize,
206    /// Compression ratio: original / compressed.
207    pub ratio: f64,
208    /// Method used.
209    pub method: EccMethod,
210}
211
212/// Aggregate statistics returned by [`EmbeddingCompressionCodec::codec_stats`].
213#[derive(Debug, Clone)]
214pub struct EccCodecStats {
215    /// Total number of compression operations logged.
216    pub total_compressed: usize,
217    /// Total bytes saved across all operations.
218    pub total_bytes_saved: usize,
219    /// Average compression ratio.
220    pub avg_ratio: f64,
221    /// Per-method breakdown: (total_ops, total_bytes_saved, avg_ratio).
222    pub per_method: HashMap<EccMethod, (usize, usize, f64)>,
223}
224
225// ---------------------------------------------------------------------------
226// Core codec struct
227// ---------------------------------------------------------------------------
228
229/// Maximum number of records kept in the compression log.
230const MAX_LOG_ENTRIES: usize = 500;
231
232/// Codec for compressing dense embedding vectors.
233///
234/// Use [`register_codec`][Self::register_codec] to create named codec configurations,
235/// then call [`compress`][Self::compress] / [`decompress`][Self::decompress] as needed.
236pub struct EmbeddingCompressionCodec {
237    /// Registry mapping id → spec.
238    registry: HashMap<EccCodecId, EccCodecSpec>,
239    /// Monotonically increasing id counter.
240    next_id: EccCodecId,
241    /// Bounded audit log (FIFO, max 500 entries).
242    log: VecDeque<EccCompressionRecord>,
243}
244
245impl Default for EmbeddingCompressionCodec {
246    fn default() -> Self {
247        Self::new()
248    }
249}
250
251impl EmbeddingCompressionCodec {
252    /// Create a new codec with an empty registry.
253    pub fn new() -> Self {
254        Self {
255            registry: HashMap::new(),
256            next_id: 1,
257            log: VecDeque::with_capacity(MAX_LOG_ENTRIES),
258        }
259    }
260
261    // -----------------------------------------------------------------------
262    // Registry
263    // -----------------------------------------------------------------------
264
265    /// Register a new codec and return its id.
266    ///
267    /// `bits` must be 4, 8, or 16. `block_size` is used by PQ / HybridPQ.
268    pub fn register_codec(
269        &mut self,
270        name: &str,
271        method: EccMethod,
272        bits: u8,
273        block_size: usize,
274    ) -> EccCodecId {
275        let id = self.next_id;
276        self.next_id = self.next_id.wrapping_add(1);
277        let spec = EccCodecSpec {
278            id,
279            name: name.to_string(),
280            method,
281            bits,
282            block_size,
283        };
284        self.registry.insert(id, spec);
285        id
286    }
287
288    /// Register a codec from a full [`EccCodecConfig`].
289    pub fn register_from_config(&mut self, config: &EccCodecConfig) -> EccCodecId {
290        self.register_codec(
291            &config.name,
292            config.method,
293            config.quantize_bits,
294            config.block_size,
295        )
296    }
297
298    /// Look up a codec spec by id.
299    pub fn get_spec(&self, id: EccCodecId) -> Option<&EccCodecSpec> {
300        self.registry.get(&id)
301    }
302
303    /// Return the number of registered codecs.
304    pub fn codec_count(&self) -> usize {
305        self.registry.len()
306    }
307
308    // -----------------------------------------------------------------------
309    // Compress / decompress
310    // -----------------------------------------------------------------------
311
312    /// Compress a single embedding vector using the specified codec.
313    pub fn compress(
314        &mut self,
315        codec_id: EccCodecId,
316        embedding: &[f64],
317    ) -> Result<EccCompressed, EccError> {
318        if embedding.is_empty() {
319            return Err(EccError::EmptyEmbedding);
320        }
321        let spec = self
322            .registry
323            .get(&codec_id)
324            .ok_or(EccError::CodecNotFound(codec_id))?
325            .clone();
326
327        validate_bits(spec.bits)?;
328
329        let min_val = embedding.iter().cloned().fold(f64::INFINITY, f64::min);
330        let max_val = embedding.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
331
332        let data = match spec.method {
333            EccMethod::ScalarQuantization => {
334                compress_scalar(embedding, spec.bits, min_val, max_val)?
335            }
336            EccMethod::ProductQuantization => {
337                compress_pq(embedding, spec.bits, spec.block_size, min_val, max_val)?
338            }
339            EccMethod::DeltaCoding => compress_delta(embedding, spec.bits, min_val, max_val)?,
340            EccMethod::RunLengthEncoding => compress_rle(embedding, spec.bits, min_val, max_val)?,
341            EccMethod::HybridPQ => {
342                compress_hybrid_pq(embedding, spec.bits, spec.block_size, min_val, max_val)?
343            }
344        };
345
346        let original_bytes = embedding.len() * 8;
347        let compressed_bytes = data.len();
348        let ratio = if compressed_bytes == 0 {
349            1.0
350        } else {
351            original_bytes as f64 / compressed_bytes as f64
352        };
353
354        self.push_log(EccCompressionRecord {
355            ts: unix_ts(),
356            original_bytes,
357            compressed_bytes,
358            ratio,
359            method: spec.method,
360        });
361
362        Ok(EccCompressed {
363            codec_id,
364            method: spec.method,
365            data,
366            original_dim: embedding.len(),
367            min_val,
368            max_val,
369        })
370    }
371
372    /// Decompress a previously compressed payload.
373    pub fn decompress(&self, compressed: &EccCompressed) -> Result<Vec<f64>, EccError> {
374        let spec = self
375            .registry
376            .get(&compressed.codec_id)
377            .ok_or(EccError::CodecNotFound(compressed.codec_id))?;
378
379        validate_bits(spec.bits)?;
380
381        match compressed.method {
382            EccMethod::ScalarQuantization => decompress_scalar(
383                &compressed.data,
384                spec.bits,
385                compressed.original_dim,
386                compressed.min_val,
387                compressed.max_val,
388            ),
389            EccMethod::ProductQuantization => decompress_pq(
390                &compressed.data,
391                spec.bits,
392                spec.block_size,
393                compressed.original_dim,
394                compressed.min_val,
395                compressed.max_val,
396            ),
397            EccMethod::DeltaCoding => decompress_delta(
398                &compressed.data,
399                spec.bits,
400                compressed.original_dim,
401                compressed.min_val,
402                compressed.max_val,
403            ),
404            EccMethod::RunLengthEncoding => decompress_rle(
405                &compressed.data,
406                spec.bits,
407                compressed.original_dim,
408                compressed.min_val,
409                compressed.max_val,
410            ),
411            EccMethod::HybridPQ => decompress_hybrid_pq(
412                &compressed.data,
413                spec.bits,
414                spec.block_size,
415                compressed.original_dim,
416                compressed.min_val,
417                compressed.max_val,
418            ),
419        }
420    }
421
422    /// Compress a batch of embeddings with the same codec.
423    pub fn compress_batch(
424        &mut self,
425        codec_id: EccCodecId,
426        embeddings: &[Vec<f64>],
427    ) -> Vec<Result<EccCompressed, EccError>> {
428        embeddings
429            .iter()
430            .map(|emb| self.compress(codec_id, emb))
431            .collect()
432    }
433
434    /// Decompress a batch of compressed payloads.
435    pub fn decompress_batch(&self, batch: &[EccCompressed]) -> Vec<Result<Vec<f64>, EccError>> {
436        batch.iter().map(|c| self.decompress(c)).collect()
437    }
438
439    // -----------------------------------------------------------------------
440    // Metrics / estimation
441    // -----------------------------------------------------------------------
442
443    /// Compute mean squared error between original and decompressed vectors.
444    pub fn reconstruction_error(original: &[f64], decompressed: &[f64]) -> f64 {
445        if original.is_empty() || decompressed.is_empty() {
446            return 0.0;
447        }
448        let len = original.len().min(decompressed.len());
449        let sum: f64 = original[..len]
450            .iter()
451            .zip(decompressed[..len].iter())
452            .map(|(a, b)| {
453                let d = a - b;
454                d * d
455            })
456            .sum();
457        sum / len as f64
458    }
459
460    /// Estimate the theoretical compression ratio for a given codec and vector dimension.
461    ///
462    /// The estimate is based on the quantization bit-width and method overhead.
463    pub fn estimate_ratio(&self, codec_id: EccCodecId, dim: usize) -> f64 {
464        let spec = match self.registry.get(&codec_id) {
465            Some(s) => s,
466            None => return 1.0,
467        };
468        if dim == 0 {
469            return 1.0;
470        }
471        let original_bits = dim * 64;
472        let quantized_bits: usize = match spec.method {
473            EccMethod::ScalarQuantization => dim * spec.bits as usize,
474            EccMethod::ProductQuantization => {
475                let blocks = dim.div_ceil(spec.block_size);
476                blocks * spec.block_size * spec.bits as usize
477            }
478            EccMethod::DeltaCoding => {
479                // delta values tend to cluster near zero — assume 70 % of bits needed
480                (dim * spec.bits as usize * 7) / 10
481            }
482            EccMethod::RunLengthEncoding => {
483                // RLE overhead: 2 bytes per run; assume 50 % unique symbols
484                let runs = (dim / 2).max(1);
485                runs * (8 + spec.bits as usize)
486            }
487            EccMethod::HybridPQ => {
488                let blocks = dim.div_ceil(spec.block_size);
489                let pq_bits = blocks * spec.block_size * spec.bits as usize;
490                // RLE on residuals reduces by ~30 %
491                (pq_bits * 7) / 10
492            }
493        };
494        // header overhead (min_val + max_val + original_dim) ≈ 3 × 8 bytes = 192 bits
495        let total_compressed_bits = quantized_bits + 192;
496        original_bits as f64 / total_compressed_bits as f64
497    }
498
499    /// Compute aggregate statistics over the compression log.
500    pub fn codec_stats(&self) -> EccCodecStats {
501        let total_compressed = self.log.len();
502        let mut total_bytes_saved: usize = 0;
503        let mut ratio_sum: f64 = 0.0;
504        let mut per_method: HashMap<EccMethod, (usize, usize, f64)> = HashMap::new();
505
506        for rec in &self.log {
507            let saved = rec.original_bytes.saturating_sub(rec.compressed_bytes);
508            total_bytes_saved += saved;
509            ratio_sum += rec.ratio;
510
511            let entry = per_method.entry(rec.method).or_insert((0, 0, 0.0));
512            entry.0 += 1;
513            entry.1 += saved;
514            entry.2 += rec.ratio;
515        }
516
517        // Normalize per-method average ratios
518        for entry in per_method.values_mut() {
519            if entry.0 > 0 {
520                entry.2 /= entry.0 as f64;
521            }
522        }
523
524        let avg_ratio = if total_compressed > 0 {
525            ratio_sum / total_compressed as f64
526        } else {
527            1.0
528        };
529
530        EccCodecStats {
531            total_compressed,
532            total_bytes_saved,
533            avg_ratio,
534            per_method,
535        }
536    }
537
538    /// Return a read-only view of the compression log.
539    pub fn log_entries(&self) -> &VecDeque<EccCompressionRecord> {
540        &self.log
541    }
542
543    // -----------------------------------------------------------------------
544    // Private helpers
545    // -----------------------------------------------------------------------
546
547    fn push_log(&mut self, record: EccCompressionRecord) {
548        if self.log.len() >= MAX_LOG_ENTRIES {
549            self.log.pop_front();
550        }
551        self.log.push_back(record);
552    }
553}
554
555// ---------------------------------------------------------------------------
556// Internal helpers
557// ---------------------------------------------------------------------------
558
559fn unix_ts() -> u64 {
560    SystemTime::now()
561        .duration_since(UNIX_EPOCH)
562        .map(|d| d.as_secs())
563        .unwrap_or(0)
564}
565
566fn validate_bits(bits: u8) -> Result<(), EccError> {
567    match bits {
568        4 | 8 | 16 => Ok(()),
569        _ => Err(EccError::UnsupportedBitWidth(bits)),
570    }
571}
572
573/// Maximum representable unsigned integer for a given bit-width.
574#[inline]
575fn max_quant_val(bits: u8) -> u64 {
576    match bits {
577        4 => 15,
578        8 => 255,
579        16 => 65535,
580        _ => 255,
581    }
582}
583
584/// Encode a single f64 value in [min_val, max_val] to a quantized integer.
585#[inline]
586fn quantize_val(v: f64, min_val: f64, max_val: f64, levels: u64) -> u64 {
587    let range = max_val - min_val;
588    if range < f64::EPSILON {
589        return 0;
590    }
591    let norm = (v - min_val) / range;
592    let q = (norm * levels as f64).round() as i64;
593    q.clamp(0, levels as i64) as u64
594}
595
596/// Decode a quantized integer back to f64.
597#[inline]
598fn dequantize_val(q: u64, min_val: f64, max_val: f64, levels: u64) -> f64 {
599    if levels == 0 {
600        return min_val;
601    }
602    let norm = q as f64 / levels as f64;
603    min_val + norm * (max_val - min_val)
604}
605
606// ---------------------------------------------------------------------------
607// Scalar Quantization
608// ---------------------------------------------------------------------------
609
610fn compress_scalar(
611    embedding: &[f64],
612    bits: u8,
613    min_val: f64,
614    max_val: f64,
615) -> Result<Vec<u8>, EccError> {
616    let levels = max_quant_val(bits);
617    match bits {
618        4 => {
619            // Pack two 4-bit values per byte
620            let mut out = Vec::with_capacity(embedding.len().div_ceil(2));
621            let mut iter = embedding.iter();
622            loop {
623                match iter.next() {
624                    None => break,
625                    Some(a) => {
626                        let qa = quantize_val(*a, min_val, max_val, levels) as u8;
627                        let qb = match iter.next() {
628                            Some(b) => quantize_val(*b, min_val, max_val, levels) as u8,
629                            None => 0,
630                        };
631                        out.push((qa & 0x0F) | ((qb & 0x0F) << 4));
632                    }
633                }
634            }
635            Ok(out)
636        }
637        8 => {
638            let out: Vec<u8> = embedding
639                .iter()
640                .map(|&v| quantize_val(v, min_val, max_val, levels) as u8)
641                .collect();
642            Ok(out)
643        }
644        16 => {
645            let mut out = Vec::with_capacity(embedding.len() * 2);
646            for &v in embedding {
647                let q = quantize_val(v, min_val, max_val, levels) as u16;
648                out.extend_from_slice(&q.to_le_bytes());
649            }
650            Ok(out)
651        }
652        _ => Err(EccError::UnsupportedBitWidth(bits)),
653    }
654}
655
656fn decompress_scalar(
657    data: &[u8],
658    bits: u8,
659    original_dim: usize,
660    min_val: f64,
661    max_val: f64,
662) -> Result<Vec<f64>, EccError> {
663    let levels = max_quant_val(bits);
664    match bits {
665        4 => {
666            let mut out = Vec::with_capacity(original_dim);
667            for &byte in data {
668                let a = (byte & 0x0F) as u64;
669                let b = ((byte >> 4) & 0x0F) as u64;
670                out.push(dequantize_val(a, min_val, max_val, levels));
671                if out.len() < original_dim {
672                    out.push(dequantize_val(b, min_val, max_val, levels));
673                }
674            }
675            out.truncate(original_dim);
676            if out.len() != original_dim {
677                return Err(EccError::DimensionMismatch {
678                    expected: original_dim,
679                    got: out.len(),
680                });
681            }
682            Ok(out)
683        }
684        8 => {
685            if data.len() != original_dim {
686                return Err(EccError::CorruptData(format!(
687                    "expected {} bytes, got {}",
688                    original_dim,
689                    data.len()
690                )));
691            }
692            Ok(data
693                .iter()
694                .map(|&b| dequantize_val(b as u64, min_val, max_val, levels))
695                .collect())
696        }
697        16 => {
698            if data.len() != original_dim * 2 {
699                return Err(EccError::CorruptData(format!(
700                    "expected {} bytes, got {}",
701                    original_dim * 2,
702                    data.len()
703                )));
704            }
705            let mut out = Vec::with_capacity(original_dim);
706            for chunk in data.chunks_exact(2) {
707                let q = u16::from_le_bytes([chunk[0], chunk[1]]) as u64;
708                out.push(dequantize_val(q, min_val, max_val, levels));
709            }
710            Ok(out)
711        }
712        _ => Err(EccError::UnsupportedBitWidth(bits)),
713    }
714}
715
716// ---------------------------------------------------------------------------
717// Product Quantization
718// ---------------------------------------------------------------------------
719//
720// Each sub-block is quantized independently using its own local min/max.
721// Layout: [ num_blocks:u32 | block_size:u32 | block_0_min:f64 | block_0_max:f64 |
722//           block_0_quantized_bytes... | block_1_min:f64 | block_1_max:f64 | ... ]
723
724fn compress_pq(
725    embedding: &[f64],
726    bits: u8,
727    block_size: usize,
728    _global_min: f64,
729    _global_max: f64,
730) -> Result<Vec<u8>, EccError> {
731    let bs = block_size.max(1);
732    let num_blocks = embedding.len().div_ceil(bs);
733    let mut out: Vec<u8> = Vec::new();
734
735    // Header
736    out.extend_from_slice(&(num_blocks as u32).to_le_bytes());
737    out.extend_from_slice(&(bs as u32).to_le_bytes());
738
739    for block_idx in 0..num_blocks {
740        let start = block_idx * bs;
741        let end = (start + bs).min(embedding.len());
742        let block = &embedding[start..end];
743
744        let bmin = block.iter().cloned().fold(f64::INFINITY, f64::min);
745        let bmax = block.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
746
747        out.extend_from_slice(&bmin.to_le_bytes());
748        out.extend_from_slice(&bmax.to_le_bytes());
749
750        let qbytes = compress_scalar(block, bits, bmin, bmax)?;
751        // Store length prefix for variable-length last block
752        out.extend_from_slice(&(qbytes.len() as u32).to_le_bytes());
753        out.extend_from_slice(&qbytes);
754    }
755    Ok(out)
756}
757
758fn decompress_pq(
759    data: &[u8],
760    bits: u8,
761    block_size: usize,
762    original_dim: usize,
763    _global_min: f64,
764    _global_max: f64,
765) -> Result<Vec<f64>, EccError> {
766    let bs = block_size.max(1);
767    if data.len() < 8 {
768        return Err(EccError::CorruptData("PQ header too short".to_string()));
769    }
770
771    let num_blocks = u32::from_le_bytes(
772        data[0..4]
773            .try_into()
774            .map_err(|_| EccError::CorruptData("num_blocks".to_string()))?,
775    ) as usize;
776
777    let stored_bs = u32::from_le_bytes(
778        data[4..8]
779            .try_into()
780            .map_err(|_| EccError::CorruptData("block_size".to_string()))?,
781    ) as usize;
782
783    let _ = stored_bs; // block_size sanity — we use the caller's bs
784    let mut offset = 8usize;
785    let mut out: Vec<f64> = Vec::with_capacity(original_dim);
786
787    for block_idx in 0..num_blocks {
788        if offset + 20 > data.len() {
789            return Err(EccError::CorruptData(format!(
790                "block {} header missing",
791                block_idx
792            )));
793        }
794        let bmin = f64::from_le_bytes(
795            data[offset..offset + 8]
796                .try_into()
797                .map_err(|_| EccError::CorruptData("bmin".to_string()))?,
798        );
799        offset += 8;
800        let bmax = f64::from_le_bytes(
801            data[offset..offset + 8]
802                .try_into()
803                .map_err(|_| EccError::CorruptData("bmax".to_string()))?,
804        );
805        offset += 8;
806        let qlen = u32::from_le_bytes(
807            data[offset..offset + 4]
808                .try_into()
809                .map_err(|_| EccError::CorruptData("qlen".to_string()))?,
810        ) as usize;
811        offset += 4;
812
813        if offset + qlen > data.len() {
814            return Err(EccError::CorruptData(format!(
815                "block {} data truncated",
816                block_idx
817            )));
818        }
819        let qbytes = &data[offset..offset + qlen];
820        offset += qlen;
821
822        // Figure out dimension of this block
823        let block_start = block_idx * bs;
824        let block_dim = (original_dim - block_start).min(bs);
825
826        let block_vals = decompress_scalar(qbytes, bits, block_dim, bmin, bmax)?;
827        out.extend_from_slice(&block_vals);
828    }
829
830    out.truncate(original_dim);
831    if out.len() != original_dim {
832        return Err(EccError::DimensionMismatch {
833            expected: original_dim,
834            got: out.len(),
835        });
836    }
837    Ok(out)
838}
839
840// ---------------------------------------------------------------------------
841// Delta Coding
842// ---------------------------------------------------------------------------
843//
844// Compress: sort indices by value, compute deltas of sorted values, quantize deltas.
845// Decompress: reverse — cumsum, un-sort, dequantize.
846//
847// Layout: [ perm:u32×dim | delta_min:f64 | delta_max:f64 | quantized_deltas ]
848
849fn compress_delta(
850    embedding: &[f64],
851    bits: u8,
852    _min_val: f64,
853    _max_val: f64,
854) -> Result<Vec<u8>, EccError> {
855    let n = embedding.len();
856    // Sort indices by value
857    let mut perm: Vec<u32> = (0..n as u32).collect();
858    perm.sort_unstable_by(|&a, &b| {
859        embedding[a as usize]
860            .partial_cmp(&embedding[b as usize])
861            .unwrap_or(std::cmp::Ordering::Equal)
862    });
863
864    // Compute deltas of sorted values
865    let mut deltas: Vec<f64> = Vec::with_capacity(n);
866    let mut prev = embedding[perm[0] as usize];
867    deltas.push(prev); // first value stored as-is
868    for i in 1..n {
869        let cur = embedding[perm[i] as usize];
870        deltas.push(cur - prev);
871        prev = cur;
872    }
873
874    let dmin = deltas.iter().cloned().fold(f64::INFINITY, f64::min);
875    let dmax = deltas.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
876
877    let qbytes = compress_scalar(&deltas, bits, dmin, dmax)?;
878
879    let mut out: Vec<u8> = Vec::with_capacity(n * 4 + 16 + qbytes.len());
880    // Write permutation
881    for &p in &perm {
882        out.extend_from_slice(&p.to_le_bytes());
883    }
884    // Write delta range
885    out.extend_from_slice(&dmin.to_le_bytes());
886    out.extend_from_slice(&dmax.to_le_bytes());
887    // Write quantized deltas
888    out.extend_from_slice(&qbytes);
889
890    Ok(out)
891}
892
893fn decompress_delta(
894    data: &[u8],
895    bits: u8,
896    original_dim: usize,
897    _min_val: f64,
898    _max_val: f64,
899) -> Result<Vec<f64>, EccError> {
900    let n = original_dim;
901    let perm_bytes = n * 4;
902    if data.len() < perm_bytes + 16 {
903        return Err(EccError::CorruptData("delta header too short".to_string()));
904    }
905
906    let mut perm: Vec<u32> = Vec::with_capacity(n);
907    for i in 0..n {
908        let off = i * 4;
909        let p = u32::from_le_bytes(
910            data[off..off + 4]
911                .try_into()
912                .map_err(|_| EccError::CorruptData("perm bytes".to_string()))?,
913        );
914        perm.push(p);
915    }
916
917    let off = perm_bytes;
918    let dmin = f64::from_le_bytes(
919        data[off..off + 8]
920            .try_into()
921            .map_err(|_| EccError::CorruptData("dmin".to_string()))?,
922    );
923    let dmax = f64::from_le_bytes(
924        data[off + 8..off + 16]
925            .try_into()
926            .map_err(|_| EccError::CorruptData("dmax".to_string()))?,
927    );
928
929    let qbytes = &data[off + 16..];
930    let deltas = decompress_scalar(qbytes, bits, n, dmin, dmax)?;
931
932    // Cumulative sum to recover sorted values
933    let mut sorted_vals: Vec<f64> = Vec::with_capacity(n);
934    let mut acc = deltas[0];
935    sorted_vals.push(acc);
936    for &d in &deltas[1..] {
937        acc += d;
938        sorted_vals.push(acc);
939    }
940
941    // Un-sort: perm[i] = original index of i-th sorted value
942    let mut out = vec![0.0f64; n];
943    for (sorted_idx, &orig_idx) in perm.iter().enumerate() {
944        let idx = orig_idx as usize;
945        if idx >= n {
946            return Err(EccError::CorruptData(format!(
947                "perm index {} out of range {}",
948                idx, n
949            )));
950        }
951        out[idx] = sorted_vals[sorted_idx];
952    }
953    Ok(out)
954}
955
956// ---------------------------------------------------------------------------
957// Run-Length Encoding
958// ---------------------------------------------------------------------------
959//
960// Quantize values to N bits, then RLE-encode the quantized symbol stream.
961// Layout: [ num_runs:u32 | (run_value:u16 | run_len:u16)... ]
962// run_value stores the quantized symbol; run_len is clamped to u16::MAX.
963
964fn compress_rle(
965    embedding: &[f64],
966    bits: u8,
967    min_val: f64,
968    max_val: f64,
969) -> Result<Vec<u8>, EccError> {
970    let levels = max_quant_val(bits);
971    let quantized: Vec<u16> = embedding
972        .iter()
973        .map(|&v| quantize_val(v, min_val, max_val, levels) as u16)
974        .collect();
975
976    // Build runs
977    let mut runs: Vec<(u16, u16)> = Vec::new(); // (value, count)
978    let mut i = 0;
979    while i < quantized.len() {
980        let val = quantized[i];
981        let mut count = 1u16;
982        while (i + count as usize) < quantized.len()
983            && quantized[i + count as usize] == val
984            && count < u16::MAX
985        {
986            count += 1;
987        }
988        runs.push((val, count));
989        i += count as usize;
990    }
991
992    let mut out = Vec::with_capacity(4 + runs.len() * 4);
993    out.extend_from_slice(&(runs.len() as u32).to_le_bytes());
994    for (val, cnt) in &runs {
995        out.extend_from_slice(&val.to_le_bytes());
996        out.extend_from_slice(&cnt.to_le_bytes());
997    }
998    Ok(out)
999}
1000
1001fn decompress_rle(
1002    data: &[u8],
1003    bits: u8,
1004    original_dim: usize,
1005    min_val: f64,
1006    max_val: f64,
1007) -> Result<Vec<f64>, EccError> {
1008    if data.len() < 4 {
1009        return Err(EccError::CorruptData("RLE header too short".to_string()));
1010    }
1011    let num_runs = u32::from_le_bytes(
1012        data[0..4]
1013            .try_into()
1014            .map_err(|_| EccError::CorruptData("num_runs".to_string()))?,
1015    ) as usize;
1016
1017    if data.len() < 4 + num_runs * 4 {
1018        return Err(EccError::CorruptData("RLE data truncated".to_string()));
1019    }
1020
1021    // Use the same quantization levels as compress_rle
1022    let levels: u64 = max_quant_val(bits);
1023    let mut out: Vec<f64> = Vec::with_capacity(original_dim);
1024    let mut offset = 4usize;
1025    for _ in 0..num_runs {
1026        let val = u16::from_le_bytes(
1027            data[offset..offset + 2]
1028                .try_into()
1029                .map_err(|_| EccError::CorruptData("run_val".to_string()))?,
1030        ) as u64;
1031        let cnt = u16::from_le_bytes(
1032            data[offset + 2..offset + 4]
1033                .try_into()
1034                .map_err(|_| EccError::CorruptData("run_cnt".to_string()))?,
1035        ) as usize;
1036        offset += 4;
1037
1038        let decoded = dequantize_val(val, min_val, max_val, levels);
1039        for _ in 0..cnt {
1040            out.push(decoded);
1041        }
1042    }
1043
1044    out.truncate(original_dim);
1045    if out.len() != original_dim {
1046        return Err(EccError::DimensionMismatch {
1047            expected: original_dim,
1048            got: out.len(),
1049        });
1050    }
1051    Ok(out)
1052}
1053
1054// ---------------------------------------------------------------------------
1055// HybridPQ — PQ per block + RLE on quantized residual symbols
1056// ---------------------------------------------------------------------------
1057//
1058// Layout: [ num_blocks:u32 | block_size:u32 |
1059//           ( bmin:f64 | bmax:f64 | rle_len:u32 | rle_bytes... ) × num_blocks ]
1060
1061fn compress_hybrid_pq(
1062    embedding: &[f64],
1063    bits: u8,
1064    block_size: usize,
1065    _global_min: f64,
1066    _global_max: f64,
1067) -> Result<Vec<u8>, EccError> {
1068    let bs = block_size.max(1);
1069    let num_blocks = embedding.len().div_ceil(bs);
1070    let mut out: Vec<u8> = Vec::new();
1071
1072    out.extend_from_slice(&(num_blocks as u32).to_le_bytes());
1073    out.extend_from_slice(&(bs as u32).to_le_bytes());
1074
1075    for block_idx in 0..num_blocks {
1076        let start = block_idx * bs;
1077        let end = (start + bs).min(embedding.len());
1078        let block = &embedding[start..end];
1079
1080        let bmin = block.iter().cloned().fold(f64::INFINITY, f64::min);
1081        let bmax = block.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
1082
1083        out.extend_from_slice(&bmin.to_le_bytes());
1084        out.extend_from_slice(&bmax.to_le_bytes());
1085
1086        let rle_bytes = compress_rle(block, bits, bmin, bmax)?;
1087        out.extend_from_slice(&(rle_bytes.len() as u32).to_le_bytes());
1088        out.extend_from_slice(&rle_bytes);
1089    }
1090    Ok(out)
1091}
1092
1093fn decompress_hybrid_pq(
1094    data: &[u8],
1095    bits: u8,
1096    block_size: usize,
1097    original_dim: usize,
1098    _global_min: f64,
1099    _global_max: f64,
1100) -> Result<Vec<f64>, EccError> {
1101    let bs = block_size.max(1);
1102    if data.len() < 8 {
1103        return Err(EccError::CorruptData(
1104            "HybridPQ header too short".to_string(),
1105        ));
1106    }
1107
1108    let num_blocks = u32::from_le_bytes(
1109        data[0..4]
1110            .try_into()
1111            .map_err(|_| EccError::CorruptData("num_blocks".to_string()))?,
1112    ) as usize;
1113
1114    let _stored_bs = u32::from_le_bytes(
1115        data[4..8]
1116            .try_into()
1117            .map_err(|_| EccError::CorruptData("block_size".to_string()))?,
1118    );
1119
1120    let mut offset = 8usize;
1121    let mut out: Vec<f64> = Vec::with_capacity(original_dim);
1122
1123    for block_idx in 0..num_blocks {
1124        if offset + 20 > data.len() {
1125            return Err(EccError::CorruptData(format!(
1126                "HybridPQ block {} header missing",
1127                block_idx
1128            )));
1129        }
1130
1131        let bmin = f64::from_le_bytes(
1132            data[offset..offset + 8]
1133                .try_into()
1134                .map_err(|_| EccError::CorruptData("bmin".to_string()))?,
1135        );
1136        offset += 8;
1137        let bmax = f64::from_le_bytes(
1138            data[offset..offset + 8]
1139                .try_into()
1140                .map_err(|_| EccError::CorruptData("bmax".to_string()))?,
1141        );
1142        offset += 8;
1143        let rlen = u32::from_le_bytes(
1144            data[offset..offset + 4]
1145                .try_into()
1146                .map_err(|_| EccError::CorruptData("rlen".to_string()))?,
1147        ) as usize;
1148        offset += 4;
1149
1150        if offset + rlen > data.len() {
1151            return Err(EccError::CorruptData(format!(
1152                "HybridPQ block {} data truncated",
1153                block_idx
1154            )));
1155        }
1156        let rle_bytes = &data[offset..offset + rlen];
1157        offset += rlen;
1158
1159        let block_start = block_idx * bs;
1160        let block_dim = (original_dim - block_start).min(bs);
1161        let block_vals = decompress_rle(rle_bytes, bits, block_dim, bmin, bmax)?;
1162        out.extend_from_slice(&block_vals);
1163    }
1164
1165    out.truncate(original_dim);
1166    if out.len() != original_dim {
1167        return Err(EccError::DimensionMismatch {
1168            expected: original_dim,
1169            got: out.len(),
1170        });
1171    }
1172    Ok(out)
1173}
1174
1175// ---------------------------------------------------------------------------
1176// Tests
1177// ---------------------------------------------------------------------------
1178
1179#[cfg(test)]
1180mod tests {
1181    use super::*;
1182
1183    // Helpers ---------------------------------------------------------------
1184
1185    fn make_rand_embedding(dim: usize, seed: u64) -> Vec<f64> {
1186        let mut state = seed;
1187        (0..dim)
1188            .map(|_| {
1189                let r = xorshift64(&mut state);
1190                // map to [-1.0, 1.0]
1191                (r as f64 / u64::MAX as f64) * 2.0 - 1.0
1192            })
1193            .collect()
1194    }
1195
1196    fn make_codec_sq(bits: u8) -> (EmbeddingCompressionCodec, EccCodecId) {
1197        let mut c = EmbeddingCompressionCodec::new();
1198        let id = c.register_codec("sq", EccMethod::ScalarQuantization, bits, 8);
1199        (c, id)
1200    }
1201
1202    fn roundtrip_mse(codec: &mut EmbeddingCompressionCodec, id: EccCodecId, emb: &[f64]) -> f64 {
1203        let comp = codec.compress(id, emb).expect("compress failed");
1204        let decomp = codec.decompress(&comp).expect("decompress failed");
1205        EmbeddingCompressionCodec::reconstruction_error(emb, &decomp)
1206    }
1207
1208    // Registry tests --------------------------------------------------------
1209
1210    #[test]
1211    fn test_register_codec_returns_unique_ids() {
1212        let mut c = EmbeddingCompressionCodec::new();
1213        let id1 = c.register_codec("a", EccMethod::ScalarQuantization, 8, 8);
1214        let id2 = c.register_codec("b", EccMethod::ProductQuantization, 8, 8);
1215        assert_ne!(id1, id2);
1216    }
1217
1218    #[test]
1219    fn test_codec_count() {
1220        let mut c = EmbeddingCompressionCodec::new();
1221        assert_eq!(c.codec_count(), 0);
1222        c.register_codec("x", EccMethod::ScalarQuantization, 8, 8);
1223        assert_eq!(c.codec_count(), 1);
1224        c.register_codec("y", EccMethod::DeltaCoding, 8, 8);
1225        assert_eq!(c.codec_count(), 2);
1226    }
1227
1228    #[test]
1229    fn test_get_spec_returns_correct_spec() {
1230        let mut c = EmbeddingCompressionCodec::new();
1231        let id = c.register_codec("mycodec", EccMethod::RunLengthEncoding, 16, 32);
1232        let spec = c.get_spec(id).expect("spec not found");
1233        assert_eq!(spec.name, "mycodec");
1234        assert_eq!(spec.method, EccMethod::RunLengthEncoding);
1235        assert_eq!(spec.bits, 16);
1236        assert_eq!(spec.block_size, 32);
1237    }
1238
1239    #[test]
1240    fn test_get_spec_unknown_id_returns_none() {
1241        let c = EmbeddingCompressionCodec::new();
1242        assert!(c.get_spec(999).is_none());
1243    }
1244
1245    #[test]
1246    fn test_register_from_config() {
1247        let mut c = EmbeddingCompressionCodec::new();
1248        let cfg = EccCodecConfig {
1249            name: "cfg-test".to_string(),
1250            method: EccMethod::HybridPQ,
1251            quantize_bits: 4,
1252            use_delta_coding: false,
1253            block_size: 4,
1254        };
1255        let id = c.register_from_config(&cfg);
1256        let spec = c.get_spec(id).expect("no spec");
1257        assert_eq!(spec.bits, 4);
1258        assert_eq!(spec.method, EccMethod::HybridPQ);
1259    }
1260
1261    // Error handling --------------------------------------------------------
1262
1263    #[test]
1264    fn test_compress_unknown_codec_id_returns_err() {
1265        let mut c = EmbeddingCompressionCodec::new();
1266        let err = c.compress(42, &[0.1, 0.2]).unwrap_err();
1267        assert!(matches!(err, EccError::CodecNotFound(42)));
1268    }
1269
1270    #[test]
1271    fn test_compress_empty_embedding_returns_err() {
1272        let mut c = EmbeddingCompressionCodec::new();
1273        let id = c.register_codec("x", EccMethod::ScalarQuantization, 8, 8);
1274        let err = c.compress(id, &[]).unwrap_err();
1275        assert!(matches!(err, EccError::EmptyEmbedding));
1276    }
1277
1278    #[test]
1279    fn test_decompress_unknown_codec_id_returns_err() {
1280        let c = EmbeddingCompressionCodec::new();
1281        let bad = EccCompressed {
1282            codec_id: 99,
1283            method: EccMethod::ScalarQuantization,
1284            data: vec![0u8; 4],
1285            original_dim: 4,
1286            min_val: 0.0,
1287            max_val: 1.0,
1288        };
1289        assert!(c.decompress(&bad).is_err());
1290    }
1291
1292    #[test]
1293    fn test_unsupported_bit_width_returns_err() {
1294        let mut c = EmbeddingCompressionCodec::new();
1295        let id = c.register_codec("x", EccMethod::ScalarQuantization, 7, 8);
1296        let err = c.compress(id, &[0.1, 0.2, 0.3]).unwrap_err();
1297        assert!(matches!(err, EccError::UnsupportedBitWidth(7)));
1298    }
1299
1300    // Scalar Quantization roundtrips ----------------------------------------
1301
1302    #[test]
1303    fn test_sq8_roundtrip_basic() {
1304        let (mut c, id) = make_codec_sq(8);
1305        let emb = vec![0.0, 0.25, 0.5, 0.75, 1.0];
1306        let mse = roundtrip_mse(&mut c, id, &emb);
1307        assert!(mse < 1e-4, "mse={mse}");
1308    }
1309
1310    #[test]
1311    fn test_sq16_roundtrip_high_precision() {
1312        let (mut c, id) = make_codec_sq(16);
1313        let emb = make_rand_embedding(128, 42);
1314        let mse = roundtrip_mse(&mut c, id, &emb);
1315        assert!(mse < 1e-7, "mse={mse}");
1316    }
1317
1318    #[test]
1319    fn test_sq4_roundtrip_low_precision() {
1320        let (mut c, id) = make_codec_sq(4);
1321        let emb = make_rand_embedding(64, 7);
1322        let mse = roundtrip_mse(&mut c, id, &emb);
1323        // 4-bit is lossy; allow generous tolerance
1324        assert!(mse < 0.05, "mse={mse}");
1325    }
1326
1327    #[test]
1328    fn test_sq8_all_same_values() {
1329        let (mut c, id) = make_codec_sq(8);
1330        let emb = vec![0.5f64; 32];
1331        let comp = c.compress(id, &emb).expect("compress");
1332        let decomp = c.decompress(&comp).expect("decompress");
1333        assert_eq!(decomp.len(), 32);
1334        for &v in &decomp {
1335            // When all values are identical the range is zero → all quantize to 0
1336            assert!((v - 0.5).abs() < 1.0);
1337        }
1338    }
1339
1340    #[test]
1341    fn test_sq8_negative_values() {
1342        let (mut c, id) = make_codec_sq(8);
1343        let emb = vec![-1.0, -0.5, 0.0, 0.5, 1.0];
1344        let mse = roundtrip_mse(&mut c, id, &emb);
1345        assert!(mse < 1e-4, "mse={mse}");
1346    }
1347
1348    #[test]
1349    fn test_sq8_single_element() {
1350        let (mut c, id) = make_codec_sq(8);
1351        let emb = vec![std::f64::consts::PI];
1352        let comp = c.compress(id, &emb).expect("compress");
1353        let decomp = c.decompress(&comp).expect("decompress");
1354        assert_eq!(decomp.len(), 1);
1355    }
1356
1357    #[test]
1358    fn test_sq8_odd_dimension() {
1359        let (mut c, id) = make_codec_sq(8);
1360        let emb = make_rand_embedding(13, 99);
1361        let comp = c.compress(id, &emb).expect("compress");
1362        let decomp = c.decompress(&comp).expect("decompress");
1363        assert_eq!(decomp.len(), 13);
1364    }
1365
1366    #[test]
1367    fn test_sq4_byte_packing_odd_dim() {
1368        let (mut c, id) = make_codec_sq(4);
1369        let emb = make_rand_embedding(7, 11);
1370        let comp = c.compress(id, &emb).expect("compress");
1371        // 7 values → 4 bytes packed
1372        assert_eq!(comp.data.len(), 4);
1373        let decomp = c.decompress(&comp).expect("decompress");
1374        assert_eq!(decomp.len(), 7);
1375    }
1376
1377    // Product Quantization roundtrips ---------------------------------------
1378
1379    #[test]
1380    fn test_pq8_roundtrip() {
1381        let mut c = EmbeddingCompressionCodec::new();
1382        let id = c.register_codec("pq8", EccMethod::ProductQuantization, 8, 8);
1383        let emb = make_rand_embedding(64, 1234);
1384        let mse = roundtrip_mse(&mut c, id, &emb);
1385        assert!(mse < 1e-4, "mse={mse}");
1386    }
1387
1388    #[test]
1389    fn test_pq_block_size_1() {
1390        let mut c = EmbeddingCompressionCodec::new();
1391        let id = c.register_codec("pq1", EccMethod::ProductQuantization, 8, 1);
1392        let emb = make_rand_embedding(16, 55);
1393        let mse = roundtrip_mse(&mut c, id, &emb);
1394        assert!(mse < 1e-4, "mse={mse}");
1395    }
1396
1397    #[test]
1398    fn test_pq_non_divisible_dim() {
1399        let mut c = EmbeddingCompressionCodec::new();
1400        let id = c.register_codec("pq-nd", EccMethod::ProductQuantization, 8, 8);
1401        let emb = make_rand_embedding(20, 77);
1402        let comp = c.compress(id, &emb).expect("compress");
1403        let decomp = c.decompress(&comp).expect("decompress");
1404        assert_eq!(decomp.len(), 20);
1405    }
1406
1407    #[test]
1408    fn test_pq16_roundtrip_high_precision() {
1409        let mut c = EmbeddingCompressionCodec::new();
1410        let id = c.register_codec("pq16", EccMethod::ProductQuantization, 16, 16);
1411        let emb = make_rand_embedding(128, 9876);
1412        let mse = roundtrip_mse(&mut c, id, &emb);
1413        assert!(mse < 1e-7, "mse={mse}");
1414    }
1415
1416    // Delta Coding roundtrips -----------------------------------------------
1417
1418    #[test]
1419    fn test_delta_roundtrip_basic() {
1420        let mut c = EmbeddingCompressionCodec::new();
1421        let id = c.register_codec("dc8", EccMethod::DeltaCoding, 8, 8);
1422        let emb = make_rand_embedding(64, 321);
1423        let mse = roundtrip_mse(&mut c, id, &emb);
1424        assert!(mse < 1e-3, "mse={mse}");
1425    }
1426
1427    #[test]
1428    fn test_delta_single_element() {
1429        let mut c = EmbeddingCompressionCodec::new();
1430        let id = c.register_codec("dc8", EccMethod::DeltaCoding, 8, 8);
1431        let emb = vec![0.42f64];
1432        let comp = c.compress(id, &emb).expect("compress");
1433        let decomp = c.decompress(&comp).expect("decompress");
1434        assert_eq!(decomp.len(), 1);
1435    }
1436
1437    #[test]
1438    fn test_delta_monotone_input() {
1439        let mut c = EmbeddingCompressionCodec::new();
1440        let id = c.register_codec("dc8", EccMethod::DeltaCoding, 8, 8);
1441        let emb: Vec<f64> = (0..32).map(|i| i as f64 / 32.0).collect();
1442        let mse = roundtrip_mse(&mut c, id, &emb);
1443        assert!(mse < 1e-3, "mse={mse}");
1444    }
1445
1446    #[test]
1447    fn test_delta_16bit_precision() {
1448        let mut c = EmbeddingCompressionCodec::new();
1449        let id = c.register_codec("dc16", EccMethod::DeltaCoding, 16, 8);
1450        let emb = make_rand_embedding(32, 42);
1451        let mse = roundtrip_mse(&mut c, id, &emb);
1452        assert!(mse < 1e-6, "mse={mse}");
1453    }
1454
1455    // RLE roundtrips --------------------------------------------------------
1456
1457    #[test]
1458    fn test_rle_roundtrip_basic() {
1459        let mut c = EmbeddingCompressionCodec::new();
1460        let id = c.register_codec("rle8", EccMethod::RunLengthEncoding, 8, 8);
1461        let emb = make_rand_embedding(64, 555);
1462        let mse = roundtrip_mse(&mut c, id, &emb);
1463        // RLE uses 8-bit quantization (255 levels); allow ~1e-4 MSE
1464        assert!(mse < 1e-4, "mse={mse}");
1465    }
1466
1467    #[test]
1468    fn test_rle_with_repeated_values() {
1469        let mut c = EmbeddingCompressionCodec::new();
1470        let id = c.register_codec("rle8", EccMethod::RunLengthEncoding, 8, 8);
1471        let mut emb = vec![0.5f64; 50];
1472        emb.extend(vec![0.1f64; 50]);
1473        let comp = c
1474            .compress(id, &emb)
1475            .expect("test: compress with repeated values");
1476        let decomp = c
1477            .decompress(&comp)
1478            .expect("test: decompress with repeated values");
1479        assert_eq!(decomp.len(), 100);
1480        // only 2 runs
1481        // num_runs=2 → 4 + 2*4 = 12 bytes
1482        assert_eq!(comp.data.len(), 12);
1483    }
1484
1485    #[test]
1486    fn test_rle_single_value() {
1487        let mut c = EmbeddingCompressionCodec::new();
1488        let id = c.register_codec("rle8", EccMethod::RunLengthEncoding, 8, 8);
1489        let emb = vec![0.7f64];
1490        let comp = c.compress(id, &emb).expect("test: compress single value");
1491        let decomp = c.decompress(&comp).expect("test: decompress single value");
1492        assert_eq!(decomp.len(), 1);
1493    }
1494
1495    // HybridPQ roundtrips ---------------------------------------------------
1496
1497    #[test]
1498    fn test_hybridpq_roundtrip_basic() {
1499        let mut c = EmbeddingCompressionCodec::new();
1500        let id = c.register_codec("hpq8", EccMethod::HybridPQ, 8, 8);
1501        let emb = make_rand_embedding(64, 777);
1502        let mse = roundtrip_mse(&mut c, id, &emb);
1503        // HybridPQ uses 8-bit RLE per block; allow ~1e-4 MSE
1504        assert!(mse < 1e-4, "mse={mse}");
1505    }
1506
1507    #[test]
1508    fn test_hybridpq_non_divisible_dim() {
1509        let mut c = EmbeddingCompressionCodec::new();
1510        let id = c.register_codec("hpq8", EccMethod::HybridPQ, 8, 8);
1511        let emb = make_rand_embedding(17, 888);
1512        let comp = c
1513            .compress(id, &emb)
1514            .expect("test: compress hybridpq repeated");
1515        let decomp = c
1516            .decompress(&comp)
1517            .expect("test: decompress hybridpq repeated");
1518        assert_eq!(decomp.len(), 17);
1519    }
1520
1521    #[test]
1522    fn test_hybridpq_repeated_values_compression() {
1523        let mut c = EmbeddingCompressionCodec::new();
1524        let id = c.register_codec("hpq8", EccMethod::HybridPQ, 8, 8);
1525        let emb = vec![0.3f64; 64];
1526        let comp_hpq = c
1527            .compress(id, &emb)
1528            .expect("test: compress hybridpq repeated");
1529        // Should still roundtrip correctly
1530        let decomp = c
1531            .decompress(&comp_hpq)
1532            .expect("test: decompress hybridpq repeated");
1533        assert_eq!(decomp.len(), 64);
1534    }
1535
1536    // Batch operations ------------------------------------------------------
1537
1538    #[test]
1539    fn test_compress_batch_all_succeed() {
1540        let mut c = EmbeddingCompressionCodec::new();
1541        let id = c.register_codec("sq8", EccMethod::ScalarQuantization, 8, 8);
1542        let batch: Vec<Vec<f64>> = (0..5).map(|i| make_rand_embedding(16, i + 100)).collect();
1543        let results = c.compress_batch(id, &batch);
1544        assert_eq!(results.len(), 5);
1545        for r in &results {
1546            assert!(r.is_ok());
1547        }
1548    }
1549
1550    #[test]
1551    fn test_decompress_batch_all_succeed() {
1552        let mut c = EmbeddingCompressionCodec::new();
1553        let id = c.register_codec("sq8", EccMethod::ScalarQuantization, 8, 8);
1554        let batch: Vec<Vec<f64>> = (0..3).map(|i| make_rand_embedding(32, i + 200)).collect();
1555        let compressed: Vec<EccCompressed> = batch
1556            .iter()
1557            .map(|emb| c.compress(id, emb).expect("test: compress in batch"))
1558            .collect();
1559        let decompressed = c.decompress_batch(&compressed);
1560        assert_eq!(decompressed.len(), 3);
1561        for (orig, decomp_res) in batch.iter().zip(decompressed.iter()) {
1562            let decomp = decomp_res.as_ref().expect("decompress failed");
1563            assert_eq!(decomp.len(), orig.len());
1564        }
1565    }
1566
1567    #[test]
1568    fn test_compress_batch_empty_batch() {
1569        let mut c = EmbeddingCompressionCodec::new();
1570        let id = c.register_codec("sq8", EccMethod::ScalarQuantization, 8, 8);
1571        let results = c.compress_batch(id, &[]);
1572        assert!(results.is_empty());
1573    }
1574
1575    #[test]
1576    fn test_decompress_batch_empty() {
1577        let c = EmbeddingCompressionCodec::new();
1578        let result = c.decompress_batch(&[]);
1579        assert!(result.is_empty());
1580    }
1581
1582    // Reconstruction error --------------------------------------------------
1583
1584    #[test]
1585    fn test_reconstruction_error_identical_vectors() {
1586        let v = vec![0.1, 0.2, 0.3];
1587        let mse = EmbeddingCompressionCodec::reconstruction_error(&v, &v);
1588        assert!(mse.abs() < f64::EPSILON);
1589    }
1590
1591    #[test]
1592    fn test_reconstruction_error_known_value() {
1593        let a = vec![0.0f64, 0.0, 0.0];
1594        let b = vec![1.0f64, 1.0, 1.0];
1595        let mse = EmbeddingCompressionCodec::reconstruction_error(&a, &b);
1596        assert!((mse - 1.0).abs() < f64::EPSILON);
1597    }
1598
1599    #[test]
1600    fn test_reconstruction_error_empty_returns_zero() {
1601        let mse = EmbeddingCompressionCodec::reconstruction_error(&[], &[]);
1602        assert_eq!(mse, 0.0);
1603    }
1604
1605    #[test]
1606    fn test_reconstruction_error_different_lengths_uses_min() {
1607        let a = vec![0.0f64, 0.0, 0.0, 0.0];
1608        let b = vec![1.0f64, 1.0];
1609        let mse = EmbeddingCompressionCodec::reconstruction_error(&a, &b);
1610        // Only first 2 elements compared: (0-1)^2 * 2 / 2 = 1.0
1611        assert!((mse - 1.0).abs() < f64::EPSILON);
1612    }
1613
1614    // Estimate ratio --------------------------------------------------------
1615
1616    #[test]
1617    fn test_estimate_ratio_sq8_better_than_one() {
1618        let mut c = EmbeddingCompressionCodec::new();
1619        let id = c.register_codec("sq8", EccMethod::ScalarQuantization, 8, 8);
1620        let ratio = c.estimate_ratio(id, 512);
1621        assert!(ratio > 1.0, "ratio={ratio}");
1622    }
1623
1624    #[test]
1625    fn test_estimate_ratio_sq16_less_than_sq8() {
1626        let mut c = EmbeddingCompressionCodec::new();
1627        let id8 = c.register_codec("sq8", EccMethod::ScalarQuantization, 8, 8);
1628        let id16 = c.register_codec("sq16", EccMethod::ScalarQuantization, 16, 8);
1629        let r8 = c.estimate_ratio(id8, 256);
1630        let r16 = c.estimate_ratio(id16, 256);
1631        assert!(r8 > r16, "r8={r8}, r16={r16}");
1632    }
1633
1634    #[test]
1635    fn test_estimate_ratio_unknown_codec() {
1636        let c = EmbeddingCompressionCodec::new();
1637        let ratio = c.estimate_ratio(999, 128);
1638        assert_eq!(ratio, 1.0);
1639    }
1640
1641    #[test]
1642    fn test_estimate_ratio_zero_dim() {
1643        let mut c = EmbeddingCompressionCodec::new();
1644        let id = c.register_codec("sq8", EccMethod::ScalarQuantization, 8, 8);
1645        let ratio = c.estimate_ratio(id, 0);
1646        assert_eq!(ratio, 1.0);
1647    }
1648
1649    #[test]
1650    fn test_estimate_ratio_hybrid_pq_less_than_pq() {
1651        let mut c = EmbeddingCompressionCodec::new();
1652        let id_pq = c.register_codec("pq8", EccMethod::ProductQuantization, 8, 8);
1653        let id_hpq = c.register_codec("hpq8", EccMethod::HybridPQ, 8, 8);
1654        let rpq = c.estimate_ratio(id_pq, 128);
1655        let rhpq = c.estimate_ratio(id_hpq, 128);
1656        // HybridPQ claims ~30% better (0.7× factor on PQ bits)
1657        assert!(rhpq > rpq, "rpq={rpq} rhpq={rhpq}");
1658    }
1659
1660    // Codec stats -----------------------------------------------------------
1661
1662    #[test]
1663    fn test_codec_stats_empty_log() {
1664        let c = EmbeddingCompressionCodec::new();
1665        let stats = c.codec_stats();
1666        assert_eq!(stats.total_compressed, 0);
1667        assert_eq!(stats.avg_ratio, 1.0);
1668        assert!(stats.per_method.is_empty());
1669    }
1670
1671    #[test]
1672    fn test_codec_stats_after_compressions() {
1673        let mut c = EmbeddingCompressionCodec::new();
1674        let id = c.register_codec("sq8", EccMethod::ScalarQuantization, 8, 8);
1675        let emb = make_rand_embedding(128, 11);
1676        c.compress(id, &emb).expect("test: compress for stats");
1677        c.compress(id, &emb)
1678            .expect("test: compress for stats second");
1679        let stats = c.codec_stats();
1680        assert_eq!(stats.total_compressed, 2);
1681        assert!(stats.avg_ratio > 1.0);
1682        assert!(stats
1683            .per_method
1684            .contains_key(&EccMethod::ScalarQuantization));
1685    }
1686
1687    #[test]
1688    fn test_codec_stats_per_method_breakdown() {
1689        let mut c = EmbeddingCompressionCodec::new();
1690        let id_sq = c.register_codec("sq8", EccMethod::ScalarQuantization, 8, 8);
1691        let id_rle = c.register_codec("rle8", EccMethod::RunLengthEncoding, 8, 8);
1692        let emb = make_rand_embedding(64, 22);
1693        c.compress(id_sq, &emb)
1694            .expect("test: compress sq8 for per-method stats");
1695        c.compress(id_rle, &emb)
1696            .expect("test: compress rle for per-method stats");
1697        let stats = c.codec_stats();
1698        assert_eq!(stats.total_compressed, 2);
1699        let (sq_ops, _, _) = stats.per_method[&EccMethod::ScalarQuantization];
1700        let (rle_ops, _, _) = stats.per_method[&EccMethod::RunLengthEncoding];
1701        assert_eq!(sq_ops, 1);
1702        assert_eq!(rle_ops, 1);
1703    }
1704
1705    // Log bounding ----------------------------------------------------------
1706
1707    #[test]
1708    fn test_log_bounded_to_500() {
1709        let mut c = EmbeddingCompressionCodec::new();
1710        let id = c.register_codec("sq8", EccMethod::ScalarQuantization, 8, 8);
1711        let emb = vec![0.5f64; 8];
1712        for _ in 0..600 {
1713            c.compress(id, &emb)
1714                .expect("test: compress for log bound test");
1715        }
1716        assert_eq!(c.log_entries().len(), 500);
1717    }
1718
1719    // Compression payload fields --------------------------------------------
1720
1721    #[test]
1722    fn test_compressed_fields_populated() {
1723        let mut c = EmbeddingCompressionCodec::new();
1724        let id = c.register_codec("sq8", EccMethod::ScalarQuantization, 8, 8);
1725        let emb = vec![0.1, 0.5, 0.9];
1726        let comp = c
1727            .compress(id, &emb)
1728            .expect("test: compress for fields test");
1729        assert_eq!(comp.codec_id, id);
1730        assert_eq!(comp.method, EccMethod::ScalarQuantization);
1731        assert_eq!(comp.original_dim, 3);
1732        assert!((comp.min_val - 0.1).abs() < 1e-10);
1733        assert!((comp.max_val - 0.9).abs() < 1e-10);
1734    }
1735
1736    #[test]
1737    fn test_compressed_data_non_empty() {
1738        let mut c = EmbeddingCompressionCodec::new();
1739        let id = c.register_codec("sq8", EccMethod::ScalarQuantization, 8, 8);
1740        let emb = make_rand_embedding(32, 33);
1741        let comp = c
1742            .compress(id, &emb)
1743            .expect("test: compress for non-empty data check");
1744        assert!(!comp.data.is_empty());
1745    }
1746
1747    // Method display --------------------------------------------------------
1748
1749    #[test]
1750    fn test_ecc_method_display() {
1751        assert_eq!(
1752            EccMethod::ScalarQuantization.to_string(),
1753            "ScalarQuantization"
1754        );
1755        assert_eq!(
1756            EccMethod::ProductQuantization.to_string(),
1757            "ProductQuantization"
1758        );
1759        assert_eq!(EccMethod::DeltaCoding.to_string(), "DeltaCoding");
1760        assert_eq!(
1761            EccMethod::RunLengthEncoding.to_string(),
1762            "RunLengthEncoding"
1763        );
1764        assert_eq!(EccMethod::HybridPQ.to_string(), "HybridPQ");
1765    }
1766
1767    // Default codec config --------------------------------------------------
1768
1769    #[test]
1770    fn test_default_config() {
1771        let cfg = EccCodecConfig::default();
1772        assert_eq!(cfg.quantize_bits, 8);
1773        assert_eq!(cfg.method, EccMethod::ScalarQuantization);
1774        assert!(!cfg.use_delta_coding);
1775        assert_eq!(cfg.block_size, 8);
1776    }
1777
1778    // xorshift64 PRNG -------------------------------------------------------
1779
1780    #[test]
1781    fn test_xorshift64_nondeterministic_zero_free() {
1782        let mut state: u64 = 12345;
1783        let mut any_nonzero = false;
1784        for _ in 0..100 {
1785            let v = xorshift64(&mut state);
1786            if v != 0 {
1787                any_nonzero = true;
1788            }
1789        }
1790        assert!(any_nonzero);
1791    }
1792
1793    #[test]
1794    fn test_xorshift64_produces_different_values() {
1795        let mut state: u64 = 99999;
1796        let a = xorshift64(&mut state);
1797        let b = xorshift64(&mut state);
1798        assert_ne!(a, b);
1799    }
1800
1801    // Additional edge cases -------------------------------------------------
1802
1803    #[test]
1804    fn test_sq8_dim_1024() {
1805        let (mut c, id) = make_codec_sq(8);
1806        let emb = make_rand_embedding(1024, 2024);
1807        let mse = roundtrip_mse(&mut c, id, &emb);
1808        assert!(mse < 1e-4, "mse={mse}");
1809    }
1810
1811    #[test]
1812    fn test_pq_large_block_exceeds_dim() {
1813        let mut c = EmbeddingCompressionCodec::new();
1814        let id = c.register_codec("pq-lb", EccMethod::ProductQuantization, 8, 128);
1815        let emb = make_rand_embedding(16, 42); // dim < block_size
1816        let comp = c
1817            .compress(id, &emb)
1818            .expect("test: compress should succeed when block size exceeds dim");
1819        let decomp = c
1820            .decompress(&comp)
1821            .expect("test: decompress should succeed when block size exceeds dim");
1822        assert_eq!(decomp.len(), 16);
1823    }
1824
1825    #[test]
1826    fn test_all_methods_roundtrip_dim_128() {
1827        let methods = [
1828            EccMethod::ScalarQuantization,
1829            EccMethod::ProductQuantization,
1830            EccMethod::DeltaCoding,
1831            EccMethod::RunLengthEncoding,
1832            EccMethod::HybridPQ,
1833        ];
1834        let emb = make_rand_embedding(128, 1111);
1835        for method in &methods {
1836            let mut c = EmbeddingCompressionCodec::new();
1837            let id = c.register_codec("test", *method, 8, 8);
1838            let comp = c
1839                .compress(id, &emb)
1840                .unwrap_or_else(|_| panic!("compress {method}"));
1841            let decomp = c
1842                .decompress(&comp)
1843                .unwrap_or_else(|_| panic!("decompress {method}"));
1844            assert_eq!(decomp.len(), 128, "dim mismatch for {method}");
1845        }
1846    }
1847
1848    #[test]
1849    fn test_compressed_size_smaller_than_original_sq8() {
1850        let (mut c, id) = make_codec_sq(8);
1851        let emb = make_rand_embedding(256, 42);
1852        let comp = c
1853            .compress(id, &emb)
1854            .expect("test: compress sq8 should succeed");
1855        let original_bytes = 256 * 8;
1856        assert!(
1857            comp.data.len() < original_bytes,
1858            "data.len()={}",
1859            comp.data.len()
1860        );
1861    }
1862
1863    #[test]
1864    fn test_compressed_size_smaller_than_original_sq16() {
1865        let (mut c, id) = make_codec_sq(16);
1866        let emb = make_rand_embedding(256, 42);
1867        let comp = c
1868            .compress(id, &emb)
1869            .expect("test: compress sq16 should succeed");
1870        let original_bytes = 256 * 8;
1871        assert!(comp.data.len() < original_bytes);
1872    }
1873
1874    #[test]
1875    fn test_hybridpq_with_block_size_4() {
1876        let mut c = EmbeddingCompressionCodec::new();
1877        let id = c.register_codec("hpq4", EccMethod::HybridPQ, 8, 4);
1878        let emb = make_rand_embedding(32, 444);
1879        let mse = roundtrip_mse(&mut c, id, &emb);
1880        // HybridPQ/8bit: allow ~1e-4 MSE
1881        assert!(mse < 1e-4, "mse={mse}");
1882    }
1883
1884    #[test]
1885    fn test_delta_coding_all_negative() {
1886        let mut c = EmbeddingCompressionCodec::new();
1887        let id = c.register_codec("dc8", EccMethod::DeltaCoding, 8, 8);
1888        let emb: Vec<f64> = (0..32).map(|i| -(i as f64) / 32.0).collect();
1889        let mse = roundtrip_mse(&mut c, id, &emb);
1890        assert!(mse < 1e-3, "mse={mse}");
1891    }
1892
1893    #[test]
1894    fn test_codec_stats_total_bytes_saved() {
1895        let mut c = EmbeddingCompressionCodec::new();
1896        let id = c.register_codec("sq8", EccMethod::ScalarQuantization, 8, 8);
1897        let emb = make_rand_embedding(256, 7);
1898        c.compress(id, &emb)
1899            .expect("test: compress should succeed for stats check");
1900        let stats = c.codec_stats();
1901        // 256 * 8 bytes original vs 256 bytes sq8 → should save 1792 bytes
1902        assert!(stats.total_bytes_saved > 0);
1903    }
1904
1905    #[test]
1906    fn test_sq8_compress_data_length() {
1907        let (mut c, id) = make_codec_sq(8);
1908        let emb = make_rand_embedding(128, 10);
1909        let comp = c
1910            .compress(id, &emb)
1911            .expect("test: compress sq8 should succeed");
1912        // SQ8: 1 byte per element
1913        assert_eq!(comp.data.len(), 128);
1914    }
1915
1916    #[test]
1917    fn test_sq16_compress_data_length() {
1918        let (mut c, id) = make_codec_sq(16);
1919        let emb = make_rand_embedding(64, 20);
1920        let comp = c
1921            .compress(id, &emb)
1922            .expect("test: compress sq16 should succeed");
1923        // SQ16: 2 bytes per element
1924        assert_eq!(comp.data.len(), 128);
1925    }
1926
1927    #[test]
1928    fn test_sq4_compress_data_length_even_dim() {
1929        let (mut c, id) = make_codec_sq(4);
1930        let emb = make_rand_embedding(64, 30);
1931        let comp = c
1932            .compress(id, &emb)
1933            .expect("test: compress sq4 should succeed");
1934        // SQ4: 2 values per byte → 32 bytes
1935        assert_eq!(comp.data.len(), 32);
1936    }
1937
1938    #[test]
1939    fn test_ecc_error_clone_and_partial_eq() {
1940        let e = EccError::EmptyEmbedding;
1941        assert_eq!(e.clone(), EccError::EmptyEmbedding);
1942    }
1943
1944    #[test]
1945    fn test_ecc_method_hash_map_key() {
1946        let mut map: HashMap<EccMethod, usize> = HashMap::new();
1947        map.insert(EccMethod::ScalarQuantization, 1);
1948        map.insert(EccMethod::HybridPQ, 2);
1949        assert_eq!(
1950            *map.get(&EccMethod::ScalarQuantization)
1951                .expect("test: ScalarQuantization should be in map"),
1952            1
1953        );
1954    }
1955}