haematite 0.7.0

//! Leaf enumeration and spine construction (CHUNKING-POLICY.md §2.2).
//!
//! Grouping is a pure function of the policy applied to the ordered entry/child
//! list, so the root is history-independent at every level.
//!
//! - Leaves group by `boundary_after`, plus (under v2) a `boundary_before` split
//!   that isolates oversized entries as singletons.
//! - Internal levels group by `boundary_after`; under v2 a group closes only
//!   once it holds >= 2 entries (the min-2 rule), which halves the level
//!   unconditionally, so [`build_spine`]'s collapse-all fallback is unreachable
//!   under v2 and retained only for v1's all-boundary degenerate case.

use crate::store::NodeStore;

use super::super::cursor::{TreeError, load_node};
use super::super::node::{Hash, InternalNode, LeafNode, Node};
use super::super::policy::TreePolicy;
use super::{ChildRef, Entry};

/// Walk the tree left-to-right collecting every leaf reference (separator key +
/// hash) in key order. Reads only refs, not leaf contents — policy-free.
pub(super) fn collect_leaf_refs<S: NodeStore + ?Sized>(
    store: &S,
    root_hash: Hash,
) -> Result<Vec<ChildRef>, TreeError> {
    let height = tree_height(store, root_hash)?;
    if height == 0 {
        // The root is a leaf; its separator is unknown without loading it, but the
        // leftmost leaf is never chosen by separator comparison, so an empty
        // separator is safe.
        return Ok(vec![(Vec::new(), root_hash)]);
    }

    let mut leaves = Vec::new();
    collect_leaf_refs_inner(store, root_hash, height, &mut leaves)?;
    Ok(leaves)
}

/// Number of internal levels above the leaves (0 when the root is a leaf).
fn tree_height<S: NodeStore + ?Sized>(store: &S, root_hash: Hash) -> Result<usize, TreeError> {
    let mut height = 0;
    let mut hash = root_hash;
    loop {
        match &*load_node(store, hash)? {
            Node::Leaf(_leaf) => return Ok(height),
            Node::Internal(internal) => {
                let Some((_separator, child_hash)) = internal.children().first() else {
                    return Err(TreeError::InvalidNode);
                };
                height = height.saturating_add(1);
                hash = *child_hash;
            }
        }
    }
}

fn collect_leaf_refs_inner<S: NodeStore + ?Sized>(
    store: &S,
    hash: Hash,
    levels_above_leaves: usize,
    leaves: &mut Vec<ChildRef>,
) -> Result<(), TreeError> {
    let node = load_node(store, hash)?;
    let Node::Internal(internal) = &*node else {
        return Err(TreeError::InvalidNode);
    };

    if levels_above_leaves <= 1 {
        leaves.extend_from_slice(internal.children());
        return Ok(());
    }

    for (_separator, child_hash) in internal.children() {
        collect_leaf_refs_inner(store, *child_hash, levels_above_leaves - 1, leaves)?;
    }
    Ok(())
}

pub(super) fn store_leaf_replacements<S: NodeStore + ?Sized>(
    store: &mut S,
    entries: Vec<Entry>,
    policy: TreePolicy,
) -> Result<Vec<ChildRef>, TreeError> {
    if entries.is_empty() {
        return Ok(Vec::new());
    }

    let mut replacements = Vec::new();
    for chunk in group_leaf_entries(entries, policy) {
        replacements.push(store_leaf(store, chunk)?);
    }
    Ok(replacements)
}

/// Split a flat ordered entry list into leaf chunks.
///
/// A chunk closes AFTER an entry whose `boundary_after` fires, and (under v2)
/// BEFORE an entry whose `boundary_before` fires — so an oversized entry, which
/// fires both, is always a singleton. Under v1 `boundary_before` is always
/// false, so this reduces exactly to the pre-v2 `boundary_after`-only split.
fn group_leaf_entries(entries: Vec<Entry>, policy: TreePolicy) -> Vec<Vec<Entry>> {
    let mut chunks = Vec::new();
    let mut current: Vec<Entry> = Vec::new();

    for (key, value) in entries {
        if !current.is_empty() && policy.leaf_boundary_before(key.len(), value.len()) {
            chunks.push(std::mem::take(&mut current));
        }
        let closes = policy.leaf_boundary_after(key.as_slice(), value.len());
        current.push((key, value));
        if closes {
            chunks.push(std::mem::take(&mut current));
        }
    }

    if !current.is_empty() {
        chunks.push(current);
    }

    chunks
}

