haematite 0.6.2

// BRANCH-004: Snapshot pruning and unreachable node reclamation.
// LEDGER A1: prune consumes the durable branch-ref denominator (§6(c)).

use std::collections::{HashMap, HashSet};
use std::fmt;
use std::sync::Arc;

use crate::store::{DeleteNode, NodeStore};
use crate::tree::{Hash, Node};

use super::refstore::BranchRefStore;
use super::registry::BranchRegistry;
use super::snapshot::{SnapshotError, SnapshotRegistry};

/// Summary of one caller-initiated pruning run.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PruneReport {
    /// Number of content-addressed tree nodes deleted from the store.
    pub node_count: usize,
    /// Total logical bytes reclaimed, measured as each deleted node's
    /// serialised size before deletion.
    pub bytes_reclaimed: usize,
}

/// Errors raised while removing a snapshot and reclaiming unreferenced nodes.
#[derive(Debug)]
pub enum PruneError {
    /// No named snapshot exists for the requested prune target.
    UnknownSnapshot { name: String },
    /// The snapshot registry could not persist or load the requested mutation.
    SnapshotRegistry(SnapshotError),
    /// A root or child hash referenced by the graph walk was absent from the store.
    MissingNode { hash: Hash },
    /// The node store returned an error while reading a referenced hash.
    StoreRead { hash: Hash },
    /// The node store returned an error while deleting an unreferenced hash.
    NodeDelete { hash: Hash },
}

impl fmt::Display for PruneError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::UnknownSnapshot { name } => write!(formatter, "unknown snapshot: {name}"),
            Self::SnapshotRegistry(error) => write!(formatter, "snapshot registry error: {error}"),
            Self::MissingNode { hash } => write!(formatter, "missing tree node {hash}"),
            Self::StoreRead { hash } => write!(formatter, "failed to read tree node {hash}"),
            Self::NodeDelete { hash } => write!(formatter, "failed to delete tree node {hash}"),
        }
    }
}

impl std::error::Error for PruneError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            Self::SnapshotRegistry(error) => Some(error),
            Self::UnknownSnapshot { .. }
            | Self::MissingNode { .. }
            | Self::StoreRead { .. }
            | Self::NodeDelete { .. } => None,
        }
    }
}

impl From<SnapshotError> for PruneError {
    fn from(error: SnapshotError) -> Self {
        Self::SnapshotRegistry(error)
    }
}

/// Removes `name` from the snapshot registry and reclaims tree nodes that are
/// reachable from the removed snapshot root and from no live root.
///
/// Live roots are the union of three sets (BRANCH-COMMIT-PATH.md §6(c)): every
/// active branch root registered in `branches`, the durable branch-ref pin set
/// `refs.protected_roots()` (every named branch's fork anchors and committed
/// heads — restart-proof with zero startup ceremony), and the roots of the
/// named snapshots that remain after removal. The commit log is intentionally
/// not accepted by this API, so pruning cannot mutate its append-only contents.
///
/// # Commit/prune exclusivity (§6(b), §14.4)
///
/// `commit_branch` takes the node store as `&mut S`; prune borrows it here as
/// `&S`. Safe Rust therefore cannot run the two concurrently against one
/// store: every prune executes entirely-before or entirely-after each
/// `commit_branch` call. That whole-call exclusivity is what closes the
/// materialisation gap — the window in which a commit's freshly written nodes
/// are reachable from no live root (and, under content-address dedup, may be
/// byte-identical to nodes reachable only from the snapshot being pruned).
/// Entirely-before, the commit's nodes do not exist yet and any dedup casualty
/// is simply rewritten by the commit's own `put`; entirely-after, the new head
/// was registered and durably recorded before `commit_branch` returned, so it
/// is in this run's live set. Caveat: the argument is carried by the borrow
/// checker, so a store wrapper with interior mutability (one that hands out
/// `&S` while writes proceed) voids it and must itself fence prune against
/// commit — the commit/prune `RwLock` successor recorded in §14.4 for the
/// A2/A6 shared-store work.
///
/// The reclaim set is computed *before* the snapshot is removed from the
/// registry: a graph-walk failure (a missing or unreadable node) therefore
/// leaves the snapshot in place and the whole operation retryable, rather than
/// dropping the registry entry and orphaning its nodes. The snapshot is removed
/// only once a complete deletion set has been collected; node deletion follows.
pub fn prune<S>(
    store: &S,
    branches: &BranchRegistry,
    refs: &BranchRefStore,
    snapshots: &mut SnapshotRegistry,
    name: &str,
) -> Result<PruneReport, PruneError>
where
    S: NodeStore + DeleteNode + ?Sized,
{
    // Resolve the target root without removing the snapshot yet, so a failed
    // walk below is fully recoverable (the snapshot is still named).
    let removed_root = snapshots
        .get(name)
        .ok_or_else(|| PruneError::UnknownSnapshot {
            name: name.to_owned(),
        })?;

    let removed_reachable = collect_reachable(store, [removed_root])?;
    let live_roots = live_roots_excluding(branches, refs, snapshots, name);
    let live_reachable = collect_reachable(store, live_roots)?;

    let unreferenced = unreferenced_nodes(removed_reachable, &live_reachable);
    let report = PruneReport {
        node_count: unreferenced.len(),
        bytes_reclaimed: unreferenced.iter().map(|(_hash, bytes)| *bytes).sum(),
    };

    // Both walks succeeded: commit the registry removal, then delete the
    // now-confirmed unreferenced nodes. `&mut SnapshotRegistry` gives exclusive
    // access, so the entry cannot have vanished since the `get` above.
    snapshots
        .remove(name)?
        .ok_or_else(|| PruneError::UnknownSnapshot {
            name: name.to_owned(),
        })?;

    for (hash, _bytes) in unreferenced {
        store
            .delete(&hash)
            .map_err(|_error| PruneError::NodeDelete { hash })?;
    }

    Ok(report)
}

