synor/crates/synor-contract-test/src/mock_storage.rs

//! Mock storage implementation for contract testing.
//!
//! Provides isolated, in-memory storage for testing contracts without
//! affecting any persistent state.

use std::collections::HashMap;

use parking_lot::RwLock;
use synor_types::Hash256;
use synor_vm::{
    storage::{ContractStorage, StorageKey, StorageValue},
    ContractId, StorageChange,
};

/// Mock storage for isolated contract testing.
///
/// Provides a thread-safe, in-memory storage implementation with
/// additional features useful for testing:
/// - Snapshots for rollback
/// - Storage inspection and iteration
/// - Change tracking
///
/// # Example
///
/// ```rust,ignore
/// let mut storage = MockStorage::new();
///
/// // Set some values
/// storage.set(&contract_id, key, value);
///
/// // Take a snapshot
/// let snapshot = storage.snapshot();
///
/// // Make more changes
/// storage.set(&contract_id, key2, value2);
///
/// // Restore to snapshot
/// storage.restore(snapshot);
/// ```
#[derive(Debug, Default)]
pub struct MockStorage {
    /// Storage data by contract.
    data: RwLock<HashMap<ContractId, HashMap<StorageKey, StorageValue>>>,
    /// Pending changes (for transaction-like behavior).
    pending: RwLock<HashMap<ContractId, HashMap<StorageKey, Option<StorageValue>>>>,
    /// Change history for inspection.
    history: RwLock<Vec<StorageChange>>,
    /// Whether to track history.
    track_history: bool,
}

impl MockStorage {
    /// Creates a new empty mock storage.
    pub fn new() -> Self {
        MockStorage {
            data: RwLock::new(HashMap::new()),
            pending: RwLock::new(HashMap::new()),
            history: RwLock::new(Vec::new()),
            track_history: true,
        }
    }

    /// Creates mock storage with history tracking disabled.
    pub fn without_history() -> Self {
        MockStorage {
            data: RwLock::new(HashMap::new()),
            pending: RwLock::new(HashMap::new()),
            history: RwLock::new(Vec::new()),
            track_history: false,
        }
    }

    /// Creates mock storage with pre-set values.
    pub fn with_data(data: HashMap<ContractId, HashMap<StorageKey, StorageValue>>) -> Self {
        MockStorage {
            data: RwLock::new(data),
            pending: RwLock::new(HashMap::new()),
            history: RwLock::new(Vec::new()),
            track_history: true,
        }
    }

    /// Gets all keys for a contract.
    pub fn keys(&self, contract: &ContractId) -> Vec<StorageKey> {
        self.data
            .read()
            .get(contract)
            .map(|m| m.keys().copied().collect())
            .unwrap_or_default()
    }

    /// Gets all key-value pairs for a contract.
    pub fn entries(&self, contract: &ContractId) -> Vec<(StorageKey, StorageValue)> {
        self.data
            .read()
            .get(contract)
            .map(|m| m.iter().map(|(k, v)| (*k, v.clone())).collect())
            .unwrap_or_default()
    }

    /// Gets the number of storage entries for a contract.
    pub fn len(&self, contract: &ContractId) -> usize {
        self.data.read().get(contract).map(|m| m.len()).unwrap_or(0)
    }

    /// Checks if a contract has any storage.
    pub fn is_empty(&self, contract: &ContractId) -> bool {
        self.len(contract) == 0
    }

    /// Checks if any contract has storage.
    pub fn is_completely_empty(&self) -> bool {
        self.data.read().is_empty()
    }

    /// Clears all storage for a contract.
    pub fn clear_contract(&self, contract: &ContractId) {
        self.data.write().remove(contract);
        self.pending.write().remove(contract);
    }

    /// Clears all storage.
    pub fn clear_all(&self) {
        self.data.write().clear();
        self.pending.write().clear();
        self.history.write().clear();
    }

    /// Takes a snapshot of the current storage state.
    pub fn snapshot(&self) -> StorageSnapshot {
        StorageSnapshot {
            data: self.data.read().clone(),
        }
    }

    /// Restores storage to a previous snapshot.
    pub fn restore(&self, snapshot: StorageSnapshot) {
        *self.data.write() = snapshot.data;
        self.pending.write().clear();
    }

    /// Gets the change history.
    pub fn get_history(&self) -> Vec<StorageChange> {
        self.history.read().clone()
    }

    /// Clears the change history.
    pub fn clear_history(&self) {
        self.history.write().clear();
    }

