haematite 0.6.2

use super::{BranchError, BranchHandle, DEFAULT_SHARD_ID, ShardId};
use crate::store::{MemoryStore, NodeStore};
use crate::tree::{Cursor, Hash, LeafNode, Node, TreeError, insert};
use crate::wal::{LookupResult, WalBuffer};
use std::cell::Cell;
use std::convert::Infallible;
use std::fmt::Debug;
use std::sync::Arc;

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

fn empty_root(store: &mut MemoryStore) -> Result<Hash, TreeError> {
    let leaf = LeafNode::new(Vec::new())?;
    Ok(store.put(&Node::Leaf(leaf)))
}

fn build_tree(store: &mut MemoryStore, entries: &[(&[u8], &[u8])]) -> Result<Hash, TreeError> {
    let mut root = empty_root(store)?;
    for (key, value) in entries {
        root = insert(store, root, key, value)?;
    }
    Ok(root)
}

fn assert_debug_clone<T: Debug + Clone>() {}

/// Peek into a shard's buffer through the same state mutex the handle uses.
fn buffer_lookup(
    branch: &BranchHandle,
    shard_id: ShardId,
    key: &[u8],
) -> Result<LookupResult, BranchError> {
    let state = branch.lock_state(shard_id)?;
    Ok(state.buffer.get(key))
}

#[test]
fn branch_handle_records_fork_point_current_root_and_empty_buffer() -> Result<(), BranchError> {
    let root = hash(1);
    let branch = BranchHandle::new(root);

    assert_eq!(branch.fork_point(), root);
    assert_eq!(branch.current_root(), root);
    assert_eq!(branch.primary_shard(), DEFAULT_SHARD_ID);
    assert!(!branch.shard_is_dirty(DEFAULT_SHARD_ID)?);
    assert!(!branch.is_named());
    assert_debug_clone::<BranchHandle>();
    Ok(())
}

#[test]
fn branch_handle_clone_shares_the_same_branch_buffer() -> Result<(), BranchError> {
    let branch = BranchHandle::new(hash(2));
    let clone = branch.clone();

    branch.put(DEFAULT_SHARD_ID, b"key", b"branch")?;

    assert_eq!(
        buffer_lookup(&clone, DEFAULT_SHARD_ID, b"key")?,
        LookupResult::BufferedValue(b"branch".to_vec())
    );
    assert!(clone.shard_is_dirty(DEFAULT_SHARD_ID)?);
    Ok(())
}

#[test]
fn branch_put_does_not_affect_parent_tree() -> Result<(), BranchError> {
    let mut store = MemoryStore::new();
    let root = build_tree(&mut store, &[(b"parent", b"value")])?;
    let branch = BranchHandle::new(root);

    branch.put(DEFAULT_SHARD_ID, b"branch", b"only")?;

    assert_eq!(Cursor::new(&store, root).get(b"branch")?, None);
    assert_eq!(
        branch.get(DEFAULT_SHARD_ID, &store, b"branch")?,
        Some(b"only".to_vec())
    );
    Ok(())
}

#[test]
fn branch_does_not_observe_parent_writes_after_fork() -> Result<(), BranchError> {
    let mut store = MemoryStore::new();
    let root = build_tree(&mut store, &[(b"a", b"1")])?;
    let branch = BranchHandle::new(root);

    let parent_root = insert(&mut store, root, b"b", b"2")?;
    assert_eq!(
        Cursor::new(&store, parent_root).get(b"b")?,
        Some(b"2".to_vec())
    );

    assert_eq!(branch.get(DEFAULT_SHARD_ID, &store, b"b")?, None);
    Ok(())
}

#[test]
fn branch_get_checks_buffer_before_shared_tree() -> Result<(), BranchError> {
    let mut store = MemoryStore::new();
    let root = build_tree(&mut store, &[(b"key", b"tree"), (b"other", b"tree")])?;
    let branch = BranchHandle::new(root);

    branch.put(DEFAULT_SHARD_ID, b"key", b"buffer")?;

    assert_eq!(
        branch.get(DEFAULT_SHARD_ID, &store, b"key")?,
        Some(b"buffer".to_vec())
    );
    assert_eq!(
        branch.get(DEFAULT_SHARD_ID, &store, b"other")?,
        Some(b"tree".to_vec())
    );
    Ok(())
}

