haematite 0.6.1

use crate::tree::Hash;

use super::commit::BranchCommitError;
use super::handle::{BranchError, BranchHandle, ShardId};
use super::lifecycle::create_branch_record;
use super::policy::{BranchKind, check_fork_into};
use super::refstore::BranchRefStore;
use super::registry::BranchRegistry;
use super::snapshot::Timestamp;

/// Fork a single-shard database or shard at its current committed root hash.
///
/// The operation records `root` as both the fork point and the current branch
/// root, then allocates one empty branch-local WAL buffer. It accepts no node
/// store and therefore cannot read or copy tree nodes; work is constant with
/// respect to the number of database entries.
#[must_use]
pub fn fork(root: Hash) -> BranchHandle {
    BranchHandle::new(root)
}

/// Fork a single-shard database or shard and register its live root.
///
/// The root remains registered until the last clone of the returned
/// [`BranchHandle`] is dropped. The caller vouches that `root` is safe to
/// anchor to (typically the database's own current committed root); to fork
/// an arbitrary historical hash with the anchor-protection check, use
/// [`fork_at`](super::fork_at) instead.
#[must_use]
pub fn fork_registered(root: Hash, registry: &BranchRegistry) -> BranchHandle {
    let branch = fork(root);
    let guard = registry.register_roots(anchor_role_pins(&branch));
    branch.with_registry_guard(guard)
}

/// Fork a logical database with one independent root per shard.
///
/// Each supplied `(shard_id, root)` records that shard's fork point and allocates
/// a distinct empty WAL buffer. The function coordinates shard state inside one
/// [`BranchHandle`] but performs no cross-shard locking or tree access.
pub fn fork_shards<I>(roots: I) -> Result<BranchHandle, BranchError>
where
    I: IntoIterator<Item = (ShardId, Hash)>,
{
    BranchHandle::from_shard_roots(roots)
}

/// Fork multiple shards and register every shard root as live.
///
/// Registering all shard roots keeps future pruning conservative: any tree root
/// referenced by any live branch shard remains visible through
/// [`BranchRegistry::live_roots`](super::BranchRegistry::live_roots).
pub fn fork_shards_registered<I>(
    roots: I,
    registry: &BranchRegistry,
) -> Result<BranchHandle, BranchError>
where
    I: IntoIterator<Item = (ShardId, Hash)>,
{
    let branch = fork_shards(roots)?;
    let guard = registry.register_roots(anchor_role_pins(&branch));
    Ok(branch.with_registry_guard(guard))
}

/// Create a durable named branch with an explicit immutable kind.
///
/// Direct creation has no parent lineage. Use [`fork_from_with_kind`] for Work
/// branches that must inherit a Namespace boundary from their parent.
pub fn create_branch_with_kind<I>(
    name: &str,
    roots: I,
    kind: BranchKind,
    refs: &mut BranchRefStore,
    registry: &BranchRegistry,
    timestamp: Timestamp,
) -> Result<BranchHandle, BranchCommitError>
where
    I: IntoIterator<Item = (ShardId, Hash)>,
{
    create_branch_record(name, roots, kind, None, refs, registry, timestamp)
}

/// Durable fork from a named parent's committed heads with an explicit kind.
///
/// Work children inherit the parent's resolved Namespace lineage. Callers do
/// not provide lineage: both source policy state and inherited destination
/// state come from the parent's durable ref record. Namespace children are
/// direct-create-only and are refused from every parent before roots are read
/// or a destination record is installed.
pub fn fork_from_with_kind(
    parent: &BranchHandle,
    child_name: &str,
    child_kind: BranchKind,
    refs: &mut BranchRefStore,
    registry: &BranchRegistry,
    timestamp: Timestamp,
) -> Result<BranchHandle, BranchCommitError> {
    if parent.registry_guard().is_none() {
        return Err(BranchError::UnregisteredParent.into());
    }
    let binding = parent
        .binding_snapshot()
        .ok_or(BranchCommitError::UnnamedBranch)?;
    let parent_record = refs.cas_check(&binding.name, binding.created, binding.seq)?;
    let source = parent_record.name.clone();
    let source_kind = parent_record.kind;
    let source_lineage = parent_record
        .resolved_namespace_lineage()
        .map(str::to_owned);
    let destination_lineage = match child_kind {
        BranchKind::Namespace => Some(child_name.to_owned()),
        BranchKind::Work => source_lineage.clone(),
    };

    check_fork_into(
        &source,
        source_kind,
        source_lineage.as_deref(),
        child_name,
        child_kind,
        destination_lineage.as_deref(),
    )?;

    // Keep the parent's shard locks held through child pin registration and
    // record installation: releasing after the read could let a concurrent
    // volatile parent commit unpin the chosen anchor before the child owns it.
    let states = parent.lock_states_ascending()?;
    let heads: Vec<(ShardId, Hash)> = states
        .iter()
        .map(|(shard_id, state)| (*shard_id, state.current_root))
        .collect();
    let stored_lineage = match child_kind {
        BranchKind::Namespace => None,
        BranchKind::Work => destination_lineage,
    };
    let child = create_branch_record(
        child_name,
        heads,
        child_kind,
        stored_lineage,
        refs,
        registry,
        timestamp,
    )?;
    drop(states);
    Ok(child)
}

