sim_incremental_core/

snapshot.rs

//! Bounded graph snapshot records.

use std::collections::{BTreeMap, BTreeSet};

use crate::{
    BudgetKind, FingerprintValue, IncrementalEngine, IncrementalError, Observation,
    ObservationKind, QueryResult, Revision, SnapshotBudgets, SnapshotError, state::Node,
};

/// A deterministic snapshot of memoized graph state.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct GraphSnapshot<K, V> {
    /// Snapshot nodes in stable key order.
    pub nodes: Vec<SnapshotNode<K, V>>,
}

impl<K, V> GraphSnapshot<K, V> {
    /// Creates a graph snapshot from already-ordered nodes.
    #[must_use]
    pub fn new(nodes: Vec<SnapshotNode<K, V>>) -> Self {
        Self { nodes }
    }
}

/// One memo node inside a graph snapshot.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct SnapshotNode<K, V> {
    /// The query key.
    pub key: K,
    /// The memo revision.
    pub revision: Revision,
    /// Whether the memo needs verification before reuse.
    pub dirty: bool,
    /// The memoized value, when one exists.
    pub value: Option<V>,
    /// Dependency observations captured during the last execution, without
    /// process-local cutoff fingerprints.
    pub dependencies: Vec<SnapshotObservation<K>>,
}

/// Durable dependency observation with no process-local fingerprint authority.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct SnapshotObservation<K> {
    /// Observed dependency key.
    pub key: K,
    /// Kind of observation made.
    pub kind: ObservationKind,
    /// Revision observed at snapshot time.
    pub revision: Revision,
}

/// Summary of a snapshot restore operation.
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct RestoreReport {
    /// Number of nodes restored.
    pub nodes: usize,
    /// Number of nodes marked dirty to recover from partial or stale snapshot
    /// contents.
    pub recovered_dirty: usize,
}

impl<K, V> IncrementalEngine<K, V>
where
    K: Ord + Clone,
    V: Clone + Eq + FingerprintValue,
{
    /// Exports memo graph state reachable from `roots`.
    pub fn snapshot<I>(
        &mut self,
        roots: I,
        budgets: SnapshotBudgets,
    ) -> QueryResult<K, GraphSnapshot<K, V>>
    where
        I: IntoIterator<Item = K>,
    {
        let mut pending = roots.into_iter().collect::<BTreeSet<_>>();
        let Some(root) = pending.iter().next().cloned() else {
            return Ok(GraphSnapshot::new(Vec::new()));
        };
        let mut seen = BTreeSet::new();
        let mut nodes = Vec::new();
        let mut edges = 0_usize;
        while let Some(key) = pending.iter().next().cloned() {
            pending.remove(&key);
            if !seen.insert(key.clone()) {
                continue;
            }
            if nodes.len().saturating_add(1) > budgets.max_nodes {
                let token = self.alloc_continuation(root);
                return Err(IncrementalError::BudgetExceeded {
                    kind: BudgetKind::Output,
                    limit: budgets.max_nodes,
                    consumed: nodes.len().saturating_add(1),
                    continuation: Some(token),
                });
            }
            let node = self
                .nodes
                .get(&key)
                .ok_or_else(|| IncrementalError::UnknownQuery { key: key.clone() })?;
            edges = edges.saturating_add(node.dependencies.len());
            if edges > budgets.max_edges {
                let token = self.alloc_continuation(root);
                return Err(IncrementalError::BudgetExceeded {
                    kind: BudgetKind::Output,
                    limit: budgets.max_edges,
                    consumed: edges,
                    continuation: Some(token),
                });
            }
            for observation in &node.dependencies {
                if matches!(observation.kind(), ObservationKind::Read) {
                    pending.insert(observation.key().clone());
                }
            }
            nodes.push(SnapshotNode {
                key,
                revision: node.revision,
                dirty: node.dirty,
                value: node.value.clone(),
                dependencies: node
                    .dependencies
                    .iter()
                    .map(|observation| SnapshotObservation {
                        key: observation.key().clone(),
                        kind: observation.kind().clone(),
                        revision: observation.revision(),
                    })
                    .collect(),
            });
        }
        Ok(GraphSnapshot::new(nodes))
    }

    /// Restores memo graph state and rebuilds reverse dependency edges.
    pub fn restore_snapshot(
        &mut self,
        snapshot: GraphSnapshot<K, V>,
    ) -> Result<RestoreReport, SnapshotError<K>> {
        let mut keys = BTreeSet::new();
        for node in &snapshot.nodes {
            if !keys.insert(node.key.clone()) {
                return Err(SnapshotError::DuplicateNode {
                    key: node.key.clone(),
                });
            }
        }

        self.nodes.clear();
        self.reverse.clear();
        let mut recovered_dirty = 0_usize;
        let mut max_revision = self.next_revision.saturating_sub(1);
        let restored_fingerprints = snapshot
            .nodes
            .iter()
            .filter_map(|node| {
                node.value
                    .as_ref()
                    .map(|value| (node.key.clone(), value.incremental_fingerprint()))
            })
            .collect::<BTreeMap<_, _>>();
        for snapshot_node in snapshot.nodes {
            max_revision = max_revision.max(snapshot_node.revision.get());
            // Snapshot revisions are historical hints. Recheck every restored
            // value through its registered query before it can be reused.
            let dirty = true;
            let recovered = !snapshot_node.dirty;
            let fingerprint = snapshot_node
                .value
                .as_ref()
                .map(FingerprintValue::incremental_fingerprint);
            if recovered {
                recovered_dirty += 1;
            }
            restore_external_revisions(&mut self.source_revisions, &snapshot_node);
            let dependencies = snapshot_node
                .dependencies
                .into_iter()
                .map(|observation| match observation.kind {
                    ObservationKind::Read => Observation::read(
                        observation.key.clone(),
                        observation.revision,
                        restored_fingerprints
                            .get(&observation.key)
                            .copied()
                            .unwrap_or_else(|| crate::ValueFingerprint::new(0)),
                    ),
                    kind => Observation::new(observation.key, kind, observation.revision, None),
                })
                .collect();
            self.nodes.insert(
                snapshot_node.key,
                Node {
                    revision: snapshot_node.revision,
                    dirty,
                    value: snapshot_node.value,
                    fingerprint,
                    dependencies,
                },
            );
        }
        self.next_revision = self.next_revision.max(max_revision.saturating_add(1));
        self.rebuild_reverse();
        Ok(RestoreReport {
            nodes: self.nodes.len(),
            recovered_dirty,
        })
    }

    fn rebuild_reverse(&mut self) {
        self.reverse.clear();
        for (key, node) in &self.nodes {
            for observation in &node.dependencies {
                self.reverse
                    .entry(observation.key().clone())
                    .or_default()
                    .insert(key.clone());
            }
        }
    }
}

fn restore_external_revisions<K, V>(
    source_revisions: &mut BTreeMap<K, Revision>,
    snapshot_node: &SnapshotNode<K, V>,
) where
    K: Ord + Clone,
{
    for observation in &snapshot_node.dependencies {
        if matches!(observation.kind, ObservationKind::Read) {
            continue;
        }
        source_revisions
            .entry(observation.key.clone())
            .and_modify(|revision| {
                if observation.revision > *revision {
                    *revision = observation.revision;
                }
            })
            .or_insert(observation.revision);
    }
}