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

//! Shard state management.
//!
//! Each shard maintains its own Merkle state tree for accounts and storage.

use std::collections::HashMap;
use std::sync::Arc;

use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use synor_types::Hash256;

use crate::{ShardError, ShardId, ShardResult};

/// Individual shard state with Merkle tree.
#[derive(Clone, Debug)]
pub struct ShardState {
    /// Shard identifier.
    pub shard_id: ShardId,
    /// State root hash.
    pub state_root: Hash256,
    /// Account states (simplified - production would use Merkle Patricia Trie).
    accounts: HashMap<Hash256, AccountState>,
    /// Block height within this shard.
    pub block_height: u64,
    /// Last finalized block hash.
    pub last_finalized: Hash256,
}

/// Account state within a shard.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct AccountState {
    /// Account address.
    pub address: Hash256,
    /// Balance in smallest unit.
    pub balance: u128,
    /// Account nonce (transaction count).
    pub nonce: u64,
    /// Storage root for contract accounts.
    pub storage_root: Hash256,
    /// Code hash for contract accounts.
    pub code_hash: Hash256,
}

impl Default for AccountState {
    fn default() -> Self {
        Self {
            address: Hash256::from_bytes([0u8; 32]),
            balance: 0,
            nonce: 0,
            storage_root: Hash256::from_bytes([0u8; 32]),
            code_hash: Hash256::from_bytes([0u8; 32]),
        }
    }
}

impl ShardState {
    /// Creates a new empty shard state.
    pub fn new(shard_id: ShardId) -> Self {
        Self {
            shard_id,
            state_root: Hash256::from_bytes([0u8; 32]),
            accounts: HashMap::new(),
            block_height: 0,
            last_finalized: Hash256::from_bytes([0u8; 32]),
        }
    }

    /// Gets an account state.
    pub fn get_account(&self, address: &Hash256) -> Option<&AccountState> {
        self.accounts.get(address)
    }

    /// Updates an account state.
    pub fn update_account(&mut self, address: Hash256, state: AccountState) {
        self.accounts.insert(address, state);
        self.update_state_root();
    }

    /// Gets the account balance.
    pub fn get_balance(&self, address: &Hash256) -> u128 {
        self.accounts.get(address).map(|a| a.balance).unwrap_or(0)
    }

    /// Transfers balance between accounts (within same shard).
    pub fn transfer(
        &mut self,
        from: &Hash256,
        to: &Hash256,
        amount: u128,
    ) -> ShardResult<()> {
        let from_balance = self.get_balance(from);
        if from_balance < amount {
            return Err(ShardError::Internal("Insufficient balance".into()));
        }

        // Update sender
        let mut from_state = self.accounts.get(from).cloned().unwrap_or_default();
        from_state.address = *from;
        from_state.balance = from_balance - amount;
        from_state.nonce += 1;
        self.accounts.insert(*from, from_state);

        // Update receiver
        let mut to_state = self.accounts.get(to).cloned().unwrap_or_default();
        to_state.address = *to;
        to_state.balance += amount;
        self.accounts.insert(*to, to_state);

        self.update_state_root();
        Ok(())
    }

    /// Updates the state root after modifications.
    fn update_state_root(&mut self) {
        // Simplified: hash all account states
        // Production would use proper Merkle Patricia Trie
        let mut hasher = blake3::Hasher::new();
        hasher.update(&self.shard_id.to_le_bytes());
        hasher.update(&self.block_height.to_le_bytes());

        let mut sorted_accounts: Vec<_> = self.accounts.iter().collect();
        sorted_accounts.sort_by_key(|(k, _)| *k);

        for (addr, state) in sorted_accounts {
            hasher.update(addr.as_bytes());
            hasher.update(&state.balance.to_le_bytes());
            hasher.update(&state.nonce.to_le_bytes());
        }

        let hash = hasher.finalize();
        self.state_root = Hash256::from_bytes(*hash.as_bytes());
    }

    /// Advances to the next block.
    pub fn advance_block(&mut self, block_hash: Hash256) {
        self.block_height += 1;
        self.last_finalized = block_hash;
        self.update_state_root();
    }

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

    /// Generates a state proof for an account.
    pub fn generate_proof(&self, address: &Hash256) -> StateProof {
        StateProof {
            shard_id: self.shard_id,
            state_root: self.state_root,
            address: *address,
            account: self.accounts.get(address).cloned(),
            // Simplified: production would include Merkle path
            merkle_path: vec![],
        }
    }
}

/// Merkle state proof for cross-shard verification.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct StateProof {
    /// Shard the proof is from.
    pub shard_id: ShardId,
    /// State root at time of proof.
    pub state_root: Hash256,
    /// Account address being proved.
    pub address: Hash256,
    /// Account state (None if doesn't exist).
    pub account: Option<AccountState>,
    /// Merkle path from account to root.
    pub merkle_path: Vec<Hash256>,
}

impl StateProof {
    /// Verifies the proof against a known state root.
    pub fn verify(&self, expected_root: &Hash256) -> bool {
        // Simplified verification
        // Production would verify full Merkle path
        &self.state_root == expected_root
    }
}

/// Manages state across all shards.
pub struct ShardStateManager {
    /// Per-shard states.
    shards: Arc<RwLock<HashMap<ShardId, ShardState>>>,
    /// Number of shards.
    num_shards: u16,
}

