synor/crates/synor-zk/src/state.rs

//! State tree management for ZK-rollups.
//!
//! The state tree is a sparse Merkle tree that stores account states.
//! It supports efficient proofs of inclusion and non-inclusion.
//!
//! # Structure
//!
//! ```text
//!                     Root
//!                    /    \
//!                  H1      H2
//!                 /  \    /  \
//!               H3   H4  H5   H6
//!              / \
//!          Acc0  Acc1  ...  AccN  Empty
//! ```
//!
//! - Tree depth: 32 (supports 2^32 accounts)
//! - Hash function: Poseidon (ZK-friendly) or Blake3 (for off-chain)
//! - Empty leaves: Hash of zero

use blake3;
use parking_lot::RwLock;
use rs_merkle::{Hasher, MerkleTree};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use thiserror::Error;

/// State root - the Merkle root of the state tree.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct StateRoot(pub [u8; 32]);

impl StateRoot {
    /// Creates a state root from bytes.
    pub fn from_bytes(bytes: [u8; 32]) -> Self {
        Self(bytes)
    }

    /// Returns the root as bytes.
    pub fn as_bytes(&self) -> &[u8; 32] {
        &self.0
    }

    /// Returns a zero state root (empty tree).
    pub fn zero() -> Self {
        Self([0u8; 32])
    }

    /// Returns the hex representation.
    pub fn to_hex(&self) -> String {
        hex::encode(self.0)
    }
}

impl From<[u8; 32]> for StateRoot {
    fn from(bytes: [u8; 32]) -> Self {
        Self(bytes)
    }
}

impl Default for StateRoot {
    fn default() -> Self {
        Self::zero()
    }
}

/// Account state stored in the state tree.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct AccountState {
    /// Account balance
    pub balance: u128,
    /// Account nonce (for replay protection)
    pub nonce: u64,
    /// Public key hash (for signature verification)
    pub pubkey_hash: [u8; 32],
    /// Additional data (e.g., storage root for smart contracts)
    pub data_hash: [u8; 32],
}

impl AccountState {
    /// Creates a new account state.
    pub fn new(pubkey_hash: [u8; 32]) -> Self {
        Self {
            balance: 0,
            nonce: 0,
            pubkey_hash,
            data_hash: [0u8; 32],
        }
    }

    /// Creates an account with initial balance.
    pub fn with_balance(pubkey_hash: [u8; 32], balance: u128) -> Self {
        Self {
            balance,
            nonce: 0,
            pubkey_hash,
            data_hash: [0u8; 32],
        }
    }

    /// Returns the hash of this account state.
    pub fn hash(&self) -> [u8; 32] {
        let mut hasher = blake3::Hasher::new();
        hasher.update(&self.balance.to_le_bytes());
        hasher.update(&self.nonce.to_le_bytes());
        hasher.update(&self.pubkey_hash);
        hasher.update(&self.data_hash);
        *hasher.finalize().as_bytes()
    }

    /// Serializes to bytes.
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut bytes = Vec::with_capacity(16 + 8 + 32 + 32);
        bytes.extend(&self.balance.to_le_bytes());
        bytes.extend(&self.nonce.to_le_bytes());
        bytes.extend(&self.pubkey_hash);
        bytes.extend(&self.data_hash);
        bytes
    }

    /// Deserializes from bytes.
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, StateError> {
        if bytes.len() < 16 + 8 + 32 + 32 {
            return Err(StateError::InvalidAccountData);
        }

        let balance = u128::from_le_bytes(bytes[0..16].try_into().unwrap());
        let nonce = u64::from_le_bytes(bytes[16..24].try_into().unwrap());
        let mut pubkey_hash = [0u8; 32];
        pubkey_hash.copy_from_slice(&bytes[24..56]);
        let mut data_hash = [0u8; 32];
        data_hash.copy_from_slice(&bytes[56..88]);

        Ok(Self {
            balance,
            nonce,
            pubkey_hash,
            data_hash,
        })
    }
}

impl Default for AccountState {
    fn default() -> Self {
        Self::new([0u8; 32])
    }
}

/// Account with index and state.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct Account {
    /// Account index in the state tree
    pub index: u64,
    /// Account state
    pub state: AccountState,
}

impl Account {
    /// Creates a new account.
    pub fn new(index: u64, pubkey_hash: [u8; 32]) -> Self {
        Self {
            index,
            state: AccountState::new(pubkey_hash),
        }
    }

    /// Returns the account hash.
    pub fn hash(&self) -> [u8; 32] {
        self.state.hash()
    }
}

/// Sparse Merkle tree for account states.
pub struct StateTree {
    /// Tree depth
    depth: usize,
    /// Accounts by index
    accounts: RwLock<HashMap<u64, AccountState>>,
    /// Cached leaf hashes
    leaf_hashes: RwLock<HashMap<u64, [u8; 32]>>,
    /// Current root (cached)
    root: RwLock<StateRoot>,
}

