haematite 0.6.1

Content-addressed, branchable, actor-native storage engine
Documentation
use std::collections::{HashMap, HashSet};
use std::sync::{Arc, Mutex, MutexGuard};

use crate::tree::Hash;

/// Thread-safe registry of roots referenced by live branch handles.
///
/// The public view is a set of live root hashes for the pruner. Internally the
/// registry keeps reference counts so two live branches at the same root do not
/// let one dropped handle deregister a root still used by the other branch.
/// Registering one hash twice yields refcount 2 on that hash — the §14.1
/// anchor pattern (anchor role + head role pinned independently) falls out of
/// plain counting, with no special casing.
#[derive(Clone, Debug, Default)]
pub struct BranchRegistry {
    counts: Arc<Mutex<HashMap<Hash, usize>>>,
}

impl BranchRegistry {
    /// Create an empty active-branch registry.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Register one active branch root.
    pub fn register(&self, root: Hash) {
        register_in(&mut self.lock_counts(), root);
    }

    /// Deregister one active branch root reference.
    ///
    /// Missing roots are ignored so drop paths remain idempotent from the
    /// registry's perspective.
    pub fn deregister(&self, root: Hash) {
        deregister_in(&mut self.lock_counts(), root);
    }

    /// Atomically swap a superseded head pin for its advanced replacement.
    ///
    /// Register-new and deregister-old run under ONE counts guard (BRANCH-
    /// COMMIT-PATH.md §5 step 5, §14.1): a [`live_roots`](Self::live_roots)
    /// snapshot taken at any instant therefore contains `old` or `new` (or
    /// both), never neither — the in-memory half of the prune registration
    /// hazard. Doing this as separate `register`/`deregister` calls would
    /// reopen the window in which a prune sees neither root and reclaims
    /// nodes the advancing branch still reaches.
    ///
    /// `commit_branch` does not call this: it performs the same swap via
    /// [`advance_in`] on the counts guard it already holds (§16.1 MF1). This
    /// method is the standalone form for orchestration layers that advance a
    /// pin outside a commit.
    pub fn advance(&self, old: Hash, new: Hash) {
        advance_in(&mut self.lock_counts(), old, new);
    }

    /// Return every root currently referenced by at least one live branch.
    #[must_use]
    pub fn live_roots(&self) -> HashSet<Hash> {
        self.lock_counts().keys().copied().collect()
    }

    pub(crate) fn register_roots<I>(&self, roots: I) -> BranchRegistryGuard
    where
        I: IntoIterator<Item = Hash>,
    {
        let roots: Vec<Hash> = roots.into_iter().collect();
        for root in roots.iter().copied() {
            self.register(root);
        }

        BranchRegistryGuard {
            registry: self.clone(),
            roots: Mutex::new(roots),
        }
    }

    /// Register `roots` only if every distinct root is already live here or
    /// `is_protected` vouches for it — the §16.4 Q3 hard-refusal substrate.
    ///
    /// Membership check and registration happen under the ONE counts guard:
    /// a root that is live only through another handle's pin cannot be
    /// deregistered by that handle's drop in a gap between our check and our
    /// register, so a `fork_at` that returns `Ok` is pinned with no window in
    /// which prune could observe the anchor unprotected. `is_protected` MUST
    /// NOT touch this registry (it is consulted with the counts lock held);
    /// the blessed caller passes closures over pre-built lookup sets.
    ///
    /// On refusal returns the offending root; nothing is registered.
    #[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
    pub(crate) fn register_roots_protected<I, F>(
        &self,
        roots: I,
        is_protected: F,
    ) -> Result<BranchRegistryGuard, Hash>
    where
        I: IntoIterator<Item = Hash>,
        F: Fn(&Hash) -> bool,
    {
        let roots: Vec<Hash> = roots.into_iter().collect();
        let mut counts = self.lock_counts();
        for root in &roots {
            if !counts.contains_key(root) && !is_protected(root) {
                return Err(*root);
            }
        }
        for root in roots.iter().copied() {
            register_in(&mut counts, root);
        }
        drop(counts);

        Ok(BranchRegistryGuard {
            registry: self.clone(),
            roots: Mutex::new(roots),
        })
    }

