cp_sync/

merkle.rs

//! Merkle tree for cognitive state

use cp_core::Result;

/// Merkle tree for computing state roots
pub struct MerkleTree {
    /// Root hash of the tree
    root: Option<[u8; 32]>,
    /// Leaf hashes
    leaves: Vec<[u8; 32]>,
}

impl MerkleTree {
    /// Create an empty Merkle tree
    pub fn new() -> Self {
        Self {
            root: None,
            leaves: Vec::new(),
        }
    }

    /// Add a leaf to the tree
    pub fn add_leaf(&mut self, data: &[u8]) {
        let hash = blake3::hash(data);
        self.leaves.push(*hash.as_bytes());
        self.root = None; // Invalidate cached root
    }

    /// Compute the root hash
    pub fn root(&mut self) -> [u8; 32] {
        if let Some(root) = self.root {
            return root;
        }

        if self.leaves.is_empty() {
            return [0u8; 32];
        }

        let root = self.compute_root(&self.leaves);
        self.root = Some(root);
        root
    }

    /// Recursively compute root from leaves
    fn compute_root(&self, hashes: &[[u8; 32]]) -> [u8; 32] {
        if hashes.is_empty() {
            return [0u8; 32];
        }
        if hashes.len() == 1 {
            return hashes[0];
        }

        let mut next_level = Vec::with_capacity((hashes.len() + 1) / 2);

        for chunk in hashes.chunks(2) {
            let mut hasher = blake3::Hasher::new();
            hasher.update(&chunk[0]);
            if chunk.len() > 1 {
                hasher.update(&chunk[1]);
            } else {
                // Duplicate last hash if odd number
                hasher.update(&chunk[0]);
            }
            next_level.push(*hasher.finalize().as_bytes());
        }

        self.compute_root(&next_level)
    }

    /// Get number of leaves
    pub fn len(&self) -> usize {
        self.leaves.len()
    }

    /// Check if tree is empty
    pub fn is_empty(&self) -> bool {
        self.leaves.is_empty()
    }

    /// Clear all leaves
    pub fn clear(&mut self) {
        self.leaves.clear();
        self.root = None;
    }

    /// Generate a Merkle inclusion proof for the leaf at the given index.
    ///
    /// Returns a vector of sibling hashes from leaf to root. Each entry is
    /// (sibling_hash, is_left) where is_left indicates the sibling is on the
    /// left side of the pair.
    pub fn generate_proof(&mut self, leaf_index: usize) -> Option<Vec<([u8; 32], bool)>> {
        if leaf_index >= self.leaves.len() {
            return None;
        }

        // Ensure root is computed (builds the tree)
        let _ = self.root();

        let mut proof = Vec::new();
        let mut level = self.leaves.clone();
        let mut index = leaf_index;

        while level.len() > 1 {
            let sibling_index = if index % 2 == 0 {
                // Even index: sibling is to the right
                if index + 1 < level.len() {
                    index + 1
                } else {
                    index // odd number of nodes, duplicate self
                }
            } else {
                // Odd index: sibling is to the left
                index - 1
            };

            let is_left = sibling_index < index;
            proof.push((level[sibling_index], is_left));

            // Build next level
            let mut next_level = Vec::with_capacity((level.len() + 1) / 2);
            for chunk in level.chunks(2) {
                let mut hasher = blake3::Hasher::new();
                hasher.update(&chunk[0]);
                if chunk.len() > 1 {
                    hasher.update(&chunk[1]);
                } else {
                    hasher.update(&chunk[0]);
                }
                next_level.push(*hasher.finalize().as_bytes());
            }

            index /= 2;
            level = next_level;
        }

        Some(proof)
    }

    /// Verify a Merkle inclusion proof for a given leaf hash against a root.
    ///
    /// `proof` is a sequence of (sibling_hash, is_left) from leaf to root.
    pub fn verify_proof(
        leaf_hash: &[u8; 32],
        proof: &[([u8; 32], bool)],
        expected_root: &[u8; 32],
    ) -> bool {
        let mut current = *leaf_hash;

        for (sibling, is_left) in proof {
            let mut hasher = blake3::Hasher::new();
            if *is_left {
                hasher.update(sibling);
                hasher.update(&current);
            } else {
                hasher.update(&current);
                hasher.update(sibling);
            }
            current = *hasher.finalize().as_bytes();
        }

        current == *expected_root
    }
}