#[test]
fn branch_delete_shadows_shared_tree_value() -> Result<(), BranchError> {
    let mut store = MemoryStore::new();
    let root = build_tree(&mut store, &[(b"key", b"tree")])?;
    let branch = BranchHandle::new(root);

    branch.delete(DEFAULT_SHARD_ID, b"key")?;

    assert_eq!(branch.get(DEFAULT_SHARD_ID, &store, b"key")?, None);
    Ok(())
}

#[test]
fn per_shard_roots_and_buffers_are_independent() -> Result<(), BranchError> {
    let shard_three_root = hash(3);
    let shard_five_root = hash(5);
    let branch = BranchHandle::from_shard_roots([(3, shard_three_root), (5, shard_five_root)])?;

    assert_eq!(branch.shard_count(), 2);
    assert_eq!(branch.primary_shard(), 3);
    assert_eq!(branch.shard_fork_point(3), Some(shard_three_root));
    assert_eq!(branch.shard_fork_point(5), Some(shard_five_root));
    let shard_three = branch
        .shard(3)
        .ok_or(BranchError::UnknownShard { shard_id: 3 })?;
    let shard_five = branch
        .shard(5)
        .ok_or(BranchError::UnknownShard { shard_id: 5 })?;
    assert!(!Arc::ptr_eq(&shard_three.state, &shard_five.state));

    branch.put(3, b"same-key", b"shard-three")?;

    assert_eq!(
        buffer_lookup(&branch, 3, b"same-key")?,
        LookupResult::BufferedValue(b"shard-three".to_vec())
    );
    assert_eq!(
        buffer_lookup(&branch, 5, b"same-key")?,
        LookupResult::NotBuffered
    );
    assert!(branch.shard_is_dirty(3)?);
    assert!(!branch.shard_is_dirty(5)?);
    Ok(())
}

#[test]
fn branch_routes_gets_to_the_requested_shard() -> Result<(), BranchError> {
    let mut store = MemoryStore::new();
    let shard_three_root = build_tree(&mut store, &[(b"key", b"three")])?;
    let shard_five_root = build_tree(&mut store, &[(b"key", b"five")])?;
    let branch = BranchHandle::from_shard_roots([(3, shard_three_root), (5, shard_five_root)])?;

    branch.put(3, b"buffered", b"three-buffer")?;
    branch.put(5, b"buffered", b"five-buffer")?;

    assert_eq!(branch.get(3, &store, b"key")?, Some(b"three".to_vec()));
    assert_eq!(branch.get(5, &store, b"key")?, Some(b"five".to_vec()));
    assert_eq!(
        branch.get(3, &store, b"buffered")?,
        Some(b"three-buffer".to_vec())
    );
    assert_eq!(
        branch.get(5, &store, b"buffered")?,
        Some(b"five-buffer".to_vec())
    );
    Ok(())
}

#[test]
fn empty_and_duplicate_shard_construction_errors() {
    assert!(matches!(
        BranchHandle::from_shard_roots([]),
        Err(BranchError::NoShards)
    ));
    assert!(matches!(
        BranchHandle::from_shard_roots([(7, hash(7)), (7, hash(8))]),
        Err(BranchError::DuplicateShard { shard_id: 7 })
    ));
}

#[test]
fn shard_is_dirty_on_unknown_shard_is_a_typed_error() {
    let branch = BranchHandle::new(hash(6));

    assert!(matches!(
        branch.shard_is_dirty(9),
        Err(BranchError::UnknownShard { shard_id: 9 })
    ));
}

#[test]
fn clones_observe_root_advance_and_buffer_reset_atomically() -> Result<(), BranchError> {
    // Clone coherence (BRANCH-COMMIT-PATH.md §3): a commit's step-6 swap —
    // root advance plus buffer reset — happens as one write under one state
    // guard, so a sibling clone sees the whole new generation at once.
    let mut store = MemoryStore::new();
    let old_root = build_tree(&mut store, &[(b"key", b"old")])?;
    let new_root = build_tree(&mut store, &[(b"key", b"new")])?;
    let branch = BranchHandle::new(old_root);
    let clone = branch.clone();

    branch.put(DEFAULT_SHARD_ID, b"key", b"buffered")?;
    assert_eq!(
        clone.get(DEFAULT_SHARD_ID, &store, b"key")?,
        Some(b"buffered".to_vec())
    );
    assert!(clone.shard_is_dirty(DEFAULT_SHARD_ID)?);

    {
        let mut state = branch.lock_state(DEFAULT_SHARD_ID)?;
        state.current_root = new_root;
        state.buffer = WalBuffer::new();
    }

    assert_eq!(clone.shard_current_root(DEFAULT_SHARD_ID), Some(new_root));
    assert_eq!(clone.current_root(), new_root);
    assert!(!clone.shard_is_dirty(DEFAULT_SHARD_ID)?);
    assert_eq!(
        clone.get(DEFAULT_SHARD_ID, &store, b"key")?,
        Some(b"new".to_vec())
    );
    // Only the advanceable half moved: the fork point is immutable.
    assert_eq!(clone.fork_point(), old_root);
    Ok(())
}