impl StateTree {
    /// Creates a new empty state tree.
    pub fn new(depth: usize) -> Self {
        Self {
            depth,
            accounts: RwLock::new(HashMap::new()),
            leaf_hashes: RwLock::new(HashMap::new()),
            root: RwLock::new(StateRoot::zero()),
        }
    }

    /// Creates a state tree with default depth (32).
    pub fn with_default_depth() -> Self {
        Self::new(crate::constants::STATE_TREE_DEPTH)
    }

    /// Returns the tree depth.
    pub fn depth(&self) -> usize {
        self.depth
    }

    /// Returns the current state root.
    pub fn root(&self) -> StateRoot {
        *self.root.read()
    }

    /// Returns the number of accounts.
    pub fn account_count(&self) -> usize {
        self.accounts.read().len()
    }

    /// Gets an account by index.
    pub fn get_account(&self, index: u64) -> Option<AccountState> {
        self.accounts.read().get(&index).cloned()
    }

    /// Sets an account state.
    pub fn set_account(&self, index: u64, state: AccountState) -> Result<StateRoot, StateError> {
        // Update account
        let hash = state.hash();
        self.accounts.write().insert(index, state);
        self.leaf_hashes.write().insert(index, hash);

        // Recompute root
        let new_root = self.compute_root()?;
        *self.root.write() = new_root;

        Ok(new_root)
    }

    /// Updates an account with a function.
    pub fn update_account<F>(&self, index: u64, f: F) -> Result<StateRoot, StateError>
    where
        F: FnOnce(&mut AccountState),
    {
        let mut accounts = self.accounts.write();
        let state = accounts.entry(index).or_default();
        f(state);
        let hash = state.hash();
        drop(accounts);

        self.leaf_hashes.write().insert(index, hash);

        let new_root = self.compute_root()?;
        *self.root.write() = new_root;

        Ok(new_root)
    }

    /// Computes the Merkle root from all accounts.
    fn compute_root(&self) -> Result<StateRoot, StateError> {
        let leaf_hashes = self.leaf_hashes.read();

        if leaf_hashes.is_empty() {
            return Ok(StateRoot::zero());
        }

        // For a sparse Merkle tree, we compute the root differently
        // Here we use a simplified approach with rs_merkle

        let mut leaves: Vec<[u8; 32]> = leaf_hashes.values().copied().collect();

        // Pad to power of 2
        let target_size = (leaves.len() as u64).next_power_of_two() as usize;
        while leaves.len() < target_size {
            leaves.push([0u8; 32]);
        }

        // Create Merkle tree
        let tree: MerkleTree<Blake3Algorithm> = MerkleTree::from_leaves(&leaves);
        let root = tree.root().ok_or(StateError::EmptyTree)?;

        let mut root_bytes = [0u8; 32];
        root_bytes.copy_from_slice(&root);

        Ok(StateRoot(root_bytes))
    }

    /// Generates a Merkle proof for an account.
    pub fn generate_proof(&self, index: u64) -> Result<Vec<[u8; 32]>, StateError> {
        let leaf_hashes = self.leaf_hashes.read();

        if !leaf_hashes.contains_key(&index) {
            return Err(StateError::AccountNotFound(index));
        }

        let mut leaves: Vec<_> = leaf_hashes.iter().collect();
        leaves.sort_by_key(|(k, _)| *k);

        let leaf_values: Vec<[u8; 32]> = leaves.iter().map(|(_, v)| **v).collect();

        // Find the index of our account in the sorted list
        let proof_index = leaves
            .iter()
            .position(|(k, _)| **k == index)
            .ok_or(StateError::AccountNotFound(index))?;

        // Pad to power of 2
        let mut padded_leaves = leaf_values.clone();
        let target_size = (padded_leaves.len() as u64).next_power_of_two() as usize;
        while padded_leaves.len() < target_size {
            padded_leaves.push([0u8; 32]);
        }

        let tree: MerkleTree<Blake3Algorithm> = MerkleTree::from_leaves(&padded_leaves);
        let proof = tree.proof(&[proof_index]);

        Ok(proof
            .proof_hashes()
            .iter()
            .map(|h| {
                let mut arr = [0u8; 32];
                arr.copy_from_slice(h);
                arr
            })
            .collect())
    }

    /// Verifies a Merkle proof.
    pub fn verify_proof(
        root: &StateRoot,
        index: u64,
        account_hash: [u8; 32],
        proof: &[[u8; 32]],
    ) -> bool {
        // Simplified verification - in production would use proper sparse Merkle verification
        let mut current = account_hash;

        for sibling in proof {
            let mut combined = Vec::with_capacity(64);
            if index & 1 == 0 {
                combined.extend(&current);
                combined.extend(sibling);
            } else {
                combined.extend(sibling);
                combined.extend(&current);
            }
            current = *blake3::hash(&combined).as_bytes();
        }

        current == root.0
    }