    /// Poison-tolerant counts lock, exposed crate-wide for `commit_branch`'s
    /// step 1 (§16.1 MF1): the commit path holds this guard from before the
    /// durable install to the end, so its step-5 registry dance is a pure
    /// write into an already-held guard — nothing fallible after the commit
    /// point. Tier 3 (last) in the §16.1 lock order.
    pub(crate) fn lock_counts(&self) -> MutexGuard<'_, HashMap<Hash, usize>> {
        match self.counts.lock() {
            Ok(guard) => guard,
            Err(poisoned) => poisoned.into_inner(),
        }
    }
}

/// [`BranchRegistry::advance`]'s body, operating on an already-held counts
/// guard so `commit_branch` (which holds the guard across the §5 commit point
/// per §16.1 MF1) performs the same one-lock swap without re-locking.
pub(crate) fn advance_in(counts: &mut HashMap<Hash, usize>, old: Hash, new: Hash) {
    // Register-new strictly before deregister-old inside the guard: even
    // if a future refactor splits this lock, the ordering stays leak-safe
    // (transient over-pin) rather than corruption-prone (transient gap).
    register_in(counts, new);
    deregister_in(counts, old);
}

fn register_in(counts: &mut HashMap<Hash, usize>, root: Hash) {
    counts
        .entry(root)
        .and_modify(|count| *count = count.saturating_add(1))
        .or_insert(1);
}

fn deregister_in(counts: &mut HashMap<Hash, usize>, root: Hash) {
    if let Some(count) = counts.get_mut(&root) {
        if *count > 1 {
            *count -= 1;
        } else {
            counts.remove(&root);
        }
    }
}

/// Drop guard that keeps registry roots live until the last handle clone is gone.
///
/// The pinned root list is mutable behind a mutex so a commit can retarget a
/// pin from a superseded head to its advanced replacement via
/// [`replace`](Self::replace): drop then deregisters the CURRENT pins, not the
/// ones captured at registration (BRANCH-COMMIT-PATH.md §5 step 5).
#[derive(Debug)]
pub struct BranchRegistryGuard {
    registry: BranchRegistry,
    roots: Mutex<Vec<Hash>>,
}

impl BranchRegistryGuard {
    /// Retarget one pinned occurrence of `old` to `new`.
    ///
    /// Exactly one occurrence moves, mirroring the single
    /// [`BranchRegistry::advance`] call it accompanies — a refcount-2 anchor
    /// (§14.1) keeps its second, anchor-role pin untouched. If `old` is not
    /// pinned (a blessed caller never does this), `new` is appended instead:
    /// the advanced root must always be released on drop, and an extra
    /// deregister of `old` elsewhere is idempotently ignored by the registry —
    /// leak-safe in every interleaving. The blessed caller is `commit_branch`
    /// step 5 (A1 stage 3).
    pub(crate) fn replace(&self, old: Hash, new: Hash) {
        let mut roots = self.lock_roots();
        if let Some(slot) = roots.iter_mut().find(|slot| **slot == old) {
            *slot = new;
        } else {
            roots.push(new);
        }
    }

    /// Poison-tolerant pin-list lock: the list is only ever mutated as a
    /// single slot write or push under the guard, so a poisoned mutex still
    /// holds a valid list, and drop MUST still release the pins.
    fn lock_roots(&self) -> MutexGuard<'_, Vec<Hash>> {
        match self.roots.lock() {
            Ok(guard) => guard,
            Err(poisoned) => poisoned.into_inner(),
        }
    }
}

impl Drop for BranchRegistryGuard {
    fn drop(&mut self) {
        let roots: Vec<Hash> = self.lock_roots().clone();
        for root in roots {
            self.registry.deregister(root);
        }
    }
}

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