xenith-core 0.1.0

use crate::{Result, StateKey, StateValue, XenithError};
use async_trait::async_trait;
use std::collections::HashMap;
use std::sync::{Arc, Mutex};

/// Optional on-chain location metadata for a [`StateKey`].
///
/// Used by `SyncEngine` with [`crate::ReadStrategy::Quorum`] to know
/// which EVM contract address and storage slot to read when verifying on-chain
/// agreement across multiple chains.
#[derive(Clone, Debug, Default)]
pub struct KeyMetadata {
    /// EVM contract address whose storage slot holds this state.
    pub address: Option<[u8; 20]>,
    /// 32-byte storage slot index within `address`.
    pub slot: Option<[u8; 32]>,
}

/// Pluggable persistence layer for synced state.
///
/// Implementations must be `Send + Sync`; use `Arc<dyn StateStore>` for shared access.
///
/// # Example
///
/// ```rust,no_run
/// use xenith_core::{InMemoryStore, StateStore, StateKey};
/// use std::sync::Arc;
///
/// # async fn example() {
/// let store: Arc<dyn StateStore> = Arc::new(InMemoryStore::default());
/// let keys = store.list_prefix("uniswap").await.unwrap();
/// # }
/// ```
#[async_trait]
pub trait StateStore: Send + Sync {
    /// Retrieve the value for `key`, or `None` if it has not been set.
    async fn get(&self, key: &StateKey) -> Result<Option<StateValue>>;

    /// Insert or overwrite the value for `key`.
    async fn set(&self, key: &StateKey, value: StateValue) -> Result<()>;

    /// Remove the value for `key`. No-op if the key does not exist.
    async fn delete(&self, key: &StateKey) -> Result<()>;

    /// Return all keys whose string representation starts with `prefix`.
    async fn list_prefix(&self, prefix: &str) -> Result<Vec<StateKey>>;

    /// Retrieve the [`KeyMetadata`] for `key`, or `None` if not set.
    async fn get_metadata(&self, key: &StateKey) -> Result<Option<KeyMetadata>>;

    /// Store [`KeyMetadata`] for `key`.
    async fn set_metadata(&self, key: &StateKey, meta: KeyMetadata) -> Result<()>;
}

/// In-process, heap-backed implementation of [`StateStore`].
///
/// Suitable for testing and single-process use. Backed by a `Mutex<HashMap>`
/// wrapped in `Arc` so it is cheaply cloneable across tasks.
///
/// # Example
///
/// ```
/// use xenith_core::{InMemoryStore, StateStore, StateKey, StateValue, StateVersion, ChainId};
/// use bytes::Bytes;
///
/// # #[tokio::main(flavor = "current_thread")]
/// # async fn main() {
/// let store = InMemoryStore::default();
/// let key = StateKey::new("proto", "pool", "0x1");
/// let val = StateValue {
///     data: Bytes::new(),
///     version: StateVersion { timestamp_ms: 1, sequence: 0, source_chain: 1 },
///     updated_at: 0,
///     source_chain: ChainId::from(1),
/// };
/// store.set(&key, val.clone()).await.unwrap();
/// assert_eq!(store.get(&key).await.unwrap(), Some(val));
/// # }
/// ```
#[derive(Clone, Default)]
pub struct InMemoryStore {
    inner: Arc<Mutex<HashMap<String, StateValue>>>,
    meta: Arc<Mutex<HashMap<String, KeyMetadata>>>,
}

#[async_trait]
impl StateStore for InMemoryStore {
    async fn get(&self, key: &StateKey) -> Result<Option<StateValue>> {
        let map = self
            .inner
            .lock()
            .map_err(|_| XenithError::StoreError("lock poisoned".into()))?;
        Ok(map.get(key.as_ref()).cloned())
    }

    async fn set(&self, key: &StateKey, value: StateValue) -> Result<()> {
        let mut map = self
            .inner
            .lock()
            .map_err(|_| XenithError::StoreError("lock poisoned".into()))?;
        map.insert(key.as_ref().to_owned(), value);
        Ok(())
    }