/// Each fork anchor pinned TWICE — the §14.1 refcount-2 pattern shared by
/// every registered-handle constructor: one pin in the anchor role (released
/// only on guard drop; `merge` needs the divergence point for the handle's
/// whole life) and one in the current-head role (the pin `commit_branch`'s
/// step-5 `advance` retargets). With refcount 2 from birth, the first commit
/// needs no special case: advancing the head role leaves the anchor role's
/// count holding the anchor live.
fn anchor_role_pins(branch: &BranchHandle) -> impl Iterator<Item = Hash> + '_ {
    branch.fork_points().flat_map(|anchor| [anchor, anchor])
}

#[cfg(test)]
mod tests {
    use super::{fork, fork_registered, fork_shards, fork_shards_registered};
    use crate::branch::{BranchError, BranchRegistry};
    use crate::tree::Hash;

    fn hash(byte: u8) -> Hash {
        Hash::from_bytes([byte; 32])
    }

    #[test]
    fn fork_records_current_root_as_fork_point_and_current_root() {
        let root = hash(1);

        let branch = fork(root);

        assert_eq!(branch.fork_point(), root);
        assert_eq!(branch.current_root(), root);
    }

    #[test]
    fn fork_allocates_one_empty_buffer_for_single_shard() -> Result<(), BranchError> {
        let branch = fork(hash(2));

        assert_eq!(branch.shard_count(), 1);
        assert!(!branch.shard_is_dirty(branch.primary_shard())?);
        Ok(())
    }

    #[test]
    fn fork_shards_allocates_one_empty_buffer_per_shard() -> Result<(), BranchError> {
        let branch = fork_shards([(3, hash(3)), (5, hash(5)), (8, hash(8))])?;

        assert_eq!(branch.shard_count(), 3);
        for shard_id in [3, 5, 8] {
            assert!(!branch.shard_is_dirty(shard_id)?);
        }
        Ok(())
    }

    #[test]
    fn fork_shards_preserves_each_independent_root() -> Result<(), BranchError> {
        let shard_three_root = hash(3);
        let shard_five_root = hash(5);

        let branch = fork_shards([(3, shard_three_root), (5, shard_five_root)])?;

        assert_eq!(branch.shard_fork_point(3), Some(shard_three_root));
        assert_eq!(branch.shard_current_root(3), Some(shard_three_root));
        assert_eq!(branch.shard_fork_point(5), Some(shard_five_root));
        assert_eq!(branch.shard_current_root(5), Some(shard_five_root));
        Ok(())
    }

    #[test]
    fn fork_single_shard_has_no_tree_store_to_read_or_copy() {
        let roots_for_large_and_small_trees_are_both_just_hashes = [hash(9), hash(10)];

        let first = fork(roots_for_large_and_small_trees_are_both_just_hashes[0]);
        let second = fork(roots_for_large_and_small_trees_are_both_just_hashes[1]);

        assert_eq!(first.shard_count(), 1);
        assert_eq!(second.shard_count(), 1);
    }

    #[test]
    fn fork_registered_registers_root_until_drop() {
        let registry = BranchRegistry::new();
        let root = hash(11);

        let branch = fork_registered(root, &registry);
        assert!(registry.live_roots().contains(&root));

        drop(branch);
        assert!(!registry.live_roots().contains(&root));
    }

    #[test]
    fn fork_shards_registered_registers_all_shard_roots() -> Result<(), BranchError> {
        let registry = BranchRegistry::new();
        let shard_three_root = hash(12);
        let shard_five_root = hash(13);

        let branch =
            fork_shards_registered([(3, shard_three_root), (5, shard_five_root)], &registry)?;
        let live = registry.live_roots();
        assert!(live.contains(&shard_three_root));
        assert!(live.contains(&shard_five_root));

        drop(branch);
        assert!(registry.live_roots().is_empty());
        Ok(())
    }
}