impl Default for MerkleTree {
    fn default() -> Self {
        Self::new()
    }
}

/// Compute a Merkle root from a set of items
#[allow(dead_code)]
pub fn compute_merkle_root<T, F>(items: &[T], hash_fn: F) -> Result<[u8; 32]>
where
    F: Fn(&T) -> [u8; 32],
{
    let mut tree = MerkleTree::new();
    for item in items {
        let hash = hash_fn(item);
        tree.leaves.push(hash);
    }
    Ok(tree.root())
}

#[cfg(test)]
mod tests {
    use super::*;
    use cp_core::{Chunk, CognitiveDiff, Document, Edge, EdgeKind, Embedding, Hlc};
    use uuid::Uuid;

    // Helper to create a simple embedding from f32 vector
    fn create_test_embedding(chunk_id: Uuid) -> Embedding {
        // Create a simple f32 vector for testing
        let vector: Vec<f32> = vec![0.0; 1536];
        Embedding::new(chunk_id, &vector, [0u8; 32], 0)
    }

    #[test]
    fn test_empty_tree() {
        let mut tree = MerkleTree::new();
        assert_eq!(tree.root(), [0u8; 32]);
    }

    #[test]
    fn test_single_leaf() {
        let mut tree = MerkleTree::new();
        tree.add_leaf(b"hello");

        let expected = blake3::hash(b"hello");
        assert_eq!(tree.root(), *expected.as_bytes());
    }

    #[test]
    fn test_two_leaves() {
        let mut tree = MerkleTree::new();
        tree.add_leaf(b"hello");
        tree.add_leaf(b"world");

        let root = tree.root();
        assert_ne!(root, [0u8; 32]);
    }

    #[test]
    fn test_deterministic() {
        let mut tree1 = MerkleTree::new();
        tree1.add_leaf(b"a");
        tree1.add_leaf(b"b");
        tree1.add_leaf(b"c");

        let mut tree2 = MerkleTree::new();
        tree2.add_leaf(b"a");
        tree2.add_leaf(b"b");
        tree2.add_leaf(b"c");

        assert_eq!(tree1.root(), tree2.root());
    }

    // Additional comprehensive tests for MerkleTree

    #[test]
    fn test_merkle_three_leaves() {
        let mut tree = MerkleTree::new();
        tree.add_leaf(b"first");
        tree.add_leaf(b"second");
        tree.add_leaf(b"third");

        let root = tree.root();
        assert_ne!(root, [0u8; 32]);
        assert_eq!(tree.len(), 3);
    }

    #[test]
    fn test_merkle_four_leaves() {
        let mut tree = MerkleTree::new();
        tree.add_leaf(b"one");
        tree.add_leaf(b"two");
        tree.add_leaf(b"three");
        tree.add_leaf(b"four");

        let root = tree.root();
        assert_ne!(root, [0u8; 32]);
        assert_eq!(tree.len(), 4);
    }

    #[test]
    fn test_merkle_many_leaves() {
        let mut tree = MerkleTree::new();
        for i in 0..100 {
            tree.add_leaf(format!("item{}", i).as_bytes());
        }

        let root = tree.root();
        assert_ne!(root, [0u8; 32]);
        assert_eq!(tree.len(), 100);
        assert!(!tree.is_empty());
    }

    #[test]
    fn test_merkle_clear() {
        let mut tree = MerkleTree::new();
        tree.add_leaf(b"test");
        assert_eq!(tree.len(), 1);

        tree.clear();
        assert_eq!(tree.len(), 0);
        assert!(tree.is_empty());
        assert_eq!(tree.root(), [0u8; 32]);
    }

    #[test]
    fn test_merkle_root_caching() {
        let mut tree = MerkleTree::new();
        tree.add_leaf(b"data");

        let root1 = tree.root();
        let root2 = tree.root();

        // Should return cached root
        assert_eq!(root1, root2);
    }

    #[test]
    fn test_merkle_root_invalidated_on_new_leaf() {
        let mut tree = MerkleTree::new();
        tree.add_leaf(b"first");

        let root1 = tree.root();

        // Adding new leaf should invalidate cache
        tree.add_leaf(b"second");

        let root2 = tree.root();

        assert_ne!(root1, root2);
    }

