kiromi-ai-memory 0.2.2

// SPDX-License-Identifier: Apache-2.0 OR MIT
//! In-memory storage backend. Test-only — never exposed in release CLIs.

use std::collections::BTreeMap;

use async_trait::async_trait;
use bytes::Bytes;
use parking_lot::RwLock;

use crate::error::{Error, Result};
use crate::storage::{Storage, StorageEntry, StorageKey};

/// Process-local in-memory blob store. Cloneable; clones share the same map.
#[derive(Clone, Default, Debug)]
pub struct InMemoryBackend {
    inner: std::sync::Arc<RwLock<BTreeMap<String, Bytes>>>,
}

impl InMemoryBackend {
    /// Empty backend.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }
}

#[async_trait]
impl Storage for InMemoryBackend {
    fn id(&self) -> String {
        "in-memory".to_string()
    }

    async fn put(&self, key: &StorageKey, body: Bytes) -> Result<()> {
        self.inner.write().insert(key.as_str().to_string(), body);
        Ok(())
    }

    async fn get(&self, key: &StorageKey) -> Result<Bytes> {
        self.inner
            .read()
            .get(key.as_str())
            .cloned()
            .ok_or_else(|| Error::storage(format!("not found: {key}"), io_not_found(key.as_str())))
    }

    async fn exists(&self, key: &StorageKey) -> Result<bool> {
        Ok(self.inner.read().contains_key(key.as_str()))
    }

    async fn delete(&self, key: &StorageKey) -> Result<()> {
        self.inner.write().remove(key.as_str());
        Ok(())
    }

    async fn list_prefix(&self, prefix: &StorageKey) -> Result<Vec<StorageEntry>> {
        let p = prefix.as_str();
        let guard = self.inner.read();
        Ok(guard
            .iter()
            .filter(|(k, _)| k.starts_with(p))
            .map(|(k, v)| StorageEntry {
                key: StorageKey::new(k.clone()),
                size: u64::try_from(v.len()).unwrap_or(u64::MAX),
            })
            .collect())
    }
}

fn io_not_found(key: &str) -> std::io::Error {
    std::io::Error::new(std::io::ErrorKind::NotFound, key.to_string())
}

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

    #[tokio::test]
    async fn put_get_exists_delete() {
        let s = InMemoryBackend::new();
        let k = StorageKey::new("a/b.md");
        assert!(!s.exists(&k).await.unwrap());
        s.put(&k, Bytes::from_static(b"hi")).await.unwrap();
        assert!(s.exists(&k).await.unwrap());
        assert_eq!(s.get(&k).await.unwrap(), Bytes::from_static(b"hi"));
        s.delete(&k).await.unwrap();
        assert!(!s.exists(&k).await.unwrap());
    }

    #[tokio::test]
    async fn rename_within_default_impl_moves_blob() {
        let s = InMemoryBackend::new();
        s.put(&StorageKey::new("a"), Bytes::from_static(b"hi"))
            .await
            .unwrap();
        s.rename_within(&StorageKey::new("a"), &StorageKey::new("b/c"))
            .await
            .unwrap();
        assert!(!s.exists(&StorageKey::new("a")).await.unwrap());
        assert_eq!(
            s.get(&StorageKey::new("b/c")).await.unwrap(),
            Bytes::from_static(b"hi")
        );
    }

    #[tokio::test]
    async fn rename_within_missing_src_is_noop() {
        let s = InMemoryBackend::new();
        s.rename_within(&StorageKey::new("a"), &StorageKey::new("b"))
            .await
            .unwrap();
    }

    #[tokio::test]
    async fn rename_within_same_keys_is_noop() {
        let s = InMemoryBackend::new();
        s.put(&StorageKey::new("a"), Bytes::from_static(b"hi"))
            .await
            .unwrap();
        s.rename_within(&StorageKey::new("a"), &StorageKey::new("a"))
            .await
            .unwrap();
        assert_eq!(
            s.get(&StorageKey::new("a")).await.unwrap(),
            Bytes::from_static(b"hi")
        );
    }

    #[tokio::test]
    async fn list_prefix_filters() {
        let s = InMemoryBackend::new();
        s.put(&StorageKey::new("a/x"), Bytes::from_static(b"1"))
            .await
            .unwrap();
        s.put(&StorageKey::new("a/y"), Bytes::from_static(b"22"))
            .await
            .unwrap();
        s.put(&StorageKey::new("b/z"), Bytes::from_static(b"333"))
            .await
            .unwrap();
        let mut got: Vec<_> = s
            .list_prefix(&StorageKey::new("a/"))
            .await
            .unwrap()
            .into_iter()
            .map(|e| (e.key.as_str().to_string(), e.size))
            .collect();
        got.sort();
        assert_eq!(got, vec![("a/x".to_string(), 1), ("a/y".to_string(), 2),]);
    }
}