    async fn delete(&self, key: &StateKey) -> Result<()> {
        let mut map = self
            .inner
            .lock()
            .map_err(|_| XenithError::StoreError("lock poisoned".into()))?;
        map.remove(key.as_ref());
        Ok(())
    }

    async fn list_prefix(&self, prefix: &str) -> Result<Vec<StateKey>> {
        let map = self
            .inner
            .lock()
            .map_err(|_| XenithError::StoreError("lock poisoned".into()))?;
        let mut keys: Vec<StateKey> = map
            .keys()
            .filter(|k| k.starts_with(prefix))
            .map(|k| StateKey::from_raw(k.clone()))
            .collect();
        keys.sort_by(|a, b| a.as_ref().cmp(b.as_ref()));
        Ok(keys)
    }

    async fn get_metadata(&self, key: &StateKey) -> Result<Option<KeyMetadata>> {
        let map = self
            .meta
            .lock()
            .map_err(|_| XenithError::StoreError("lock poisoned".into()))?;
        Ok(map.get(key.as_ref()).cloned())
    }

    async fn set_metadata(&self, key: &StateKey, meta: KeyMetadata) -> Result<()> {
        let mut map = self
            .meta
            .lock()
            .map_err(|_| XenithError::StoreError("lock poisoned".into()))?;
        map.insert(key.as_ref().to_owned(), meta);
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ChainId;
    use bytes::Bytes;

    fn val(ts: u64) -> StateValue {
        use crate::StateVersion;
        StateValue {
            data: Bytes::from_static(b"x"),
            version: StateVersion {
                timestamp_ms: ts,
                sequence: 0,
                source_chain: 1,
            },
            updated_at: 0,
            source_chain: ChainId(1),
        }
    }

    #[tokio::test]
    async fn set_then_get_returns_value() {
        let store = InMemoryStore::default();
        let key = StateKey::new("proto", "pool", "0x1");
        store.set(&key, val(1)).await.unwrap();
        assert_eq!(store.get(&key).await.unwrap(), Some(val(1)));
    }

    #[tokio::test]
    async fn get_missing_key_returns_none() {
        let store = InMemoryStore::default();
        let key = StateKey::new("proto", "pool", "missing");
        assert_eq!(store.get(&key).await.unwrap(), None);
    }

    #[tokio::test]
    async fn delete_removes_key() {
        let store = InMemoryStore::default();
        let key = StateKey::new("proto", "pool", "0x2");
        store.set(&key, val(1)).await.unwrap();
        store.delete(&key).await.unwrap();
        assert_eq!(store.get(&key).await.unwrap(), None);
    }

    #[tokio::test]
    async fn delete_missing_key_is_noop() {
        let store = InMemoryStore::default();
        let key = StateKey::new("proto", "pool", "ghost");
        store.delete(&key).await.unwrap(); // must not error
    }

    #[tokio::test]
    async fn list_prefix_returns_matching_keys() {
        let store = InMemoryStore::default();
        let a = StateKey::new("uniswap", "pool", "0xaaa");
        let b = StateKey::new("uniswap", "pool", "0xbbb");
        let other = StateKey::new("aave", "reserve", "0xaaa");
        store.set(&a, val(1)).await.unwrap();
        store.set(&b, val(2)).await.unwrap();
        store.set(&other, val(3)).await.unwrap();

        let keys = store.list_prefix("uniswap").await.unwrap();
        assert_eq!(keys.len(), 2);
        assert!(keys.contains(&a));
        assert!(keys.contains(&b));
        assert!(!keys.contains(&other));
    }

    #[tokio::test]
    async fn list_prefix_empty_when_no_match() {
        let store = InMemoryStore::default();
        store
            .set(&StateKey::new("aave", "x", "1"), val(1))
            .await
            .unwrap();
        assert!(store.list_prefix("uniswap").await.unwrap().is_empty());
    }
}