Skip to main content

diskann_record/save/
mod.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6//! # Saving Records to Disk
7//!
8//! This module provides the writer-side of the framework. User types implement [`Save`]
9//! (or, for primitive-like leaves, [`Saveable`]) and obtain a [`Context`] from which they
10//! request side-car artifact writers and assemble a [`Record`] of named fields.
11//!
12//! The generic entry point is `save`; `save_to_disk` (available under the `disk`
13//! feature) is the disk-backed convenience wrapper that serializes a value into a
14//! caller-chosen directory plus a manifest path.
15//!
16//! # Building Records
17//!
18//! The [`save_fields!`](crate::save_fields) macro is the idiomatic way to build a record
19//! from a struct or destructured enum variant. It handles per-field error context and
20//! invokes [`Saveable::save`] on each value.
21//!
22//! # Side-Car Artifacts
23//!
24//! Binary blobs (e.g. vector buffers) are written to side-car files via
25//! [`Context::write`], which returns a [`Writer`]. The handle returned by
26//! [`Writer::finish`](context::Writer::finish) can be embedded into the record as a
27//! [`Handle`]; it serializes as a `$handle` reference and is rehydrated on the load side.
28pub use crate::value::{Handle, Keys, Record, Value, Versioned};
29
30mod context;
31pub use context::{Context, Writer};
32pub(crate) use context::{SaveContext, WriterInner, delegate_write_and_seek};
33
34mod error;
35pub use error::{Error, Result};
36
37use crate::Version;
38
39/// Serialize `x` into `context`, returning the context's committed output.
40///
41/// This is the generic entry point: it threads a [`Context`] borrowing `context` through
42/// the value's [`Saveable::save`] impl, then commits the resulting manifest via
43/// [`SaveContext::finish`]. `save_to_disk` is the disk-backed convenience wrapper
44/// available under the `disk` feature.
45///
46/// # Errors
47///
48/// Returns [`Error`] if a user impl returns an error or if the context fails to commit
49/// the manifest.
50pub(crate) fn save<T, C>(x: &T, context: C) -> Result<C::Output>
51where
52    T: Saveable,
53    C: SaveContext,
54{
55    let value = x.save(Context::new(&context))?;
56    context.finish(value)
57}
58
59/// Serialize `x` to disk.
60///
61/// The manifest (a JSON document) is written atomically to `metadata`; any side-car
62/// artifacts the type's [`Save::save`] impl creates via [`Context::write`] are written
63/// into `dir`.
64///
65/// # Errors
66///
67/// Returns [`Error`] if the directory cannot be written to, if the manifest cannot be
68/// serialized, or if a user impl returns an error.
69#[cfg(feature = "disk")]
70pub fn save_to_disk<T>(
71    x: &T,
72    dir: impl AsRef<std::path::Path>,
73    metadata: impl AsRef<std::path::Path>,
74) -> Result<()>
75where
76    T: Saveable,
77{
78    let context =
79        crate::backend::disk::DiskSaveContext::new(dir.as_ref().into(), metadata.as_ref().into())?;
80    save(x, context)
81}
82
83/// Implemented by user types that map to a versioned [`Record`].
84///
85/// This is the primary trait for structured user types. A [`Save`] impl describes the
86/// versioned schema of `Self`: its associated [`VERSION`](Self::VERSION) is attached to
87/// the [`Record`] produced by [`Self::save`](Self::save).
88///
89/// # Enums
90///
91/// Enum types are encoded by returning a [`Record`] with a single user key whose name
92/// is the variant tag and whose value is the variant's payload (frequently
93/// [`Value::Null`] for unit variants). See the crate-level docs for a worked example.
94pub trait Save {
95    /// The schema version attached to records produced by this impl.
96    ///
97    /// Loaders compare this against the version stored in the manifest to decide
98    /// between [`Load::load`](crate::load::Load::load) and
99    /// [`Load::load_legacy`](crate::load::Load::load_legacy).
100    const VERSION: Version;
101
102    /// Serialize `self` into a [`Record`].
103    ///
104    /// Use the supplied [`Context`] to request side-car artifact writers. Use the
105    /// [`save_fields!`](crate::save_fields) macro to populate the record.
106    fn save(&self, context: Context<'_>) -> Result<Record<'_>>;
107}
108
109/// Implemented by any value that can be written into a [`Value`].
110///
111/// This is the bottom of the trait hierarchy and is implemented for:
112///
113/// * Primitive numeric types (signed, unsigned, floats, `NonZero*`).
114/// * [`bool`], [`str`], [`String`], and [`Handle`].
115/// * [`Option<T>`] (serializes `None` as [`Value::Null`]).
116/// * `&[T]` and [`Vec<T>`] (serialize as [`Value::Array`]).
117/// * Any `T: Save` (wraps the produced record in [`Value::Object`] with the type's
118///   [`Save::VERSION`]).
119///
120/// Most user types should implement [`Save`] (which gets a [`Saveable`] impl for free
121/// via the blanket below) rather than [`Saveable`] directly.
122pub trait Saveable {
123    /// Serialize `self` into a [`Value`].
124    fn save(&self, context: Context<'_>) -> Result<Value<'_>>;
125}
126
127impl<T> Saveable for T
128where
129    T: Save,
130{
131    fn save(&self, context: Context<'_>) -> Result<Value<'_>> {
132        let record = self.save(context)?;
133        Ok(record.into_value(T::VERSION))
134    }
135}
136
137//////////////////
138// Random Stuff //
139//////////////////
140
141impl<T> Saveable for [T]
142where
143    T: Saveable,
144{
145    fn save(&self, context: Context<'_>) -> Result<Value<'_>> {
146        let values: Result<Vec<_>> = self.iter().map(|t| t.save(context.clone())).collect();
147        values.map(Value::Array)
148    }
149}
150
151impl<T> Saveable for Vec<T>
152where
153    T: Saveable,
154{
155    fn save(&self, context: Context<'_>) -> Result<Value<'_>> {
156        self.as_slice().save(context)
157    }
158}
159
160impl Saveable for str {
161    fn save(&self, _: Context<'_>) -> Result<Value<'_>> {
162        Ok(Value::String(self.into()))
163    }
164}
165
166impl Saveable for String {
167    fn save(&self, _: Context<'_>) -> Result<Value<'_>> {
168        Ok(Value::String(self.as_str().into()))
169    }
170}
171
172impl Saveable for Handle {
173    fn save(&self, _: Context<'_>) -> Result<Value<'_>> {
174        Ok(Value::Handle(self.clone()))
175    }
176}
177
178impl Saveable for bool {
179    fn save(&self, _: Context<'_>) -> Result<Value<'_>> {
180        Ok(Value::Bool(*self))
181    }
182}
183
184impl<T> Saveable for Option<T>
185where
186    T: Saveable,
187{
188    fn save(&self, context: Context<'_>) -> Result<Value<'_>> {
189        match self {
190            None => Ok(Value::Null),
191            Some(t) => t.save(context),
192        }
193    }
194}
195
196macro_rules! save_number {
197    ($T:ty) => {
198        impl Saveable for $T {
199            fn save(&self, _: Context<'_>) -> Result<Value<'_>> {
200                Ok(Value::Number((*self).into()))
201            }
202        }
203    };
204    ($($Ts:ty),+ $(,)?) => {
205        $(save_number!($Ts);)+
206    }
207}
208
209save_number!(usize, u64, u32, u16, u8, i64, i32, i16, i8, f32, f64);
210
211// NonZero* primitives serialize as their inner numeric type. Loaders reject zero.
212macro_rules! save_nonzero {
213    ($T:ty) => {
214        impl Saveable for $T {
215            fn save(&self, _: Context<'_>) -> Result<Value<'_>> {
216                Ok(Value::Number(self.get().into()))
217            }
218        }
219    };
220    ($($Ts:ty),+ $(,)?) => {
221        $(save_nonzero!($Ts);)+
222    }
223}
224
225save_nonzero!(
226    std::num::NonZeroU32,
227    std::num::NonZeroU64,
228    std::num::NonZeroUsize
229);
230
231#[derive(Debug, Clone, Copy)]
232#[doc(hidden)]
233pub struct Serializing(pub &'static str);
234
235impl std::fmt::Display for Serializing {
236    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
237        write!(f, "while serializing field \"{}\"", self.0)
238    }
239}
240
241/// Build a [`Record`] from a list of fields.
242///
243/// Two forms are supported:
244///
245/// * `save_fields!(self, context, [a, b, c])` reads each field as `self.a`,
246///   `self.b`, etc. Use this from `Save::save` for plain structs.
247/// * `save_fields!(context, [a, b, c])` reads each field from a local binding of
248///   the same name. Use this inside enum match arms where the variant's payload
249///   has already been destructured into local bindings. Those bindings are
250///   assumed to be references (which is automatic when matching against `&self`);
251///   for an owned local, take a reference explicitly first.
252#[macro_export]
253macro_rules! save_fields {
254    ($me:ident, $context:ident, [$($field:ident),+ $(,)?]) => {{
255        $crate::save::Record::from_iter(
256            [
257                $(
258                    (
259                        ::std::borrow::Cow::Borrowed(stringify!($field)),
260                        <_ as $crate::save::Saveable>::save(
261                            &$me.$field,
262                            $context.clone()
263                        ).map_err(|err| {
264                            err.context($crate::save::Serializing(stringify!($field)))
265                        })?
266                    ),
267                )+
268            ]
269        )
270    }};
271    ($context:ident, [$($field:ident),+ $(,)?]) => {{
272        $crate::save::Record::from_iter(
273            [
274                $(
275                    (
276                        ::std::borrow::Cow::Borrowed(stringify!($field)),
277                        <_ as $crate::save::Saveable>::save(
278                            $field,
279                            $context.clone()
280                        ).map_err(|err| {
281                            err.context($crate::save::Serializing(stringify!($field)))
282                        })?
283                    ),
284                )+
285            ]
286        )
287    }};
288}