#[test]
fn concurrent_clone_never_observes_a_mixed_generation() -> Result<(), BranchError> {
    // Two coherent generations: A = (root_a, buffer holding "gen") and
    // B = (root_b, empty buffer). A flipper thread swaps between them under
    // the state guard while this thread samples (root, buffer) pairs through
    // the same guard. Observing root_a with an empty buffer, or root_b with
    // the marker still buffered, would be exactly the torn read §3's
    // one-mutex-per-shard structure exists to make unrepresentable.
    let root_a = hash(41);
    let root_b = hash(42);
    let branch = BranchHandle::new(root_a);
    branch.put(DEFAULT_SHARD_ID, b"gen", b"a")?;
    let clone = branch.clone();

    let flipper = std::thread::spawn(move || -> Result<(), BranchError> {
        for _ in 0..2_000 {
            let mut state = branch.lock_state(DEFAULT_SHARD_ID)?;
            if state.current_root == root_a {
                state.current_root = root_b;
                state.buffer = WalBuffer::new();
            } else {
                state.current_root = root_a;
                let mut buffer = WalBuffer::new();
                buffer.put(b"gen", b"a");
                state.buffer = buffer;
            }
        }
        Ok(())
    });

    for _ in 0..2_000 {
        let generation = {
            let state = clone.lock_state(DEFAULT_SHARD_ID)?;
            (state.current_root, state.buffer.get(b"gen"))
        };
        let coherent_a = generation == (root_a, LookupResult::BufferedValue(b"a".to_vec()));
        let coherent_b = generation == (root_b, LookupResult::NotBuffered);
        assert!(
            coherent_a || coherent_b,
            "clone observed a mixed root/buffer generation"
        );
    }

    match flipper.join() {
        Ok(result) => result,
        Err(_panic) => Err(BranchError::BufferPoisoned {
            shard_id: DEFAULT_SHARD_ID,
        }),
    }
}

#[derive(Debug)]
struct CountingStore {
    root: Hash,
    node: Node,
    gets: Cell<usize>,
}

impl CountingStore {
    fn new(node: Node) -> Self {
        let root = node.hash();
        Self {
            root,
            node,
            gets: Cell::new(0),
        }
    }
}

impl NodeStore for CountingStore {
    type Error = Infallible;

    fn get(&self, hash: &Hash) -> Result<Option<std::sync::Arc<Node>>, Self::Error> {
        self.gets.set(self.gets.get().saturating_add(1));
        if *hash == self.root {
            Ok(Some(std::sync::Arc::new(self.node.clone())))
        } else {
            Ok(None)
        }
    }

    fn put(&mut self, node: &Node) -> Result<Hash, Self::Error> {
        Ok(node.hash())
    }
}

#[test]
fn branch_buffer_hit_does_not_read_tree_nodes() -> Result<(), TreeError> {
    let leaf = LeafNode::new(vec![(b"key".to_vec(), b"tree".to_vec())])?;
    let store = CountingStore::new(Node::Leaf(leaf));
    let branch = BranchHandle::new(store.root);

    branch
        .put(DEFAULT_SHARD_ID, b"key", b"buffer")
        .map_err(branch_error_to_tree_error)?;
    assert_eq!(
        branch
            .get(DEFAULT_SHARD_ID, &store, b"key")
            .map_err(branch_error_to_tree_error)?,
        Some(b"buffer".to_vec())
    );
    assert_eq!(store.gets.get(), 0);
    Ok(())
}

fn branch_error_to_tree_error(error: BranchError) -> TreeError {
    match error {
        BranchError::Tree(error) => error,
        BranchError::NoShards
        | BranchError::DuplicateShard { .. }
        | BranchError::UnknownShard { .. }
        | BranchError::BufferPoisoned { .. }
        | BranchError::UnregisteredParent => TreeError::InvalidNode,
    }
}