synor/crates/synor-vm/src/storage.rs

//! Contract storage abstraction.
//!
//! Provides a key-value storage interface for contracts with support for:
//! - Merkle proofs
//! - Storage rent (optional)
//! - Efficient updates

use std::collections::HashMap;

use borsh::{BorshDeserialize, BorshSerialize};
use synor_types::Hash256;

use crate::ContractId;

/// Storage key (256-bit).
#[derive(
    Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, BorshSerialize, BorshDeserialize,
)]
pub struct StorageKey(pub [u8; 32]);

impl StorageKey {
    /// Creates a new storage key.
    pub fn new(bytes: [u8; 32]) -> Self {
        StorageKey(bytes)
    }

    /// Creates from a hash.
    pub fn from_hash(hash: &Hash256) -> Self {
        StorageKey(*hash.as_bytes())
    }

    /// Creates from a string key (hashed).
    pub fn from_string_key(key: &str) -> Self {
        let hash: [u8; 32] = blake3::hash(key.as_bytes()).into();
        StorageKey(hash)
    }

    /// Creates from bytes (padded or hashed).
    pub fn from_bytes(bytes: &[u8]) -> Self {
        if bytes.len() == 32 {
            let mut arr = [0u8; 32];
            arr.copy_from_slice(bytes);
            StorageKey(arr)
        } else if bytes.len() < 32 {
            let mut arr = [0u8; 32];
            arr[..bytes.len()].copy_from_slice(bytes);
            StorageKey(arr)
        } else {
            let hash: [u8; 32] = blake3::hash(bytes).into();
            StorageKey(hash)
        }
    }

    /// Creates a key for a map entry (key + index).
    pub fn map_key(base: &StorageKey, index: &[u8]) -> Self {
        let mut input = Vec::with_capacity(32 + index.len());
        input.extend_from_slice(&base.0);
        input.extend_from_slice(index);
        let hash: [u8; 32] = blake3::hash(&input).into();
        StorageKey(hash)
    }

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

/// Storage value (variable length).
#[derive(Clone, Debug, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
pub struct StorageValue(pub Vec<u8>);

impl StorageValue {
    /// Creates a new storage value.
    pub fn new(data: Vec<u8>) -> Self {
        StorageValue(data)
    }

    /// Creates from a u64.
    pub fn from_u64(value: u64) -> Self {
        StorageValue(value.to_le_bytes().to_vec())
    }

    /// Creates from a u128.
    pub fn from_u128(value: u128) -> Self {
        StorageValue(value.to_le_bytes().to_vec())
    }

    /// Parses as u64.
    pub fn as_u64(&self) -> Option<u64> {
        if self.0.len() == 8 {
            let mut arr = [0u8; 8];
            arr.copy_from_slice(&self.0);
            Some(u64::from_le_bytes(arr))
        } else {
            None
        }
    }

    /// Parses as u128.
    pub fn as_u128(&self) -> Option<u128> {
        if self.0.len() == 16 {
            let mut arr = [0u8; 16];
            arr.copy_from_slice(&self.0);
            Some(u128::from_le_bytes(arr))
        } else {
            None
        }
    }

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

    /// Returns the length.
    pub fn len(&self) -> usize {
        self.0.len()
    }

    /// Checks if empty.
    pub fn is_empty(&self) -> bool {
        self.0.is_empty()
    }
}

/// Storage access type for metering.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum StorageAccess {
    /// Read an existing value.
    Read,
    /// Write a new value.
    Write,
    /// Update an existing value.
    Update,
    /// Delete a value.
    Delete,
}

/// Contract storage interface.
pub trait ContractStorage: Send + Sync {
    /// Gets a value from storage.
    fn get(&self, contract: &ContractId, key: &StorageKey) -> Option<StorageValue>;

    /// Sets a value in storage.
    fn set(&mut self, contract: &ContractId, key: StorageKey, value: StorageValue);

    /// Deletes a value from storage.
    fn delete(&mut self, contract: &ContractId, key: &StorageKey) -> Option<StorageValue>;

    /// Checks if a key exists.
    fn contains(&self, contract: &ContractId, key: &StorageKey) -> bool;

    /// Gets the storage root hash for a contract.
    fn root(&self, contract: &ContractId) -> Hash256;

    /// Commits pending changes.
    fn commit(&mut self);

    /// Rolls back pending changes.
    fn rollback(&mut self);
}

/// In-memory storage implementation.
#[derive(Clone, Debug, Default)]
pub struct MemoryStorage {
    /// Storage data by contract.
    data: HashMap<ContractId, HashMap<StorageKey, StorageValue>>,
    /// Pending changes (for rollback).
    pending: HashMap<ContractId, HashMap<StorageKey, Option<StorageValue>>>,
}

impl MemoryStorage {
    /// Creates a new in-memory storage.
    pub fn new() -> Self {
        MemoryStorage {
            data: HashMap::new(),
            pending: HashMap::new(),
        }
    }

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

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

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

