storage-engine-playground/crates/query-storage/src/memory.rs

//! In-memory backend, keyed by relation name. Always available.

use std::collections::HashMap;

use query_ops::value::Value;

use crate::{Storage, StorageError};

/// In-memory backend, useful as the default in tests and as a correctness
/// oracle for other backends.
#[derive(Debug, Default)]
pub struct MemoryStorage {
    relations: HashMap<String, MemoryRelation>,
}

#[derive(Debug)]
struct MemoryRelation {
    arity: usize,
    rows: Vec<Vec<Value>>,
}

impl MemoryStorage {
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }
}

impl Storage for MemoryStorage {
    fn create_relation(&mut self, name: &str, arity: usize) -> Result<(), StorageError> {
        if self.relations.contains_key(name) {
            return Err(StorageError::RelationExists(name.to_string()));
        }
        self.relations.insert(
            name.to_string(),
            MemoryRelation {
                arity,
                rows: Vec::new(),
            },
        );
        Ok(())
    }

    fn arity(&self, name: &str) -> Result<usize, StorageError> {
        self.relations
            .get(name)
            .map(|r| r.arity)
            .ok_or_else(|| StorageError::RelationNotFound(name.to_string()))
    }

    fn scan(&self, name: &str) -> Result<Vec<Vec<Value>>, StorageError> {
        self.relations
            .get(name)
            .map(|r| r.rows.clone())
            .ok_or_else(|| StorageError::RelationNotFound(name.to_string()))
    }

    fn insert(&mut self, name: &str, row: Vec<Value>) -> Result<(), StorageError> {
        let relation = self
            .relations
            .get_mut(name)
            .ok_or_else(|| StorageError::RelationNotFound(name.to_string()))?;
        if row.len() != relation.arity {
            return Err(StorageError::ArityMismatch {
                expected: relation.arity,
                got: row.len(),
            });
        }
        relation.rows.push(row);
        Ok(())
    }
}

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

    fn i(x: i64) -> Value {
        Value::Int(x)
    }

    #[test]
    fn create_insert_scan_roundtrip() -> Result<(), StorageError> {
        let mut storage = MemoryStorage::new();
        storage.create_relation("edge", 2)?;
        storage.insert("edge", vec![i(1), i(2)])?;
        storage.insert("edge", vec![i(2), i(3)])?;
        let rows = storage.scan("edge")?;
        assert_eq!(rows, vec![vec![i(1), i(2)], vec![i(2), i(3)]]);
        Ok(())
    }

    #[test]
    fn duplicate_create_returns_err() -> Result<(), StorageError> {
        let mut storage = MemoryStorage::new();
        storage.create_relation("edge", 2)?;
        assert!(matches!(
            storage.create_relation("edge", 2),
            Err(StorageError::RelationExists(_))
        ));
        Ok(())
    }

    #[test]
    fn scan_unknown_relation_returns_err() {
        let storage = MemoryStorage::new();
        assert!(matches!(
            storage.scan("missing"),
            Err(StorageError::RelationNotFound(_))
        ));
    }

    #[test]
    fn arity_unknown_relation_returns_err() {
        let storage = MemoryStorage::new();
        assert!(matches!(
            storage.arity("missing"),
            Err(StorageError::RelationNotFound(_))
        ));
    }

    #[test]
    fn insert_wrong_arity_returns_err() -> Result<(), StorageError> {
        let mut storage = MemoryStorage::new();
        storage.create_relation("edge", 2)?;
        assert!(matches!(
            storage.insert("edge", vec![i(1)]),
            Err(StorageError::ArityMismatch {
                expected: 2,
                got: 1
            })
        ));
        Ok(())
    }

    #[test]
    fn scan_as_table_materializes_table() -> Result<(), StorageError> {
        let mut storage = MemoryStorage::new();
        storage.create_relation("edge", 2)?;
        storage.insert("edge", vec![i(1), i(2)])?;
        let table = scan_as_table(&storage, "edge")?;
        assert_eq!(table.arity, 2);
        assert_eq!(table.rows, vec![vec![i(1), i(2)]]);
        Ok(())
    }
}