1mod number;
86pub use number::Number;
87
88mod version;
89pub use version::Version;
90
91mod value;
92pub use value::{Handle, Keys, Record, Value, Versioned};
93
94pub mod load;
95pub mod save;
96
97mod backend;
98pub use backend::memory::{MemoryContext, MemorySaveContext};
99
100const _: () = assert!(
106 usize::BITS == 64,
107 "diskann-record requires a 64-bit target: usize/isize MUST be 64 bits wide !!",
108);
109
110#[doc(hidden)]
116pub const fn is_reserved(s: &str) -> bool {
117 if let Some(first) = s.as_bytes().first()
118 && *first == b"$"[0]
119 {
120 true
121 } else {
122 false
123 }
124}
125
126#[cfg(all(test, feature = "disk"))]
131mod tests {
132 use super::*;
133
134 use std::io::{Read, Write};
135 use std::path::{Path, PathBuf};
136
137 #[derive(Debug, PartialEq)]
138 struct Test {
139 x: String,
140 y: f32,
141 enabled: bool,
142 inner: Inner,
143 vector: Vec<u8>,
145 nickname: Option<String>,
146 absent: Option<String>,
147 }
148
149 #[derive(Debug, PartialEq)]
150 struct Inner {
151 z: usize,
152 w: Vec<i8>,
153 flags: Vec<bool>,
154 maybe_count: Option<u32>,
155 maybe_missing: Option<u32>,
156 sparse: Vec<Option<i32>>,
157 }
158
159 impl save::Save for Inner {
160 const VERSION: Version = Version::new(0, 0);
161 fn save(&self, context: save::Context<'_>) -> save::Result<save::Record<'_>> {
162 Ok(save_fields!(
163 self,
164 context,
165 [z, w, flags, maybe_count, maybe_missing, sparse]
166 ))
167 }
168 }
169
170 impl save::Save for Test {
171 const VERSION: Version = Version::new(0, 0);
172 fn save(&self, context: save::Context<'_>) -> save::Result<save::Record<'_>> {
173 let mut io = context.write(Some("auxiliary.bin"))?;
177 io.write_all(&self.vector).map_err(save::Error::new)?;
178
179 let mut record = save_fields!(self, context, [x, y, enabled, inner, nickname, absent]);
180 record.insert("vector", io.finish()?)?;
181 Ok(record)
182 }
183 }
184
185 impl load::Load<'_> for Test {
186 const VERSION: Version = Version::new(0, 0);
187 fn load(object: load::Object<'_>) -> load::Result<Self> {
188 load_fields!(
189 object,
190 [
191 x,
192 y,
193 enabled,
194 inner,
195 nickname: Option<String>,
196 absent: Option<String>,
197 vector: save::Handle,
198 ]
199 );
200
201 let mut io = object.read(&vector)?;
202 let mut vector = Vec::new();
203 io.read_to_end(&mut vector).unwrap();
204
205 Ok(Self {
206 x,
207 y,
208 enabled,
209 inner,
210 vector,
211 nickname,
212 absent,
213 })
214 }
215
216 fn load_legacy(_object: load::Object<'_>) -> load::Result<Self> {
217 panic!("nope!");
218 }
219 }
220
221 impl load::Load<'_> for Inner {
222 const VERSION: Version = Version::new(0, 0);
223 fn load(object: load::Object<'_>) -> load::Result<Self> {
224 load_fields!(
225 object,
226 [
227 z,
228 w,
229 flags,
230 maybe_count: Option<u32>,
231 maybe_missing: Option<u32>,
232 sparse: Vec<Option<i32>>,
233 ]
234 );
235 Ok(Self {
236 z,
237 w,
238 flags,
239 maybe_count,
240 maybe_missing,
241 sparse,
242 })
243 }
244
245 fn load_legacy(_object: load::Object<'_>) -> load::Result<Self> {
246 panic!("nope!");
247 }
248 }
249
250 #[test]
251 fn round_trip_uses_isolated_temp_dir() -> anyhow::Result<()> {
252 let inner = Inner {
253 z: 10,
254 w: vec![-1, -2, -3],
255 flags: vec![true, false, true],
256 maybe_count: Some(42),
257 maybe_missing: None,
258 sparse: vec![Some(1), None, Some(-3), None],
259 };
260
261 let t = Test {
262 x: "hello".into(),
263 y: 5.0,
264 enabled: true,
265 inner,
266 vector: vec![0, 1, 2, 3, 4, 5],
267 nickname: Some("friend".into()),
268 absent: None,
269 };
270
271 let temp_dir = tempfile::tempdir()?;
274 let dir = temp_dir.path();
275 let metadata = dir.join("metadata.json");
276
277 save::save_to_disk(&t, dir, &metadata)?;
278 let we_are_back: Test = load::load_from_disk(&metadata, dir)?;
279
280 assert_eq!(t, we_are_back);
281 Ok(())
282 }
283
284 #[derive(Debug, PartialEq)]
289 enum Metric {
290 L2,
291 Cosine,
292 Weighted { weights: Vec<f32> },
293 }
294
295 impl save::Save for Metric {
296 const VERSION: Version = Version::new(0, 0);
297 fn save(&self, context: save::Context<'_>) -> save::Result<save::Record<'_>> {
298 let mut record = save::Record::empty();
299 match self {
300 Self::L2 => {
301 record.insert("L2", save::Value::Null)?;
302 }
303 Self::Cosine => {
304 record.insert("Cosine", save::Value::Null)?;
305 }
306 Self::Weighted { weights } => {
307 let payload = save_fields!(context, [weights]).into_value(Self::VERSION);
308 record.insert("Weighted", payload)?;
309 }
310 }
311 Ok(record)
312 }
313 }
314
315 impl load::Load<'_> for Metric {
316 const VERSION: Version = Version::new(0, 0);
317 fn load(object: load::Object<'_>) -> load::Result<Self> {
318 match object.single_key()? {
319 "L2" => Ok(Self::L2),
320 "Cosine" => Ok(Self::Cosine),
321 "Weighted" => {
322 let inner = object
323 .child("Weighted")?
324 .as_object()
325 .ok_or(load::error::Kind::TypeMismatch)?;
326 load_fields!(inner, [weights: Vec<f32>]);
327 Ok(Self::Weighted { weights })
328 }
329 _ => Err(load::error::Kind::UnknownVariant.into()),
330 }
331 }
332 fn load_legacy(_object: load::Object<'_>) -> load::Result<Self> {
333 Err(load::error::Kind::UnknownVersion.into())
334 }
335 }
336
337 #[derive(Debug, PartialEq)]
338 struct MetricBag {
339 primary: Metric,
340 alternatives: Vec<Metric>,
341 }
342
343 impl save::Save for MetricBag {
344 const VERSION: Version = Version::new(0, 0);
345 fn save(&self, context: save::Context<'_>) -> save::Result<save::Record<'_>> {
346 Ok(save_fields!(self, context, [primary, alternatives]))
347 }
348 }
349
350 impl load::Load<'_> for MetricBag {
351 const VERSION: Version = Version::new(0, 0);
352 fn load(object: load::Object<'_>) -> load::Result<Self> {
353 load_fields!(object, [primary: Metric, alternatives: Vec<Metric>]);
354 Ok(Self {
355 primary,
356 alternatives,
357 })
358 }
359 fn load_legacy(_object: load::Object<'_>) -> load::Result<Self> {
360 panic!("nope!");
361 }
362 }
363
364 #[test]
365 fn enum_round_trip_through_disk() -> anyhow::Result<()> {
366 let bag = MetricBag {
367 primary: Metric::Weighted {
368 weights: vec![0.25, 0.5, 0.25],
369 },
370 alternatives: vec![
371 Metric::L2,
372 Metric::Cosine,
373 Metric::Weighted { weights: vec![1.0] },
374 ],
375 };
376
377 let temp_dir = tempfile::tempdir()?;
378 let dir = temp_dir.path();
379 let metadata = dir.join("metadata.json");
380
381 save::save_to_disk(&bag, dir, &metadata)?;
382 let restored: MetricBag = load::load_from_disk(&metadata, dir)?;
383
384 assert_eq!(bag, restored);
385 Ok(())
386 }
387
388 #[derive(Debug, PartialEq)]
389 struct StructShape {
390 x: i32,
391 }
392
393 impl save::Save for StructShape {
394 const VERSION: Version = Version::new(0, 0);
395 fn save(&self, context: save::Context<'_>) -> save::Result<save::Record<'_>> {
396 Ok(save_fields!(self, context, [x]))
397 }
398 }
399
400 impl load::Load<'_> for StructShape {
401 const VERSION: Version = Version::new(0, 0);
402 fn load(object: load::Object<'_>) -> load::Result<Self> {
403 load_fields!(object, [x: i32]);
404 Ok(Self { x })
405 }
406 fn load_legacy(_object: load::Object<'_>) -> load::Result<Self> {
407 panic!("nope!");
408 }
409 }
410
411 #[derive(Debug, PartialEq)]
412 enum EnumShape {
413 Only { x: i32 },
414 }
415
416 impl save::Save for EnumShape {
417 const VERSION: Version = Version::new(0, 0);
418 fn save(&self, context: save::Context<'_>) -> save::Result<save::Record<'_>> {
419 let mut record = save::Record::empty();
420 match self {
421 Self::Only { x } => {
422 let payload = save_fields!(context, [x]).into_value(Self::VERSION);
423 record.insert("Only", payload)?;
424 }
425 }
426 Ok(record)
427 }
428 }
429
430 impl load::Load<'_> for EnumShape {
431 const VERSION: Version = Version::new(0, 0);
432 fn load(object: load::Object<'_>) -> load::Result<Self> {
433 match object.single_key()? {
434 "Only" => {
435 let inner = object
436 .child("Only")?
437 .as_object()
438 .ok_or(load::error::Kind::TypeMismatch)?;
439 load_fields!(inner, [x: i32]);
440 Ok(Self::Only { x })
441 }
442 _ => Err(load::error::Kind::UnknownVariant.into()),
443 }
444 }
445 fn load_legacy(_object: load::Object<'_>) -> load::Result<Self> {
446 panic!("nope!");
447 }
448 }
449
450 #[test]
451 fn loading_enum_as_struct_is_rejected() -> anyhow::Result<()> {
452 let value = EnumShape::Only { x: 7 };
456 let temp_dir = tempfile::tempdir()?;
457 let dir = temp_dir.path();
458 let metadata = dir.join("metadata.json");
459
460 save::save_to_disk(&value, dir, &metadata)?;
461 let err = load::load_from_disk::<StructShape>(&metadata, dir)
462 .expect_err("loading enum data into a struct shape should fail");
463 let msg = format!("{err}");
464 assert!(
465 msg.contains("missing field"),
466 "expected MissingField error, got: {msg}"
467 );
468 Ok(())
469 }
470
471 #[test]
472 fn loading_struct_as_enum_is_rejected() -> anyhow::Result<()> {
473 let value = StructShape { x: 7 };
476 let temp_dir = tempfile::tempdir()?;
477 let dir = temp_dir.path();
478 let metadata = dir.join("metadata.json");
479
480 save::save_to_disk(&value, dir, &metadata)?;
481 let err = load::load_from_disk::<EnumShape>(&metadata, dir)
482 .expect_err("loading struct data into an enum shape should fail");
483 let msg = format!("{err}");
484 assert!(
485 msg.contains("unknown variant"),
486 "expected UnknownVariant error, got: {msg}"
487 );
488 Ok(())
489 }
490
491 #[derive(Debug, PartialEq)]
497 struct Upgraded {
498 count: u32,
500 scaled: u32,
502 }
503
504 impl save::Save for Upgraded {
505 const VERSION: Version = Version::new(0, 0);
507 fn save(&self, context: save::Context<'_>) -> save::Result<save::Record<'_>> {
508 Ok(save_fields!(self, context, [count]))
509 }
510 }
511
512 impl load::Load<'_> for Upgraded {
513 const VERSION: Version = Version::new(1, 0);
515 fn load(_object: load::Object<'_>) -> load::Result<Self> {
516 panic!("matching-version load must not run for a legacy record");
517 }
518 fn load_legacy(object: load::Object<'_>) -> load::Result<Self> {
519 load_fields!(object, [count: u32]);
521 Ok(Self {
522 count, scaled: count * 10, })
525 }
526 }
527
528 #[test]
529 fn legacy_record_dispatches_to_load_legacy() -> anyhow::Result<()> {
530 let temp_dir = tempfile::tempdir()?;
533 let dir = temp_dir.path();
534 let metadata = dir.join("metadata.json");
535
536 save::save_to_disk(
537 &Upgraded {
538 count: 4,
539 scaled: 0,
540 },
541 dir,
542 &metadata,
543 )?;
544 let restored: Upgraded = load::load_from_disk(&metadata, dir)?;
545
546 assert_eq!(
547 restored,
548 Upgraded {
549 count: 4,
550 scaled: 40
551 }
552 );
553 Ok(())
554 }
555
556 #[derive(Debug, PartialEq)]
559 struct NoUpgrade {
560 value: i32,
561 }
562
563 impl save::Save for NoUpgrade {
564 const VERSION: Version = Version::new(0, 0);
565 fn save(&self, context: save::Context<'_>) -> save::Result<save::Record<'_>> {
566 Ok(save_fields!(self, context, [value]))
567 }
568 }
569
570 impl load::Load<'_> for NoUpgrade {
571 const VERSION: Version = Version::new(2, 0);
572 fn load(_object: load::Object<'_>) -> load::Result<Self> {
573 panic!("matching-version load must not run for a legacy record");
574 }
575 fn load_legacy(_object: load::Object<'_>) -> load::Result<Self> {
576 if _object.version().major < 1 {
578 return Err(load::error::Kind::UnknownVersion.into());
579 }
580 panic!("should not reach this point");
581 }
582 }
583
584 #[test]
585 fn legacy_record_without_upgrade_path_is_rejected() -> anyhow::Result<()> {
586 let temp_dir = tempfile::tempdir()?;
587 let dir = temp_dir.path();
588 let metadata = dir.join("metadata.json");
589
590 save::save_to_disk(&NoUpgrade { value: 7 }, dir, &metadata)?;
591 let err = load::load_from_disk::<NoUpgrade>(&metadata, dir)
592 .expect_err("a legacy record with no upgrade path must fail");
593 let msg = format!("{err}");
594 assert!(
595 msg.contains("unknown version"),
596 "expected UnknownVersion error, got: {msg}"
597 );
598 Ok(())
599 }
600
601 #[derive(Debug, PartialEq)]
608 struct Primitives {
609 a: u16,
610 b: u32,
611 c: u64,
612 d: i16,
613 e: i32,
614 f: i64,
615 g: f64,
616 nz32: std::num::NonZeroU32,
617 nz64: std::num::NonZeroU64,
618 nzsize: std::num::NonZeroUsize,
619 }
620
621 impl save::Save for Primitives {
622 const VERSION: Version = Version::new(0, 0);
623 fn save(&self, context: save::Context<'_>) -> save::Result<save::Record<'_>> {
624 Ok(save_fields!(
625 self,
626 context,
627 [a, b, c, d, e, f, g, nz32, nz64, nzsize]
628 ))
629 }
630 }
631
632 impl load::Load<'_> for Primitives {
633 const VERSION: Version = Version::new(0, 0);
634 fn load(object: load::Object<'_>) -> load::Result<Self> {
635 load_fields!(
636 object,
637 [
638 a: u16,
639 b: u32,
640 c: u64,
641 d: i16,
642 e: i32,
643 f: i64,
644 g: f64,
645 nz32: std::num::NonZeroU32,
646 nz64: std::num::NonZeroU64,
647 nzsize: std::num::NonZeroUsize,
648 ]
649 );
650 Ok(Self {
651 a,
652 b,
653 c,
654 d,
655 e,
656 f,
657 g,
658 nz32,
659 nz64,
660 nzsize,
661 })
662 }
663 fn load_legacy(_object: load::Object<'_>) -> load::Result<Self> {
664 panic!("nope!");
665 }
666 }
667
668 #[test]
669 fn builtin_primitives_round_trip() -> anyhow::Result<()> {
670 let value = Primitives {
671 a: 4242,
672 b: 4_000_000_000,
673 c: 1 << 40,
674 d: -12345,
675 e: -2_000_000_000,
676 f: -(1 << 40),
677 g: -2.5e-9,
678 nz32: std::num::NonZeroU32::new(7).unwrap(),
679 nz64: std::num::NonZeroU64::new(1 << 50).unwrap(),
680 nzsize: std::num::NonZeroUsize::new(99).unwrap(),
681 };
682
683 let temp_dir = tempfile::tempdir()?;
684 let dir = temp_dir.path();
685 let metadata = dir.join("metadata.json");
686
687 save::save_to_disk(&value, dir, &metadata)?;
688 let restored: Primitives = load::load_from_disk(&metadata, dir)?;
689
690 assert_eq!(value, restored);
691 Ok(())
692 }
693
694 #[test]
695 fn nonzero_rejects_zero_on_load() -> anyhow::Result<()> {
696 #[derive(Debug)]
699 struct NzHolder {
700 _nz: std::num::NonZeroU32,
701 }
702
703 impl load::Load<'_> for NzHolder {
704 const VERSION: Version = Version::new(0, 0);
705 fn load(object: load::Object<'_>) -> load::Result<Self> {
706 load_fields!(object, [nz: std::num::NonZeroU32]);
707 Ok(Self { _nz: nz })
708 }
709 fn load_legacy(_object: load::Object<'_>) -> load::Result<Self> {
710 panic!("nope!");
711 }
712 }
713
714 let temp_dir = tempfile::tempdir()?;
715 let dir = temp_dir.path();
716 let metadata = dir.join("metadata.json");
717 let manifest = serde_json::json!({
718 "files": [],
719 "value": { "$version": "0.0", "nz": 0 },
720 });
721 std::fs::write(&metadata, serde_json::to_vec(&manifest)?)?;
722
723 let err = load::load_from_disk::<NzHolder>(&metadata, dir)
724 .expect_err("zero stored in a NonZero field must be rejected");
725 let msg = format!("{err}");
726 assert!(
727 msg.contains("number out of range"),
728 "expected NumberOutOfRange error, got: {msg}"
729 );
730 Ok(())
731 }
732
733 #[derive(Debug)]
741 struct HandleOnly {
742 _blob: Vec<u8>,
743 }
744
745 impl load::Load<'_> for HandleOnly {
746 const VERSION: Version = Version::new(0, 0);
747 fn load(object: load::Object<'_>) -> load::Result<Self> {
748 load_fields!(object, [blob: save::Handle]);
749 let mut io = object.read(&blob)?;
750 let mut buf = Vec::new();
751 io.read_to_end(&mut buf).map_err(load::Error::new)?;
752 Ok(Self { _blob: buf })
753 }
754 fn load_legacy(_object: load::Object<'_>) -> load::Result<Self> {
755 panic!("nope!");
756 }
757 }
758
759 fn write_handle_manifest(dir: &Path, handle_target: &str) -> std::io::Result<PathBuf> {
762 let manifest = serde_json::json!({
763 "files": [handle_target],
767 "value": {
768 "$version": "0.0",
769 "blob": { "$handle": handle_target },
770 },
771 });
772 let metadata = dir.join("metadata.json");
773 std::fs::write(&metadata, serde_json::to_vec(&manifest)?)?;
774 Ok(metadata)
775 }
776
777 #[test]
778 fn handle_with_parent_traversal_is_rejected() -> anyhow::Result<()> {
779 let temp_dir = tempfile::tempdir()?;
780 let dir = temp_dir.path();
781 let metadata = write_handle_manifest(dir, "../escape.bin")?;
782
783 let err = load::load_from_disk::<HandleOnly>(&metadata, dir)
784 .expect_err("handle escaping the manifest directory must be rejected");
785 let msg = format!("{err}");
786 assert!(
787 msg.contains("escapes the manifest directory"),
788 "expected manifest-escape rejection, got: {msg}"
789 );
790 Ok(())
791 }
792
793 #[test]
794 fn handle_with_absolute_path_is_rejected() -> anyhow::Result<()> {
795 let temp_dir = tempfile::tempdir()?;
796 let dir = temp_dir.path();
797 let absolute = if cfg!(windows) {
800 "C:\\Windows\\System32\\drivers\\etc\\hosts"
801 } else {
802 "/etc/passwd"
803 };
804 let metadata = write_handle_manifest(dir, absolute)?;
805
806 let err = load::load_from_disk::<HandleOnly>(&metadata, dir)
807 .expect_err("absolute-path handle must be rejected");
808 let msg = format!("{err}");
809 assert!(
810 msg.contains("escapes the manifest directory"),
811 "expected manifest-escape rejection, got: {msg}"
812 );
813 Ok(())
814 }
815}