Skip to main content

rill_ml/
persistence.rs

1//! Model state persistence via a versioned [`Snapshot`] envelope.
2//!
3//! Only available when the `serde` feature is enabled.
4
5use crate::error::RillError;
6
7/// The current snapshot format version.
8pub const SNAPSHOT_FORMAT_VERSION: u32 = 1;
9
10/// A versioned envelope around a serializable model state.
11///
12/// Versioning is centralized here so individual models do not need to
13/// duplicate format-version fields.
14///
15/// # Examples
16///
17/// ```
18/// # #[cfg(feature = "serde")] {
19/// use rill_ml::persistence::Snapshot;
20/// use rill_ml::stats::Mean;
21/// use rill_ml::OnlineStatistic;
22///
23/// let mut mean = Mean::new();
24/// mean.update(1.0).unwrap();
25/// mean.update(2.0).unwrap();
26///
27/// let snap = Snapshot::new(mean);
28/// let json = serde_json::to_string(&snap).unwrap();
29/// let restored: Snapshot<Mean> = serde_json::from_str(&json).unwrap();
30/// let m = restored.into_model().unwrap();
31/// assert!((m.value() - 1.5).abs() < 1e-12);
32/// # }
33/// ```
34#[derive(Debug, Clone)]
35#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
36pub struct Snapshot<T> {
37    /// The format version of this snapshot.
38    pub format_version: u32,
39    /// The model state.
40    pub model: T,
41}
42
43impl<T> Snapshot<T> {
44    /// Wrap a model in a new snapshot with the current format version.
45    pub fn new(model: T) -> Self {
46        Self {
47            format_version: SNAPSHOT_FORMAT_VERSION,
48            model,
49        }
50    }
51
52    /// Consume the snapshot and return the model, verifying the format version.
53    ///
54    /// Returns [`RillError::IncompatibleStateVersion`] if the version does not
55    /// match [`SNAPSHOT_FORMAT_VERSION`].
56    pub fn into_model(self) -> Result<T, RillError> {
57        if self.format_version != SNAPSHOT_FORMAT_VERSION {
58            return Err(RillError::IncompatibleStateVersion {
59                expected: SNAPSHOT_FORMAT_VERSION,
60                actual: self.format_version,
61            });
62        }
63        Ok(self.model)
64    }
65
66    /// Consume the snapshot, verify its format version, and run an
67    /// application-provided model-state validator before returning the model.
68    ///
69    /// The snapshot envelope can only validate its own version field because
70    /// `T` may be an application type. Use this method at trust boundaries to
71    /// enforce model-specific invariants before activating restored state.
72    ///
73    /// # Errors
74    ///
75    /// Returns [`RillError::IncompatibleStateVersion`] for a version mismatch,
76    /// or propagates the validator's error.
77    pub fn into_model_with_validation<F>(self, validate: F) -> Result<T, RillError>
78    where
79        F: FnOnce(&T) -> Result<(), RillError>,
80    {
81        let model = self.into_model()?;
82        validate(&model)?;
83        Ok(model)
84    }
85}
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90    use crate::stats::Mean;
91    #[cfg(feature = "serde")]
92    use crate::traits::OnlineStatistic;
93
94    #[cfg(feature = "serde")]
95    #[test]
96    fn snapshot_roundtrip() {
97        let mut mean = Mean::new();
98        mean.update(1.0).unwrap();
99        mean.update(2.0).unwrap();
100        let snap = Snapshot::new(mean);
101        let json = serde_json::to_string(&snap).unwrap();
102        let restored: Snapshot<Mean> = serde_json::from_str(&json).unwrap();
103        let m = restored.into_model().unwrap();
104        assert!((m.value() - 1.5).abs() < 1e-12);
105        assert_eq!(m.count(), 2);
106    }
107
108    #[test]
109    fn incompatible_version_rejected() {
110        let snap = Snapshot {
111            format_version: 999,
112            model: Mean::new(),
113        };
114        assert!(snap.into_model().is_err());
115    }
116
117    #[test]
118    fn application_validation_runs_before_activation() {
119        let snap = Snapshot::new(Mean::new());
120        let result = snap.into_model_with_validation(|_| {
121            Err(RillError::InvalidState(
122                "application check failed".to_owned(),
123            ))
124        });
125        assert!(matches!(result, Err(RillError::InvalidState(_))));
126    }
127}