heddle-objects 0.25.1

// SPDX-License-Identifier: Apache-2.0
//! In-memory object store — reference implementation and test utility.
//!
//! Enable with the `memory-backend` Cargo feature, or use it automatically in
//! `#[cfg(test)]` contexts (no feature flag needed for tests).

use std::{collections::HashMap, sync::RwLock};

use crate::{
    object::{
        Action, ActionId, AnnotatedTag, Blob, BytesTreeSource, ContentHash, OpenedTreeBody, State,
        StateAttachment, StateAttachmentId, StateId, Tree, TreeEntryReader, TreeResumeCursor,
        TreeScheme, decode_tree_delta_header, is_delta_tree, is_redacted_tree, is_streamable_tree,
    },
    store::{HeddleError, ObjectCacheControl, ObjectStore, Result, SidecarStore, codec},
    sync::RwLockExt,
};

/// A non-persistent, in-memory implementation of [`ObjectStore`].
///
/// Useful for testing and as a reference implementation for custom backends.
/// All data is lost when the store is dropped.
///
/// # Example
///
/// ```ignore
/// use cli::store::InMemoryStore;
/// use cli::{ObjectStore, Blob};
///
/// let store = InMemoryStore::new();
/// let blob = Blob::from("hello world");
/// let hash = store.put_blob(&blob).unwrap();
/// let retrieved = store.get_blob(&hash).unwrap().unwrap();
/// assert_eq!(retrieved.content(), b"hello world");
/// ```
#[derive(Default)]
pub struct InMemoryStore {
    annotated_tags: RwLock<HashMap<ContentHash, Vec<u8>>>,
    blobs: RwLock<HashMap<ContentHash, Vec<u8>>>,
    trees: RwLock<HashMap<ContentHash, Vec<u8>>>,
    states: RwLock<HashMap<StateId, Vec<u8>>>,
    state_attachments: RwLock<HashMap<StateAttachmentId, Vec<u8>>>,
    actions: RwLock<HashMap<ActionId, Vec<u8>>>,
    redactions: RwLock<HashMap<ContentHash, Vec<u8>>>,
    state_visibility: RwLock<HashMap<StateId, Vec<u8>>>,
    /// Raw HRT1 partial projections keyed by the canonical tree hash they
    /// project — a slot DISTINCT from `trees` (the full-tree slot).
    partial_trees: RwLock<HashMap<ContentHash, Vec<u8>>>,
}

impl InMemoryStore {
    /// Create a new, empty in-memory store.
    pub fn new() -> Self {
        Self::default()
    }

    fn materialized_tree(&self, hash: &ContentHash) -> Result<Option<Tree>> {
        let Some(bytes) = self.trees.read_or_poisoned().get(hash).cloned() else {
            return Ok(None);
        };
        if is_delta_tree(&bytes) {
            return Err(HeddleError::InvalidObject(
                "HDC1 anchor must be materialized; delta chains are forbidden".to_string(),
            ));
        }
        Ok(Some(codec::decode_tree_serialized_with_key(
            &bytes, *hash, None,
        )?))
    }
}

impl ObjectCacheControl for InMemoryStore {
    fn clear_recent_caches(&self) {
        // The in-memory implementation is already its own source of truth and
        // has no decoded-object cache distinct from its stored bytes.
    }
}

impl ObjectStore for InMemoryStore {
    fn get_annotated_tag(&self, hash: &ContentHash) -> Result<Option<AnnotatedTag>> {
        self.annotated_tags
            .read_or_poisoned()
            .get(hash)
            .map(|bytes| {
                AnnotatedTag::decode_current_msgpack(bytes)
                    .map_err(|error| HeddleError::InvalidObject(error.to_string()))
            })
            .transpose()
    }

    fn put_annotated_tag(&self, tag: &AnnotatedTag) -> Result<ContentHash> {
        let hash = tag.hash();
        self.annotated_tags
            .write_or_poisoned()
            .insert(hash, tag.encode_current_msgpack());
        Ok(hash)
    }

    fn list_annotated_tags(&self) -> Result<Vec<ContentHash>> {
        Ok(self
            .annotated_tags
            .read_or_poisoned()
            .keys()
            .copied()
            .collect())
    }

    fn get_blob(&self, hash: &ContentHash) -> Result<Option<Blob>> {
        Ok(self
            .blobs
            .read_or_poisoned()
            .get(hash)
            .map(|v| Blob::new(v.clone())))
    }

    fn put_blob(&self, blob: &Blob) -> Result<ContentHash> {
        let hash = blob.hash();
        self.blobs
            .write_or_poisoned()
            .insert(hash, blob.content().to_vec());
        Ok(hash)
    }

    fn has_blob(&self, hash: &ContentHash) -> Result<bool> {
        Ok(self.blobs.read_or_poisoned().contains_key(hash))
    }

    fn blob_size(&self, hash: &ContentHash) -> Result<Option<u64>> {
        // InMemoryStore keeps raw uncompressed bytes — the length of
        // the stored buffer is the blob size, no header parsing needed.
        Ok(self
            .blobs
            .read_or_poisoned()
            .get(hash)
            .map(|v| v.len() as u64))
    }

