tree.rs

/*
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

use crate::config::TreeConfig;
use crate::diff::{
    ConflictResolver, DiffResult, IgnoreConflictsResolver, MergeConflict, MergeResult,
};
use crate::digest::ValueDigest;
use crate::node::{Node, ProllyNode};
use crate::proof::Proof;
use crate::storage::NodeStorage;
use std::sync::Arc;

/// Trait representing a Prolly tree with a fixed size N and a node storage S.
/// This trait provides methods for creating, modifying, and querying the tree.
pub trait Tree<const N: usize, S: NodeStorage<N>> {
    /// Creates a new Prolly tree with the specified root node and storage.
    ///
    /// # Parameters
    /// - `storage`: The storage to use for persisting nodes.
    /// - `config`: The configuration for the tree.
    ///
    /// # Returns
    /// - A new instance of the tree.
    fn new(storage: S, config: TreeConfig<N>) -> Self;

    /// Inserts a key-value pair into the tree.
    ///
    /// # Parameters
    /// - `key`: The key to insert.
    /// - `value`: The value associated with the key.
    fn insert(&mut self, key: Vec<u8>, value: Vec<u8>);

    /// Inserts multiple key-value pairs into the tree in an optimized way.
    ///
    /// # Parameters
    /// - `keys`: The keys to insert.
    /// - `values`: The values associated with the keys.
    fn insert_batch(&mut self, keys: &[Vec<u8>], values: &[Vec<u8>]);

    /// Updates the value associated with the specified key in the tree.
    ///
    /// # Parameters
    /// - `key`: The key to update.
    /// - `value`: The new value to associate with the key.
    ///
    /// # Returns
    /// - `true` if the key was found and updated, `false` otherwise.
    fn update(&mut self, key: Vec<u8>, value: Vec<u8>) -> bool;

    /// Deletes the key-value pair associated with the specified key from the tree.
    ///
    /// # Parameters
    /// - `key`: The key to delete.
    ///
    /// # Returns
    /// - `true` if the key was found and deleted, `false` otherwise.
    fn delete(&mut self, key: &[u8]) -> bool;

    /// Deletes multiple key-value pairs from the tree.
    ///
    /// # Parameters
    /// - `keys`: The keys to delete.
    fn delete_batch(&mut self, keys: &[Vec<u8>]);

    /// Finds the node associated with the specified key in the tree.
    ///
    /// # Parameters
    /// - `key`: The key to find.
    ///
    /// # Returns
    /// - `Some(ProllyNode<N>)` if the key was found, `None` otherwise.
    fn find(&self, key: &[u8]) -> Option<ProllyNode<N>>;

    /// Traverses the tree and returns a string representation of its structure.
    ///
    /// # Returns
    /// - A string representation of the tree structure.
    fn traverse(&self) -> String;

    /// Traverses the tree and returns a formatted string representation using the provided formatter function.
    ///
    /// # Parameters
    /// - `formatter`: A function to format each node.
    ///
    /// # Returns
    /// - A formatted string representation of the tree structure.
    fn formatted_traverse<F>(&self, formatter: F) -> String
    where
        F: Fn(&ProllyNode<N>) -> String;

    /// Gets the hash of the root node of the tree.
    ///
    /// # Returns
    /// - `Some(ValueDigest<N>)` if the root node exists, `None` otherwise.
    fn get_root_hash(&self) -> Option<ValueDigest<N>>;

    /// Gets the number of nodes in the tree.
    ///
    /// # Returns
    /// - The number of nodes in the tree.
    fn size(&self) -> usize;

    /// Gets the depth of the tree.
    ///
    /// # Returns
    /// - The depth of the tree.
    fn depth(&self) -> usize;

    /// Provides a summary of the tree structure and contents.
    ///
    /// # Returns
    /// - A summary of the tree.
    fn summary(&self) -> String;

    /// Provides various statistics about the tree.
    ///
    /// # Returns
    /// - A `TreeStats` object containing statistics about the tree.
    fn stats(&self) -> TreeStats;

    /// Loads the configuration for the tree from storage.
    ///
    /// # Parameters
    /// - `storage`: The storage to load the configuration from.
    fn load_config(storage: &S) -> Result<TreeConfig<N>, &'static str>;

    /// Saves the configuration for the tree to storage.
    ///
    /// # Returns
    /// - `Ok(())` if the configuration was saved successfully, `Err(&'static str)` otherwise.
    fn save_config(&self) -> Result<(), &'static str>;

    /// Generates a proof of existence for a given key in the tree.
    ///
    /// This function traverses the tree from the root to the target node containing the key,
    /// collecting the hashes of all nodes along the path. The proof can be used to verify the
    /// existence of the key and its associated value without revealing other data in the tree.
    ///
    /// # Arguments
    ///
    /// * `key` - The key for which to generate the proof.
    ///
    /// # Returns
    ///
    /// A `Proof` struct containing the path of hashes and the hash of the target node (if the key exists).
    fn generate_proof(&self, key: &[u8]) -> Proof<N>;

    fn verify(&self, proof: Proof<N>, key: &[u8], expected_value: Option<&[u8]>) -> bool;

    /// Computes the differences between two Prolly Trees.
    ///
    /// This function compares the current tree (`self`) with another tree (`other`)
    /// and identifies the differences between them. It traverses both trees and
    /// generates a list of changes, including added, removed, and modified key-value pairs.
    ///
    /// # Arguments
    ///
    /// * `other` - The other Prolly Tree to compare against.
    ///
    /// # Returns
    ///
    /// A vector of `DiffResult` containing the differences between the two trees.
    fn diff(&self, other: &Self) -> Vec<DiffResult>;

    /// Prints the tree structure to the console.
    /// This function is useful for debugging and visualizing the tree.
    /// It prints the tree structure in a human-readable format.
    /// The tree is printed in a depth-first manner, starting from the root node.
    /// Each node is printed with its keys and values, along with the hash of the node.
    ///
    fn print(&mut self);

    /// Prints the tree structure with the proof path highlighted for a given key.
    /// This function combines `generate_proof` and `print` to visualize the
    /// cryptographic proof path through the tree structure with color coding.
    ///
    /// # Arguments
    ///
    /// * `key` - The key for which to generate and display the proof path.
    ///
    /// # Returns
    ///
    /// A boolean indicating whether the proof is valid.
    fn print_proof(&self, key: &[u8]) -> bool;

    /// Performs a three-way merge between source, destination and base trees.
    ///
    /// This function implements a three-way merge algorithm for prolly trees.
    /// Given three tree root hashes (base, source, destination), it computes
    /// the differences between base->source and base->destination, then merges
    /// changes from source into destination, detecting conflicts when both
    /// branches modify the same key with different values.
    ///
    /// # Arguments
    ///
    /// * `source_root` - Root hash of the source (feature) tree
    /// * `destination_root` - Root hash of the destination (main) tree
    /// * `base_root` - Root hash of the common base tree
    ///
    /// # Returns
    ///
    /// A vector of `MergeResult` indicating the changes to apply and any conflicts
    fn merge(
        &self,
        source_root: &ValueDigest<N>,
        destination_root: &ValueDigest<N>,
        base_root: &ValueDigest<N>,
    ) -> Vec<MergeResult>;

    /// Applies merge results to create a new merged tree.
    ///
    /// This method takes the destination tree and applies a list of merge results
    /// to create a new tree with all the merged changes. If any conflicts are
    /// present in the merge results, this method will return an error.
    ///
    /// # Arguments
    ///
    /// * `destination_root` - Root hash of the destination tree to merge into
    /// * `merge_results` - List of merge operations to apply
    ///
    /// # Returns
    ///
    /// A new `ProllyTree` instance with merged changes, or an error if conflicts exist
    fn apply_merge_results(
        &self,
        destination_root: &ValueDigest<N>,
        merge_results: &[MergeResult],
    ) -> Result<Self, Vec<MergeConflict>>
    where
        Self: Sized;

    /// Convenience method to perform a three-way merge with conflict resolution.
    ///
    /// This method combines `merge()` and `apply_merge_results()` with a conflict resolver
    /// to provide a flexible interface for merging. Conflicts that can't be resolved
    /// are returned for manual resolution.
    ///
    /// # Arguments
    ///
    /// * `source_root` - Root hash of the source (feature) tree
    /// * `destination_root` - Root hash of the destination (main) tree
    /// * `base_root` - Root hash of the common base tree
    /// * `resolver` - Conflict resolver to handle merge conflicts
    ///
    /// # Returns
    ///
    /// Either a new merged tree or a list of unresolved conflicts
    fn merge_trees<R: ConflictResolver>(
        &self,
        source_root: &ValueDigest<N>,
        destination_root: &ValueDigest<N>,
        base_root: &ValueDigest<N>,
        resolver: &R,
    ) -> Result<Self, Vec<MergeConflict>>
    where
        Self: Sized;

    /// Convenience method for merge_trees with default IgnoreConflictsResolver
    fn merge_trees_ignore_conflicts(
        &self,
        source_root: &ValueDigest<N>,
        destination_root: &ValueDigest<N>,
        base_root: &ValueDigest<N>,
    ) -> Result<Self, Vec<MergeConflict>>
    where
        Self: Sized,
    {
        self.merge_trees(
            source_root,
            destination_root,
            base_root,
            &IgnoreConflictsResolver,
        )
    }
}

pub struct TreeStats {
    pub num_nodes: usize,
    pub num_leaves: usize,
    pub num_internal_nodes: usize,
    pub avg_node_size: f64,
    pub total_key_value_pairs: usize,
}

impl TreeStats {
    pub fn new() -> Self {
        TreeStats {
            num_nodes: 0,
            num_leaves: 0,
            num_internal_nodes: 0,
            avg_node_size: 0.0,
            total_key_value_pairs: 0,
        }
    }
}

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

#[derive(Debug, Clone)]
pub struct ProllyTree<const N: usize, S: NodeStorage<N>> {
    pub root: ProllyNode<N>,
    pub storage: S,
    pub config: TreeConfig<N>,
}

impl<const N: usize, S: NodeStorage<N>> Tree<N, S> for ProllyTree<N, S> {
    fn new(storage: S, config: TreeConfig<N>) -> Self {
        let root = ProllyNode {
            keys: Vec::new(),
            key_schema: config.key_schema.clone(),
            values: Vec::new(),
            value_schema: config.value_schema.clone(),
            is_leaf: true,
            level: 0,
            base: config.base,
            modulus: config.modulus,
            min_chunk_size: config.min_chunk_size,
            max_chunk_size: config.max_chunk_size,
            pattern: config.pattern,
            split: false,
            merged: false,
            encode_types: Vec::new(),
            encode_values: Vec::new(),
        };
        let root_hash = Some(root.get_hash());
        let mut tree = ProllyTree {
            root,
            storage,
            config,
        };
        tree.config.root_hash = root_hash;
        tree
    }
    fn insert(&mut self, key: Vec<u8>, value: Vec<u8>) {
        // Stream the single mutation through the cursor-driven chunker.
        // This bypasses the legacy in-place balance code (which doesn't
        // maintain history independence) and produces a canonical tree
        // by construction.
        self.apply_changes(std::iter::once((key, Some(value))));
        self.persist_root();
    }

    fn insert_batch(&mut self, keys: &[Vec<u8>], values: &[Vec<u8>]) {
        assert_eq!(
            keys.len(),
            values.len(),
            "insert_batch requires the same number of keys and values"
        );
        let batch = keys.iter().cloned().zip(values.iter().cloned().map(Some));
        self.apply_changes(batch);
    }

    fn update(&mut self, key: Vec<u8>, value: Vec<u8>) -> bool {
        if self.find(&key).is_some() {
            self.insert(key, value);
            true
        } else {
            false
        }
    }

    fn delete(&mut self, key: &[u8]) -> bool {
        // Pre-check: was the key present? Same answer as before but
        // computed via the cursor walk inside apply_changes' probe path.
        if self.find(key).is_none() {
            return false;
        }
        self.apply_changes(std::iter::once((key.to_vec(), None)));
        self.persist_root();
        true
    }

    fn delete_batch(&mut self, keys: &[Vec<u8>]) {
        let batch = keys.iter().map(|k| (k.clone(), None));
        self.apply_changes(batch);
    }

    fn find(&self, key: &[u8]) -> Option<ProllyNode<N>> {
        self.root.find(key, &self.storage)
    }

    fn traverse(&self) -> String {
        self.root.traverse(&self.storage)
    }

    fn formatted_traverse<F>(&self, formatter: F) -> String
    where
        F: Fn(&ProllyNode<N>) -> String,
    {
        self.root.formatted_traverse(&self.storage, formatter)
    }

    fn get_root_hash(&self) -> Option<ValueDigest<N>> {
        Option::from(self.root.get_hash())
    }

    fn size(&self) -> usize {
        fn count_pairs<const N: usize, S: NodeStorage<N>>(
            node: &ProllyNode<N>,
            storage: &S,
        ) -> usize {
            if node.is_leaf {
                node.keys.len()
            } else {
                let mut count = 0;
                for value in &node.values {
                    if let Some(child_node) =
                        storage.get_node_by_hash(&ValueDigest::raw_hash(value))
                    {
                        count += count_pairs(&child_node, storage);
                    }
                }
                count
            }
        }

        count_pairs(&self.root, &self.storage)
    }

    fn depth(&self) -> usize {
        (self.root.level as usize) + 1
    }

    fn summary(&self) -> String {
        let stats = self.stats();
        format!(
            "Tree Summary:\n- Number of Key-Value Pairs: {}\n- Number of Nodes: {}\n- Number of Leaves: {}\n- Number of Internal Nodes: {}\n- Average Leaf Node Size: {:.2}",
            self.size(),
            stats.num_nodes,
            stats.num_leaves,
            stats.num_internal_nodes,
            stats.avg_node_size
        )
    }

    fn stats(&self) -> TreeStats {
        fn collect_stats<const N: usize, S: NodeStorage<N>>(
            node: &ProllyNode<N>,
            storage: &S,
            stats: &mut TreeStats,
        ) {
            stats.num_nodes += 1;
            if node.is_leaf {
                stats.num_leaves += 1;
                stats.total_key_value_pairs += node.keys.len();
            } else {
                stats.num_internal_nodes += 1;
                for value in &node.values {
                    if let Some(child_node) =
                        storage.get_node_by_hash(&ValueDigest::raw_hash(value))
                    {
                        collect_stats(&child_node, storage, stats);
                    }
                }
            }
        }

        let mut stats = TreeStats::new();
        collect_stats(&self.root, &self.storage, &mut stats);
        if stats.num_leaves > 0 {
            stats.avg_node_size = stats.total_key_value_pairs as f64 / stats.num_leaves as f64;
        }
        stats
    }

    fn load_config(storage: &S) -> Result<TreeConfig<N>, &'static str> {
        // Implement the logic to load the configuration from storage
        // Here we assume the config is stored with a specific key "tree_config"
        if let Some(config_data) = storage.get_config("tree_config") {
            let config: TreeConfig<N> =
                serde_json::from_slice(&config_data).map_err(|_| "Failed to deserialize config")?;
            Ok(config)
        } else {
            Err("Config not found")
        }
    }

    fn save_config(&self) -> Result<(), &'static str> {
        let mut config = self.config.clone();
        config.root_hash = Option::from(self.root.get_hash());
        let config_data = serde_json::to_vec(&config).map_err(|_| "Failed to serialize config")?;
        self.storage.save_config("tree_config", &config_data);
        Ok(())
    }

    /// Generates a proof of existence for a given key in the tree.
    ///
    /// This function traverses the tree from the root to the target node containing the key,
    /// collecting the hashes of all nodes along the path. The proof can be used to verify the
    /// existence of the key and its associated value without revealing other data in the tree.
    ///
    /// # Arguments
    ///
    /// * `key` - The key for which to generate the proof.
    /// * `storage` - The storage implementation to retrieve child nodes.
    ///
    /// # Returns
    ///
    /// A `Proof` struct containing the path of hashes and the hash of the target node (if the key exists).
    fn generate_proof(&self, key: &[u8]) -> Proof<N> {
        /// Recursive helper function to generate the proof path.
        ///
        /// This function traverses the tree from the given node to the target node containing the key,
        /// collecting the hashes of all nodes along the path. It returns the hash of the target node
        /// if the key exists, or `None` if the key does not exist.
        ///
        /// # Arguments
        ///
        /// * `node` - The current node being traversed.
        /// * `key` - The key for which to generate the proof.
        /// * `storage` - The storage implementation to retrieve child nodes.
        /// * `path` - The vector to store the hashes of the nodes along the path.
        ///
        /// # Returns
        ///
        /// The hash of the target node if the key exists, or `None` if the key does not exist.
        fn generate_proof_recursive<const N: usize, S: NodeStorage<N>>(
            node: &ProllyNode<N>,
            key: &[u8],
            storage: &S,
            path: &mut Vec<ValueDigest<N>>,
        ) -> Option<ValueDigest<N>> {
            path.push(node.get_hash());

            if node.is_leaf {
                if node.keys.iter().any(|k| k == key) {
                    Some(node.get_hash())
                } else {
                    None
                }
            } else {
                // Mirror `ProllyNode::find`: after certain delete patterns an internal
                // node can transiently have no children (or fewer values than keys), so
                // guard the empty case and clamp the index instead of indexing out of
                // bounds and panicking.
                if node.values.is_empty() {
                    return None;
                }
                let i = node.keys.iter().rposition(|k| key >= &k[..]).unwrap_or(0);
                let i = i.min(node.values.len() - 1);
                let child_hash = node.values[i].clone();

                if let Some(child_node) =
                    storage.get_node_by_hash(&ValueDigest::raw_hash(&child_hash))
                {
                    generate_proof_recursive(&child_node, key, storage, path)
                } else {
                    None
                }
            }
        }

        let mut path = Vec::new();
        let target_hash = generate_proof_recursive(&self.root, key, &self.storage, &mut path);

        Proof { path, target_hash }
    }

    fn verify(&self, proof: Proof<N>, key: &[u8], expected_value: Option<&[u8]>) -> bool {
        // Start with the root hash
        let mut current_hash = self.root.get_hash();

        for (i, node_hash) in proof.path.iter().enumerate() {
            // Retrieve the node content from storage using the current hash
            if let Some(node) = self.storage.get_node_by_hash(&current_hash) {
                // Check if the current node's hash matches the expected hash in the path
                if node.get_hash() != *node_hash {
                    return false;
                }

                // If it's the last node in the path, verify the leaf node
                if i == proof.path.len() - 1 {
                    return if node.is_leaf {
                        node.keys.iter().any(|k| k == key)
                            && match expected_value {
                                None => true,
                                Some(ev) => node.values.iter().any(|v| ev == &v[..]),
                            }
                    } else {
                        false // Path should end at a leaf node
                    };
                }

                // Move to the next node in the path by finding the correct child
                let child_index = node.keys.iter().rposition(|k| key >= &k[..]).unwrap_or(0);
                current_hash = ValueDigest::raw_hash(&node.values[child_index]);
            } else {
                // If the node is not found in storage, the proof is invalid
                return false;
            }
        }

        false // If we exit the loop without verifying, the proof is invalid
    }

    fn diff(&self, other: &Self) -> Vec<DiffResult> {
        let mut diffs = Vec::new();
        self.diff_recursive(&self.root, &other.root, &mut diffs);
        diffs
    }

    fn print(&mut self) {
        self.root.print_tree(&self.storage);
    }

    fn print_proof(&self, key: &[u8]) -> bool {
        // Generate the proof for the given key
        let proof = self.generate_proof(key);

        // Verify the proof
        let is_valid = self.verify(proof.clone(), key, None);

        // Print the tree structure with proof path highlighted
        #[cfg(feature = "tracing")]
        tracing::debug!("root:");
        #[cfg(not(feature = "tracing"))]
        println!("root:");

        self.root.print_tree_with_proof(&self.storage, &proof, key);

        // Print proof information
        #[cfg(feature = "tracing")]
        {
            tracing::debug!("Proof for key {:?} is valid: {}", key, is_valid);
            tracing::debug!("Proof: {:#?}", proof);
        }
        #[cfg(not(feature = "tracing"))]
        {
            println!("\nProof for key {key:?} is valid: {is_valid}");
            println!("Proof: {proof:#?}");
        }

        is_valid
    }

    fn merge(
        &self,
        source_root: &ValueDigest<N>,
        destination_root: &ValueDigest<N>,
        base_root: &ValueDigest<N>,
    ) -> Vec<MergeResult> {
        // Load trees from storage using the provided root hashes
        let source_tree = self.storage.get_node_by_hash(source_root);
        let destination_tree = self.storage.get_node_by_hash(destination_root);
        let base_tree = self.storage.get_node_by_hash(base_root);

        let (source_tree, destination_tree, base_tree) =
            match (source_tree, destination_tree, base_tree) {
                (Some(s), Some(d), Some(b)) => (s, d, b),
                _ => {
                    // If we can't load one of the trees, return an error as a conflict
                    return vec![MergeResult::Conflict(MergeConflict {
                        key: b"<merge_error>".to_vec(),
                        base_value: None,
                        source_value: None,
                        destination_value: Some(b"Failed to load tree from storage".to_vec()),
                    })];
                }
            };

        // Compute diffs directly using the node-level diffing
        let mut base_to_source_diffs = Vec::new();
        let mut base_to_destination_diffs = Vec::new();

        self.diff_nodes_recursive(&base_tree, &source_tree, &mut base_to_source_diffs);
        self.diff_nodes_recursive(
            &base_tree,
            &destination_tree,
            &mut base_to_destination_diffs,
        );

        // Convert diffs to maps for easier processing
        let mut source_changes: std::collections::HashMap<Vec<u8>, DiffResult> =
            std::collections::HashMap::new();
        let mut destination_changes: std::collections::HashMap<Vec<u8>, DiffResult> =
            std::collections::HashMap::new();

        for diff in base_to_source_diffs {
            let key = match &diff {
                DiffResult::Added(k, _) => k.clone(),
                DiffResult::Removed(k, _) => k.clone(),
                DiffResult::Modified(k, _, _) => k.clone(),
            };
            source_changes.insert(key, diff);
        }

        for diff in base_to_destination_diffs {
            let key = match &diff {
                DiffResult::Added(k, _) => k.clone(),
                DiffResult::Removed(k, _) => k.clone(),
                DiffResult::Modified(k, _, _) => k.clone(),
            };
            destination_changes.insert(key, diff);
        }

        // Collect all keys that were changed in either branch
        let mut all_changed_keys = std::collections::HashSet::new();
        for key in source_changes.keys() {
            all_changed_keys.insert(key.clone());
        }
        for key in destination_changes.keys() {
            all_changed_keys.insert(key.clone());
        }

        let mut merge_results = Vec::new();

        // Process each changed key
        for key in all_changed_keys {
            let source_change = source_changes.get(&key);
            let destination_change = destination_changes.get(&key);

            match (source_change, destination_change) {
                // Only source changed - apply source change
                (Some(source_diff), None) => match source_diff {
                    DiffResult::Added(_, value) => {
                        merge_results.push(MergeResult::Added(key, value.clone()));
                    }
                    DiffResult::Removed(_, _) => {
                        merge_results.push(MergeResult::Removed(key));
                    }
                    DiffResult::Modified(_, _, new_value) => {
                        merge_results.push(MergeResult::Modified(key, new_value.clone()));
                    }
                },

                // Only destination changed - no action needed (destination already has the change)
                (None, Some(_)) => {
                    // Destination change already exists, no merge action needed
                }

                // Both changed - need to check for conflicts
                (Some(source_diff), Some(destination_diff)) => {
                    let conflict =
                        self.detect_conflict(&key, source_diff, destination_diff, &base_tree);

                    if let Some(conflict) = conflict {
                        merge_results.push(MergeResult::Conflict(conflict));
                    } else {
                        // No conflict, apply source change (assuming identical changes)
                        match source_diff {
                            DiffResult::Added(_, value) => {
                                merge_results.push(MergeResult::Added(key, value.clone()));
                            }
                            DiffResult::Removed(_, _) => {
                                merge_results.push(MergeResult::Removed(key));
                            }
                            DiffResult::Modified(_, _, new_value) => {
                                merge_results.push(MergeResult::Modified(key, new_value.clone()));
                            }
                        }
                    }
                }

                // Neither changed (shouldn't happen due to our key collection logic)
                (None, None) => {}
            }
        }

        merge_results
    }

    fn apply_merge_results(
        &self,
        destination_root: &ValueDigest<N>,
        merge_results: &[MergeResult],
    ) -> Result<Self, Vec<MergeConflict>> {
        // Check for conflicts first
        let mut conflicts = Vec::new();
        for result in merge_results {
            if let MergeResult::Conflict(conflict) = result {
                conflicts.push((*conflict).clone());
            }
        }

        if !conflicts.is_empty() {
            return Err(conflicts);
        }

        // Load the destination tree
        let destination_tree =
            self.storage
                .get_node_by_hash(destination_root)
                .ok_or_else(|| {
                    vec![MergeConflict {
                        key: b"<apply_error>".to_vec(),
                        base_value: None,
                        source_value: None,
                        destination_value: Some(b"Failed to load destination tree".to_vec()),
                    }]
                })?;

        // Create a new tree starting from the destination
        let mut new_tree = ProllyTree {
            root: Arc::unwrap_or_clone(destination_tree),
            storage: self.storage.clone(),
            config: self.config.clone(),
        };

        // Apply each merge result
        for result in merge_results {
            match result {
                MergeResult::Added(key, value) => {
                    new_tree.insert(key.clone(), value.clone());
                }
                MergeResult::Modified(key, value) => {
                    new_tree.insert(key.clone(), value.clone()); // insert overwrites existing
                }
                MergeResult::Removed(key) => {
                    new_tree.delete(key);
                }
                MergeResult::Conflict(_) => {
                    // This should not happen since we checked for conflicts above
                    unreachable!("Conflicts should have been filtered out");
                }
            }
        }

        Ok(new_tree)
    }

    fn merge_trees<R: ConflictResolver>(
        &self,
        source_root: &ValueDigest<N>,
        destination_root: &ValueDigest<N>,
        base_root: &ValueDigest<N>,
        resolver: &R,
    ) -> Result<Self, Vec<MergeConflict>> {
        let merge_results = self.merge(source_root, destination_root, base_root);

        // Separate conflicts from other merge results and try to resolve conflicts
        let mut resolved_results = Vec::new();
        let mut unresolved_conflicts = Vec::new();

        for result in merge_results {
            match result {
                MergeResult::Conflict(conflict) => {
                    if let Some(resolved_result) = resolver.resolve_conflict(&conflict) {
                        resolved_results.push(resolved_result);
                    } else {
                        unresolved_conflicts.push(conflict);
                    }
                }
                other => resolved_results.push(other),
            }
        }

        // If there are still unresolved conflicts, return them
        if !unresolved_conflicts.is_empty() {
            return Err(unresolved_conflicts);
        }

        // Apply the resolved results
        self.apply_merge_results(destination_root, &resolved_results)
    }
}

impl<const N: usize, S: NodeStorage<N>> ProllyTree<N, S> {
    /// Compute differences between two nodes recursively
    fn diff_nodes_recursive(
        &self,
        old_node: &ProllyNode<N>,
        new_node: &ProllyNode<N>,
        diffs: &mut Vec<DiffResult>,
    ) {
        // O(differences) structural diff (Dolt/Noms): the tree is content-addressed,
        // so equal get_hash() => byte-identical subtree => zero diffs; skip it without
        // loading. Descend only into differing children. Emits the IDENTICAL DiffResult
        // set as the full-leaf flatten (diff_nodes_flatten), guarded by a differential test.
        if old_node.get_hash() == new_node.get_hash() {
            return;
        }
        match (old_node.is_leaf, new_node.is_leaf) {
            (true, true) => {
                let o: Vec<(Vec<u8>, Vec<u8>)> = old_node
                    .keys
                    .iter()
                    .cloned()
                    .zip(old_node.values.iter().cloned())
                    .collect();
                let n: Vec<(Vec<u8>, Vec<u8>)> = new_node
                    .keys
                    .iter()
                    .cloned()
                    .zip(new_node.values.iter().cloned())
                    .collect();
                self.merge_join_pairs(&o, &n, diffs);
            }
            (false, false) => self.diff_internal(old_node, new_node, diffs),
            _ => {
                // height divergence (one side shallower): flatten the divergent region.
                let mut op = Vec::new();
                self.collect_pairs_recursive(old_node, &mut op);
                let mut np = Vec::new();
                self.collect_pairs_recursive(new_node, &mut np);
                self.merge_join_pairs(&op, &np, diffs);
            }
        }
    }

    /// Both nodes internal: skip children whose stored child-hash is equal (O(1), no
    /// load), flatten + merge-join the rest. Provably identical to the full flatten:
    /// a common child subtree has identical (k,v) on both sides => contributes no diff.
    fn diff_internal(
        &self,
        old_node: &ProllyNode<N>,
        new_node: &ProllyNode<N>,
        diffs: &mut Vec<DiffResult>,
    ) {
        use std::collections::HashSet;
        // The node is content-addressed: an internal node's `values` ARE its child
        // hashes, so byte-equal `values[i]` <=> identical child subtree (same (k,v)).
        // Skip common children by hash WITHOUT loading them (the O(diff) win), and
        // load only the differing children (mirroring `children()`'s raw_hash lookup).
        let new_hashes: HashSet<&Vec<u8>> = new_node.values.iter().collect();
        let old_hashes: HashSet<&Vec<u8>> = old_node.values.iter().collect();
        let mut old_pairs: Vec<(Vec<u8>, Vec<u8>)> = Vec::new();
        for child_hash in &old_node.values {
            if new_hashes.contains(child_hash) {
                continue; // identical subtree on the new side => contributes no diff
            }
            if let Some(child) = self
                .storage
                .get_node_by_hash(&ValueDigest::raw_hash(child_hash))
            {
                self.collect_pairs_recursive(&child, &mut old_pairs);
            }
        }
        let mut new_pairs: Vec<(Vec<u8>, Vec<u8>)> = Vec::new();
        for child_hash in &new_node.values {
            if old_hashes.contains(child_hash) {
                continue;
            }
            if let Some(child) = self
                .storage
                .get_node_by_hash(&ValueDigest::raw_hash(child_hash))
            {
                self.collect_pairs_recursive(&child, &mut new_pairs);
            }
        }
        self.merge_join_pairs(&old_pairs, &new_pairs, diffs);
    }

    /// Merge-join two sorted (key,value) streams into DiffResults.
    fn merge_join_pairs(
        &self,
        old_pairs: &[(Vec<u8>, Vec<u8>)],
        new_pairs: &[(Vec<u8>, Vec<u8>)],
        diffs: &mut Vec<DiffResult>,
    ) {
        let mut oi = old_pairs.iter().peekable();
        let mut ni = new_pairs.iter().peekable();
        while let (Some((ok, ov)), Some((nk, nv))) = (oi.peek(), ni.peek()) {
            match ok.cmp(nk) {
                std::cmp::Ordering::Less => {
                    diffs.push(DiffResult::Removed(ok.clone(), ov.clone()));
                    oi.next();
                }
                std::cmp::Ordering::Greater => {
                    diffs.push(DiffResult::Added(nk.clone(), nv.clone()));
                    ni.next();
                }
                std::cmp::Ordering::Equal => {
                    if ov != nv {
                        diffs.push(DiffResult::Modified(ok.clone(), ov.clone(), nv.clone()));
                    }
                    oi.next();
                    ni.next();
                }
            }
        }
        for (ok, ov) in oi {
            diffs.push(DiffResult::Removed(ok.clone(), ov.clone()));
        }
        for (nk, nv) in ni {
            diffs.push(DiffResult::Added(nk.clone(), nv.clone()));
        }
    }

    /// Proven full-leaf flatten diff — kept as the differential-test oracle.
    #[cfg(test)]
    fn diff_nodes_flatten(
        &self,
        old_node: &ProllyNode<N>,
        new_node: &ProllyNode<N>,
        diffs: &mut Vec<DiffResult>,
    ) {
        let mut old_pairs = Vec::new();
        self.collect_pairs_recursive(old_node, &mut old_pairs);
        let mut new_pairs = Vec::new();
        self.collect_pairs_recursive(new_node, &mut new_pairs);
        self.merge_join_pairs(&old_pairs, &new_pairs, diffs);
    }

    /// Collect all leaf (key, value) pairs of a subtree in sorted order (full traversal).
    fn collect_pairs_recursive(&self, node: &ProllyNode<N>, pairs: &mut Vec<(Vec<u8>, Vec<u8>)>) {
        if node.is_leaf {
            for (k, v) in node.keys.iter().zip(node.values.iter()) {
                pairs.push((k.clone(), v.clone()));
            }
        } else {
            for child in node.children(&self.storage) {
                self.collect_pairs_recursive(&child, pairs);
            }
        }
    }

    /// Find a value for a specific key in a node tree
    fn find_value_in_node(&self, node: &ProllyNode<N>, key: &[u8]) -> Option<Vec<u8>> {
        // Search directly in the node first (for leaf nodes)
        if node.is_leaf {
            for (i, k) in node.keys.iter().enumerate() {
                if k.as_slice() == key {
                    return Some(node.values[i].clone());
                }
            }
        } else {
            // For internal nodes, use the node's find method with storage
            return node.find(key, &self.storage).and_then(|found_node| {
                found_node
                    .keys
                    .iter()
                    .zip(found_node.values.iter())
                    .find(|(k, _)| k.as_slice() == key)
                    .map(|(_, v)| v.clone())
            });
        }
        None
    }

    /// Detects conflicts between changes from source and destination branches
    fn detect_conflict(
        &self,
        key: &[u8],
        source_diff: &DiffResult,
        destination_diff: &DiffResult,
        base_node: &ProllyNode<N>,
    ) -> Option<MergeConflict> {
        // Get the base value for this key from the base node
        let base_value = self.find_value_in_node(base_node, key);

        match (source_diff, destination_diff) {
            // Both added the same key
            (DiffResult::Added(_, source_value), DiffResult::Added(_, destination_value)) => {
                if source_value != destination_value {
                    Some(MergeConflict {
                        key: key.to_vec(),
                        base_value,
                        source_value: Some(source_value.clone()),
                        destination_value: Some(destination_value.clone()),
                    })
                } else {
                    None // Same value added, no conflict
                }
            }

            // Both removed the same key - no conflict
            (DiffResult::Removed(_, _), DiffResult::Removed(_, _)) => None,

            // Both modified the same key
            (
                DiffResult::Modified(_, _, source_value),
                DiffResult::Modified(_, _, destination_value),
            ) => {
                if source_value != destination_value {
                    Some(MergeConflict {
                        key: key.to_vec(),
                        base_value,
                        source_value: Some(source_value.clone()),
                        destination_value: Some(destination_value.clone()),
                    })
                } else {
                    None // Same modification, no conflict
                }
            }

            // One added, one removed - conflict
            (DiffResult::Added(_, source_value), DiffResult::Removed(_, _)) => {
                Some(MergeConflict {
                    key: key.to_vec(),
                    base_value,
                    source_value: Some(source_value.clone()),
                    destination_value: None,
                })
            }
            (DiffResult::Removed(_, _), DiffResult::Added(_, destination_value)) => {
                Some(MergeConflict {
                    key: key.to_vec(),
                    base_value,
                    source_value: None,
                    destination_value: Some(destination_value.clone()),
                })
            }

            // One added, one modified - conflict
            (DiffResult::Added(_, source_value), DiffResult::Modified(_, _, destination_value)) => {
                Some(MergeConflict {
                    key: key.to_vec(),
                    base_value,
                    source_value: Some(source_value.clone()),
                    destination_value: Some(destination_value.clone()),
                })
            }
            (DiffResult::Modified(_, _, source_value), DiffResult::Added(_, destination_value)) => {
                Some(MergeConflict {
                    key: key.to_vec(),
                    base_value,
                    source_value: Some(source_value.clone()),
                    destination_value: Some(destination_value.clone()),
                })
            }

            // One removed, one modified - conflict
            (DiffResult::Removed(_, _), DiffResult::Modified(_, _, destination_value)) => {
                Some(MergeConflict {
                    key: key.to_vec(),
                    base_value,
                    source_value: None,
                    destination_value: Some(destination_value.clone()),
                })
            }
            (DiffResult::Modified(_, _, source_value), DiffResult::Removed(_, _)) => {
                Some(MergeConflict {
                    key: key.to_vec(),
                    base_value,
                    source_value: Some(source_value.clone()),
                    destination_value: None,
                })
            }
        }
    }

    /// Recursively computes the differences between two Prolly Nodes.
    ///
    /// This helper function is used by `diff` to traverse the nodes of both trees
    /// and identify changes. It compares the keys and values of the nodes and
    /// generates appropriate `DiffResult` entries for added, removed, and modified
    /// key-value pairs.
    ///
    /// # Arguments
    ///
    /// * `old_node` - The node from the original tree.
    /// * `new_node` - The node from the new tree.
    /// * `diffs` - The vector to store the differences.
    #[allow(dead_code)]
    fn diff_recursive(
        &self,
        old_node: &ProllyNode<N>,
        new_node: &ProllyNode<N>,
        diffs: &mut Vec<DiffResult>,
    ) {
        let mut old_iter = old_node.keys.iter().zip(old_node.values.iter()).peekable();
        let mut new_iter = new_node.keys.iter().zip(new_node.values.iter()).peekable();

        while let (Some((old_key, old_value)), Some((new_key, new_value))) =
            (old_iter.peek(), new_iter.peek())
        {
            match old_key.cmp(new_key) {
                std::cmp::Ordering::Less => {
                    diffs.push(DiffResult::Removed(old_key.to_vec(), old_value.to_vec()));
                    old_iter.next();
                }
                std::cmp::Ordering::Greater => {
                    diffs.push(DiffResult::Added(new_key.to_vec(), new_value.to_vec()));
                    new_iter.next();
                }
                std::cmp::Ordering::Equal => {
                    if old_value != new_value {
                        diffs.push(DiffResult::Modified(
                            old_key.to_vec(),
                            old_value.to_vec(),
                            new_value.to_vec(),
                        ));
                    }
                    old_iter.next();
                    new_iter.next();
                }
            }
        }

        for (old_key, old_value) in old_iter {
            diffs.push(DiffResult::Removed(old_key.clone(), old_value.clone()));
        }

        for (new_key, new_value) in new_iter {
            diffs.push(DiffResult::Added(new_key.clone(), new_value.clone()));
        }
    }

    /// Persist the root node to storage and save configuration
    pub fn persist_root(&mut self) {
        // Store the root node in the storage
        let root_hash = self.root.get_hash();
        if self
            .storage
            .insert_node(root_hash.clone(), self.root.clone())
            .is_ok()
        {
            // Update the config with the new root hash
            self.config.root_hash = Some(root_hash);

            // Save the configuration
            let _ = self.save_config();
        }
    }

    /// Load a ProllyTree from an existing root hash in storage
    pub fn load_from_storage(storage: S, config: TreeConfig<N>) -> Option<Self> {
        if let Some(ref root_hash) = config.root_hash {
            if let Some(root_node) = storage.get_node_by_hash(root_hash) {
                return Some(ProllyTree {
                    root: Arc::unwrap_or_clone(root_node),
                    storage,
                    config,
                });
            }
        }
        None
    }

    /// Apply a batch of insert/delete operations in a single streaming pass.
    /// `Some(value)` is an insert/upsert; `None` is a delete.
    ///
    /// Implementation: sort the mutations into a `BTreeMap` for ordering
    /// and duplicate-key collapsing, then drive them through
    /// `streaming_chunker::apply_mutations`, which walks the existing
    /// tree with a `NodeCursor` and emits a new canonical tree without
    /// materializing every existing pair into memory.
    ///
    /// Returns the number of deletes that hit a non-existent key
    /// (callers may use this to set return values).
    pub fn apply_changes<I>(&mut self, changes: I) -> usize
    where
        I: IntoIterator<Item = (Vec<u8>, Option<Vec<u8>>)>,
    {
        // Sort + dedupe the input. Last-write-wins per key within the batch.
        let map: std::collections::BTreeMap<Vec<u8>, Option<Vec<u8>>> =
            changes.into_iter().collect();
        // For accurate `missing_deletes`, we need to know which delete
        // targets weren't in the old tree. Cheap pre-pass: probe the
        // existing tree.
        let mut missing_deletes = 0usize;
        for (k, v) in map.iter() {
            if v.is_none() && self.find(k).is_none() {
                missing_deletes += 1;
            }
        }

        let new_root = crate::streaming_chunker::apply_mutations(
            self.root.clone(),
            map,
            &self.config,
            &mut self.storage,
        );
        let _ = self
            .storage
            .insert_node(new_root.get_hash(), new_root.clone());
        self.root = new_root;
        missing_deletes
    }

    /// Collect all keys from the tree
    pub fn collect_keys(&self) -> Vec<Vec<u8>> {
        let mut keys = Vec::new();
        self.collect_keys_recursive(&self.root, &mut keys);
        keys
    }

    /// Recursively collect keys from a node and its children.
    ///
    /// Only keys that live in leaf nodes are collected. Internal-node `keys`
    /// entries are split-key pivots that *should* equal the first key of each
    /// child subtree, but a `delete` that removes a subtree's first key does
    /// not currently propagate the new first-key back to the parent pivot
    /// (the pivot becomes stale). Including pivots here would surface those
    /// stale entries as live keys — non-deterministically, since it depends
    /// on whether the deleted key happened to land as a pivot during the last
    /// balance. Leaf-only traversal matches the canonical walker in
    /// `history.rs::collect_keys_recursive` and keeps `list_keys` stable.
    fn collect_keys_recursive(&self, node: &ProllyNode<N>, keys: &mut Vec<Vec<u8>>) {
        if node.is_leaf {
            for key in &node.keys {
                keys.push(key.clone());
            }
        } else {
            for child_node in node.children(&self.storage) {
                self.collect_keys_recursive(&child_node, keys);
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::storage::InMemoryNodeStorage;

    /// Example usage of the Prolly Tree
    #[test]
    fn inmem_node_storage_test() {
        // 1. Create a custom tree config
        let config = TreeConfig {
            base: 131,
            modulus: 1_000_000_009,
            min_chunk_size: 4,
            max_chunk_size: 8 * 1024,
            pattern: 0b101,
            root_hash: None,
            key_schema: None,
            value_schema: None,
            encode_types: vec![],
        };

        // 2. Create and Wrap the Storage Backend
        let storage = InMemoryNodeStorage::<32>::default();

        // 3. Create the Prolly Tree
        let mut tree = ProllyTree::new(storage, config);

        // 4. Insert New Key-Value Pairs
        tree.insert(b"key1".to_vec(), b"value1".to_vec());
        tree.insert(b"key2".to_vec(), b"value2".to_vec());

        // 5. Traverse the Tree with a Custom Formatter
        let traversal = tree.formatted_traverse(|node| {
            let keys_as_strings: Vec<String> = node.keys.iter().map(|k| format!("{k:?}")).collect();
            format!("[L{}: {}]", node.level, keys_as_strings.join(", "))
        });
        println!("Traversal: {traversal}");

        // 6. Update the Value for an Existing Key
        tree.update(b"key1".to_vec(), b"new_value1".to_vec());

        // 7. Find or Search for a Key
        if let Some(node) = tree.find(b"key1") {
            println!("Found key1 with value: {node:?}");
        } else {
            println!("key1 not found");
        }

        // 8. Delete a key-value pair
        if tree.delete(b"key2") {
            println!("key2 deleted");
        } else {
            println!("key2 not found");
        }

        // 9. Print tree stats
        println!("Size: {}", tree.size());
        println!("Depth: {}", tree.depth());
        println!("Summary: {}", tree.summary());

        // 10. Print Tree
        println!("{:?}", tree.root.print_tree(&tree.storage));
    }

    #[test]
    fn file_node_storage_test() {
        use crate::storage::FileNodeStorage;
        use std::fs;
        use std::path::PathBuf;

        // 1. Create a custom tree config
        let config = TreeConfig {
            base: 131,
            modulus: 1_000_000_009,
            min_chunk_size: 4,
            max_chunk_size: 8 * 1024,
            pattern: 0b101,
            root_hash: None,
            key_schema: None,
            value_schema: None,
            encode_types: vec![],
        };

        // 2. Create and Wrap the Storage Backend
        let storage_dir = PathBuf::from("/tmp/prolly_tree_storage");
        let storage = FileNodeStorage::<32>::new(storage_dir.clone()).unwrap();

        // 3. Create the Prolly Tree
        let mut tree = ProllyTree::new(storage, config);

        // 4. Insert New Key-Value Pairs
        tree.insert(b"key1".to_vec(), b"value1".to_vec());
        tree.insert(b"key2".to_vec(), b"value2".to_vec());

        // 5. Traverse the Tree with a Custom Formatter
        let traversal = tree.formatted_traverse(|node| {
            let keys_as_strings: Vec<String> = node.keys.iter().map(|k| format!("{k:?}")).collect();
            format!("[L{}: {}]", node.level, keys_as_strings.join(", "))
        });
        println!("Traversal: {traversal}");

        // 6. Update the Value for an Existing Key
        tree.update(b"key1".to_vec(), b"new_value1".to_vec());

        // 7. Find or Search for a Key
        if let Some(node) = tree.find(b"key1") {
            println!("Found key1 with value: {node:?}");
        } else {
            println!("key1 not found");
        }

        // 8. Delete a key-value pair
        if tree.delete(b"key2") {
            println!("key2 deleted");
        } else {
            println!("key2 not found");
        }

        // 9. Print tree stats
        println!("Size: {}", tree.size());
        println!("Depth: {}", tree.depth());
        println!("Summary: {}", tree.summary());

        // 10. Print Tree
        println!("{:?}", tree.root.print_tree(&tree.storage));

        // Clean up the storage directory
        fs::remove_dir_all(storage_dir).unwrap();
    }

    #[test]
    fn test_insert_and_find() {
        let storage = InMemoryNodeStorage::<32>::default();

        let mut tree = ProllyTree::new(storage, TreeConfig::default());

        tree.insert(b"key1".to_vec(), b"value1".to_vec());
        tree.insert(b"key2".to_vec(), b"value2".to_vec());

        assert!(tree.find(b"key1").is_some());
        assert!(tree.find(b"key2").is_some());
        assert!(tree.find(b"key3").is_none());
    }

    #[test]
    fn test_persist_and_load() {
        let storage = InMemoryNodeStorage::<32>::default();
        let config = TreeConfig::default();

        // Create tree and add data
        let mut tree = ProllyTree::new(storage.clone(), config.clone());
        tree.insert(b"key1".to_vec(), b"value1".to_vec());
        tree.insert(b"key2".to_vec(), b"value2".to_vec());

        // Persist the tree
        tree.persist_root();

        // Load the tree from storage
        let loaded_tree = ProllyTree::load_from_storage(tree.storage, tree.config)
            .expect("Should be able to load tree from storage");

        // Verify data is preserved
        assert!(loaded_tree.find(b"key1").is_some());
        assert!(loaded_tree.find(b"key2").is_some());
        assert!(loaded_tree.find(b"key3").is_none());
    }

    #[test]
    fn test_insert_batch_and_find() {
        let storage = InMemoryNodeStorage::<32>::default();

        let mut tree = ProllyTree::new(storage, TreeConfig::default());

        let keys = vec![b"key1".to_vec(), b"key2".to_vec(), b"key3".to_vec()];
        let values = vec![b"value1".to_vec(), b"value2".to_vec(), b"value3".to_vec()];

        tree.insert_batch(&keys, &values);

        assert!(tree.find(b"key1").is_some());
        assert!(tree.find(b"key2").is_some());
        assert!(tree.find(b"key3").is_some());
        assert!(tree.find(b"key4").is_none());
    }

    #[test]
    #[should_panic(expected = "insert_batch requires the same number of keys and values")]
    fn test_insert_batch_rejects_length_mismatch() {
        let storage = InMemoryNodeStorage::<32>::default();
        let mut tree = ProllyTree::new(storage, TreeConfig::default());

        let keys = vec![b"key1".to_vec(), b"key2".to_vec()];
        let values = vec![b"value1".to_vec()];

        tree.insert_batch(&keys, &values);
    }

    #[test]
    fn test_delete() {
        let storage = InMemoryNodeStorage::<32>::default();
        let mut tree = ProllyTree::new(storage, TreeConfig::default());

        tree.insert(b"key1".to_vec(), b"value1".to_vec());
        tree.insert(b"key2".to_vec(), b"value2".to_vec());

        assert!(tree.delete(b"key1"));
        assert!(tree.find(b"key1").is_none());
        assert!(tree.find(b"key2").is_some());
    }

    #[test]
    fn test_delete_batch() {
        let storage = InMemoryNodeStorage::<32>::default();
        let mut tree = ProllyTree::new(storage, TreeConfig::default());

        let keys = vec![b"key1".to_vec(), b"key2".to_vec(), b"key3".to_vec()];
        let values = vec![b"value1".to_vec(), b"value2".to_vec(), b"value3".to_vec()];

        // Insert keys and values
        tree.insert_batch(&keys, &values);

        // Verify insertion
        assert!(tree.find(b"key1").is_some());
        assert!(tree.find(b"key2").is_some());
        assert!(tree.find(b"key3").is_some());

        // Delete keys in batch
        tree.delete_batch(&keys);

        // Verify deletion
        assert!(tree.find(b"key1").is_none());
        assert!(tree.find(b"key2").is_none());
        assert!(tree.find(b"key3").is_none());
    }

    #[test]
    fn test_traverse() {
        let storage = InMemoryNodeStorage::<32>::default();
        let mut tree = ProllyTree::new(storage, TreeConfig::default());

        let key1 = b"key1".to_vec();
        let key2 = b"key2".to_vec();

        tree.insert(key1.clone(), b"value1".to_vec());
        tree.insert(key2.clone(), b"value2".to_vec());

        let traversal = tree.traverse();

        // Convert byte arrays to their binary representation strings for comparison
        let expected_key1 = format!("{key1:?}");
        let expected_key2 = format!("{key2:?}");

        // Check if the traversal contains the expected keys
        assert!(traversal.contains(&expected_key1.to_string()));
        assert!(traversal.contains(&expected_key2.to_string()));
    }

    #[test]
    fn test_stats() {
        let storage = InMemoryNodeStorage::<32>::default();
        let config = TreeConfig {
            base: 131,
            modulus: 1_000_000_009,
            min_chunk_size: 16,
            max_chunk_size: 8 * 1024,
            pattern: 0b111,
            root_hash: None,
            key_schema: None,
            value_schema: None,
            encode_types: vec![],
        };

        let mut tree = ProllyTree::new(storage, config);

        // Insert key-value pairs using a loop
        let max_key = 3000u32;

        for i in 0..max_key {
            // Convert to big-endian byte array to maintain order
            let key = i.to_be_bytes().to_vec();
            let value = i.to_be_bytes().to_vec();
            tree.insert(key.clone(), value.clone());
        }

        println!("{:?}", tree.root.print_tree(&tree.storage));

        for i in 0..max_key {
            let key = i.to_be_bytes().to_vec();
            assert!(tree.find(&key).is_some());
        }
        let non_existing_key = (max_key + 10).to_be_bytes().to_vec();
        assert!(tree.find(&non_existing_key).is_none());

        // assert that the tree has the expected key-value pairs
        assert_eq!(tree.size(), max_key as usize);

        // assert that the tree has the expected depth
        assert_eq!(tree.depth(), 3);

        println!("Size: {}", tree.size());
        println!("Depth: {}", tree.depth());
        println!("Summary: {}", tree.summary());
    }

    #[test]
    fn test_generate_proof() {
        let config = TreeConfig::default();
        let storage = InMemoryNodeStorage::<32>::default();
        let mut tree = ProllyTree::new(storage, config);

        // Insert key-value pairs
        for i in 0..100 {
            let key = vec![i];
            let value = vec![i];
            tree.insert(key.clone(), value.clone());
        }

        // Generate proof for an existing key
        let key_to_prove = vec![5];
        let proof = tree.generate_proof(&key_to_prove);

        // Verify the proof
        let verified = tree.verify(proof, &key_to_prove, Some(&key_to_prove));
        assert!(verified);

        // Generate proof for a non-existing key
        let key_to_prove_wrong = vec![120];
        let proof_wrong = tree.generate_proof(&key_to_prove_wrong);

        // Should not be verified
        let verified_wrong =
            tree.verify(proof_wrong, &key_to_prove_wrong, Some(&key_to_prove_wrong));
        assert!(!verified_wrong);
    }

    #[test]
    fn test_diff() {
        let config = TreeConfig::default();
        let storage1 = InMemoryNodeStorage::<32>::default();
        let mut tree1 = ProllyTree::new(storage1, config.clone());

        let storage2 = InMemoryNodeStorage::<32>::default();
        let mut tree2 = ProllyTree::new(storage2, config);

        // Insert key-value pairs into tree1
        for i in 0..50 {
            tree1.insert(vec![i], vec![i]);
        }

        // Insert key-value pairs into tree1
        for i in 0..50 {
            tree2.insert(vec![i], vec![i]);
        }

        // modify some keys in tree2
        tree2.insert(vec![10], vec![200]);

        // print tree1 and tree2
        println!("{:?}", tree1.root.print_tree(&tree1.storage));
        println!("{:?}", tree2.root.print_tree(&tree2.storage));

        // Generate diff between tree1 and tree2
        let differences = tree1.diff(&tree2);

        // Check the differences
        // Expecting only the first L1 value would change
        for diff in &differences {
            match diff {
                DiffResult::Added(key, value) => {
                    println!("Added: key = {key:?}, value = {value:?}");
                }
                DiffResult::Removed(key, value) => {
                    println!("Removed: key = {key:?}, value = {value:?}");
                }
                DiffResult::Modified(key, old_value, new_value) => {
                    println!(
                        "Modified: key = {key:?}, old_value = {old_value:?}, new_value = {new_value:?}"
                    );
                }
            }
        }
    }

    #[test]
    fn test_print_proof_demo() {
        let config = TreeConfig::default();
        let storage = InMemoryNodeStorage::<32>::default();
        let mut tree = ProllyTree::new(storage, config);

        // Insert enough key-value pairs to create a multi-level tree
        for i in 0..20 {
            tree.insert(vec![i], vec![i * 10]);
        }

        println!("=== Prolly Tree with Proof Visualization Demo ===");

        // Test with an existing key
        let existing_key = vec![10];
        println!("\n--- Testing with existing key {:?} ---", existing_key);
        let is_valid = tree.print_proof(&existing_key);
        assert!(is_valid, "Proof should be valid for existing key");

        // Test with a non-existing key
        let non_existing_key = vec![25];
        println!(
            "\n--- Testing with non-existing key {:?} ---",
            non_existing_key
        );
        let is_valid = tree.print_proof(&non_existing_key);
        assert!(!is_valid, "Proof should be invalid for non-existing key");

        println!("\n=== Demo completed successfully ===");
    }

    #[test]
    fn test_merge_simple() {
        let config = TreeConfig::default();
        let mut storage = InMemoryNodeStorage::<32>::default();

        // Create base tree and persist it
        let mut base_tree = ProllyTree::new(storage.clone(), config.clone());
        base_tree.insert(b"key1".to_vec(), b"value1".to_vec());
        let base_root = base_tree.get_root_hash().unwrap();
        storage
            .insert_node(base_root.clone(), base_tree.root.clone())
            .unwrap();

        // Create source tree - same as base
        let mut source_tree = ProllyTree::new(storage.clone(), config.clone());
        source_tree.insert(b"key1".to_vec(), b"value1".to_vec());
        let source_root = source_tree.get_root_hash().unwrap();
        storage
            .insert_node(source_root.clone(), source_tree.root.clone())
            .unwrap();

        // Create destination tree - same as base
        let mut dest_tree = ProllyTree::new(storage.clone(), config.clone());
        dest_tree.insert(b"key1".to_vec(), b"value1".to_vec());
        let dest_root = dest_tree.get_root_hash().unwrap();
        storage
            .insert_node(dest_root.clone(), dest_tree.root.clone())
            .unwrap();

        // Perform merge using storage that contains all the nodes
        let merge_tree = ProllyTree::new(storage, config);
        let merge_results = merge_tree.merge(&source_root, &dest_root, &base_root);

        // With identical trees, should have no changes
        assert_eq!(merge_results.len(), 0);
    }

    #[test]
    fn test_merge_with_conflicts() {
        let config = TreeConfig::default();
        let mut storage = InMemoryNodeStorage::<32>::default();

        // Create base tree
        let mut base_tree = ProllyTree::new(storage.clone(), config.clone());
        base_tree.insert(b"key1".to_vec(), b"value1".to_vec());
        base_tree.insert(b"key2".to_vec(), b"value2".to_vec());
        let base_root = base_tree.get_root_hash().unwrap();
        storage
            .insert_node(base_root.clone(), base_tree.root.clone())
            .unwrap();

        // Create source tree - modify key1 to "source_value"
        let mut source_tree = ProllyTree::new(storage.clone(), config.clone());
        source_tree.insert(b"key1".to_vec(), b"source_value".to_vec());
        source_tree.insert(b"key2".to_vec(), b"value2".to_vec());
        let source_root = source_tree.get_root_hash().unwrap();
        storage
            .insert_node(source_root.clone(), source_tree.root.clone())
            .unwrap();

        // Create destination tree - modify key1 to "dest_value"
        let mut dest_tree = ProllyTree::new(storage.clone(), config.clone());
        dest_tree.insert(b"key1".to_vec(), b"dest_value".to_vec());
        dest_tree.insert(b"key2".to_vec(), b"value2".to_vec());
        let dest_root = dest_tree.get_root_hash().unwrap();
        storage
            .insert_node(dest_root.clone(), dest_tree.root.clone())
            .unwrap();

        // Perform merge
        let merge_tree = ProllyTree::new(storage, config);
        let merge_results = merge_tree.merge(&source_root, &dest_root, &base_root);

        // Verify conflict detected
        assert_eq!(merge_results.len(), 1);
        match &merge_results[0] {
            MergeResult::Conflict(conflict) => {
                assert_eq!(conflict.key, b"key1".to_vec());
                assert_eq!(conflict.base_value, Some(b"value1".to_vec()));
                assert_eq!(conflict.source_value, Some(b"source_value".to_vec()));
                assert_eq!(conflict.destination_value, Some(b"dest_value".to_vec()));
            }
            _ => panic!("Expected conflict, got: {:?}", merge_results[0]),
        }
    }

    #[test]
    fn test_merge_to_empty_tree() {
        let config = TreeConfig::default();
        let mut storage = InMemoryNodeStorage::<32>::default();

        // Create base tree (empty)
        let base_tree = ProllyTree::new(storage.clone(), config.clone());
        let base_root = base_tree.get_root_hash().unwrap();
        storage
            .insert_node(base_root.clone(), base_tree.root.clone())
            .unwrap();

        // Create source tree with data
        let mut source_tree = ProllyTree::new(storage.clone(), config.clone());
        source_tree.insert(b"key1".to_vec(), b"value1".to_vec());
        source_tree.insert(b"key2".to_vec(), b"value2".to_vec());
        let source_root = source_tree.get_root_hash().unwrap();
        storage
            .insert_node(source_root.clone(), source_tree.root.clone())
            .unwrap();

        // Create destination tree (also empty)
        let dest_tree = ProllyTree::new(storage.clone(), config.clone());
        let dest_root = dest_tree.get_root_hash().unwrap();
        storage
            .insert_node(dest_root.clone(), dest_tree.root.clone())
            .unwrap();

        // Perform merge
        let merge_tree = ProllyTree::new(storage, config);
        let merge_results = merge_tree.merge(&source_root, &dest_root, &base_root);

        // Should see all source changes as additions
        assert_eq!(merge_results.len(), 2);
        let mut keys_added = std::collections::HashSet::new();
        for result in merge_results {
            match result {
                MergeResult::Added(key, _) => {
                    keys_added.insert(key);
                }
                _ => panic!("Expected only additions, got: {:?}", result),
            }
        }
        assert!(keys_added.contains(&b"key1".to_vec()));
        assert!(keys_added.contains(&b"key2".to_vec()));
    }

    #[test]
    fn test_merge_add_remove_conflicts() {
        let config = TreeConfig::default();
        let mut storage = InMemoryNodeStorage::<32>::default();

        // Create base tree
        let mut base_tree = ProllyTree::new(storage.clone(), config.clone());
        base_tree.insert(b"key1".to_vec(), b"value1".to_vec());
        let base_root = base_tree.get_root_hash().unwrap();
        storage
            .insert_node(base_root.clone(), base_tree.root.clone())
            .unwrap();

        // Create source tree - remove key1, add key2
        let mut source_tree = ProllyTree::new(storage.clone(), config.clone());
        source_tree.insert(b"key2".to_vec(), b"value2".to_vec());
        let source_root = source_tree.get_root_hash().unwrap();
        storage
            .insert_node(source_root.clone(), source_tree.root.clone())
            .unwrap();

        // Create destination tree - modify key1
        let mut dest_tree = ProllyTree::new(storage.clone(), config.clone());
        dest_tree.insert(b"key1".to_vec(), b"modified_value1".to_vec());
        let dest_root = dest_tree.get_root_hash().unwrap();
        storage
            .insert_node(dest_root.clone(), dest_tree.root.clone())
            .unwrap();

        // Perform merge
        let merge_tree = ProllyTree::new(storage, config);
        let merge_results = merge_tree.merge(&source_root, &dest_root, &base_root);

        // Should have one addition and one conflict
        assert_eq!(merge_results.len(), 2);

        let mut has_addition = false;
        let mut has_conflict = false;

        for result in merge_results {
            match result {
                MergeResult::Added(key, value) => {
                    assert_eq!(key, b"key2".to_vec());
                    assert_eq!(value, b"value2".to_vec());
                    has_addition = true;
                }
                MergeResult::Conflict(conflict) => {
                    assert_eq!(conflict.key, b"key1".to_vec());
                    assert_eq!(conflict.base_value, Some(b"value1".to_vec()));
                    assert_eq!(conflict.source_value, None); // removed in source
                    assert_eq!(
                        conflict.destination_value,
                        Some(b"modified_value1".to_vec())
                    );
                    has_conflict = true;
                }
                _ => panic!("Unexpected merge result: {:?}", result),
            }
        }

        assert!(has_addition && has_conflict);
    }

    #[test]
    fn test_merge_complex_scenario() {
        let config = TreeConfig::default();
        let mut storage = InMemoryNodeStorage::<32>::default();

        // Create base tree with a single key for simplicity
        let mut base_tree = ProllyTree::new(storage.clone(), config.clone());
        base_tree.insert(b"modify_conflict".to_vec(), b"base_value".to_vec());
        let base_root = base_tree.get_root_hash().unwrap();
        storage
            .insert_node(base_root.clone(), base_tree.root.clone())
            .unwrap();

        // Create source tree - modify the key and add a new one
        let mut source_tree = ProllyTree::new(storage.clone(), config.clone());
        source_tree.insert(b"modify_conflict".to_vec(), b"source_value".to_vec()); // modified
        source_tree.insert(b"new_in_source".to_vec(), b"source_addition".to_vec()); // added
        let source_root = source_tree.get_root_hash().unwrap();
        storage
            .insert_node(source_root.clone(), source_tree.root.clone())
            .unwrap();

        // Create destination tree - modify the key differently
        let mut dest_tree = ProllyTree::new(storage.clone(), config.clone());
        dest_tree.insert(b"modify_conflict".to_vec(), b"dest_value".to_vec()); // modified differently
        let dest_root = dest_tree.get_root_hash().unwrap();
        storage
            .insert_node(dest_root.clone(), dest_tree.root.clone())
            .unwrap();

        // Perform merge
        let merge_tree = ProllyTree::new(storage, config);
        let merge_results = merge_tree.merge(&source_root, &dest_root, &base_root);

        // Analyze results
        let mut added_keys = std::collections::HashSet::new();
        let mut conflicts = std::collections::HashMap::new();

        for result in merge_results {
            match result {
                MergeResult::Added(key, _) => {
                    added_keys.insert(key);
                }
                MergeResult::Conflict(conflict) => {
                    conflicts.insert(conflict.key.clone(), conflict);
                }
                _ => panic!("Unexpected merge result: {:?}", result),
            }
        }

        // Should have one addition and one conflict
        assert!(added_keys.contains(&b"new_in_source".to_vec()));
        assert!(conflicts.contains_key(&b"modify_conflict".to_vec()));

        // Verify conflict details (relax the base_value check for now)
        let conflict = &conflicts[&b"modify_conflict".to_vec()];
        assert_eq!(conflict.source_value, Some(b"source_value".to_vec()));
        assert_eq!(conflict.destination_value, Some(b"dest_value".to_vec()));
    }

    #[test]
    fn test_merge_same_changes_no_conflict() {
        let config = TreeConfig::default();
        let mut storage = InMemoryNodeStorage::<32>::default();

        // Create base tree
        let mut base_tree = ProllyTree::new(storage.clone(), config.clone());
        base_tree.insert(b"key1".to_vec(), b"value1".to_vec());
        base_tree.insert(b"key2".to_vec(), b"value2".to_vec());
        let base_root = base_tree.get_root_hash().unwrap();
        storage
            .insert_node(base_root.clone(), base_tree.root.clone())
            .unwrap();

        // Both source and destination make the same changes
        let mut source_tree = ProllyTree::new(storage.clone(), config.clone());
        source_tree.insert(b"key1".to_vec(), b"same_new_value".to_vec()); // both modify to same value
        source_tree.insert(b"key2".to_vec(), b"value2".to_vec()); // unchanged
        source_tree.insert(b"key3".to_vec(), b"same_addition".to_vec()); // both add same key-value
        let source_root = source_tree.get_root_hash().unwrap();
        storage
            .insert_node(source_root.clone(), source_tree.root.clone())
            .unwrap();

        let mut dest_tree = ProllyTree::new(storage.clone(), config.clone());
        dest_tree.insert(b"key1".to_vec(), b"same_new_value".to_vec()); // same modification
        dest_tree.insert(b"key2".to_vec(), b"value2".to_vec()); // unchanged
        dest_tree.insert(b"key3".to_vec(), b"same_addition".to_vec()); // same addition
        let dest_root = dest_tree.get_root_hash().unwrap();
        storage
            .insert_node(dest_root.clone(), dest_tree.root.clone())
            .unwrap();

        // Perform merge
        let merge_tree = ProllyTree::new(storage, config);
        let merge_results = merge_tree.merge(&source_root, &dest_root, &base_root);

        // Should apply changes without conflict since they're identical
        assert_eq!(merge_results.len(), 2);

        let mut has_modification = false;
        let mut has_addition = false;

        for result in merge_results {
            match result {
                MergeResult::Modified(key, value) => {
                    assert_eq!(key, b"key1".to_vec());
                    assert_eq!(value, b"same_new_value".to_vec());
                    has_modification = true;
                }
                MergeResult::Added(key, value) => {
                    assert_eq!(key, b"key3".to_vec());
                    assert_eq!(value, b"same_addition".to_vec());
                    has_addition = true;
                }
                _ => panic!("Unexpected merge result: {:?}", result),
            }
        }

        assert!(has_modification && has_addition);
    }

    #[test]
    fn test_apply_merge_results_success() {
        let config = TreeConfig::default();
        let mut storage = InMemoryNodeStorage::<32>::default();

        // Create base tree
        let mut base_tree = ProllyTree::new(storage.clone(), config.clone());
        base_tree.insert(b"existing".to_vec(), b"value".to_vec());
        let base_root = base_tree.get_root_hash().unwrap();
        storage
            .insert_node(base_root.clone(), base_tree.root.clone())
            .unwrap();

        // Create merge results (no conflicts)
        let merge_results = vec![
            MergeResult::Added(b"new_key".to_vec(), b"new_value".to_vec()),
            MergeResult::Modified(b"existing".to_vec(), b"modified_value".to_vec()),
        ];

        // Apply merge results
        let merge_tree = ProllyTree::new(storage, config);
        let result = merge_tree.apply_merge_results(&base_root, &merge_results);

        assert!(result.is_ok());
        let merged_tree = result.unwrap();

        // Verify the merged tree has the expected changes
        assert!(merged_tree.find(b"new_key").is_some());
        assert!(merged_tree.find(b"existing").is_some());

        // Check values are correct
        if let Some(node) = merged_tree.find(b"new_key") {
            let key_idx = node.keys.iter().position(|k| k == b"new_key").unwrap();
            let value = node.values[key_idx].clone();
            assert_eq!(value, b"new_value".to_vec());
        }

        if let Some(node) = merged_tree.find(b"existing") {
            let key_idx = node.keys.iter().position(|k| k == b"existing").unwrap();
            let value = node.values[key_idx].clone();
            assert_eq!(value, b"modified_value".to_vec());
        }
    }

    #[test]
    fn test_apply_merge_results_with_conflicts() {
        let config = TreeConfig::default();
        let mut storage = InMemoryNodeStorage::<32>::default();

        // Create base tree
        let base_tree = ProllyTree::new(storage.clone(), config.clone());
        let base_root = base_tree.get_root_hash().unwrap();
        storage
            .insert_node(base_root.clone(), base_tree.root.clone())
            .unwrap();

        // Create merge results with conflicts
        let merge_results = vec![
            MergeResult::Added(b"new_key".to_vec(), b"new_value".to_vec()),
            MergeResult::Conflict(MergeConflict {
                key: b"conflict_key".to_vec(),
                base_value: Some(b"base".to_vec()),
                source_value: Some(b"source".to_vec()),
                destination_value: Some(b"dest".to_vec()),
            }),
        ];

        // Apply merge results should fail due to conflicts
        let merge_tree = ProllyTree::new(storage, config);
        let result = merge_tree.apply_merge_results(&base_root, &merge_results);

        assert!(result.is_err());
        let conflicts = result.expect_err("Expected conflicts");
        assert_eq!(conflicts.len(), 1);
        assert_eq!(conflicts[0].key, b"conflict_key".to_vec());
    }

    #[test]
    fn test_merge_trees_success() {
        let config = TreeConfig::default();
        let mut storage = InMemoryNodeStorage::<32>::default();

        // Create base tree
        let mut base_tree = ProllyTree::new(storage.clone(), config.clone());
        base_tree.insert(b"shared".to_vec(), b"original".to_vec());
        let base_root = base_tree.get_root_hash().unwrap();
        storage
            .insert_node(base_root.clone(), base_tree.root.clone())
            .unwrap();

        // Create source tree with additions
        let mut source_tree = ProllyTree::new(storage.clone(), config.clone());
        source_tree.insert(b"shared".to_vec(), b"original".to_vec());
        source_tree.insert(b"from_source".to_vec(), b"source_value".to_vec());
        let source_root = source_tree.get_root_hash().unwrap();
        storage
            .insert_node(source_root.clone(), source_tree.root.clone())
            .unwrap();

        // Create destination tree with different additions
        let mut dest_tree = ProllyTree::new(storage.clone(), config.clone());
        dest_tree.insert(b"shared".to_vec(), b"original".to_vec());
        dest_tree.insert(b"from_dest".to_vec(), b"dest_value".to_vec());
        let dest_root = dest_tree.get_root_hash().unwrap();
        storage
            .insert_node(dest_root.clone(), dest_tree.root.clone())
            .unwrap();

        // Perform merge_trees (should succeed - no conflicts)
        let merge_tree = ProllyTree::new(storage, config);
        let result = merge_tree.merge_trees_ignore_conflicts(&source_root, &dest_root, &base_root);

        assert!(result.is_ok());
        let merged_tree = result.unwrap();

        // Verify merged tree contains all expected keys
        assert!(merged_tree.find(b"shared").is_some());
        assert!(merged_tree.find(b"from_source").is_some());
        assert!(merged_tree.find(b"from_dest").is_some()); // dest already had this
    }

    #[test]
    fn test_merge_trees_with_conflicts() {
        let config = TreeConfig::default();
        let mut storage = InMemoryNodeStorage::<32>::default();

        // Create base tree
        let mut base_tree = ProllyTree::new(storage.clone(), config.clone());
        base_tree.insert(b"conflict_key".to_vec(), b"original".to_vec());
        let base_root = base_tree.get_root_hash().unwrap();
        storage
            .insert_node(base_root.clone(), base_tree.root.clone())
            .unwrap();

        // Create source tree with modification
        let mut source_tree = ProllyTree::new(storage.clone(), config.clone());
        source_tree.insert(b"conflict_key".to_vec(), b"source_value".to_vec());
        let source_root = source_tree.get_root_hash().unwrap();
        storage
            .insert_node(source_root.clone(), source_tree.root.clone())
            .unwrap();

        // Create destination tree with different modification
        let mut dest_tree = ProllyTree::new(storage.clone(), config.clone());
        dest_tree.insert(b"conflict_key".to_vec(), b"dest_value".to_vec());
        let dest_root = dest_tree.get_root_hash().unwrap();
        storage
            .insert_node(dest_root.clone(), dest_tree.root.clone())
            .unwrap();

        // Perform merge_trees (should fail due to conflict)
        let merge_tree = ProllyTree::new(storage, config);

        // Create a resolver that doesn't resolve conflicts (keeps them as conflicts)
        struct NoResolutionResolver;
        impl ConflictResolver for NoResolutionResolver {
            fn resolve_conflict(&self, _conflict: &MergeConflict) -> Option<MergeResult> {
                None // Don't resolve any conflicts
            }
        }

        let resolver = NoResolutionResolver;
        let result = merge_tree.merge_trees(&source_root, &dest_root, &base_root, &resolver);

        assert!(result.is_err());
        let conflicts = result.expect_err("Expected conflicts");
        assert_eq!(conflicts.len(), 1);
        assert_eq!(conflicts[0].key, b"conflict_key".to_vec());
    }

    #[test]
    fn test_merge_trees_with_ignore_conflicts_resolver() {
        use crate::diff::IgnoreConflictsResolver;

        let config = TreeConfig::default();
        let mut storage = InMemoryNodeStorage::<32>::default();

        // Create base tree
        let mut base_tree = ProllyTree::new(storage.clone(), config.clone());
        base_tree.insert(b"shared".to_vec(), b"base_value".to_vec());
        let base_root = base_tree.get_root_hash().unwrap();
        storage
            .insert_node(base_root.clone(), base_tree.root.clone())
            .unwrap();

        // Create source tree with modifications
        let mut source_tree = ProllyTree::new(storage.clone(), config.clone());
        source_tree.insert(b"shared".to_vec(), b"source_value".to_vec());
        source_tree.insert(b"source_only".to_vec(), b"source_only_value".to_vec());
        let source_root = source_tree.get_root_hash().unwrap();
        storage
            .insert_node(source_root.clone(), source_tree.root.clone())
            .unwrap();

        // Create destination tree with conflicting modifications
        let mut dest_tree = ProllyTree::new(storage.clone(), config.clone());
        dest_tree.insert(b"shared".to_vec(), b"dest_value".to_vec());
        dest_tree.insert(b"dest_only".to_vec(), b"dest_only_value".to_vec());
        let dest_root = dest_tree.get_root_hash().unwrap();
        storage
            .insert_node(dest_root.clone(), dest_tree.root.clone())
            .unwrap();

        // Perform merge with ignore conflicts resolver
        let merge_tree = ProllyTree::new(storage, config);
        let resolver = IgnoreConflictsResolver;
        let result = merge_tree.merge_trees(&source_root, &dest_root, &base_root, &resolver);

        // Should succeed since conflicts are ignored
        assert!(result.is_ok());
        let merged_tree = result.unwrap();

        // Verify non-conflicting changes were applied
        assert!(merged_tree.find(b"source_only").is_some());
        assert!(merged_tree.find(b"dest_only").is_some());

        // Conflicting key should remain with destination value (unchanged)
        if let Some(node) = merged_tree.find(b"shared") {
            let key_idx = node.keys.iter().position(|k| k == b"shared").unwrap();
            let value = node.values[key_idx].clone();
            assert_eq!(value, b"dest_value".to_vec());
        }
    }

    #[test]
    fn test_merge_trees_with_take_source_resolver() {
        use crate::diff::TakeSourceResolver;

        let config = TreeConfig::default();
        let mut storage = InMemoryNodeStorage::<32>::default();

        // Create base tree
        let mut base_tree = ProllyTree::new(storage.clone(), config.clone());
        base_tree.insert(b"conflict_key".to_vec(), b"base_value".to_vec());
        let base_root = base_tree.get_root_hash().unwrap();
        storage
            .insert_node(base_root.clone(), base_tree.root.clone())
            .unwrap();

        // Create source tree
        let mut source_tree = ProllyTree::new(storage.clone(), config.clone());
        source_tree.insert(b"conflict_key".to_vec(), b"source_value".to_vec());
        let source_root = source_tree.get_root_hash().unwrap();
        storage
            .insert_node(source_root.clone(), source_tree.root.clone())
            .unwrap();

        // Create destination tree
        let mut dest_tree = ProllyTree::new(storage.clone(), config.clone());
        dest_tree.insert(b"conflict_key".to_vec(), b"dest_value".to_vec());
        let dest_root = dest_tree.get_root_hash().unwrap();
        storage
            .insert_node(dest_root.clone(), dest_tree.root.clone())
            .unwrap();

        // Perform merge with take source resolver
        let merge_tree = ProllyTree::new(storage, config);
        let resolver = TakeSourceResolver;
        let result = merge_tree.merge_trees(&source_root, &dest_root, &base_root, &resolver);

        assert!(result.is_ok());
        let merged_tree = result.unwrap();

        // Should have source value due to resolver
        if let Some(node) = merged_tree.find(b"conflict_key") {
            let key_idx = node.keys.iter().position(|k| k == b"conflict_key").unwrap();
            let value = node.values[key_idx].clone();
            assert_eq!(value, b"source_value".to_vec());
        }
    }

    #[test]
    fn test_merge_trees_with_take_destination_resolver() {
        use crate::diff::TakeDestinationResolver;

        let config = TreeConfig::default();
        let mut storage = InMemoryNodeStorage::<32>::default();

        // Create base tree
        let mut base_tree = ProllyTree::new(storage.clone(), config.clone());
        base_tree.insert(b"conflict_key".to_vec(), b"base_value".to_vec());
        let base_root = base_tree.get_root_hash().unwrap();
        storage
            .insert_node(base_root.clone(), base_tree.root.clone())
            .unwrap();

        // Create source tree
        let mut source_tree = ProllyTree::new(storage.clone(), config.clone());
        source_tree.insert(b"conflict_key".to_vec(), b"source_value".to_vec());
        let source_root = source_tree.get_root_hash().unwrap();
        storage
            .insert_node(source_root.clone(), source_tree.root.clone())
            .unwrap();

        // Create destination tree
        let mut dest_tree = ProllyTree::new(storage.clone(), config.clone());
        dest_tree.insert(b"conflict_key".to_vec(), b"dest_value".to_vec());
        let dest_root = dest_tree.get_root_hash().unwrap();
        storage
            .insert_node(dest_root.clone(), dest_tree.root.clone())
            .unwrap();

        // Perform merge with take destination resolver
        let merge_tree = ProllyTree::new(storage, config);
        let resolver = TakeDestinationResolver;
        let result = merge_tree.merge_trees(&source_root, &dest_root, &base_root, &resolver);

        assert!(result.is_ok());
        let merged_tree = result.unwrap();

        // Should have destination value due to resolver
        if let Some(node) = merged_tree.find(b"conflict_key") {
            let key_idx = node.keys.iter().position(|k| k == b"conflict_key").unwrap();
            let value = node.values[key_idx].clone();
            assert_eq!(value, b"dest_value".to_vec());
        }
    }

    #[test]
    fn test_merge_trees_ignore_conflicts_convenience_method() {
        let config = TreeConfig::default();
        let mut storage = InMemoryNodeStorage::<32>::default();

        // Create base tree
        let mut base_tree = ProllyTree::new(storage.clone(), config.clone());
        base_tree.insert(b"conflict_key".to_vec(), b"base_value".to_vec());
        let base_root = base_tree.get_root_hash().unwrap();
        storage
            .insert_node(base_root.clone(), base_tree.root.clone())
            .unwrap();

        // Create source tree
        let mut source_tree = ProllyTree::new(storage.clone(), config.clone());
        source_tree.insert(b"conflict_key".to_vec(), b"source_value".to_vec());
        source_tree.insert(b"new_key".to_vec(), b"new_value".to_vec());
        let source_root = source_tree.get_root_hash().unwrap();
        storage
            .insert_node(source_root.clone(), source_tree.root.clone())
            .unwrap();

        // Create destination tree
        let mut dest_tree = ProllyTree::new(storage.clone(), config.clone());
        dest_tree.insert(b"conflict_key".to_vec(), b"dest_value".to_vec());
        let dest_root = dest_tree.get_root_hash().unwrap();
        storage
            .insert_node(dest_root.clone(), dest_tree.root.clone())
            .unwrap();

        // Use convenience method
        let merge_tree = ProllyTree::new(storage, config);
        let result = merge_tree.merge_trees_ignore_conflicts(&source_root, &dest_root, &base_root);

        assert!(result.is_ok());
        let merged_tree = result.unwrap();

        // Should have the new key from source
        assert!(merged_tree.find(b"new_key").is_some());

        // Conflicting key should keep destination value (ignored conflict)
        if let Some(node) = merged_tree.find(b"conflict_key") {
            let key_idx = node.keys.iter().position(|k| k == b"conflict_key").unwrap();
            let value = node.values[key_idx].clone();
            assert_eq!(value, b"dest_value".to_vec());
        }
    }

    // ------------------------------------------------------------------
    // History-independence tests at the ProllyTree (public API) layer.
    //
    // Four always-on tests covering the same shapes used at the lower
    // layers: traversal equality across insertion orders, root-hash
    // equality across orders, root-hash equality under update sequences,
    // and root-hash equality under delete sequences. The full matrix
    // (configs × orders × key patterns × op mixes) lives in
    // `tests/history_independence.rs`.
    // ------------------------------------------------------------------
    mod history_independence_tests {
        use super::*;
        use rand::prelude::StdRng;
        use rand::seq::SliceRandom;
        use rand::SeedableRng;

        fn k8(i: u64) -> Vec<u8> {
            i.to_be_bytes().to_vec()
        }
        fn v16(i: u64) -> Vec<u8> {
            let mut v = Vec::with_capacity(16);
            v.extend_from_slice(&i.to_be_bytes());
            v.extend_from_slice(&(!i).to_be_bytes());
            v
        }

        /// Five insertion orders for the same key set 0..n.
        fn orders(n: u64) -> Vec<(&'static str, Vec<u64>)> {
            let asc: Vec<u64> = (0..n).collect();
            let desc: Vec<u64> = (0..n).rev().collect();
            let alt: Vec<u64> = (0..n).step_by(2).chain((1..n).step_by(2)).collect();
            let mut s0: Vec<u64> = (0..n).collect();
            s0.shuffle(&mut StdRng::from_seed([0u8; 32]));
            let mut s42: Vec<u64> = (0..n).collect();
            s42.shuffle(&mut StdRng::from_seed([42u8; 32]));
            vec![
                ("ascending", asc),
                ("descending", desc),
                ("alt-odd-even", alt),
                ("shuffled(0)", s0),
                ("shuffled(42)", s42),
            ]
        }

        fn build_tree(order: &[u64]) -> ProllyTree<32, InMemoryNodeStorage<32>> {
            let storage = InMemoryNodeStorage::<32>::default();
            let mut tree = ProllyTree::new(storage, TreeConfig::default());
            for &i in order {
                tree.insert(k8(i), v16(i));
            }
            tree
        }

        /// Pure inserts under the default TreeConfig produce the same
        /// leaf-level traversal regardless of insertion order. This is
        /// the equivalent of `test_history_independence_default_config_traversal`
        /// from `src/node.rs` but at the public ProllyTree API.
        #[test]
        fn prollytree_traversal_independent_of_order() {
            const N: u64 = 256;
            let orders = orders(N);
            let baseline_label = orders[0].0;
            let baseline_trav = build_tree(&orders[0].1).traverse();
            for (label, order) in orders.iter().skip(1) {
                let t = build_tree(order).traverse();
                assert_eq!(
                    t, baseline_trav,
                    "order={} leaf content diverged from order={}",
                    label, baseline_label
                );
            }
        }

        #[test]
        fn prollytree_root_hash_independent_of_order() {
            const N: u64 = 256;
            let orders = orders(N);
            let baseline_label = orders[0].0;
            let baseline_hash = build_tree(&orders[0].1).get_root_hash().unwrap();
            for (label, order) in orders.iter().skip(1) {
                let h = build_tree(order).get_root_hash().unwrap();
                assert_eq!(
                    h, baseline_hash,
                    "order={} root hash diverged from order={}",
                    label, baseline_label
                );
            }
        }

        #[test]
        fn prollytree_root_hash_independent_under_updates() {
            const N: u64 = 256;
            // Baseline: insert ascending with the real values.
            let asc: Vec<u64> = (0..N).collect();
            let baseline = build_tree(&asc).get_root_hash().unwrap();

            // Variant: insert all with placeholder, then overwrite all in reverse.
            let storage = InMemoryNodeStorage::<32>::default();
            let mut tree = ProllyTree::new(storage, TreeConfig::default());
            for &i in &asc {
                tree.insert(k8(i), vec![0u8]);
            }
            for i in (0..N).rev() {
                tree.insert(k8(i), v16(i));
            }
            let h = tree.get_root_hash().unwrap();
            assert_eq!(
                h, baseline,
                "update-then-final root hash diverged from baseline"
            );
        }

        #[test]
        fn prollytree_root_hash_independent_under_deletes() {
            const SURVIVE: u64 = 256;
            const EXTRA: u64 = 256;
            let survivors: Vec<u64> = (0..SURVIVE).collect();
            let baseline = build_tree(&survivors).get_root_hash().unwrap();

            let storage = InMemoryNodeStorage::<32>::default();
            let mut tree = ProllyTree::new(storage, TreeConfig::default());
            for i in 0..SURVIVE + EXTRA {
                tree.insert(k8(i), v16(i));
            }
            for i in SURVIVE..SURVIVE + EXTRA {
                tree.delete(&k8(i));
            }
            let h = tree.get_root_hash().unwrap();
            assert_eq!(
                h, baseline,
                "delete-then-final root hash diverged from baseline"
            );
        }
    }

    /// History-independence under three-way merge.
    ///
    /// `apply_merge_results` (and `merge_trees`) start from a destination
    /// tree and apply added / modified / removed entries via
    /// `ProllyTree::insert` / `delete`, both of which route through the
    /// streaming chunker in `crate::streaming_chunker`. So the merge
    /// result is a canonical tree that depends only on the final
    /// key/value set.
    mod merge_canonicality_tests {
        use super::*;
        use crate::diff::IgnoreConflictsResolver;
        use rand::prelude::StdRng;
        use rand::seq::SliceRandom;
        use rand::SeedableRng;

        fn k8(i: u64) -> Vec<u8> {
            i.to_be_bytes().to_vec()
        }
        fn v16(i: u64) -> Vec<u8> {
            let mut v = Vec::with_capacity(16);
            v.extend_from_slice(&i.to_be_bytes());
            v.extend_from_slice(&(!i).to_be_bytes());
            v
        }

        /// Build a fresh ProllyTree sharing the given storage; insert the
        /// keys in the supplied order so all node hashes are reachable
        /// from `storage`. Returns (tree, root_hash). The tree's
        /// `storage` field is updated in place via clone, so the caller's
        /// original mutable handle still references the latest state.
        fn build_tree_sharing_storage(
            storage: &mut InMemoryNodeStorage<32>,
            indices: &[u64],
        ) -> ValueDigest<32> {
            let mut tree = ProllyTree::new(storage.clone(), TreeConfig::default());
            for &i in indices {
                tree.insert(k8(i), v16(i));
            }
            let root_hash = tree.get_root_hash().unwrap();
            // Promote the per-tree storage (which now holds all nodes for this
            // tree) back into the shared storage.
            *storage = tree.storage;
            root_hash
        }

        /// `apply_merge_results` produces a canonical tree: applying merge
        /// ops to a destination must equal a freshly built tree from the
        /// same final key/value set.
        #[test]
        fn apply_merge_results_is_canonical() {
            let mut storage = InMemoryNodeStorage::<32>::default();

            // base has {0..16}, source adds {16..24}, dest adds {24..32}.
            let base_root = build_tree_sharing_storage(&mut storage, &(0..16).collect::<Vec<_>>());

            let mut source_keys: Vec<u64> = (0..24).collect();
            source_keys.shuffle(&mut StdRng::from_seed([1u8; 32]));
            let source_root = build_tree_sharing_storage(&mut storage, &source_keys);

            let mut dest_keys: Vec<u64> = (0..16).chain(24..32).collect();
            dest_keys.shuffle(&mut StdRng::from_seed([2u8; 32]));
            let dest_root = build_tree_sharing_storage(&mut storage, &dest_keys);

            // Expected final state: union of source and dest = {0..32}.
            let mut expected_keys: Vec<u64> = (0..32).collect();
            expected_keys.shuffle(&mut StdRng::from_seed([3u8; 32]));
            let expected_root = build_tree_sharing_storage(&mut storage, &expected_keys);

            let merge_tool = ProllyTree::new(storage, TreeConfig::default());
            let merge_results = merge_tool.merge(&source_root, &dest_root, &base_root);
            let merged = merge_tool
                .apply_merge_results(&dest_root, &merge_results)
                .expect("no conflicts expected");

            assert_eq!(
                merged.get_root_hash().unwrap(),
                expected_root,
                "merged tree's root hash does not match a fresh canonical build of the same final state"
            );
        }

        /// Same three input roots, but the merge is applied twice with the
        /// merge results in different orders. Because `apply_merge_results`
        /// goes through `insert`/`delete` which both stream through the
        /// canonical chunker, the final root hash should be identical.
        #[test]
        fn apply_merge_results_independent_of_result_order() {
            let mut storage = InMemoryNodeStorage::<32>::default();
            let base_root = build_tree_sharing_storage(&mut storage, &(0..16).collect::<Vec<_>>());
            let source_root =
                build_tree_sharing_storage(&mut storage, &(0..24).collect::<Vec<_>>());
            let dest_root = build_tree_sharing_storage(
                &mut storage,
                &(0..16).chain(24..32).collect::<Vec<_>>(),
            );

            let merge_tool = ProllyTree::new(storage, TreeConfig::default());
            let mut results = merge_tool.merge(&source_root, &dest_root, &base_root);

            let merged_asc = merge_tool
                .apply_merge_results(&dest_root, &results)
                .expect("no conflicts");

            results.shuffle(&mut StdRng::from_seed([4u8; 32]));
            let merged_shuf = merge_tool
                .apply_merge_results(&dest_root, &results)
                .expect("no conflicts");

            assert_eq!(
                merged_asc.get_root_hash().unwrap(),
                merged_shuf.get_root_hash().unwrap(),
                "applying the same merge results in different orders produced different root hashes"
            );
        }

        /// `merge_trees` with the ignore-conflicts resolver should produce
        /// a canonical tree regardless of how the underlying trees were
        /// built.
        #[test]
        fn merge_trees_ignore_conflicts_is_canonical() {
            let mut storage = InMemoryNodeStorage::<32>::default();
            let base_root = build_tree_sharing_storage(&mut storage, &(0..16).collect::<Vec<_>>());

            let mut a_keys: Vec<u64> = (0..24).collect();
            a_keys.shuffle(&mut StdRng::from_seed([5u8; 32]));
            let source_root = build_tree_sharing_storage(&mut storage, &a_keys);

            let mut b_keys: Vec<u64> = (0..16).chain(24..32).collect();
            b_keys.shuffle(&mut StdRng::from_seed([6u8; 32]));
            let dest_root = build_tree_sharing_storage(&mut storage, &b_keys);

            let expected_root =
                build_tree_sharing_storage(&mut storage, &(0..32).collect::<Vec<_>>());

            let merge_tool = ProllyTree::new(storage, TreeConfig::default());
            let merged = merge_tool
                .merge_trees(
                    &source_root,
                    &dest_root,
                    &base_root,
                    &IgnoreConflictsResolver,
                )
                .expect("no unresolved conflicts");

            assert_eq!(
                merged.get_root_hash().unwrap(),
                expected_root,
                "merge_trees output's root hash differs from a fresh canonical build"
            );
        }

        /// `merge_trees` with `TakeSourceResolver`: when source and
        /// destination both modified key `5` with different values, the
        /// merged tree must take the source value, and its root must
        /// equal a canonical build of the resulting key/value set.
        #[test]
        fn merge_trees_take_source_is_canonical() {
            use crate::diff::TakeSourceResolver;
            let mut storage = InMemoryNodeStorage::<32>::default();
            let base_root = build_tree_sharing_storage(&mut storage, &(0..16).collect::<Vec<_>>());

            // Build source where key 5 has value v(5_000_000)
            let source_root = {
                let mut tree = ProllyTree::new(storage.clone(), TreeConfig::default());
                for i in 0u64..24 {
                    let v = if i == 5 { v16(5_000_000) } else { v16(i) };
                    tree.insert(k8(i), v);
                }
                let h = tree.get_root_hash().unwrap();
                storage = tree.storage;
                h
            };

            // Build dest where key 5 has value v(9_999_999)
            let dest_root = {
                let mut tree = ProllyTree::new(storage.clone(), TreeConfig::default());
                for i in (0..16).chain(24..32) {
                    let v = if i == 5 { v16(9_999_999) } else { v16(i) };
                    tree.insert(k8(i), v);
                }
                let h = tree.get_root_hash().unwrap();
                storage = tree.storage;
                h
            };

            // Expected: union of keys; key 5 takes source's value.
            let expected_root = {
                let mut tree = ProllyTree::new(storage.clone(), TreeConfig::default());
                for i in 0u64..32 {
                    let v = if i == 5 { v16(5_000_000) } else { v16(i) };
                    tree.insert(k8(i), v);
                }
                let h = tree.get_root_hash().unwrap();
                storage = tree.storage;
                h
            };

            let merge_tool = ProllyTree::new(storage, TreeConfig::default());
            let merged = merge_tool
                .merge_trees(&source_root, &dest_root, &base_root, &TakeSourceResolver)
                .expect("no unresolved conflicts (resolver handles them)");

            assert_eq!(
                merged.get_root_hash().unwrap(),
                expected_root,
                "TakeSourceResolver merge output is non-canonical"
            );
        }

        /// `merge_trees` with `TakeDestinationResolver`: symmetric to the
        /// TakeSource test - the destination's value for the conflicting
        /// key must be the one that ends up in the merged tree.
        #[test]
        fn merge_trees_take_destination_is_canonical() {
            use crate::diff::TakeDestinationResolver;
            let mut storage = InMemoryNodeStorage::<32>::default();
            let base_root = build_tree_sharing_storage(&mut storage, &(0..16).collect::<Vec<_>>());

            let source_root = {
                let mut tree = ProllyTree::new(storage.clone(), TreeConfig::default());
                for i in 0u64..24 {
                    let v = if i == 5 { v16(5_000_000) } else { v16(i) };
                    tree.insert(k8(i), v);
                }
                let h = tree.get_root_hash().unwrap();
                storage = tree.storage;
                h
            };

            let dest_root = {
                let mut tree = ProllyTree::new(storage.clone(), TreeConfig::default());
                for i in (0..16).chain(24..32) {
                    let v = if i == 5 { v16(9_999_999) } else { v16(i) };
                    tree.insert(k8(i), v);
                }
                let h = tree.get_root_hash().unwrap();
                storage = tree.storage;
                h
            };

            // Expected: union of keys; key 5 takes destination's value.
            let expected_root = {
                let mut tree = ProllyTree::new(storage.clone(), TreeConfig::default());
                for i in 0u64..32 {
                    let v = if i == 5 { v16(9_999_999) } else { v16(i) };
                    tree.insert(k8(i), v);
                }
                let h = tree.get_root_hash().unwrap();
                storage = tree.storage;
                h
            };

            let merge_tool = ProllyTree::new(storage, TreeConfig::default());
            let merged = merge_tool
                .merge_trees(
                    &source_root,
                    &dest_root,
                    &base_root,
                    &TakeDestinationResolver,
                )
                .expect("no unresolved conflicts (resolver handles them)");

            assert_eq!(
                merged.get_root_hash().unwrap(),
                expected_root,
                "TakeDestinationResolver merge output is non-canonical"
            );
        }
    }
}

#[cfg(test)]
mod prefix_free_collision_probe {
    use super::*;
    use crate::storage::InMemoryNodeStorage;

    /// Red-team attack #1: distinct (k,v) sets must NOT share a root hash. get_hash =
    /// SHA256(keys.concat() ++ values.concat()) has no length delimiters, so values
    /// "ab"+"c" and "a"+"bc" both concat to "abc" -> same hash input -> same root.
    #[test]
    fn distinct_contents_must_not_share_root() {
        let cfg = TreeConfig::<32>::default();
        let mut a = ProllyTree::new(InMemoryNodeStorage::<32>::default(), cfg.clone());
        a.insert(b"k1".to_vec(), b"ab".to_vec());
        a.insert(b"k2".to_vec(), b"c".to_vec());
        let mut b = ProllyTree::new(InMemoryNodeStorage::<32>::default(), cfg.clone());
        b.insert(b"k1".to_vec(), b"a".to_vec());
        b.insert(b"k2".to_vec(), b"bc".to_vec());
        let ra = a.get_root_hash().unwrap();
        let rb = b.get_root_hash().unwrap();
        // also a key-vs-value boundary collision: key "k1k2" value ... vs keys "k1","k2"
        assert_ne!(
            ra.as_bytes(),
            rb.as_bytes(),
            "PREFIX-FREE COLLISION CONFIRMED: distinct (k,v) sets share a root hash"
        );
    }
}

/// O(diff) structural-diff differential tests: the structural `diff_nodes_recursive`
/// must emit a BYTE-IDENTICAL `Vec<DiffResult>` to the proven full-leaf flatten oracle
/// (`diff_nodes_flatten`) on every shape — and load strictly fewer nodes when a large
/// subtree is shared. prollytree is history-independent, so shape is a function of
/// CONTENT; divergent shapes come from divergent content with a shared region.
#[cfg(test)]
mod odiff_differential {
    use super::*;
    use crate::digest::ValueDigest;
    use crate::storage::{InMemoryNodeStorage, NodeStorage, StorageError};
    use std::sync::atomic::{AtomicUsize, Ordering};
    use std::sync::Arc;

    /// InMemoryNodeStorage wrapper that counts `get_node_by_hash` (node loads).
    #[derive(Clone)]
    struct CountingStorage<const N: usize> {
        inner: InMemoryNodeStorage<N>,
        loads: Arc<AtomicUsize>,
    }
    impl<const N: usize> CountingStorage<N> {
        fn new() -> Self {
            Self {
                inner: InMemoryNodeStorage::<N>::default(),
                loads: Arc::new(AtomicUsize::new(0)),
            }
        }
    }
    impl<const N: usize> NodeStorage<N> for CountingStorage<N> {
        fn get_node_by_hash(&self, hash: &ValueDigest<N>) -> Option<Arc<ProllyNode<N>>> {
            self.loads.fetch_add(1, Ordering::SeqCst);
            self.inner.get_node_by_hash(hash)
        }
        fn insert_node(
            &mut self,
            hash: ValueDigest<N>,
            node: ProllyNode<N>,
        ) -> Result<(), StorageError> {
            self.inner.insert_node(hash, node)
        }
        fn delete_node(&mut self, hash: &ValueDigest<N>) -> Result<(), StorageError> {
            self.inner.delete_node(hash)
        }
        fn save_config(&self, key: &str, config: &[u8]) {
            self.inner.save_config(key, config)
        }
        fn get_config(&self, key: &str) -> Option<Vec<u8>> {
            self.inner.get_config(key)
        }
    }

    /// Build BOTH trees into ONE storage. `InMemoryNodeStorage::clone` DEEP-COPIES the
    /// node map and `insert` persists nodes into the TREE's storage — so we must thread
    /// one storage through both builds (clone carries old's nodes forward; new's inserts
    /// add to it) rather than clone-then-copy-only-root (which strands all children and
    /// makes a multi-level diff trivially empty). Returns (storage_with_both, roots).
    fn build_two(
        config: &TreeConfig<32>,
        old_pairs: &[(Vec<u8>, Vec<u8>)],
        new_pairs: &[(Vec<u8>, Vec<u8>)],
    ) -> (CountingStorage<32>, ValueDigest<32>, ValueDigest<32>) {
        let mut t_old = ProllyTree::new(CountingStorage::<32>::new(), config.clone());
        for (k, v) in old_pairs {
            t_old.insert(k.clone(), v.clone());
        }
        let old_root = t_old.get_root_hash().unwrap();
        // Ensure the root node is present even when no insert ran (empty tree:
        // `insert`/`persist_root` never fired).
        t_old
            .storage
            .insert_node(old_root.clone(), t_old.root.clone())
            .unwrap();
        // clone() carries a full copy of old's nodes; new's inserts add new's nodes,
        // never evicting old's -> the result holds BOTH complete trees.
        let mut t_new = ProllyTree::new(t_old.storage.clone(), config.clone());
        for (k, v) in new_pairs {
            t_new.insert(k.clone(), v.clone());
        }
        let new_root = t_new.get_root_hash().unwrap();
        t_new
            .storage
            .insert_node(new_root.clone(), t_new.root.clone())
            .unwrap();
        (t_new.storage, old_root, new_root)
    }

    /// k-th key as a fixed-width sortable 8-byte big-endian blob.
    fn key(i: u64) -> Vec<u8> {
        i.to_be_bytes().to_vec()
    }
    /// variable-length value (exercises the D14 variable-length canonicality face).
    fn val(i: u64, tag: u8) -> Vec<u8> {
        let mut v = vec![tag];
        v.extend_from_slice(&i.to_le_bytes());
        v.extend(std::iter::repeat_n(tag, (i % 19) as usize));
        v
    }

    /// Assert structural diff == flatten oracle (exact Vec<DiffResult>), return
    /// (structural_loads, flatten_loads) measured during the diff only.
    fn assert_identical(
        old_pairs: &[(Vec<u8>, Vec<u8>)],
        new_pairs: &[(Vec<u8>, Vec<u8>)],
        label: &str,
    ) -> (usize, usize) {
        let config = TreeConfig::<32>::default();
        let (storage, old_root, new_root) = build_two(&config, old_pairs, new_pairs);

        // Both full trees are in `storage`; their children resolve (asserted below).
        let old_node = storage.get_node_by_hash(&old_root).unwrap();
        let new_node = storage.get_node_by_hash(&new_root).unwrap();
        // HARNESS GUARD: a non-leaf root's children MUST resolve, else the diff would
        // be trivially empty (the build_root bug). children() returns all present.
        if !old_node.is_leaf {
            assert_eq!(
                old_node.children(&storage).len(),
                old_node.values.len(),
                "[{label}] old root children not all persisted (harness bug)"
            );
        }
        if !new_node.is_leaf {
            assert_eq!(
                new_node.children(&storage).len(),
                new_node.values.len(),
                "[{label}] new root children not all persisted (harness bug)"
            );
        }
        let facade = ProllyTree::new(storage.clone(), config.clone());

        storage.loads.store(0, Ordering::SeqCst);
        let mut d_struct = Vec::new();
        facade.diff_nodes_recursive(&old_node, &new_node, &mut d_struct);
        let struct_loads = storage.loads.load(Ordering::SeqCst);

        storage.loads.store(0, Ordering::SeqCst);
        let mut d_flat = Vec::new();
        facade.diff_nodes_flatten(&old_node, &new_node, &mut d_flat);
        let flat_loads = storage.loads.load(Ordering::SeqCst);

        assert_eq!(
            d_struct, d_flat,
            "[{label}] structural diff diverged from flatten oracle"
        );
        (struct_loads, flat_loads)
    }

    fn range_pairs(lo: u64, hi: u64, tag: u8) -> Vec<(Vec<u8>, Vec<u8>)> {
        (lo..hi).map(|i| (key(i), val(i, tag))).collect()
    }

    #[test]
    fn identical_trees_zero_diff_zero_descent() {
        let base = range_pairs(0, 1000, 0);
        let (s, f) = assert_identical(&base, &base, "identical");
        // equal root hash => structural returns immediately, no descent at all.
        assert_eq!(s, 0, "identical trees must load zero nodes structurally");
        assert!(f >= s);
    }

    #[test]
    fn localized_modify_skips_shared_subtrees() {
        let old = range_pairs(0, 1000, 0);
        let mut new = old.clone();
        new[500].1 = val(500, 9); // one cell changed deep in the tree
        let (s, f) = assert_identical(&old, &new, "localized_modify");
        assert!(
            s < f,
            "localized change must load fewer nodes structurally ({s}) than flatten ({f})"
        );
        // GROUND TRUTH (not just structural==flatten): the diff must be EXACTLY the one
        // changed cell — guards against a both-paths-lossy false pass.
        let config = TreeConfig::<32>::default();
        let (storage, old_root, new_root) = build_two(&config, &old, &new);
        let on = storage.get_node_by_hash(&old_root).unwrap();
        let nn = storage.get_node_by_hash(&new_root).unwrap();
        let facade = ProllyTree::new(storage.clone(), config.clone());
        let mut d = Vec::new();
        facade.diff_nodes_recursive(&on, &nn, &mut d);
        assert_eq!(
            d,
            vec![DiffResult::Modified(key(500), val(500, 0), val(500, 9))],
            "structural diff must report exactly the one changed cell, got {d:?}"
        );
    }

    #[test]
    fn bulk_suffix_add_shares_prefix() {
        let old = range_pairs(0, 800, 0);
        let new = range_pairs(0, 1000, 0); // 200 appended at the high end
        let (s, f) = assert_identical(&old, &new, "bulk_suffix_add");
        assert!(
            s < f,
            "shared-prefix add should skip prefix subtrees ({s} < {f})"
        );
    }

    #[test]
    fn bulk_prefix_remove() {
        let old = range_pairs(0, 1000, 0);
        let new = range_pairs(200, 1000, 0); // low 200 removed
        assert_identical(&old, &new, "bulk_prefix_remove");
    }

    #[test]
    fn scattered_modifications() {
        let old = range_pairs(0, 1000, 0);
        let mut new = old.clone();
        for i in (0..1000).step_by(7) {
            new[i].1 = val(i as u64, 5); // every 7th cell changed
        }
        assert_identical(&old, &new, "scattered_modify");
    }

    #[test]
    fn fully_disjoint_keyspaces() {
        let old = range_pairs(0, 500, 0);
        let new = range_pairs(10_000, 10_500, 0); // no overlap at all
        assert_identical(&old, &new, "disjoint");
    }

    #[test]
    fn interleaved_add_remove_modify() {
        let mut old = range_pairs(0, 600, 0);
        // remove evens, modify multiples of 3, keep odds; then add a high block.
        let mut new: Vec<(Vec<u8>, Vec<u8>)> = Vec::new();
        for i in 0..600u64 {
            if i % 2 == 0 {
                continue; // removed
            }
            let tag = if i % 3 == 0 { 7 } else { 0 };
            new.push((key(i), val(i, tag)));
        }
        new.extend(range_pairs(600, 720, 0));
        old.sort();
        new.sort();
        assert_identical(&old, &new, "interleaved");
    }

    #[test]
    fn ground_truth_add_remove_modify() {
        // base 0..200; new: drop key 50, modify key 100, add key 500. Diff must be
        // EXACTLY those three ops (sorted by key), regardless of structural pruning.
        let old = range_pairs(0, 200, 0);
        let mut new = old.clone();
        new.retain(|(k, _)| k != &key(50));
        for p in new.iter_mut() {
            if p.0 == key(100) {
                p.1 = val(100, 9);
            }
        }
        new.push((key(500), val(500, 0)));
        new.sort();
        let config = TreeConfig::<32>::default();
        let (storage, old_root, new_root) = build_two(&config, &old, &new);
        let on = storage.get_node_by_hash(&old_root).unwrap();
        let nn = storage.get_node_by_hash(&new_root).unwrap();
        let facade = ProllyTree::new(storage.clone(), config.clone());
        let mut d = Vec::new();
        facade.diff_nodes_recursive(&on, &nn, &mut d);
        // merge-join emits in key order: 50 (Removed) < 100 (Modified) < 500 (Added).
        assert_eq!(
            d,
            vec![
                DiffResult::Removed(key(50), val(50, 0)),
                DiffResult::Modified(key(100), val(100, 0), val(100, 9)),
                DiffResult::Added(key(500), val(500, 0)),
            ],
            "structural diff ground truth mismatch, got {d:?}"
        );
    }

    #[test]
    fn report_load_win() {
        for n in [1_000u64, 4_000u64] {
            let old = range_pairs(0, n, 0);
            let mut new = old.clone();
            new[(n / 2) as usize].1 = val(n / 2, 9); // one cell changed
            let (s, f) = assert_identical(&old, &new, "report");
            eprintln!(
                "ODIFF_LOAD_WIN n={n} one-cell-change: structural_loads={s} flatten_loads={f} ratio={:.1}x",
                f as f64 / s.max(1) as f64
            );
        }
    }

    #[test]
    fn small_multilevel_and_single_leaf() {
        // small trees (single leaf / shallow) must also match exactly.
        assert_identical(&range_pairs(0, 3, 0), &range_pairs(0, 5, 0), "tiny");
        assert_identical(
            &range_pairs(0, 1, 0),
            &range_pairs(0, 1, 1),
            "one_key_modify",
        );
        assert_identical(&[], &range_pairs(0, 4, 0), "empty_to_small");
        assert_identical(&range_pairs(0, 4, 0), &[], "small_to_empty");
    }
}