haematite 0.7.0

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),
        }
    }

    /// Cross each protected anchor into a working root and register BOTH under
    /// the ONE counts lock — the §4.2 step-5 fork crossing / §16.4 Q3 substrate.
    ///
    /// 1. Every anchor must be already live here or vouched by `is_protected`;
    ///    else the offending anchor is returned and NOTHING is built or
    ///    registered. A root live only through another handle's pin cannot be
    ///    deregistered by that handle's drop in a gap between the check and the
    ///    register — check and register share this one guard — so a crossing 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 lock held); the caller passes a closure over
    ///    pre-built lookup sets.
    /// 2. `build` runs UNDER the counts lock — it writes the rebuilt working-root
    ///    nodes to the store. Holding the lock across the build is what makes the
    ///    rebuilt root prune-safe: a concurrent prune takes this SAME lock to
    ///    snapshot [`live_roots`](Self::live_roots), so it can never observe a
    ///    state where the working-root nodes exist but the working root is not
    ///    yet pinned — a byte-identical dedup collision with a prunable snapshot
    ///    therefore cannot reclaim them.
    /// 3. Each `(anchor, working)` is registered — anchor role (lineage, held to
    ///    guard drop) + head role (the pin `commit_branch` retargets). A no-op
    ///    crossing (working == anchor) reduces to the §14.1 refcount-2 pattern.
    ///
    /// `build` receives the checked anchors and returns the working roots in
    /// anchor order (contract: exactly `anchors.len()` of them); `is_protected`
    /// MUST NOT touch this registry.
    #[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
    pub(crate) fn cross_and_register<F, E>(
        &self,
        anchors: &[Hash],
        is_protected: impl Fn(&Hash) -> bool,
        build: F,
    ) -> Result<(Vec<Hash>, BranchRegistryGuard), CrossRegisterError<E>>
    where
        F: FnOnce(&[Hash]) -> Result<Vec<Hash>, E>,
    {
        let mut counts = self.lock_counts();
        for anchor in anchors {
            if !counts.contains_key(anchor) && !is_protected(anchor) {
                return Err(CrossRegisterError::UnprotectedAnchor(*anchor));
            }
        }
        let working = build(anchors).map_err(CrossRegisterError::Build)?;
        // Contract: exactly one working root per anchor. A short `build` would
        // otherwise leave later anchors UNREGISTERED (a prune hazard) via the
        // zip below, and an over-long one would drop roots silently. Refuse
        // LOUDLY here — BEFORE any `register_in` — so a mismatch registers
        // NOTHING (no partial pins); the guard is dropped by the early return.
        if working.len() != anchors.len() {
            return Err(CrossRegisterError::WorkingRootCountMismatch {
                expected: anchors.len(),
                got: working.len(),
            });
        }
        let mut pins = Vec::with_capacity(anchors.len().saturating_mul(2));
        for (anchor, work) in anchors.iter().zip(working.iter()) {
            register_in(&mut counts, *anchor);
            register_in(&mut counts, *work);
            pins.push(*anchor);
            pins.push(*work);
        }
        drop(counts);
        Ok((
            working,
            BranchRegistryGuard {
                registry: self.clone(),
                roots: Mutex::new(pins),
            },
        ))
    }

    /// 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(),
        }
    }
}

/// Failure of [`BranchRegistry::cross_and_register`]: an anchor that no live pin
/// or durable/named record vouches for (nothing built or registered), a failure
/// of the working-root build (the v2 rebuild), or a `build` that violated its
/// one-working-root-per-anchor contract (a would-be silent under/over-pin,
/// refused BEFORE any registration).
#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
pub(crate) enum CrossRegisterError<E> {
    UnprotectedAnchor(Hash),
    Build(E),
    WorkingRootCountMismatch { expected: usize, got: usize },
}

/// [`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;