    /// Gets pending changes without committing them.
    pub fn pending_changes(&self) -> Vec<StorageChange> {
        let pending = self.pending.read();
        let data = self.data.read();
        let mut changes = Vec::new();

        for (contract, pending_changes) in pending.iter() {
            for (key, new_value) in pending_changes {
                let old_value = data.get(contract).and_then(|m| m.get(key)).cloned();

                changes.push(StorageChange {
                    contract: *contract,
                    key: *key,
                    old_value,
                    new_value: new_value.clone(),
                });
            }
        }

        changes
    }

    /// Takes pending changes without committing.
    pub fn take_pending_changes(&self) -> Vec<StorageChange> {
        let changes = self.pending_changes();
        self.pending.write().clear();
        changes
    }

    /// Sets a value directly (bypasses pending/commit flow).
    pub fn set_direct(&self, contract: &ContractId, key: StorageKey, value: StorageValue) {
        self.data
            .write()
            .entry(*contract)
            .or_default()
            .insert(key, value);
    }

    /// Deletes a value directly.
    pub fn delete_direct(&self, contract: &ContractId, key: &StorageKey) -> Option<StorageValue> {
        self.data
            .write()
            .get_mut(contract)
            .and_then(|m| m.remove(key))
    }

    /// Gets a value directly (ignores pending changes).
    pub fn get_direct(&self, contract: &ContractId, key: &StorageKey) -> Option<StorageValue> {
        self.data
            .read()
            .get(contract)
            .and_then(|m| m.get(key))
            .cloned()
    }

    /// Iterates over all contracts.
    pub fn contracts(&self) -> Vec<ContractId> {
        self.data.read().keys().copied().collect()
    }

    /// Gets statistics about storage usage.
    pub fn stats(&self) -> StorageStats {
        let data = self.data.read();
        let mut total_entries = 0;
        let mut total_bytes = 0;

        for contract_storage in data.values() {
            total_entries += contract_storage.len();
            for value in contract_storage.values() {
                total_bytes += value.len();
            }
        }

        StorageStats {
            contract_count: data.len(),
            total_entries,
            total_bytes,
            history_len: self.history.read().len(),
        }
    }
}

impl ContractStorage for MockStorage {
    fn get(&self, contract: &ContractId, key: &StorageKey) -> Option<StorageValue> {
        // Check pending first
        let pending = self.pending.read();
        if let Some(contract_pending) = pending.get(contract) {
            if let Some(pending_value) = contract_pending.get(key) {
                return pending_value.clone();
            }
        }
        drop(pending);

        // Then check committed data
        self.data
            .read()
            .get(contract)
            .and_then(|m| m.get(key))
            .cloned()
    }

    fn set(&mut self, contract: &ContractId, key: StorageKey, value: StorageValue) {
        self.pending
            .write()
            .entry(*contract)
            .or_default()
            .insert(key, Some(value));
    }

    fn delete(&mut self, contract: &ContractId, key: &StorageKey) -> Option<StorageValue> {
        let old_value = self.get(contract, key);
        self.pending
            .write()
            .entry(*contract)
            .or_default()
            .insert(*key, None);
        old_value
    }

    fn contains(&self, contract: &ContractId, key: &StorageKey) -> bool {
        self.get(contract, key).is_some()
    }

    fn root(&self, contract: &ContractId) -> Hash256 {
        let data = self.data.read();
        let mut entries: Vec<_> = data
            .get(contract)
            .map(|m| m.iter().collect())
            .unwrap_or_default();

        entries.sort_by_key(|(k, _)| *k);

        if entries.is_empty() {
            return Hash256::default();
        }

        let mut hasher_input = Vec::new();
        for (key, value) in entries {
            hasher_input.extend_from_slice(&key.0);
            hasher_input.extend_from_slice(&value.0);
        }

        Hash256::from_bytes(blake3::hash(&hasher_input).into())
    }

    fn commit(&mut self) {
        let mut pending = self.pending.write();
        let mut data = self.data.write();
        let mut history = self.history.write();

        for (contract, changes) in pending.drain() {
            let contract_data = data.entry(contract).or_default();

            for (key, new_value) in changes {
                let old_value = contract_data.get(&key).cloned();

                // Track history if enabled
                if self.track_history {
                    history.push(StorageChange {
                        contract,
                        key,
                        old_value,
                        new_value: new_value.clone(),
                    });
                }

                // Apply change
                match new_value {
                    Some(v) => {
                        contract_data.insert(key, v);
                    }
                    None => {
                        contract_data.remove(&key);
                    }
                }
            }
        }
    }

    fn rollback(&mut self) {
        self.pending.write().clear();
    }
}

// Implement Clone for MockStorage (needed for some test patterns)
impl Clone for MockStorage {
    fn clone(&self) -> Self {
        MockStorage {
            data: RwLock::new(self.data.read().clone()),
            pending: RwLock::new(self.pending.read().clone()),
            history: RwLock::new(self.history.read().clone()),
            track_history: self.track_history,
        }
    }
}