    /// Takes pending changes and returns them as StorageChange entries.
    ///
    /// This consumes the pending changes without committing them.
    /// Use this to collect storage changes for execution results.
    pub fn take_pending_changes(&mut self) -> Vec<crate::StorageChange> {
        let mut changes = Vec::new();

        for (contract, pending_changes) in self.pending.drain() {
            for (key, new_value) in pending_changes {
                // Get old value from committed data
                let old_value = self.data.get(&contract).and_then(|m| m.get(&key)).cloned();

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

        changes
    }

    /// Returns pending changes without consuming them.
    pub fn pending_changes(&self) -> Vec<crate::StorageChange> {
        let mut changes = Vec::new();

        for (contract, pending_changes) in &self.pending {
            for (key, new_value) in pending_changes {
                // Get old value from committed data
                let old_value = self.data.get(contract).and_then(|m| m.get(key)).cloned();

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

        changes
    }
}

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

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

    fn set(&mut self, contract: &ContractId, key: StorageKey, value: StorageValue) {
        self.pending
            .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
            .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 {
        // Simple merkle root of sorted key-value pairs
        let mut entries: Vec<_> = self
            .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) {
        for (contract, changes) in self.pending.drain() {
            let contract_data = self.data.entry(contract).or_default();
            for (key, value) in changes {
                match value {
                    Some(v) => {
                        contract_data.insert(key, v);
                    }
                    None => {
                        contract_data.remove(&key);
                    }
                }
            }
        }
    }

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

/// Storage with change tracking.
#[derive(Debug)]
pub struct TrackedStorage<S: ContractStorage> {
    /// Underlying storage.
    inner: S,
    /// Tracked reads.
    reads: Vec<(ContractId, StorageKey)>,
    /// Tracked writes.
    writes: Vec<(ContractId, StorageKey, Option<StorageValue>)>,
}

impl<S: ContractStorage> TrackedStorage<S> {
    /// Creates a new tracked storage wrapper.
    pub fn new(inner: S) -> Self {
        TrackedStorage {
            inner,
            reads: Vec::new(),
            writes: Vec::new(),
        }
    }

    /// Gets recorded reads.
    pub fn reads(&self) -> &[(ContractId, StorageKey)] {
        &self.reads
    }

    /// Gets recorded writes.
    pub fn writes(&self) -> &[(ContractId, StorageKey, Option<StorageValue>)] {
        &self.writes
    }

    /// Clears tracking data.
    pub fn clear_tracking(&mut self) {
        self.reads.clear();
        self.writes.clear();
    }

    /// Unwraps the inner storage.
    pub fn into_inner(self) -> S {
        self.inner
    }
}

impl<S: ContractStorage> ContractStorage for TrackedStorage<S> {
    fn get(&self, contract: &ContractId, key: &StorageKey) -> Option<StorageValue> {
        // Note: Can't track reads here due to &self, would need interior mutability
        self.inner.get(contract, key)
    }

    fn set(&mut self, contract: &ContractId, key: StorageKey, value: StorageValue) {
        self.writes.push((*contract, key, Some(value.clone())));
        self.inner.set(contract, key, value);
    }

    fn delete(&mut self, contract: &ContractId, key: &StorageKey) -> Option<StorageValue> {
        self.writes.push((*contract, *key, None));
        self.inner.delete(contract, key)
    }

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

    fn root(&self, contract: &ContractId) -> Hash256 {
        self.inner.root(contract)
    }

    fn commit(&mut self) {
        self.inner.commit();
    }

    fn rollback(&mut self) {
        self.inner.rollback();
    }
}

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

impl StorageSnapshot {
    /// Creates a snapshot from memory storage.
    pub fn from_memory(storage: &MemoryStorage) -> Self {
        StorageSnapshot {
            data: storage.data.clone(),
        }
    }

    /// Restores memory storage from snapshot.
    pub fn restore(&self, storage: &mut MemoryStorage) {
        storage.data = self.data.clone();
        storage.pending.clear();
    }
}

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

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

    #[test]
    fn test_storage_key() {
        let key1 = StorageKey::from_string_key("balance");
        let key2 = StorageKey::from_string_key("balance");
        assert_eq!(key1, key2);

        let key3 = StorageKey::from_string_key("other");
        assert_ne!(key1, key3);
    }

    #[test]
    fn test_storage_value() {
        let value = StorageValue::from_u64(12345);
        assert_eq!(value.as_u64(), Some(12345));

        let value = StorageValue::from_u128(u128::MAX);
        assert_eq!(value.as_u128(), Some(u128::MAX));
    }

    #[test]
    fn test_map_key() {
        let base = StorageKey::from_string_key("balances");
        let index = [1u8; 32];

        let key1 = StorageKey::map_key(&base, &index);
        let key2 = StorageKey::map_key(&base, &[2u8; 32]);

        assert_ne!(key1, key2);
    }

    #[test]
    fn test_memory_storage() {
        let mut storage = MemoryStorage::new();
        let contract = test_contract_id();
        let key = StorageKey::from_string_key("test");

        // Write
        storage.set(&contract, key, StorageValue::from_u64(42));

        // Read pending
        assert_eq!(storage.get(&contract, &key).unwrap().as_u64(), Some(42));

        // Commit
        storage.commit();

        // Read committed
        assert_eq!(storage.get(&contract, &key).unwrap().as_u64(), Some(42));

        // Delete
        storage.delete(&contract, &key);
        storage.commit();
        assert!(storage.get(&contract, &key).is_none());
    }

    #[test]
    fn test_storage_rollback() {
        let mut storage = MemoryStorage::new();
        let contract = test_contract_id();
        let key = StorageKey::from_string_key("test");

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

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

        // Should still be 42
        assert_eq!(storage.get(&contract, &key).unwrap().as_u64(), Some(42));
    }

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

        let root1 = storage.root(&contract);

        storage.set(
            &contract,
            StorageKey::from_string_key("a"),
            StorageValue::from_u64(1),
        );
        storage.commit();

        let root2 = storage.root(&contract);
        assert_ne!(root1, root2);

        storage.set(
            &contract,
            StorageKey::from_string_key("b"),
            StorageValue::from_u64(2),
        );
        storage.commit();

        let root3 = storage.root(&contract);
        assert_ne!(root2, root3);
    }
}