/// Build exactly one internal level from a flat list of child refs. Every group
/// becomes exactly one internal node — including a single-child group — so all
/// children stay at the same depth and the tree stays height-balanced.
///
/// A whole input of length 1 passes through unwrapped: the spine reducing to a
/// single subtree, not a new level.
fn store_internal_replacements<S: NodeStore + ?Sized>(
    store: &mut S,
    children: Vec<ChildRef>,
    policy: TreePolicy,
) -> Result<Vec<ChildRef>, TreeError> {
    match children.len() {
        0 => Ok(Vec::new()),
        1 => Ok(children),
        _ => {
            let mut replacements = Vec::new();
            for chunk in group_internal_children(children, policy) {
                replacements.push(store_internal(store, chunk)?);
            }
            Ok(replacements)
        }
    }
}

/// Group internal children. A group closes at a `boundary_after` separator; under
/// v2's min-2 rule it closes only once the group holds >= 2 entries (so every
/// non-final group has >= 2 children ⇒ the level halves unconditionally). Under
/// v1 there is no min-2 gate — the pre-v2 `boundary_after`-only split.
fn group_internal_children(children: Vec<ChildRef>, policy: TreePolicy) -> Vec<Vec<ChildRef>> {
    let min_two = policy.is_v2();
    let mut chunks = Vec::new();
    let mut current: Vec<ChildRef> = Vec::new();

    for (separator, hash) in children {
        let closes = policy.internal_boundary_after(separator.as_slice());
        current.push((separator, hash));
        if closes && (!min_two || current.len() >= 2) {
            chunks.push(std::mem::take(&mut current));
        }
    }

    if !current.is_empty() {
        chunks.push(current);
    }

    chunks
}

fn store_leaf<S: NodeStore + ?Sized>(
    store: &mut S,
    entries: Vec<Entry>,
) -> Result<ChildRef, TreeError> {
    let separator = first_entry_key(entries.as_slice())?;
    let leaf = LeafNode::new(entries)?;
    let hash = store
        .put(&Node::Leaf(leaf))
        .map_err(|_| TreeError::InvalidNode)?;
    Ok((separator, hash))
}

fn store_internal<S: NodeStore + ?Sized>(
    store: &mut S,
    children: Vec<ChildRef>,
) -> Result<ChildRef, TreeError> {
    let separator = first_child_key(children.as_slice())?;
    let internal = InternalNode::new(children)?;
    let hash = store
        .put(&Node::Internal(internal))
        .map_err(|_| TreeError::InvalidNode)?;
    Ok((separator, hash))
}

/// Build the new root from the complete ordered leaf list. An empty list
/// collapses to the canonical empty leaf; otherwise the spine is rebuilt level
/// by level, so the root depends only on the leaf set.
pub(super) fn finish_root<S: NodeStore + ?Sized>(
    store: &mut S,
    leaves: Vec<ChildRef>,
    policy: TreePolicy,
) -> Result<Hash, TreeError> {
    if leaves.is_empty() {
        return store_empty_leaf(store);
    }

    let (_separator, hash) = build_spine(store, leaves, policy)?;
    Ok(hash)
}

/// Collapse a flat list of child refs into a single root by repeatedly building
/// one internal level at a time until a single node remains.
///
/// Termination: each pass either shrinks the list (some group had > 1 element)
/// or — the v1 all-boundary degenerate case where every separator is a boundary
/// and no group can hold two — collapses the whole list into one internal node
/// (v1's own termination guarantee, `mutate.rs` collapse-all, retired from v2).
/// Under v2 the min-2 rule guarantees a strict shrink for every input of length
/// > 1, so the collapse arm is never taken.
fn build_spine<S: NodeStore + ?Sized>(
    store: &mut S,
    mut children: Vec<ChildRef>,
    policy: TreePolicy,
) -> Result<ChildRef, TreeError> {
    loop {
        match children.len() {
            0 => return Err(TreeError::InvalidNode),
            1 => return children.into_iter().next().ok_or(TreeError::InvalidNode),
            len => {
                let next = store_internal_replacements(store, children, policy)?;
                if next.len() < len {
                    children = next;
                } else {
                    return store_internal(store, next);
                }
            }
        }
    }
}

fn store_empty_leaf<S: NodeStore + ?Sized>(store: &mut S) -> Result<Hash, TreeError> {
    let leaf = LeafNode::new(Vec::new())?;
    store
        .put(&Node::Leaf(leaf))
        .map_err(|_| TreeError::InvalidNode)
}

fn first_entry_key(entries: &[Entry]) -> Result<Vec<u8>, TreeError> {
    entries
        .first()
        .map(|(key, _value)| key.clone())
        .ok_or(TreeError::InvalidNode)
}

fn first_child_key(children: &[ChildRef]) -> Result<Vec<u8>, TreeError> {
    children
        .first()
        .map(|(key, _hash)| key.clone())
        .ok_or(TreeError::InvalidNode)
}