    fn list_blobs(&self) -> Result<Vec<ContentHash>> {
        Ok(self.blobs.read_or_poisoned().keys().copied().collect())
    }

    fn get_tree(&self, hash: &ContentHash) -> Result<Option<Tree>> {
        let Some(bytes) = self.trees.read_or_poisoned().get(hash).cloned() else {
            return Ok(None);
        };
        let anchor = if is_delta_tree(&bytes) {
            let header = decode_tree_delta_header(&bytes)?;
            self.materialized_tree(&header.anchor)?
        } else {
            None
        };
        Ok(Some(codec::decode_tree_serialized_with_key(
            &bytes,
            *hash,
            anchor.as_ref(),
        )?))
    }

    fn open_tree(
        &self,
        tree_id: &ContentHash,
        cursor: Option<&TreeResumeCursor>,
    ) -> Result<Option<TreeEntryReader<OpenedTreeBody>>> {
        let Some(body) = self.trees.read_or_poisoned().get(tree_id).cloned() else {
            return Ok(None);
        };
        let body = if is_streamable_tree(&body) {
            body
        } else {
            self.get_tree(tree_id)?
                .ok_or_else(|| HeddleError::NotFound(format!("tree {tree_id}")))?
                .encode_lean()?
        };
        Ok(Some(TreeEntryReader::open(
            OpenedTreeBody::Bytes(BytesTreeSource::sequential_verify(body)),
            *tree_id,
            cursor,
        )?))
    }

    fn get_tree_serialized(&self, hash: &ContentHash) -> Result<Option<Vec<u8>>> {
        Ok(self.trees.read_or_poisoned().get(hash).cloned())
    }

    fn put_tree(&self, tree: &Tree) -> Result<ContentHash> {
        let hash = tree.hash();
        // A V4 salted tree is stored as its full HSR1 canonical body; V3 trees
        // use the cheap HLR1 lean anchor as before.
        let body = match tree.scheme() {
            TreeScheme::V4Salted => tree.encode_canonical()?,
            TreeScheme::V3Flat => tree.encode_lean()?,
        };
        self.trees.write_or_poisoned().insert(hash, body);
        // Full tree backfilled: drop any lingering redacted projection so the
        // DERIVED partial marker clears (no auto-backfill). Idempotent.
        self.partial_trees.write_or_poisoned().remove(&hash);
        Ok(hash)
    }

    fn put_tree_serialized(&self, data: &[u8], hash: ContentHash) -> Result<ContentHash> {
        // Route an HRT1 redacted projection to the partial slot (monotone)
        // rather than through the full-tree decoder, which refuses it.
        if is_redacted_tree(data) {
            self.put_partial_tree(&hash, data)?;
            return Ok(hash);
        }
        let anchor = if is_delta_tree(data) {
            let header = decode_tree_delta_header(data)?;
            self.materialized_tree(&header.anchor)?
        } else {
            None
        };
        let tree = codec::decode_tree_serialized_with_key(data, hash, anchor.as_ref())?;
        self.trees.write_or_poisoned().insert(hash, data.to_vec());
        // Full tree backfilled: drop any lingering redacted projection.
        self.partial_trees.write_or_poisoned().remove(&hash);
        Ok(tree.hash())
    }

    fn has_tree(&self, hash: &ContentHash) -> Result<bool> {
        Ok(self.trees.read_or_poisoned().contains_key(hash))
    }

    fn list_trees(&self) -> Result<Vec<ContentHash>> {
        Ok(self.trees.read_or_poisoned().keys().copied().collect())
    }

    fn has_partial_tree(&self, hash: &ContentHash) -> Result<bool> {
        Ok(self.partial_trees.read_or_poisoned().contains_key(hash))
    }

    fn get_partial_tree_bytes(&self, hash: &ContentHash) -> Result<Option<Vec<u8>>> {
        Ok(self.partial_trees.read_or_poisoned().get(hash).cloned())
    }

    fn put_partial_tree_bytes(&self, hash: &ContentHash, bytes: &[u8]) -> Result<()> {
        self.partial_trees
            .write_or_poisoned()
            .insert(*hash, bytes.to_vec());
        Ok(())
    }

    fn list_partial_trees(&self) -> Result<Vec<ContentHash>> {
        Ok(self
            .partial_trees
            .read_or_poisoned()
            .keys()
            .copied()
            .collect())
    }

    fn remove_partial_tree(&self, hash: &ContentHash) -> Result<()> {
        self.partial_trees.write_or_poisoned().remove(hash);
        Ok(())
    }

    fn get_state(&self, id: &StateId) -> Result<Option<State>> {
        match self.states.read_or_poisoned().get(id) {
            Some(bytes) => {
                let mut state = State::decode_current_msgpack(bytes)?;
                if !state.accepts_stored_id(id) {
                    return Err(crate::error::HeddleError::InvalidObject(format!(
                        "state id mismatch: requested {id}, computed {}",
                        state.id()
                    )));
                }
                state.state_id = *id;
                Ok(Some(state))
            }
            None => Ok(None),
        }
    }

