Skip to main content

dbx_core/storage/
memory_wos.rs

1//! In-memory WOS backend using BTreeMap
2//!
3//! Provides fast O(log n) operations with proven stability
4
5use crate::error::DbxResult;
6use crate::storage::StorageBackend;
7use std::collections::{BTreeMap, HashMap};
8use std::ops::RangeBounds;
9use std::sync::RwLock;
10
11/// In-memory WOS backend using BTreeMap
12pub struct InMemoryWosBackend {
13    #[allow(clippy::type_complexity)]
14    tables: RwLock<HashMap<String, BTreeMap<Vec<u8>, Vec<u8>>>>,
15}
16
17impl InMemoryWosBackend {
18    /// Create a new in-memory WOS backend
19    pub fn new() -> Self {
20        Self {
21            tables: RwLock::new(HashMap::new()),
22        }
23    }
24}
25
26impl Default for InMemoryWosBackend {
27    fn default() -> Self {
28        Self::new()
29    }
30}
31
32impl StorageBackend for InMemoryWosBackend {
33    fn insert(&self, table: &str, key: &[u8], value: &[u8]) -> DbxResult<()> {
34        let mut tables = self.tables.write().unwrap();
35        tables
36            .entry(table.to_string())
37            .or_default()
38            .insert(key.to_vec(), value.to_vec());
39        Ok(())
40    }
41
42    fn get(&self, table: &str, key: &[u8]) -> DbxResult<Option<Vec<u8>>> {
43        let tables = self.tables.read().unwrap();
44        Ok(tables.get(table).and_then(|map| map.get(key).cloned()))
45    }
46
47    fn delete(&self, table: &str, key: &[u8]) -> DbxResult<bool> {
48        let mut tables = self.tables.write().unwrap();
49        if let Some(map) = tables.get_mut(table) {
50            Ok(map.remove(key).is_some())
51        } else {
52            Ok(false)
53        }
54    }
55
56    fn scan<R: RangeBounds<Vec<u8>> + Clone>(
57        &self,
58        table: &str,
59        range: R,
60    ) -> DbxResult<Vec<(Vec<u8>, Vec<u8>)>> {
61        let tables = self.tables.read().unwrap();
62        if let Some(map) = tables.get(table) {
63            Ok(map
64                .range(range)
65                .map(|(k, v)| (k.clone(), v.clone()))
66                .collect())
67        } else {
68            Ok(Vec::new())
69        }
70    }
71
72    fn scan_one<R: RangeBounds<Vec<u8>> + Clone>(
73        &self,
74        table: &str,
75        range: R,
76    ) -> DbxResult<Option<(Vec<u8>, Vec<u8>)>> {
77        let tables = self.tables.read().unwrap();
78        if let Some(map) = tables.get(table) {
79            Ok(map.range(range).map(|(k, v)| (k.clone(), v.clone())).next())
80        } else {
81            Ok(None)
82        }
83    }
84
85    fn flush(&self) -> DbxResult<()> {
86        // No-op for in-memory backend
87        Ok(())
88    }
89
90    fn count(&self, table: &str) -> DbxResult<usize> {
91        let tables = self.tables.read().unwrap();
92        Ok(tables.get(table).map(|m| m.len()).unwrap_or(0))
93    }
94
95    fn table_names(&self) -> DbxResult<Vec<String>> {
96        let tables = self.tables.read().unwrap();
97        Ok(tables.keys().cloned().collect())
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104
105    #[test]
106    fn test_insert_and_get() {
107        let backend = InMemoryWosBackend::new();
108        backend.insert("test", b"key1", b"value1").unwrap();
109
110        let result = backend.get("test", b"key1").unwrap();
111        assert_eq!(result, Some(b"value1".to_vec()));
112    }
113
114    #[test]
115    fn test_delete() {
116        let backend = InMemoryWosBackend::new();
117        backend.insert("test", b"key1", b"value1").unwrap();
118
119        assert!(backend.delete("test", b"key1").unwrap());
120        assert_eq!(backend.get("test", b"key1").unwrap(), None);
121    }
122
123    #[test]
124    fn test_scan() {
125        let backend = InMemoryWosBackend::new();
126        backend.insert("test", b"key1", b"value1").unwrap();
127        backend.insert("test", b"key2", b"value2").unwrap();
128        backend.insert("test", b"key3", b"value3").unwrap();
129
130        let results = backend
131            .scan("test", b"key1".to_vec()..b"key3".to_vec())
132            .unwrap();
133        assert_eq!(results.len(), 2);
134    }
135
136    #[test]
137    fn test_count() {
138        let backend = InMemoryWosBackend::new();
139        backend.insert("test", b"key1", b"value1").unwrap();
140        backend.insert("test", b"key2", b"value2").unwrap();
141
142        assert_eq!(backend.count("test").unwrap(), 2);
143    }
144}