/// Live roots that pruning `excluded` must preserve: every active branch root
/// registered in `branches`, every fork anchor and committed head durably
/// recorded in `refs`, and the roots of every named snapshot other than
/// `excluded` (the snapshot being pruned, which is excluded by name since it
/// has not been removed from the registry yet).
///
/// Contract (BRANCH-COMMIT-PATH.md §6(b)): `prune` protects exactly the nodes
/// reachable from this set, and registration and durable recording happen
/// inside `commit_branch` before it returns — therefore, at every point at
/// which prune can legally observe the store, every root any completed commit
/// references is in the live set. In-memory, [`BranchRegistry::advance`] swaps
/// a superseded head for its replacement under ONE lock, so no `live_roots()`
/// snapshot ever observes neither root; durably, the ref store swaps pins in
/// one atomic file rename. Anonymous branches are protected only while a live
/// handle's guard pins them — matching their volatile contract — while named
/// branches (and their merge anchors) stay protected across restarts via
/// [`BranchRefStore::protected_roots`].
fn live_roots_excluding(
    branches: &BranchRegistry,
    refs: &BranchRefStore,
    snapshots: &SnapshotRegistry,
    excluded: &str,
) -> HashSet<Hash> {
    let mut roots = branches.live_roots();
    roots.extend(refs.protected_roots());
    roots.extend(
        snapshots
            .list_snapshots()
            .into_iter()
            .filter(|(name, _root_hash, _timestamp)| name != excluded)
            .map(|(_name, root_hash, _timestamp)| root_hash),
    );
    roots
}

fn collect_reachable<S, I>(store: &S, roots: I) -> Result<HashMap<Hash, usize>, PruneError>
where
    S: NodeStore + ?Sized,
    I: IntoIterator<Item = Hash>,
{
    let mut reachable = HashMap::new();
    let mut stack: Vec<Hash> = roots.into_iter().collect();

    while let Some(hash) = stack.pop() {
        if reachable.contains_key(&hash) {
            continue;
        }

        let node = load_node(store, hash)?;
        let serialised_len = node.serialise().len();
        if let Node::Internal(internal) = &*node {
            stack.extend(
                internal
                    .children()
                    .iter()
                    .map(|(_lower_bound, child_hash)| *child_hash),
            );
        }
        reachable.insert(hash, serialised_len);
    }

    Ok(reachable)
}

fn load_node<S>(store: &S, hash: Hash) -> Result<Arc<Node>, PruneError>
where
    S: NodeStore + ?Sized,
{
    store
        .get(&hash)
        .map_err(|_error| PruneError::StoreRead { hash })?
        .ok_or(PruneError::MissingNode { hash })
}

fn unreferenced_nodes(
    removed_reachable: HashMap<Hash, usize>,
    live_reachable: &HashMap<Hash, usize>,
) -> Vec<(Hash, usize)> {
    removed_reachable
        .into_iter()
        .filter(|(hash, _bytes)| !live_reachable.contains_key(hash))
        .collect()
}

#[cfg(test)]
#[path = "prune_tests.rs"]
mod tests;