1use std::collections::{HashMap, VecDeque};
30use std::time::{SystemTime, UNIX_EPOCH};
31
32use thiserror::Error;
33
34#[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
49pub type EccEmbeddingCompressionCodec = EmbeddingCompressionCodec;
55
56pub type EccCodecId = u32;
58
59#[derive(Debug, Error, Clone, PartialEq)]
65pub enum EccError {
66 #[error("codec id {0} not found in registry")]
68 CodecNotFound(EccCodecId),
69
70 #[error("embedding must not be empty")]
72 EmptyEmbedding,
73
74 #[error("compressed data is corrupt: {0}")]
76 CorruptData(String),
77
78 #[error("dimension mismatch: expected {expected}, got {got}")]
80 DimensionMismatch { expected: usize, got: usize },
81
82 #[error("block size {block_size} does not divide embedding dim {dim} evenly")]
84 BlockSizeMismatch { block_size: usize, dim: usize },
85
86 #[error("unsupported quantize_bits {0}: must be 4, 8, or 16")]
88 UnsupportedBitWidth(u8),
89
90 #[error("internal error: {0}")]
92 Internal(String),
93}
94
95#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
101pub enum EccMethod {
102 ScalarQuantization,
104 ProductQuantization,
106 DeltaCoding,
108 RunLengthEncoding,
110 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#[derive(Debug, Clone)]
132pub struct EccCodecConfig {
133 pub name: String,
135 pub method: EccMethod,
137 pub quantize_bits: u8,
139 pub use_delta_coding: bool,
141 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#[derive(Debug, Clone)]
159pub struct EccCodecSpec {
160 pub id: EccCodecId,
162 pub name: String,
164 pub method: EccMethod,
166 pub bits: u8,
168 pub block_size: usize,
170}
171
172#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
178pub struct EccCompressed {
179 pub codec_id: EccCodecId,
181 pub method: EccMethod,
183 pub data: Vec<u8>,
185 pub original_dim: usize,
187 pub min_val: f64,
189 pub max_val: f64,
191}
192
193#[derive(Debug, Clone)]
199pub struct EccCompressionRecord {
200 pub ts: u64,
202 pub original_bytes: usize,
204 pub compressed_bytes: usize,
206 pub ratio: f64,
208 pub method: EccMethod,
210}
211
212#[derive(Debug, Clone)]
214pub struct EccCodecStats {
215 pub total_compressed: usize,
217 pub total_bytes_saved: usize,
219 pub avg_ratio: f64,
221 pub per_method: HashMap<EccMethod, (usize, usize, f64)>,
223}
224
225const MAX_LOG_ENTRIES: usize = 500;
231
232pub struct EmbeddingCompressionCodec {
237 registry: HashMap<EccCodecId, EccCodecSpec>,
239 next_id: EccCodecId,
241 log: VecDeque<EccCompressionRecord>,
243}
244
245impl Default for EmbeddingCompressionCodec {
246 fn default() -> Self {
247 Self::new()
248 }
249}
250
251impl EmbeddingCompressionCodec {
252 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 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 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 pub fn get_spec(&self, id: EccCodecId) -> Option<&EccCodecSpec> {
300 self.registry.get(&id)
301 }
302
303 pub fn codec_count(&self) -> usize {
305 self.registry.len()
306 }
307
308 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 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 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 pub fn decompress_batch(&self, batch: &[EccCompressed]) -> Vec<Result<Vec<f64>, EccError>> {
436 batch.iter().map(|c| self.decompress(c)).collect()
437 }
438
439 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 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 (dim * spec.bits as usize * 7) / 10
481 }
482 EccMethod::RunLengthEncoding => {
483 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 (pq_bits * 7) / 10
492 }
493 };
494 let total_compressed_bits = quantized_bits + 192;
496 original_bits as f64 / total_compressed_bits as f64
497 }
498
499 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 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 pub fn log_entries(&self) -> &VecDeque<EccCompressionRecord> {
540 &self.log
541 }
542
543 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
555fn 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#[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#[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#[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
606fn 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 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
716fn 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 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 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; 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 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
840fn 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 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 let mut deltas: Vec<f64> = Vec::with_capacity(n);
866 let mut prev = embedding[perm[0] as usize];
867 deltas.push(prev); 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 for &p in &perm {
882 out.extend_from_slice(&p.to_le_bytes());
883 }
884 out.extend_from_slice(&dmin.to_le_bytes());
886 out.extend_from_slice(&dmax.to_le_bytes());
887 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 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 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
956fn 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 let mut runs: Vec<(u16, u16)> = Vec::new(); 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 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
1054fn 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#[cfg(test)]
1180mod tests {
1181 use super::*;
1182
1183 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 (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 #[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 #[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 #[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 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 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 assert_eq!(comp.data.len(), 4);
1373 let decomp = c.decompress(&comp).expect("decompress");
1374 assert_eq!(decomp.len(), 7);
1375 }
1376
1377 #[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 #[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 #[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 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 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 #[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 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 let decomp = c
1531 .decompress(&comp_hpq)
1532 .expect("test: decompress hybridpq repeated");
1533 assert_eq!(decomp.len(), 64);
1534 }
1535
1536 #[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 #[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 assert!((mse - 1.0).abs() < f64::EPSILON);
1612 }
1613
1614 #[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 assert!(rhpq > rpq, "rpq={rpq} rhpq={rhpq}");
1658 }
1659
1660 #[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 #[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 #[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 #[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 #[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 #[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 #[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); 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 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 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 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 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 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}