/// Snapshot of storage state for rollback.
#[derive(Clone, Debug)]
pub struct StorageSnapshot {
    data: HashMap<ContractId, HashMap<StorageKey, StorageValue>>,
}

impl StorageSnapshot {
    /// Gets the number of contracts in the snapshot.
    pub fn contract_count(&self) -> usize {
        self.data.len()
    }

    /// Gets storage entries for a contract.
    pub fn get_contract(
        &self,
        contract: &ContractId,
    ) -> Option<&HashMap<StorageKey, StorageValue>> {
        self.data.get(contract)
    }
}

/// Statistics about storage usage.
#[derive(Clone, Debug, Default)]
pub struct StorageStats {
    /// Number of contracts with storage.
    pub contract_count: usize,
    /// Total number of storage entries.
    pub total_entries: usize,
    /// Total bytes of storage values.
    pub total_bytes: usize,
    /// Length of change history.
    pub history_len: usize,
}

impl std::fmt::Display for StorageStats {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "StorageStats {{ contracts: {}, entries: {}, bytes: {}, history: {} }}",
            self.contract_count, self.total_entries, self.total_bytes, self.history_len
        )
    }
}

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

    fn test_contract_id() -> ContractId {
        ContractId::from_bytes([0x42; 32])
    }

    fn test_key() -> StorageKey {
        StorageKey::from_bytes(b"test_key")
    }

    fn test_value() -> StorageValue {
        StorageValue::from_u64(12345)
    }

    #[test]
    fn test_mock_storage_basic() {
        let mut storage = MockStorage::new();
        let contract = test_contract_id();
        let key = test_key();
        let value = test_value();

        // Set value
        storage.set(&contract, key, value.clone());

        // Read pending
        assert_eq!(storage.get(&contract, &key), Some(value.clone()));

        // Commit
        storage.commit();

        // Read committed
        assert_eq!(storage.get(&contract, &key), Some(value));
    }

    #[test]
    fn test_mock_storage_rollback() {
        let mut storage = MockStorage::new();
        let contract = test_contract_id();
        let key = test_key();

        storage.set(&contract, key, StorageValue::from_u64(100));
        storage.commit();

        storage.set(&contract, key, StorageValue::from_u64(200));
        assert_eq!(
            storage.get(&contract, &key),
            Some(StorageValue::from_u64(200))
        );

        storage.rollback();
        assert_eq!(
            storage.get(&contract, &key),
            Some(StorageValue::from_u64(100))
        );
    }

    #[test]
    fn test_mock_storage_snapshot() {
        let mut storage = MockStorage::new();
        let contract = test_contract_id();
        let key = test_key();

        storage.set(&contract, key, StorageValue::from_u64(100));
        storage.commit();

        let snapshot = storage.snapshot();

        storage.set(&contract, key, StorageValue::from_u64(200));
        storage.commit();

        assert_eq!(
            storage.get(&contract, &key),
            Some(StorageValue::from_u64(200))
        );

        storage.restore(snapshot);

        assert_eq!(
            storage.get(&contract, &key),
            Some(StorageValue::from_u64(100))
        );
    }

    #[test]
    fn test_mock_storage_history() {
        let mut storage = MockStorage::new();
        let contract = test_contract_id();
        let key = test_key();

        storage.set(&contract, key, StorageValue::from_u64(100));
        storage.commit();

        storage.set(&contract, key, StorageValue::from_u64(200));
        storage.commit();

        let history = storage.get_history();
        assert_eq!(history.len(), 2);
    }

    #[test]
    fn test_mock_storage_stats() {
        let mut storage = MockStorage::new();
        let contract = test_contract_id();

        storage.set(&contract, test_key(), test_value());
        storage.commit();

        let stats = storage.stats();
        assert_eq!(stats.contract_count, 1);
        assert_eq!(stats.total_entries, 1);
        assert!(stats.total_bytes > 0);
    }

    #[test]
    fn test_mock_storage_delete() {
        let mut storage = MockStorage::new();
        let contract = test_contract_id();
        let key = test_key();

        storage.set(&contract, key, test_value());
        storage.commit();

        assert!(storage.contains(&contract, &key));

        storage.delete(&contract, &key);
        storage.commit();

        assert!(!storage.contains(&contract, &key));
    }

    #[test]
    fn test_mock_storage_direct_access() {
        let storage = MockStorage::new();
        let contract = test_contract_id();
        let key = test_key();
        let value = test_value();

        storage.set_direct(&contract, key, value.clone());
        assert_eq!(storage.get_direct(&contract, &key), Some(value));

        storage.delete_direct(&contract, &key);
        assert_eq!(storage.get_direct(&contract, &key), None);
    }
}