    #[test]
    fn test_merkle_different_data_different_root() {
        let mut tree1 = MerkleTree::new();
        tree1.add_leaf(b"hello");
        tree1.add_leaf(b"world");

        let mut tree2 = MerkleTree::new();
        tree2.add_leaf(b"foo");
        tree2.add_leaf(b"bar");

        assert_ne!(tree1.root(), tree2.root());
    }

    #[test]
    fn test_merkle_order_matters() {
        let mut tree1 = MerkleTree::new();
        tree1.add_leaf(b"a");
        tree1.add_leaf(b"b");

        let mut tree2 = MerkleTree::new();
        tree2.add_leaf(b"b");
        tree2.add_leaf(b"a");

        // Different order should produce different root
        assert_ne!(tree1.root(), tree2.root());
    }

    #[test]
    fn test_compute_merkle_root_function() {
        let items = vec![
            [1u8; 32],
            [2u8; 32],
            [3u8; 32],
        ];

        let root = compute_merkle_root(&items, |item| *item).unwrap();
        assert_ne!(root, [0u8; 32]);
    }

    #[test]
    fn test_compute_merkle_root_empty() {
        let items: Vec<[u8; 32]> = vec![];
        let root = compute_merkle_root(&items, |item| *item).unwrap();
        assert_eq!(root, [0u8; 32]);
    }

    #[test]
    fn test_compute_merkle_root_single_item() {
        let items = vec![[42u8; 32]];
        let root = compute_merkle_root(&items, |item| *item).unwrap();
        assert_eq!(root, [42u8; 32]);
    }

    // Tests for diff-related functionality using CognitiveDiff

    #[test]
    fn test_merkle_diff_empty() {
        // Test that an empty diff produces expected results
        let diff = CognitiveDiff::empty(
            [0u8; 32],
            Uuid::nil(),
            0,
            Hlc::new(0, [0u8; 16]),
        );

        assert!(diff.is_empty());
        assert_eq!(diff.change_count(), 0);
    }

    #[test]
    fn test_merkle_diff_document_added() {
        let mut diff = CognitiveDiff::empty(
            [0u8; 32],
            Uuid::nil(),
            0,
            Hlc::new(0, [0u8; 16]),
        );

        let doc = Document::new(
            std::path::PathBuf::from("test.md"),
            b"test content",
            0,
        );
        diff.added_docs.push(doc);

        assert!(!diff.is_empty());
        assert_eq!(diff.change_count(), 1);
        assert_eq!(diff.added_docs.len(), 1);
    }

    #[test]
    fn test_merkle_diff_document_modified() {
        // Document modification is modeled as remove + add
        let mut diff = CognitiveDiff::empty(
            [0u8; 32],
            Uuid::nil(),
            0,
            Hlc::new(0, [0u8; 16]),
        );

        let old_doc_id = Uuid::new_v4();
        diff.removed_doc_ids.push(old_doc_id);

        let new_doc = Document::new(
            std::path::PathBuf::from("test.md"),
            b"updated content",
            0,
        );
        diff.added_docs.push(new_doc);

        assert!(!diff.is_empty());
        assert_eq!(diff.change_count(), 2);
    }

    #[test]
    fn test_merkle_diff_document_deleted() {
        let mut diff = CognitiveDiff::empty(
            [0u8; 32],
            Uuid::nil(),
            0,
            Hlc::new(0, [0u8; 16]),
        );

        let doc_id = Uuid::new_v4();
        diff.removed_doc_ids.push(doc_id);

        assert!(!diff.is_empty());
        assert_eq!(diff.change_count(), 1);
        assert_eq!(diff.removed_doc_ids.len(), 1);
    }

    #[test]
    fn test_merkle_diff_chunk_changes() {
        let mut diff = CognitiveDiff::empty(
            [0u8; 32],
            Uuid::nil(),
            0,
            Hlc::new(0, [0u8; 16]),
        );

        let doc_id = Uuid::new_v4();
        let chunk = Chunk::new(
            doc_id,
            "test chunk content".to_string(),
            0,
            0,
        );
        diff.added_chunks.push(chunk);

        let removed_chunk_id = Uuid::new_v4();
        diff.removed_chunk_ids.push(removed_chunk_id);

        assert_eq!(diff.change_count(), 2);
        assert_eq!(diff.added_chunks.len(), 1);
        assert_eq!(diff.removed_chunk_ids.len(), 1);
    }