    fn put_state(&self, state: &State) -> Result<()> {
        self.states
            .write_or_poisoned()
            .insert(state.id(), rmp_serde::to_vec(state)?);
        Ok(())
    }

    fn has_state(&self, id: &StateId) -> Result<bool> {
        Ok(self.states.read_or_poisoned().contains_key(id))
    }

    fn list_states(&self) -> Result<Vec<StateId>> {
        Ok(self.states.read_or_poisoned().keys().copied().collect())
    }

    fn get_state_attachment(
        &self,
        state: &StateId,
        id: &StateAttachmentId,
    ) -> Result<Option<StateAttachment>> {
        let attachment = self
            .state_attachments
            .read_or_poisoned()
            .get(id)
            .map(|bytes| StateAttachment::decode_current_msgpack(bytes))
            .transpose()?;
        Ok(attachment.filter(|attachment: &StateAttachment| attachment.state_id == *state))
    }

    fn put_state_attachment(&self, attachment: &StateAttachment) -> Result<StateAttachmentId> {
        let id = attachment.id();
        self.state_attachments
            .write_or_poisoned()
            .insert(id, attachment.encode_current_msgpack()?);
        Ok(id)
    }

    fn list_state_attachments(&self, state: &StateId) -> Result<Vec<StateAttachment>> {
        let mut attachments = Vec::new();
        for bytes in self.state_attachments.read_or_poisoned().values() {
            let attachment = StateAttachment::decode_current_msgpack(bytes)?;
            if attachment.state_id == *state {
                attachments.push(attachment);
            }
        }
        Ok(attachments)
    }

    fn get_action(&self, id: &ActionId) -> Result<Option<Action>> {
        match self.actions.read_or_poisoned().get(id) {
            Some(bytes) => {
                let action: Action = rmp_serde::from_slice(bytes)?;
                let found_id = action.compute_id();
                if found_id != *id {
                    return Err(HeddleError::InvalidObject(format!(
                        "action id mismatch: requested {}, found {}",
                        id, found_id
                    )));
                }
                Ok(Some(action))
            }
            None => Ok(None),
        }
    }

    fn put_action(&self, action: &mut Action) -> Result<ActionId> {
        let id = action.id();
        self.actions
            .write_or_poisoned()
            .insert(id, rmp_serde::to_vec(action)?);
        Ok(id)
    }

    fn list_actions(&self) -> Result<Vec<ActionId>> {
        Ok(self.actions.read_or_poisoned().keys().copied().collect())
    }
}

impl SidecarStore for InMemoryStore {
    fn has_redactions_for_blob(&self, blob: &ContentHash) -> Result<bool> {
        Ok(self.redactions.read_or_poisoned().contains_key(blob))
    }

    fn get_redactions_bytes_for_blob(&self, blob: &ContentHash) -> Result<Option<Vec<u8>>> {
        Ok(self.redactions.read_or_poisoned().get(blob).cloned())
    }

    fn put_redactions_bytes_for_blob(&self, blob: &ContentHash, bytes: &[u8]) -> Result<()> {
        self.redactions
            .write_or_poisoned()
            .insert(*blob, bytes.to_vec());
        Ok(())
    }

    fn list_blobs_with_redactions(&self) -> Result<Vec<ContentHash>> {
        Ok(self.redactions.read_or_poisoned().keys().copied().collect())
    }

    fn has_state_visibility_for_state(&self, state: &StateId) -> Result<bool> {
        Ok(self.state_visibility.read_or_poisoned().contains_key(state))
    }

    fn get_state_visibility_bytes_for_state(&self, state: &StateId) -> Result<Option<Vec<u8>>> {
        Ok(self.state_visibility.read_or_poisoned().get(state).cloned())
    }

    fn put_state_visibility_bytes_for_state(&self, state: &StateId, bytes: &[u8]) -> Result<()> {
        self.state_visibility
            .write_or_poisoned()
            .insert(*state, bytes.to_vec());
        Ok(())
    }

    fn list_states_with_visibility(&self) -> Result<Vec<StateId>> {
        Ok(self
            .state_visibility
            .read_or_poisoned()
            .keys()
            .copied()
            .collect())
    }
}

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

    /// Verify InMemoryStore satisfies the full ObjectStore compliance contract.
    #[test]
    fn test_compliance() {
        let store = InMemoryStore::new();
        crate::store::store_compliance::run_compliance_tests(&store);
    }

    /// Verify that a second put of the same blob is idempotent.
    #[test]
    fn test_blob_put_idempotent() {
        let store = InMemoryStore::new();
        let blob = Blob::from("idempotent");
        let h1 = store.put_blob(&blob).unwrap();
        let h2 = store.put_blob(&blob).unwrap();
        assert_eq!(h1, h2);
        assert_eq!(store.list_blobs().unwrap().len(), 1);
    }

    /// Verify has_blob returns false for a hash that was never stored.
    #[test]
    fn test_has_blob_unknown() {
        let store = InMemoryStore::new();
        let hash = ContentHash::compute(b"never-stored");
        assert!(!store.has_blob(&hash).unwrap());
    }
}