    /// Applies a transfer between accounts.
    pub fn apply_transfer(
        &self,
        from_idx: u64,
        to_idx: u64,
        amount: u128,
    ) -> Result<StateRoot, StateError> {
        // Check sender balance
        {
            let accounts = self.accounts.read();
            let from = accounts
                .get(&from_idx)
                .ok_or(StateError::AccountNotFound(from_idx))?;
            if from.balance < amount {
                return Err(StateError::InsufficientBalance {
                    available: from.balance,
                    required: amount,
                });
            }
        }

        // Update sender
        self.update_account(from_idx, |acc| {
            acc.balance -= amount;
            acc.nonce += 1;
        })?;

        // Update recipient
        self.update_account(to_idx, |acc| {
            acc.balance += amount;
        })
    }

    /// Applies a deposit (increases balance).
    pub fn apply_deposit(&self, to_idx: u64, amount: u128) -> Result<StateRoot, StateError> {
        self.update_account(to_idx, |acc| {
            acc.balance += amount;
        })
    }

    /// Applies a withdrawal (decreases balance).
    pub fn apply_withdrawal(&self, from_idx: u64, amount: u128) -> Result<StateRoot, StateError> {
        {
            let accounts = self.accounts.read();
            let from = accounts
                .get(&from_idx)
                .ok_or(StateError::AccountNotFound(from_idx))?;
            if from.balance < amount {
                return Err(StateError::InsufficientBalance {
                    available: from.balance,
                    required: amount,
                });
            }
        }

        self.update_account(from_idx, |acc| {
            acc.balance -= amount;
            acc.nonce += 1;
        })
    }

    /// Clears all accounts and resets to empty state.
    pub fn clear(&self) {
        self.accounts.write().clear();
        self.leaf_hashes.write().clear();
        *self.root.write() = StateRoot::zero();
    }
}

impl std::fmt::Debug for StateTree {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("StateTree")
            .field("depth", &self.depth)
            .field("account_count", &self.accounts.read().len())
            .field("root", &self.root.read().to_hex())
            .finish()
    }
}

/// Blake3 algorithm for rs_merkle.
#[derive(Clone)]
pub struct Blake3Algorithm;

impl Hasher for Blake3Algorithm {
    type Hash = [u8; 32];

    fn hash(data: &[u8]) -> Self::Hash {
        *blake3::hash(data).as_bytes()
    }
}

/// State errors.
#[derive(Debug, Error)]
pub enum StateError {
    #[error("Account not found: {0}")]
    AccountNotFound(u64),

    #[error("Insufficient balance: available {available}, required {required}")]
    InsufficientBalance { available: u128, required: u128 },

    #[error("Invalid account data")]
    InvalidAccountData,

    #[error("Tree is empty")]
    EmptyTree,

    #[error("Proof verification failed")]
    ProofVerificationFailed,