    #[test]
    fn test_merkle_diff_embedding_changes() {
        let mut diff = CognitiveDiff::empty(
            [0u8; 32],
            Uuid::nil(),
            0,
            Hlc::new(0, [0u8; 16]),
        );

        let chunk_id = Uuid::new_v4();
        let embedding = create_test_embedding(chunk_id);
        diff.added_embeddings.push(embedding);

        let removed_embedding_id = Uuid::new_v4();
        diff.removed_embedding_ids.push(removed_embedding_id);

        assert_eq!(diff.change_count(), 2);
    }

    #[test]
    fn test_merkle_diff_edge_changes() {
        let mut diff = CognitiveDiff::empty(
            [0u8; 32],
            Uuid::nil(),
            0,
            Hlc::new(0, [0u8; 16]),
        );

        let edge = Edge::new(
            Uuid::new_v4(),
            Uuid::new_v4(),
            EdgeKind::DocToChunk,
        );
        diff.added_edges.push(edge);

        let removed_edge = (
            Uuid::new_v4(),
            Uuid::new_v4(),
            EdgeKind::ChunkToEmbedding,
        );
        diff.removed_edges.push(removed_edge);

        assert_eq!(diff.change_count(), 2);
    }

    #[test]
    fn test_merkle_diff_multiple_changes() {
        let mut diff = CognitiveDiff::empty(
            [0u8; 32],
            Uuid::nil(),
            0,
            Hlc::new(0, [0u8; 16]),
        );

        // Add document
        let doc = Document::new(
            std::path::PathBuf::from("test.md"),
            b"content",
            0,
        );
        diff.added_docs.push(doc);

        // Add chunk
        let chunk = Chunk::new(
            Uuid::new_v4(),
            "content".to_string(),
            0,
            0,
        );
        diff.added_chunks.push(chunk);

        // Add embedding
        let embedding = create_test_embedding(Uuid::new_v4());
        diff.added_embeddings.push(embedding);

        // Add edge
        let edge = Edge::new(
            Uuid::new_v4(),
            Uuid::new_v4(),
            EdgeKind::DocToChunk,
        );
        diff.added_edges.push(edge);

        assert_eq!(diff.change_count(), 4);
    }

    #[test]
    fn test_merkle_diff_serialization() {
        use crate::{serialize_diff, deserialize_diff};

        let diff = CognitiveDiff::empty(
            [0u8; 32],
            Uuid::nil(),
            0,
            Hlc::new(0, [0u8; 16]),
        );

        let serialized = serialize_diff(&diff).unwrap();
        assert!(!serialized.is_empty());

        let deserialized = deserialize_diff(&serialized).unwrap();
        assert_eq!(diff.metadata, deserialized.metadata);
    }

    #[test]
    fn test_merkle_diff_serialization_with_content() {
        use crate::{serialize_diff, deserialize_diff};

        let mut diff = CognitiveDiff::empty(
            [0u8; 32],
            Uuid::nil(),
            0,
            Hlc::new(0, [0u8; 16]),
        );

        let doc = Document::new(
            std::path::PathBuf::from("test.md"),
            b"Test content",
            0,
        );
        diff.added_docs.push(doc);

        let serialized = serialize_diff(&diff).unwrap();
        let deserialized = deserialize_diff(&serialized).unwrap();

        assert_eq!(diff.added_docs.len(), deserialized.added_docs.len());
        assert_eq!(diff.added_docs[0].path, deserialized.added_docs[0].path);
    }

    #[test]
    fn test_merkle_diff_estimated_size() {
        let diff = CognitiveDiff::empty(
            [0u8; 32],
            Uuid::nil(),
            0,
            Hlc::new(0, [0u8; 16]),
        );

        let size = diff.estimated_size();
        assert!(size > 0);
    }

    #[test]
    fn test_merkle_diff_apply_order() {
        // Tests that changes are tracked in the correct order
        let mut diff = CognitiveDiff::empty(
            [0u8; 32],
            Uuid::nil(),
            0,
            Hlc::new(0, [0u8; 16]),
        );

        // Add in specific order
        for i in 0..5 {
            let doc = Document::new(
                std::path::PathBuf::from(format!("doc{}.md", i)),
                format!("content{}", i).as_bytes(),
                0,
            );
            diff.added_docs.push(doc);
        }

        assert_eq!(diff.added_docs.len(), 5);
        // Verify order is preserved
        for (i, doc) in diff.added_docs.iter().enumerate() {
            assert!(doc.path.to_string_lossy().contains(&format!("{}", i)));
        }
    }