impl ShardStateManager {
    /// Creates a new state manager with initialized shards.
    pub fn new(num_shards: u16) -> Self {
        let mut shards = HashMap::new();
        for i in 0..num_shards {
            shards.insert(i, ShardState::new(i));
        }

        Self {
            shards: Arc::new(RwLock::new(shards)),
            num_shards,
        }
    }

    /// Gets the state root for a shard.
    pub fn get_state_root(&self, shard_id: ShardId) -> Option<Hash256> {
        self.shards.read().get(&shard_id).map(|s| s.state_root)
    }

    /// Gets a shard state (read-only).
    pub fn get_shard(&self, shard_id: ShardId) -> Option<ShardState> {
        self.shards.read().get(&shard_id).cloned()
    }

    /// Updates a shard state.
    pub fn update_shard(&self, shard_id: ShardId, state: ShardState) {
        self.shards.write().insert(shard_id, state);
    }

    /// Executes a function on a shard's state.
    pub fn with_shard_mut<F, R>(&self, shard_id: ShardId, f: F) -> Option<R>
    where
        F: FnOnce(&mut ShardState) -> R,
    {
        self.shards.write().get_mut(&shard_id).map(f)
    }

    /// Splits a shard into multiple new shards (for dynamic resharding).
    pub fn split_shard(&self, shard_id: ShardId, new_shard_ids: &[ShardId]) {
        let mut shards = self.shards.write();

        if let Some(old_shard) = shards.remove(&shard_id) {
            // Distribute accounts to new shards based on address
            let mut new_shards: HashMap<ShardId, ShardState> = new_shard_ids
                .iter()
                .map(|&id| (id, ShardState::new(id)))
                .collect();

            for (addr, account) in old_shard.accounts {
                // Determine which new shard this account belongs to
                let bytes = addr.as_bytes();
                let shard_num = u16::from_le_bytes([bytes[0], bytes[1]]);
                let new_shard_id = new_shard_ids[shard_num as usize % new_shard_ids.len()];

                if let Some(shard) = new_shards.get_mut(&new_shard_id) {
                    shard.update_account(addr, account);
                }
            }

            // Add new shards
            for (id, shard) in new_shards {
                shards.insert(id, shard);
            }
        }
    }

    /// Merges multiple shards into one (for dynamic resharding).
    pub fn merge_shards(&self, shard_ids: &[ShardId], into: ShardId) {
        let mut shards = self.shards.write();
        let mut merged = ShardState::new(into);

        // Collect all accounts from shards being merged
        for &shard_id in shard_ids {
            if let Some(shard) = shards.remove(&shard_id) {
                for (addr, account) in shard.accounts {
                    merged.update_account(addr, account);
                }
            }
        }

        shards.insert(into, merged);
    }

    /// Gets all shard state roots for beacon chain commitment.
    pub fn get_all_state_roots(&self) -> Vec<(ShardId, Hash256)> {
        self.shards
            .read()
            .iter()
            .map(|(&id, state)| (id, state.state_root))
            .collect()
    }

    /// Returns the total number of accounts across all shards.
    pub fn total_accounts(&self) -> usize {
        self.shards.read().values().map(|s| s.account_count()).sum()
    }
}

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

    #[test]
    fn test_shard_state_new() {
        let state = ShardState::new(5);
        assert_eq!(state.shard_id, 5);
        assert_eq!(state.block_height, 0);
        assert_eq!(state.account_count(), 0);
    }

    #[test]
    fn test_account_update() {
        let mut state = ShardState::new(0);
        let addr = Hash256::from_bytes([1u8; 32]);

        let account = AccountState {
            address: addr,
            balance: 1000,
            nonce: 0,
            ..Default::default()
        };

        state.update_account(addr, account);
        assert_eq!(state.get_balance(&addr), 1000);
        assert_eq!(state.account_count(), 1);
    }

    #[test]
    fn test_transfer() {
        let mut state = ShardState::new(0);
        let alice = Hash256::from_bytes([1u8; 32]);
        let bob = Hash256::from_bytes([2u8; 32]);

        // Give Alice some balance
        state.update_account(
            alice,
            AccountState {
                address: alice,
                balance: 1000,
                nonce: 0,
                ..Default::default()
            },
        );

        // Transfer to Bob
        state.transfer(&alice, &bob, 300).unwrap();

        assert_eq!(state.get_balance(&alice), 700);
        assert_eq!(state.get_balance(&bob), 300);
    }

    #[test]
    fn test_state_manager() {
        let manager = ShardStateManager::new(4);

        // Check all shards initialized
        for i in 0..4 {
            assert!(manager.get_state_root(i).is_some());
        }

        // Update a shard
        manager.with_shard_mut(0, |shard| {
            let addr = Hash256::from_bytes([1u8; 32]);
            shard.update_account(
                addr,
                AccountState {
                    address: addr,
                    balance: 500,
                    ..Default::default()
                },
            );
        });

        assert_eq!(manager.total_accounts(), 1);
    }

    #[test]
    fn test_state_proof() {
        let mut state = ShardState::new(0);
        let addr = Hash256::from_bytes([1u8; 32]);

        state.update_account(
            addr,
            AccountState {
                address: addr,
                balance: 1000,
                ..Default::default()
            },
        );

        let proof = state.generate_proof(&addr);
        assert!(proof.verify(&state.state_root));
        assert_eq!(proof.account.unwrap().balance, 1000);
    }
}