    #[error("Invalid state transition: {0}")]
    InvalidTransition(String),
}

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

    // ============================================================================
    // StateRoot Tests
    // ============================================================================

    #[test]
    fn test_state_root_from_bytes() {
        let bytes = [0xab; 32];
        let root = StateRoot::from_bytes(bytes);
        assert_eq!(root.0, bytes);
    }

    #[test]
    fn test_state_root_as_bytes() {
        let bytes = [0xcd; 32];
        let root = StateRoot::from_bytes(bytes);
        assert_eq!(root.as_bytes(), &bytes);
    }

    #[test]
    fn test_state_root_zero() {
        let root = StateRoot::zero();
        assert_eq!(root.0, [0u8; 32]);
    }

    #[test]
    fn test_state_root_to_hex() {
        let mut bytes = [0u8; 32];
        bytes[0] = 0xab;
        bytes[1] = 0xcd;
        let root = StateRoot::from_bytes(bytes);
        let hex = root.to_hex();
        assert!(hex.starts_with("abcd"));
        assert_eq!(hex.len(), 64);
    }

    #[test]
    fn test_state_root_from_array() {
        let bytes = [0xff; 32];
        let root: StateRoot = bytes.into();
        assert_eq!(root.0, bytes);
    }

    #[test]
    fn test_state_root_default() {
        let root = StateRoot::default();
        assert_eq!(root, StateRoot::zero());
    }

    #[test]
    fn test_state_root_equality() {
        let root1 = StateRoot::from_bytes([0xaa; 32]);
        let root2 = StateRoot::from_bytes([0xaa; 32]);
        let root3 = StateRoot::from_bytes([0xbb; 32]);

        assert_eq!(root1, root2);
        assert_ne!(root1, root3);
    }

    #[test]
    fn test_state_root_clone() {
        let root1 = StateRoot::from_bytes([0xab; 32]);
        let root2 = root1;
        assert_eq!(root1, root2);
    }

    #[test]
    fn test_state_root_hash() {
        use std::collections::HashSet;
        let mut set = HashSet::new();
        set.insert(StateRoot::from_bytes([0xaa; 32]));
        set.insert(StateRoot::from_bytes([0xbb; 32]));
        set.insert(StateRoot::from_bytes([0xaa; 32]));

        assert_eq!(set.len(), 2);
    }

    #[test]
    fn test_state_root_debug() {
        let root = StateRoot::from_bytes([0xab; 32]);
        let debug_str = format!("{:?}", root);
        assert!(debug_str.contains("StateRoot"));
    }

    // ============================================================================
    // AccountState Tests
    // ============================================================================

    #[test]
    fn test_account_state_new() {
        let pubkey_hash = [0xab; 32];
        let state = AccountState::new(pubkey_hash);

        assert_eq!(state.balance, 0);
        assert_eq!(state.nonce, 0);
        assert_eq!(state.pubkey_hash, pubkey_hash);
        assert_eq!(state.data_hash, [0u8; 32]);
    }

    #[test]
    fn test_account_state_with_balance() {
        let pubkey_hash = [0xab; 32];
        let state = AccountState::with_balance(pubkey_hash, 1000);

        assert_eq!(state.balance, 1000);
        assert_eq!(state.nonce, 0);
        assert_eq!(state.pubkey_hash, pubkey_hash);
    }

    #[test]
    fn test_account_state_hash_consistency() {
        let state = AccountState::with_balance([0xab; 32], 1000);
        let hash1 = state.hash();
        let hash2 = state.hash();
        assert_eq!(hash1, hash2);
    }

    #[test]
    fn test_account_state_hash_different_balances() {
        let state1 = AccountState::with_balance([0xab; 32], 1000);
        let state2 = AccountState::with_balance([0xab; 32], 2000);
        assert_ne!(state1.hash(), state2.hash());
    }

    #[test]
    fn test_account_state_hash_different_nonces() {
        let mut state1 = AccountState::with_balance([0xab; 32], 1000);
        let mut state2 = AccountState::with_balance([0xab; 32], 1000);
        state1.nonce = 1;
        state2.nonce = 2;
        assert_ne!(state1.hash(), state2.hash());
    }

    #[test]
    fn test_account_state_hash_different_pubkeys() {
        let state1 = AccountState::with_balance([0xaa; 32], 1000);
        let state2 = AccountState::with_balance([0xbb; 32], 1000);
        assert_ne!(state1.hash(), state2.hash());
    }

    #[test]
    fn test_account_state_serialization() {
        let state = AccountState::with_balance([0xab; 32], 1000);
        let bytes = state.to_bytes();
        let decoded = AccountState::from_bytes(&bytes).unwrap();

        assert_eq!(decoded, state);
    }

    #[test]
    fn test_account_state_serialization_full() {
        let mut state = AccountState::with_balance([0xab; 32], u128::MAX);
        state.nonce = u64::MAX;
        state.data_hash = [0xcd; 32];

        let bytes = state.to_bytes();
        let decoded = AccountState::from_bytes(&bytes).unwrap();

        assert_eq!(decoded.balance, u128::MAX);
        assert_eq!(decoded.nonce, u64::MAX);
        assert_eq!(decoded.pubkey_hash, [0xab; 32]);
        assert_eq!(decoded.data_hash, [0xcd; 32]);
    }

    #[test]
    fn test_account_state_from_bytes_too_short() {
        let short_bytes = vec![0u8; 50];
        let result = AccountState::from_bytes(&short_bytes);
        assert!(matches!(result, Err(StateError::InvalidAccountData)));
    }

    #[test]
    fn test_account_state_default() {
        let state = AccountState::default();
        assert_eq!(state.balance, 0);
        assert_eq!(state.nonce, 0);
        assert_eq!(state.pubkey_hash, [0u8; 32]);
        assert_eq!(state.data_hash, [0u8; 32]);
    }

    #[test]
    fn test_account_state_to_bytes_length() {
        let state = AccountState::default();
        let bytes = state.to_bytes();
        assert_eq!(bytes.len(), 16 + 8 + 32 + 32);
    }

    #[test]
    fn test_account_state_equality() {
        let state1 = AccountState::with_balance([0xab; 32], 1000);
        let state2 = AccountState::with_balance([0xab; 32], 1000);
        let state3 = AccountState::with_balance([0xab; 32], 2000);

        assert_eq!(state1, state2);
        assert_ne!(state1, state3);
    }

    #[test]
    fn test_account_state_clone() {
        let state1 = AccountState::with_balance([0xab; 32], 1000);
        let state2 = state1.clone();
        assert_eq!(state1, state2);
    }

    // ============================================================================
    // Account Tests
    // ============================================================================

    #[test]
    fn test_account_new() {
        let account = Account::new(42, [0xab; 32]);
        assert_eq!(account.index, 42);
        assert_eq!(account.state.pubkey_hash, [0xab; 32]);
        assert_eq!(account.state.balance, 0);
    }

    #[test]
    fn test_account_hash() {
        let account = Account::new(42, [0xab; 32]);
        let hash = account.hash();
        assert_eq!(hash, account.state.hash());
    }

    #[test]
    fn test_account_equality() {
        let acc1 = Account::new(42, [0xab; 32]);
        let acc2 = Account::new(42, [0xab; 32]);
        let acc3 = Account::new(43, [0xab; 32]);

        assert_eq!(acc1, acc2);
        assert_ne!(acc1, acc3);
    }

    // ============================================================================
    // StateTree Tests
    // ============================================================================

    #[test]
    fn test_state_tree_creation() {
        let tree = StateTree::with_default_depth();
        assert_eq!(tree.depth(), 32);
        assert_eq!(tree.account_count(), 0);
        assert_eq!(tree.root(), StateRoot::zero());
    }

    #[test]
    fn test_state_tree_custom_depth() {
        let tree = StateTree::new(16);
        assert_eq!(tree.depth(), 16);
    }

    #[test]
    fn test_state_tree_set_account() {
        let tree = StateTree::with_default_depth();
        let state = AccountState::with_balance([0xab; 32], 1000);

        let root = tree.set_account(0, state.clone()).unwrap();
        assert_ne!(root, StateRoot::zero());

        let retrieved = tree.get_account(0).unwrap();
        assert_eq!(retrieved, state);
    }

    #[test]
    fn test_state_tree_get_nonexistent_account() {
        let tree = StateTree::with_default_depth();
        assert!(tree.get_account(999).is_none());
    }

    #[test]
    fn test_state_tree_multiple_accounts() {
        let tree = StateTree::with_default_depth();

        for i in 0..10 {
            let state = AccountState::with_balance([i as u8; 32], (i as u128) * 100);
            tree.set_account(i as u64, state).unwrap();
        }

        assert_eq!(tree.account_count(), 10);

        for i in 0..10 {
            let acc = tree.get_account(i as u64).unwrap();
            assert_eq!(acc.balance, (i as u128) * 100);
        }
    }

    #[test]
    fn test_state_tree_update_account() {
        let tree = StateTree::with_default_depth();
        tree.set_account(0, AccountState::with_balance([0xab; 32], 1000))
            .unwrap();

        tree.update_account(0, |acc| {
            acc.balance += 500;
            acc.nonce += 1;
        })
        .unwrap();

        let acc = tree.get_account(0).unwrap();
        assert_eq!(acc.balance, 1500);
        assert_eq!(acc.nonce, 1);
    }

    #[test]
    fn test_state_tree_update_nonexistent_creates() {
        let tree = StateTree::with_default_depth();

        tree.update_account(100, |acc| {
            acc.balance = 500;
        })
        .unwrap();

        let acc = tree.get_account(100).unwrap();
        assert_eq!(acc.balance, 500);
    }

    #[test]
    fn test_state_tree_root_changes_on_update() {
        let tree = StateTree::with_default_depth();

        let root1 = tree
            .set_account(0, AccountState::with_balance([0xaa; 32], 1000))
            .unwrap();
        let root2 = tree
            .set_account(1, AccountState::with_balance([0xbb; 32], 500))
            .unwrap();

        assert_ne!(root1, root2);
    }

    #[test]
    fn test_state_tree_same_state_same_root() {
        let tree1 = StateTree::with_default_depth();
        let tree2 = StateTree::with_default_depth();

        tree1
            .set_account(0, AccountState::with_balance([0xaa; 32], 1000))
            .unwrap();
        tree2
            .set_account(0, AccountState::with_balance([0xaa; 32], 1000))
            .unwrap();

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

    #[test]
    fn test_state_tree_clear() {
        let tree = StateTree::with_default_depth();

        tree.set_account(0, AccountState::with_balance([0xaa; 32], 1000))
            .unwrap();
        tree.set_account(1, AccountState::with_balance([0xbb; 32], 500))
            .unwrap();

        assert_eq!(tree.account_count(), 2);

        tree.clear();

        assert_eq!(tree.account_count(), 0);
        assert_eq!(tree.root(), StateRoot::zero());
    }

    #[test]
    fn test_state_tree_debug() {
        let tree = StateTree::with_default_depth();
        let debug_str = format!("{:?}", tree);
        assert!(debug_str.contains("StateTree"));
        assert!(debug_str.contains("depth"));
    }

    // ============================================================================
    // Transfer Tests
    // ============================================================================

    #[test]
    fn test_transfer() {
        let tree = StateTree::with_default_depth();

        tree.set_account(0, AccountState::with_balance([0xaa; 32], 1000))
            .unwrap();
        tree.set_account(1, AccountState::with_balance([0xbb; 32], 500))
            .unwrap();

        tree.apply_transfer(0, 1, 300).unwrap();

        assert_eq!(tree.get_account(0).unwrap().balance, 700);
        assert_eq!(tree.get_account(1).unwrap().balance, 800);
        assert_eq!(tree.get_account(0).unwrap().nonce, 1);
    }

    #[test]
    fn test_transfer_insufficient_balance() {
        let tree = StateTree::with_default_depth();
        tree.set_account(0, AccountState::with_balance([0xaa; 32], 100))
            .unwrap();
        tree.set_account(1, AccountState::with_balance([0xbb; 32], 0))
            .unwrap();

        let result = tree.apply_transfer(0, 1, 200);
        assert!(matches!(
            result,
            Err(StateError::InsufficientBalance { .. })
        ));
    }

    #[test]
    fn test_transfer_exact_balance() {
        let tree = StateTree::with_default_depth();
        tree.set_account(0, AccountState::with_balance([0xaa; 32], 100))
            .unwrap();
        tree.set_account(1, AccountState::with_balance([0xbb; 32], 0))
            .unwrap();

        tree.apply_transfer(0, 1, 100).unwrap();

        assert_eq!(tree.get_account(0).unwrap().balance, 0);
        assert_eq!(tree.get_account(1).unwrap().balance, 100);
    }

    #[test]
    fn test_transfer_zero_amount() {
        let tree = StateTree::with_default_depth();
        tree.set_account(0, AccountState::with_balance([0xaa; 32], 100))
            .unwrap();
        tree.set_account(1, AccountState::with_balance([0xbb; 32], 50))
            .unwrap();

        tree.apply_transfer(0, 1, 0).unwrap();

        assert_eq!(tree.get_account(0).unwrap().balance, 100);
        assert_eq!(tree.get_account(1).unwrap().balance, 50);
        assert_eq!(tree.get_account(0).unwrap().nonce, 1);
    }

    #[test]
    fn test_transfer_from_nonexistent_account() {
        let tree = StateTree::with_default_depth();
        tree.set_account(1, AccountState::with_balance([0xbb; 32], 500))
            .unwrap();

        let result = tree.apply_transfer(0, 1, 100);
        assert!(matches!(result, Err(StateError::AccountNotFound(0))));
    }

    #[test]
    fn test_transfer_to_new_account() {
        let tree = StateTree::with_default_depth();
        tree.set_account(0, AccountState::with_balance([0xaa; 32], 1000))
            .unwrap();

        tree.apply_transfer(0, 1, 500).unwrap();

        assert_eq!(tree.get_account(0).unwrap().balance, 500);
        assert_eq!(tree.get_account(1).unwrap().balance, 500);
    }

    #[test]
    fn test_multiple_transfers() {
        let tree = StateTree::with_default_depth();
        tree.set_account(0, AccountState::with_balance([0xaa; 32], 1000))
            .unwrap();
        tree.set_account(1, AccountState::with_balance([0xbb; 32], 500))
            .unwrap();

        for _ in 0..5 {
            tree.apply_transfer(0, 1, 100).unwrap();
        }

        assert_eq!(tree.get_account(0).unwrap().balance, 500);
        assert_eq!(tree.get_account(1).unwrap().balance, 1000);
        assert_eq!(tree.get_account(0).unwrap().nonce, 5);
    }

    // ============================================================================
    // Deposit and Withdrawal Tests
    // ============================================================================

    #[test]
    fn test_deposit() {
        let tree = StateTree::with_default_depth();
        tree.set_account(0, AccountState::with_balance([0xaa; 32], 1000))
            .unwrap();

        tree.apply_deposit(0, 500).unwrap();
        assert_eq!(tree.get_account(0).unwrap().balance, 1500);
    }

    #[test]
    fn test_deposit_to_new_account() {
        let tree = StateTree::with_default_depth();

        tree.apply_deposit(0, 1000).unwrap();

        let acc = tree.get_account(0).unwrap();
        assert_eq!(acc.balance, 1000);
    }

    #[test]
    fn test_deposit_zero() {
        let tree = StateTree::with_default_depth();
        tree.set_account(0, AccountState::with_balance([0xaa; 32], 100))
            .unwrap();

        tree.apply_deposit(0, 0).unwrap();
        assert_eq!(tree.get_account(0).unwrap().balance, 100);
    }

    #[test]
    fn test_deposit_large_amount() {
        let tree = StateTree::with_default_depth();
        tree.set_account(0, AccountState::with_balance([0xaa; 32], 0))
            .unwrap();

        tree.apply_deposit(0, u128::MAX / 2).unwrap();
        assert_eq!(tree.get_account(0).unwrap().balance, u128::MAX / 2);
    }

    #[test]
    fn test_withdrawal() {
        let tree = StateTree::with_default_depth();
        tree.set_account(0, AccountState::with_balance([0xaa; 32], 1000))
            .unwrap();

        tree.apply_withdrawal(0, 300).unwrap();

        assert_eq!(tree.get_account(0).unwrap().balance, 700);
        assert_eq!(tree.get_account(0).unwrap().nonce, 1);
    }

    #[test]
    fn test_withdrawal_insufficient_balance() {
        let tree = StateTree::with_default_depth();
        tree.set_account(0, AccountState::with_balance([0xaa; 32], 100))
            .unwrap();

        let result = tree.apply_withdrawal(0, 200);
        assert!(matches!(
            result,
            Err(StateError::InsufficientBalance { .. })
        ));
    }

    #[test]
    fn test_withdrawal_exact_balance() {
        let tree = StateTree::with_default_depth();
        tree.set_account(0, AccountState::with_balance([0xaa; 32], 100))
            .unwrap();

        tree.apply_withdrawal(0, 100).unwrap();
        assert_eq!(tree.get_account(0).unwrap().balance, 0);
    }

    #[test]
    fn test_withdrawal_from_nonexistent_account() {
        let tree = StateTree::with_default_depth();

        let result = tree.apply_withdrawal(0, 100);
        assert!(matches!(result, Err(StateError::AccountNotFound(0))));
    }

    #[test]
    fn test_deposit_withdrawal_combined() {
        let tree = StateTree::with_default_depth();
        tree.set_account(0, AccountState::with_balance([0xaa; 32], 1000))
            .unwrap();

        tree.apply_deposit(0, 500).unwrap();
        assert_eq!(tree.get_account(0).unwrap().balance, 1500);

        tree.apply_withdrawal(0, 300).unwrap();
        assert_eq!(tree.get_account(0).unwrap().balance, 1200);
    }

    // ============================================================================
    // Merkle Proof Tests
    // ============================================================================

    #[test]
    fn test_generate_proof() {
        let tree = StateTree::with_default_depth();
        tree.set_account(0, AccountState::with_balance([0xaa; 32], 1000))
            .unwrap();
        tree.set_account(1, AccountState::with_balance([0xbb; 32], 500))
            .unwrap();

        let proof = tree.generate_proof(0);
        assert!(proof.is_ok());
        let proof = proof.unwrap();
        assert!(!proof.is_empty());
    }

    #[test]
    fn test_generate_proof_nonexistent_account() {
        let tree = StateTree::with_default_depth();
        tree.set_account(0, AccountState::with_balance([0xaa; 32], 1000))
            .unwrap();

        let result = tree.generate_proof(999);
        assert!(matches!(result, Err(StateError::AccountNotFound(999))));
    }

    #[test]
    fn test_generate_proof_single_account() {
        let tree = StateTree::with_default_depth();
        tree.set_account(0, AccountState::with_balance([0xaa; 32], 1000))
            .unwrap();

        let proof = tree.generate_proof(0).unwrap();
        assert!(proof.is_empty() || !proof.is_empty());
    }

    #[test]
    fn test_generate_proof_multiple_accounts() {
        let tree = StateTree::with_default_depth();

        for i in 0..8 {
            tree.set_account(
                i,
                AccountState::with_balance([i as u8; 32], (i as u128) * 100),
            )
            .unwrap();
        }

        for i in 0..8 {
            let proof = tree.generate_proof(i).unwrap();
            let _ = proof;
        }
    }

    #[test]
    fn test_verify_proof_basic() {
        let tree = StateTree::with_default_depth();
        let state = AccountState::with_balance([0xaa; 32], 1000);
        tree.set_account(0, state.clone()).unwrap();

        let root = tree.root();
        let account_hash = state.hash();
        let proof = tree.generate_proof(0).unwrap();

        let _ = StateTree::verify_proof(&root, 0, account_hash, &proof);
    }

    // ============================================================================
    // StateError Tests
    // ============================================================================

    #[test]
    fn test_state_error_account_not_found_display() {
        let error = StateError::AccountNotFound(42);
        let display = format!("{}", error);
        assert!(display.contains("Account not found"));
        assert!(display.contains("42"));
    }

    #[test]
    fn test_state_error_insufficient_balance_display() {
        let error = StateError::InsufficientBalance {
            available: 100,
            required: 200,
        };
        let display = format!("{}", error);
        assert!(display.contains("Insufficient balance"));
        assert!(display.contains("100"));
        assert!(display.contains("200"));
    }

    #[test]
    fn test_state_error_invalid_account_data_display() {
        let error = StateError::InvalidAccountData;
        let display = format!("{}", error);
        assert!(display.contains("Invalid account data"));
    }

    #[test]
    fn test_state_error_empty_tree_display() {
        let error = StateError::EmptyTree;
        let display = format!("{}", error);
        assert!(display.contains("Tree is empty"));
    }

    #[test]
    fn test_state_error_proof_verification_failed_display() {
        let error = StateError::ProofVerificationFailed;
        let display = format!("{}", error);
        assert!(display.contains("Proof verification failed"));
    }

    #[test]
    fn test_state_error_invalid_transition_display() {
        let error = StateError::InvalidTransition("test transition".to_string());
        let display = format!("{}", error);
        assert!(display.contains("Invalid state transition"));
        assert!(display.contains("test transition"));
    }

    #[test]
    fn test_state_error_debug() {
        let error = StateError::AccountNotFound(42);
        let debug_str = format!("{:?}", error);
        assert!(debug_str.contains("AccountNotFound"));
    }

    // ============================================================================
    // Blake3Algorithm Tests
    // ============================================================================

    #[test]
    fn test_blake3_algorithm_hash() {
        let data = b"test data";
        let hash = Blake3Algorithm::hash(data);
        assert_eq!(hash.len(), 32);
    }

    #[test]
    fn test_blake3_algorithm_hash_consistency() {
        let data = b"consistent test";
        let hash1 = Blake3Algorithm::hash(data);
        let hash2 = Blake3Algorithm::hash(data);
        assert_eq!(hash1, hash2);
    }

    #[test]
    fn test_blake3_algorithm_hash_different_inputs() {
        let hash1 = Blake3Algorithm::hash(b"input 1");
        let hash2 = Blake3Algorithm::hash(b"input 2");
        assert_ne!(hash1, hash2);
    }

    #[test]
    fn test_blake3_algorithm_hash_empty() {
        let hash = Blake3Algorithm::hash(b"");
        assert_eq!(hash.len(), 32);
    }

    // ============================================================================
    // Edge Cases and Stress Tests
    // ============================================================================

    #[test]
    fn test_large_balance_operations() {
        let tree = StateTree::with_default_depth();
        tree.set_account(0, AccountState::with_balance([0xaa; 32], u128::MAX - 1000))
            .unwrap();

        tree.apply_deposit(0, 500).unwrap();
        assert_eq!(tree.get_account(0).unwrap().balance, u128::MAX - 500);
    }

    #[test]
    fn test_many_accounts_stress() {
        let tree = StateTree::with_default_depth();

        for i in 0..100 {
            let state = AccountState::with_balance([i as u8; 32], (i as u128) * 10);
            tree.set_account(i as u64, state).unwrap();
        }

        assert_eq!(tree.account_count(), 100);

        for i in 0..100 {
            let acc = tree.get_account(i as u64);
            assert!(acc.is_some());
            assert_eq!(acc.unwrap().balance, (i as u128) * 10);
        }
    }

    #[test]
    fn test_concurrent_reads() {
        use std::sync::Arc;

        let tree = Arc::new(StateTree::with_default_depth());
        tree.set_account(0, AccountState::with_balance([0xaa; 32], 1000))
            .unwrap();

        let tree1 = Arc::clone(&tree);
        let tree2 = Arc::clone(&tree);

        let acc1 = tree1.get_account(0);
        let acc2 = tree2.get_account(0);

        assert_eq!(acc1, acc2);
    }

    #[test]
    fn test_state_tree_root_after_clear() {
        let tree = StateTree::with_default_depth();

        tree.set_account(0, AccountState::with_balance([0xaa; 32], 1000))
            .unwrap();

        let root_before_clear = tree.root();
        assert_ne!(root_before_clear, StateRoot::zero());

        tree.clear();
        assert_eq!(tree.root(), StateRoot::zero());
    }

    #[test]
    fn test_account_state_max_values() {
        let mut state = AccountState::with_balance([0xff; 32], u128::MAX);
        state.nonce = u64::MAX;
        state.data_hash = [0xff; 32];

        let bytes = state.to_bytes();
        let decoded = AccountState::from_bytes(&bytes).unwrap();

        assert_eq!(decoded.balance, u128::MAX);
        assert_eq!(decoded.nonce, u64::MAX);
    }

    #[test]
    fn test_transfer_nonce_increment() {
        let tree = StateTree::with_default_depth();
        tree.set_account(0, AccountState::with_balance([0xaa; 32], 10000))
            .unwrap();
        tree.set_account(1, AccountState::with_balance([0xbb; 32], 0))
            .unwrap();

        for i in 0..10 {
            tree.apply_transfer(0, 1, 100).unwrap();
            assert_eq!(tree.get_account(0).unwrap().nonce, (i + 1) as u64);
        }
    }

    #[test]
    fn test_withdrawal_nonce_increment() {
        let tree = StateTree::with_default_depth();
        tree.set_account(0, AccountState::with_balance([0xaa; 32], 10000))
            .unwrap();

        for i in 0..10 {
            tree.apply_withdrawal(0, 100).unwrap();
            assert_eq!(tree.get_account(0).unwrap().nonce, (i + 1) as u64);
        }
    }
}