1use std::path::Path;
29
30use serde::{de::DeserializeOwned, Serialize};
31
32use crate::{
33 codec::{FibCodeV1, FibQuantizer},
34 profile::FibQuantProfileV1,
35 scoring::FibScorer,
36 sidecar::FibSidecarIndex,
37 FibQuantError, Result,
38};
39
40#[cfg(feature = "mmap")]
41use crate::sidecar::ScoredCandidate;
42
43pub const FILE_MAGIC: [u8; 4] = *b"FBS1";
49
50pub const FILE_VERSION: u8 = 1;
52
53const MIN_HEADER_SIZE: usize = 5;
57
58#[derive(Debug)]
69pub struct FibSidecarFileV1 {
70 pub magic: [u8; 4],
72 pub version: u8,
74 pub profile: FibQuantProfileV1,
76 pub gram_values: Vec<f32>,
78 pub entry_count: u32,
80 }
82
83fn write_len_prefixed(out: &mut Vec<u8>, data: &[u8]) {
89 let len = u32::try_from(data.len()).expect("length fits in u32");
90 out.extend_from_slice(&len.to_le_bytes());
91 out.extend_from_slice(data);
92}
93
94fn read_u32_at(buf: &[u8], offset: usize, label: &str) -> Result<(u32, usize)> {
101 if offset + 4 > buf.len() {
102 return Err(FibQuantError::CorruptPayload(format!(
103 "truncated file: cannot read {} u32 at offset {} (file len {})",
104 label,
105 offset,
106 buf.len()
107 )));
108 }
109 let val = u32::from_le_bytes([
110 buf[offset],
111 buf[offset + 1],
112 buf[offset + 2],
113 buf[offset + 3],
114 ]);
115 Ok((val, offset + 4))
116}
117
118fn read_len_prefixed<'a>(buf: &'a [u8], offset: usize, label: &str) -> Result<(&'a [u8], usize)> {
121 let (len, next) = read_u32_at(buf, offset, &format!("{} length", label))?;
122 let len = len as usize;
123 if next + len > buf.len() {
124 return Err(FibQuantError::CorruptPayload(format!(
125 "truncated file: {} data at offset {} needs {} bytes but only {} remain",
126 label,
127 next,
128 len,
129 buf.len() - next
130 )));
131 }
132 Ok((&buf[next..next + len], next + len))
133}
134
135fn read_f32_slice(
137 buf: &[u8],
138 offset: usize,
139 count: usize,
140 label: &str,
141) -> Result<(Vec<f32>, usize)> {
142 let byte_len = count.checked_mul(4).ok_or_else(|| {
143 FibQuantError::ResourceLimitExceeded(format!("{} f32 count overflow", label))
144 })?;
145 if offset + byte_len > buf.len() {
146 return Err(FibQuantError::CorruptPayload(format!(
147 "truncated file: {} needs {} bytes at offset {} but only {} remain",
148 label,
149 byte_len,
150 offset,
151 buf.len() - offset
152 )));
153 }
154 let mut values = Vec::with_capacity(count);
155 for i in 0..count {
156 let base = offset + i * 4;
157 values.push(f32::from_le_bytes([
158 buf[base],
159 buf[base + 1],
160 buf[base + 2],
161 buf[base + 3],
162 ]));
163 }
164 Ok((values, offset + byte_len))
165}
166
167pub fn save_to_file<Id>(index: &FibSidecarIndex<Id>, path: &Path) -> Result<()>
183where
184 Id: Clone + Eq + std::fmt::Debug + Serialize,
185{
186 let bytes = serialize_index(index)?;
187 std::fs::write(path, &bytes).map_err(|e| {
188 FibQuantError::CorruptPayload(format!("failed to write sidecar file {:?}: {}", path, e))
189 })
190}
191
192fn serialize_index<Id>(index: &FibSidecarIndex<Id>) -> Result<Vec<u8>>
194where
195 Id: Clone + Eq + std::fmt::Debug + Serialize,
196{
197 let profile = index.profile();
199 let profile_json = serde_json::to_vec(profile).map_err(|e| {
200 FibQuantError::CorruptPayload(format!("failed to serialize profile as JSON: {}", e))
201 })?;
202
203 let gram_values = index.scorer().gram_table().values();
205 let gram_count = gram_values.len();
206
207 let entries = index.entries();
209 let entry_count = u32::try_from(entries.len())
210 .map_err(|_| FibQuantError::ResourceLimitExceeded("entry count exceeds u32".into()))?;
211
212 let estimated_size =
214 MIN_HEADER_SIZE + 4 + profile_json.len() + 4 + gram_count * 4 + 4 + entries.len() * 40; let mut out = Vec::with_capacity(estimated_size);
216
217 out.extend_from_slice(&FILE_MAGIC);
219 out.push(FILE_VERSION);
220
221 write_len_prefixed(&mut out, &profile_json);
223
224 let gram_u32 = u32::try_from(gram_count)
226 .map_err(|_| FibQuantError::ResourceLimitExceeded("gram table size exceeds u32".into()))?;
227 out.extend_from_slice(&gram_u32.to_le_bytes());
228 for &v in gram_values {
229 out.extend_from_slice(&v.to_le_bytes());
230 }
231
232 out.extend_from_slice(&entry_count.to_le_bytes());
234
235 for (id, code) in entries {
237 let id_bytes = postcard::to_stdvec(id).map_err(|e| {
238 FibQuantError::CorruptPayload(format!("failed to serialize entry id: {}", e))
239 })?;
240 write_len_prefixed(&mut out, &id_bytes);
241 let code_bytes = code.to_compact_bytes();
242 write_len_prefixed(&mut out, &code_bytes);
243 }
244
245 Ok(out)
246}
247
248pub fn load_from_file<Id>(path: &Path) -> Result<FibSidecarIndex<Id>>
265where
266 Id: Clone + Eq + std::fmt::Debug + DeserializeOwned,
267{
268 let bytes = std::fs::read(path).map_err(|e| {
269 FibQuantError::CorruptPayload(format!("failed to read sidecar file {:?}: {}", path, e))
270 })?;
271 deserialize_index(&bytes)
272}
273
274fn deserialize_index<Id>(buf: &[u8]) -> Result<FibSidecarIndex<Id>>
276where
277 Id: Clone + Eq + std::fmt::Debug + DeserializeOwned,
278{
279 if buf.len() < MIN_HEADER_SIZE {
281 return Err(FibQuantError::CorruptPayload(format!(
282 "file too short: {} bytes (need >= {})",
283 buf.len(),
284 MIN_HEADER_SIZE
285 )));
286 }
287 if buf[0..4] != FILE_MAGIC {
288 return Err(FibQuantError::CorruptPayload(format!(
289 "bad file magic: {:?} (expected {:?})",
290 &buf[0..4],
291 FILE_MAGIC
292 )));
293 }
294 if buf[4] != FILE_VERSION {
295 return Err(FibQuantError::CorruptPayload(format!(
296 "unsupported file version {} (expected {})",
297 buf[4], FILE_VERSION
298 )));
299 }
300
301 let mut offset = MIN_HEADER_SIZE; let (profile_bytes, next) = read_len_prefixed(buf, offset, "profile")?;
305 offset = next;
306 let profile: FibQuantProfileV1 = serde_json::from_slice(profile_bytes).map_err(|e| {
307 FibQuantError::CorruptPayload(format!("failed to deserialize profile JSON: {}", e))
308 })?;
309
310 let (gram_count, next) = read_u32_at(buf, offset, "gram count")?;
312 offset = next;
313 let gram_count = gram_count as usize;
314 let (gram_values, next) = read_f32_slice(buf, offset, gram_count, "gram table")?;
315 offset = next;
316
317 let (entry_count, next) = read_u32_at(buf, offset, "entry count")?;
319 offset = next;
320 let entry_count = entry_count as usize;
321
322 let quantizer = FibQuantizer::new(profile.clone())?;
327 let scorer = FibScorer::new(quantizer)?;
328
329 let reconstructed_gram = scorer.gram_table().values();
331 if reconstructed_gram.len() != gram_values.len() {
332 return Err(FibQuantError::CorruptPayload(format!(
333 "gram table size mismatch: file has {} values, reconstructed has {}",
334 gram_values.len(),
335 reconstructed_gram.len()
336 )));
337 }
338 for (i, (a, b)) in gram_values
339 .iter()
340 .zip(reconstructed_gram.iter())
341 .enumerate()
342 {
343 if (a - b).abs() > 1e-5 {
344 return Err(FibQuantError::CorruptPayload(format!(
345 "gram table value mismatch at index {}: file={}, reconstructed={}",
346 i, a, b
347 )));
348 }
349 }
350
351 let mut index = FibSidecarIndex::<Id>::new(scorer);
353
354 for i in 0..entry_count {
356 let (id_bytes, next) = read_len_prefixed(buf, offset, &format!("entry {} id", i))?;
358 offset = next;
359 let id: Id = postcard::from_bytes(id_bytes).map_err(|e| {
360 FibQuantError::CorruptPayload(format!("failed to deserialize entry {} id: {}", i, e))
361 })?;
362
363 let (code_bytes, next) = read_len_prefixed(buf, offset, &format!("entry {} code", i))?;
365 offset = next;
366 let code = FibCodeV1::from_compact_bytes(code_bytes, &profile)?;
367 index.add(id, code);
368 }
369
370 Ok(index)
371}
372
373#[cfg(feature = "mmap")]
388pub struct MmapSidecarIndex<Id>
389where
390 Id: Clone + Eq + std::fmt::Debug,
391{
392 scorer: FibScorer,
394 entries: Vec<(Id, &'static [u8])>,
396 profile: FibQuantProfileV1,
398 _mmap: memmap2::Mmap,
400}
401
402#[cfg(feature = "mmap")]
403impl<Id> MmapSidecarIndex<Id>
404where
405 Id: Clone + Eq + std::fmt::Debug + DeserializeOwned,
406{
407 #[allow(unsafe_code)]
426 pub fn open(path: &Path) -> Result<Self> {
427 let file = std::fs::File::open(path).map_err(|e| {
429 FibQuantError::CorruptPayload(format!("failed to open sidecar file {:?}: {}", path, e))
430 })?;
431
432 let mmap = unsafe {
434 memmap2::Mmap::map(&file).map_err(|e| {
435 FibQuantError::CorruptPayload(format!(
436 "failed to mmap sidecar file {:?}: {}",
437 path, e
438 ))
439 })?
440 };
441 let buf = &mmap[..];
442
443 if buf.len() < MIN_HEADER_SIZE {
445 return Err(FibQuantError::CorruptPayload(format!(
446 "file too short: {} bytes (need >= {})",
447 buf.len(),
448 MIN_HEADER_SIZE
449 )));
450 }
451 if buf[0..4] != FILE_MAGIC {
452 return Err(FibQuantError::CorruptPayload(format!(
453 "bad file magic: {:?} (expected {:?})",
454 &buf[0..4],
455 FILE_MAGIC
456 )));
457 }
458 if buf[4] != FILE_VERSION {
459 return Err(FibQuantError::CorruptPayload(format!(
460 "unsupported file version {} (expected {})",
461 buf[4], FILE_VERSION
462 )));
463 }
464
465 let mut offset = MIN_HEADER_SIZE;
466
467 let (profile_bytes, next) = read_len_prefixed(buf, offset, "profile")?;
469 offset = next;
470 let profile: FibQuantProfileV1 = serde_json::from_slice(profile_bytes).map_err(|e| {
471 FibQuantError::CorruptPayload(format!("failed to deserialize profile JSON: {}", e))
472 })?;
473
474 let (gram_count, next) = read_u32_at(buf, offset, "gram count")?;
476 offset = next;
477 let gram_count = gram_count as usize;
478 let (gram_values, next) = read_f32_slice(buf, offset, gram_count, "gram table")?;
479 offset = next;
480
481 let (entry_count, next) = read_u32_at(buf, offset, "entry count")?;
483 offset = next;
484 let entry_count = entry_count as usize;
485
486 let quantizer = FibQuantizer::new(profile.clone())?;
488 let scorer = FibScorer::new(quantizer)?;
489
490 let reconstructed_gram = scorer.gram_table().values();
492 if reconstructed_gram.len() != gram_values.len() {
493 return Err(FibQuantError::CorruptPayload(format!(
494 "gram table size mismatch: file has {} values, reconstructed has {}",
495 gram_values.len(),
496 reconstructed_gram.len()
497 )));
498 }
499 for (i, (a, b)) in gram_values
500 .iter()
501 .zip(reconstructed_gram.iter())
502 .enumerate()
503 {
504 if (a - b).abs() > 1e-5 {
505 return Err(FibQuantError::CorruptPayload(format!(
506 "gram table value mismatch at index {}: file={}, reconstructed={}",
507 i, a, b
508 )));
509 }
510 }
511
512 let mut entries: Vec<(Id, &'static [u8])> = Vec::with_capacity(entry_count);
515 for i in 0..entry_count {
516 let (id_bytes, next) = read_len_prefixed(buf, offset, &format!("entry {} id", i))?;
518 offset = next;
519 let id: Id = postcard::from_bytes(id_bytes).map_err(|e| {
520 FibQuantError::CorruptPayload(format!(
521 "failed to deserialize entry {} id: {}",
522 i, e
523 ))
524 })?;
525
526 let (code_bytes, next) = read_len_prefixed(buf, offset, &format!("entry {} code", i))?;
528 offset = next;
529
530 let code_slice: &'static [u8] =
536 unsafe { std::slice::from_raw_parts(code_bytes.as_ptr(), code_bytes.len()) };
537 entries.push((id, code_slice));
538 }
539
540 drop(gram_values);
542
543 Ok(Self {
544 scorer,
545 entries,
546 profile,
547 _mmap: mmap,
548 })
549 }
550
551 pub fn scorer(&self) -> &FibScorer {
553 &self.scorer
554 }
555
556 pub fn profile(&self) -> &FibQuantProfileV1 {
558 &self.profile
559 }
560
561 pub fn len(&self) -> usize {
563 self.entries.len()
564 }
565
566 pub fn is_empty(&self) -> bool {
568 self.entries.is_empty()
569 }
570
571 pub fn search(
579 &self,
580 query: &[f32],
581 top_k: usize,
582 oversample: usize,
583 ) -> Result<Vec<ScoredCandidate<Id>>>
584 where
585 Id: Clone,
586 {
587 let mut scored: Vec<(usize, f32)> = Vec::with_capacity(self.entries.len());
588 for (idx, (_, code_bytes)) in self.entries.iter().enumerate() {
589 let code = FibCodeV1::from_compact_bytes(code_bytes, &self.profile)?;
590 let s = self.scorer.inner_product_estimate(query, &code)?;
591 scored.push((idx, s));
592 }
593 scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
595
596 let limit = top_k.saturating_mul(oversample.max(1)).min(scored.len());
597 let candidates = scored
598 .into_iter()
599 .take(limit)
600 .enumerate()
601 .map(|(rank, (idx, score))| {
602 let id = self.entries[idx].0.clone();
603 ScoredCandidate {
604 id,
605 approximate_score: score,
606 rank,
607 }
608 })
609 .collect();
610 Ok(candidates)
611 }
612}
613
614#[cfg(feature = "mmap")]
618pub fn load_mmap<Id>(path: &Path) -> Result<MmapSidecarIndex<Id>>
619where
620 Id: Clone + Eq + std::fmt::Debug + DeserializeOwned,
621{
622 MmapSidecarIndex::open(path)
623}
624
625#[cfg(test)]
630mod tests {
631 use super::*;
632 use crate::profile::FibQuantProfileV1;
633 use crate::{FibQuantizer, FibScorer};
634 use tempfile::tempdir;
635
636 fn build_test_scorer() -> Result<FibScorer> {
637 let mut profile = FibQuantProfileV1::paper_default(8, 2, 8, 7)?;
638 profile.training_samples = 128;
639 profile.lloyd_restarts = 1;
640 profile.lloyd_iterations = 2;
641 let quantizer = FibQuantizer::new(profile)?;
642 FibScorer::new(quantizer)
643 }
644
645 fn make_vectors(d: usize, count: usize) -> Vec<Vec<f32>> {
646 (0..count)
647 .map(|seed| {
648 (0..d)
649 .map(|i| (seed as f32 * 0.1 + i as f32 * 0.05 - 0.3).sin())
650 .collect()
651 })
652 .collect()
653 }
654
655 fn build_test_index() -> Result<FibSidecarIndex<u32>> {
656 let scorer = build_test_scorer()?;
657 let d = scorer.quantizer().profile().ambient_dim as usize;
658 let mut index: FibSidecarIndex<u32> = FibSidecarIndex::new(scorer);
659 let vectors = make_vectors(d, 16);
660 for (i, v) in vectors.iter().enumerate() {
661 let code = index.scorer().quantizer().encode(v)?;
662 index.add(i as u32, code);
663 }
664 Ok(index)
665 }
666
667 #[test]
668 fn save_and_load_produces_identical_search_results() -> Result<()> {
669 let dir = tempdir()
670 .map_err(|e| FibQuantError::CorruptPayload(format!("tempdir failed: {}", e)))?;
671 let path = dir.path().join("index.fbs1");
672
673 let original = build_test_index()?;
674 save_to_file(&original, &path)?;
675 let loaded: FibSidecarIndex<u32> = load_from_file(&path)?;
676
677 assert_eq!(loaded.len(), original.len());
678
679 let query: Vec<f32> = vec![0.5, -0.3, 0.8, -0.1, 0.2, -0.4, 0.7, -0.6];
680 let original_results = original.search(&query, 5, 2)?;
681 let loaded_results = loaded.search(&query, 5, 2)?;
682
683 assert_eq!(loaded_results.len(), original_results.len());
684 for (orig, load) in original_results.iter().zip(loaded_results.iter()) {
685 assert_eq!(orig.id, load.id, "id mismatch at rank {}", orig.rank);
686 assert_eq!(orig.rank, load.rank, "rank mismatch");
687 assert!(
688 (orig.approximate_score - load.approximate_score).abs() < 1e-5,
689 "score mismatch at rank {}: {} vs {}",
690 orig.rank,
691 orig.approximate_score,
692 load.approximate_score
693 );
694 }
695 Ok(())
696 }
697
698 #[test]
699 fn file_magic_detection() -> Result<()> {
700 let dir = tempdir()
701 .map_err(|e| FibQuantError::CorruptPayload(format!("tempdir failed: {}", e)))?;
702 let path = dir.path().join("magic_test.fbs1");
703
704 let index = build_test_index()?;
705 save_to_file(&index, &path)?;
706
707 let bytes = std::fs::read(&path)
708 .map_err(|e| FibQuantError::CorruptPayload(format!("read failed: {}", e)))?;
709
710 assert_eq!(&bytes[0..4], b"FBS1", "file should start with FBS1 magic");
712 assert_eq!(bytes[4], FILE_VERSION, "version byte should be 1");
713
714 let mut bad = bytes.clone();
716 bad[0] = b'X';
717 let result: Result<FibSidecarIndex<u32>> = deserialize_index(&bad);
718 assert!(result.is_err(), "bad magic should cause load failure");
719
720 Ok(())
721 }
722
723 #[test]
724 fn truncation_rejection() -> Result<()> {
725 let dir = tempdir()
726 .map_err(|e| FibQuantError::CorruptPayload(format!("tempdir failed: {}", e)))?;
727 let path = dir.path().join("trunc_test.fbs1");
728
729 let index = build_test_index()?;
730 save_to_file(&index, &path)?;
731
732 let bytes = std::fs::read(&path)
733 .map_err(|e| FibQuantError::CorruptPayload(format!("read failed: {}", e)))?;
734
735 let truncated = &bytes[..MIN_HEADER_SIZE + 2];
737 let result: Result<FibSidecarIndex<u32>> = deserialize_index(truncated);
738 assert!(result.is_err(), "truncated file should be rejected");
739
740 let result: Result<FibSidecarIndex<u32>> = deserialize_index(&[]);
742 assert!(result.is_err(), "empty file should be rejected");
743
744 Ok(())
745 }
746
747 #[test]
748 fn empty_index_roundtrips() -> Result<()> {
749 let dir = tempdir()
750 .map_err(|e| FibQuantError::CorruptPayload(format!("tempdir failed: {}", e)))?;
751 let path = dir.path().join("empty.fbs1");
752
753 let scorer = build_test_scorer()?;
754 let index: FibSidecarIndex<u32> = FibSidecarIndex::new(scorer);
755 save_to_file(&index, &path)?;
756
757 let loaded: FibSidecarIndex<u32> = load_from_file(&path)?;
758 assert!(loaded.is_empty());
759 assert_eq!(loaded.len(), 0);
760
761 let d = loaded.scorer().quantizer().profile().ambient_dim as usize;
762 let query = vec![0.0f32; d];
763 let results = loaded.search(&query, 5, 1)?;
764 assert!(results.is_empty());
765 Ok(())
766 }
767
768 #[test]
769 fn string_id_roundtrips() -> Result<()> {
770 let dir = tempdir()
771 .map_err(|e| FibQuantError::CorruptPayload(format!("tempdir failed: {}", e)))?;
772 let path = dir.path().join("string_ids.fbs1");
773
774 let scorer = build_test_scorer()?;
775 let d = scorer.quantizer().profile().ambient_dim as usize;
776 let mut index: FibSidecarIndex<String> = FibSidecarIndex::new(scorer);
777 let vectors = make_vectors(d, 8);
778 for (i, v) in vectors.iter().enumerate() {
779 let code = index.scorer().quantizer().encode(v)?;
780 index.add(format!("vec_{}", i), code);
781 }
782 save_to_file(&index, &path)?;
783
784 let loaded: FibSidecarIndex<String> = load_from_file(&path)?;
785 assert_eq!(loaded.len(), 8);
786
787 let query: Vec<f32> = vec![0.5, -0.3, 0.8, -0.1, 0.2, -0.4, 0.7, -0.6];
788 let results = loaded.search(&query, 4, 1)?;
789 assert_eq!(results.len(), 4);
790 for r in &results {
791 assert!(r.id.starts_with("vec_"), "id should be string vec_N");
792 }
793 Ok(())
794 }
795
796 #[cfg(feature = "mmap")]
797 #[test]
798 fn mmap_search_matches_loaded_search() -> Result<()> {
799 let dir = tempdir()
800 .map_err(|e| FibQuantError::CorruptPayload(format!("tempdir failed: {}", e)))?;
801 let path = dir.path().join("mmap_test.fbs1");
802
803 let original = build_test_index()?;
804 save_to_file(&original, &path)?;
805
806 let loaded: FibSidecarIndex<u32> = load_from_file(&path)?;
807 let mmap_index: MmapSidecarIndex<u32> = load_mmap(&path)?;
808
809 assert_eq!(mmap_index.len(), loaded.len());
810
811 let query: Vec<f32> = vec![0.5, -0.3, 0.8, -0.1, 0.2, -0.4, 0.7, -0.6];
812 let loaded_results = loaded.search(&query, 5, 2)?;
813 let mmap_results = mmap_index.search(&query, 5, 2)?;
814
815 assert_eq!(loaded_results.len(), mmap_results.len());
816 for (a, b) in loaded_results.iter().zip(mmap_results.iter()) {
817 assert_eq!(a.id, b.id, "id mismatch at rank {}", a.rank);
818 assert_eq!(a.rank, b.rank);
819 assert!(
820 (a.approximate_score - b.approximate_score).abs() < 1e-5,
821 "score mismatch at rank {}: {} vs {}",
822 a.rank,
823 a.approximate_score,
824 b.approximate_score
825 );
826 }
827 Ok(())
828 }
829
830 #[cfg(feature = "mmap")]
831 #[test]
832 fn mmap_empty_index_works() -> Result<()> {
833 let dir = tempdir()
834 .map_err(|e| FibQuantError::CorruptPayload(format!("tempdir failed: {}", e)))?;
835 let path = dir.path().join("mmap_empty.fbs1");
836
837 let scorer = build_test_scorer()?;
838 let index: FibSidecarIndex<u32> = FibSidecarIndex::new(scorer);
839 save_to_file(&index, &path)?;
840
841 let mmap_index: MmapSidecarIndex<u32> = load_mmap(&path)?;
842 assert!(mmap_index.is_empty());
843
844 let d = mmap_index.profile().ambient_dim as usize;
845 let query = vec![0.0f32; d];
846 let results = mmap_index.search(&query, 5, 1)?;
847 assert!(results.is_empty());
848 Ok(())
849 }
850}