    #[test]
    fn test_merkle_diff_idempotent() {
        // Test that adding the same item twice creates a diff that would
        // contain both additions (idempotency is handled at application time)
        let mut diff = CognitiveDiff::empty(
            [0u8; 32],
            Uuid::nil(),
            0,
            Hlc::new(0, [0u8; 16]),
        );

        let doc = Document::new(
            std::path::PathBuf::from("test.md"),
            b"content",
            0,
        );

        // Add same document type twice (in real use, this would be deduplicated at apply time)
        diff.added_docs.push(doc.clone());
        diff.added_docs.push(doc);

        assert_eq!(diff.change_count(), 2);
    }

    #[test]
    fn test_merkle_comparison_using_blake3() {
        // Test using Merkle tree to compare two states
        let mut old_tree = MerkleTree::new();
        old_tree.add_leaf(b"document1");
        old_tree.add_leaf(b"document2");

        let mut new_tree = MerkleTree::new();
        new_tree.add_leaf(b"document1");
        new_tree.add_leaf(b"document2");
        new_tree.add_leaf(b"document3"); // New document

        let old_root = old_tree.root();
        let new_root = new_tree.root();

        // Roots should be different because states are different
        assert_ne!(old_root, new_root);
    }

    #[test]
    fn test_merkle_proof_single_leaf() {
        let mut tree = MerkleTree::new();
        tree.add_leaf(b"only");

        let root = tree.root();
        let proof = tree.generate_proof(0).unwrap();
        assert!(proof.is_empty(), "Single leaf proof should be empty");

        let leaf_hash = blake3::hash(b"only");
        assert!(MerkleTree::verify_proof(leaf_hash.as_bytes(), &proof, &root));
    }

    #[test]
    fn test_merkle_proof_two_leaves() {
        let mut tree = MerkleTree::new();
        tree.add_leaf(b"left");
        tree.add_leaf(b"right");

        let root = tree.root();

        // Proof for leaf 0
        let proof0 = tree.generate_proof(0).unwrap();
        let leaf0 = blake3::hash(b"left");
        assert!(MerkleTree::verify_proof(leaf0.as_bytes(), &proof0, &root));

        // Proof for leaf 1
        let proof1 = tree.generate_proof(1).unwrap();
        let leaf1 = blake3::hash(b"right");
        assert!(MerkleTree::verify_proof(leaf1.as_bytes(), &proof1, &root));
    }

    #[test]
    fn test_merkle_proof_many_leaves() {
        let mut tree = MerkleTree::new();
        let data: Vec<Vec<u8>> = (0..17).map(|i| format!("leaf_{}", i).into_bytes()).collect();
        for d in &data {
            tree.add_leaf(d);
        }

        let root = tree.root();

        for (i, d) in data.iter().enumerate() {
            let proof = tree.generate_proof(i).unwrap();
            let leaf_hash = blake3::hash(d);
            assert!(
                MerkleTree::verify_proof(leaf_hash.as_bytes(), &proof, &root),
                "Proof failed for leaf {}",
                i
            );
        }
    }

    #[test]
    fn test_merkle_proof_wrong_leaf_fails() {
        let mut tree = MerkleTree::new();
        tree.add_leaf(b"a");
        tree.add_leaf(b"b");
        tree.add_leaf(b"c");

        let root = tree.root();

        let proof = tree.generate_proof(0).unwrap();
        let wrong_leaf = blake3::hash(b"wrong");
        assert!(!MerkleTree::verify_proof(wrong_leaf.as_bytes(), &proof, &root));
    }

    #[test]
    fn test_merkle_proof_out_of_bounds() {
        let mut tree = MerkleTree::new();
        tree.add_leaf(b"a");
        assert!(tree.generate_proof(1).is_none());
        assert!(tree.generate_proof(100).is_none());
    }

    #[test]
    fn test_merkle_prove_change_detection() {
        // Test that Merkle tree can detect changes
        let mut tree1 = MerkleTree::new();
        tree1.add_leaf(b"a");
        tree1.add_leaf(b"b");

        let mut tree2 = MerkleTree::new();
        tree2.add_leaf(b"a");
        tree2.add_leaf(b"c"); // Changed from b to c

        assert_ne!(tree1.root(), tree2.root());
    }
}