aztec_pxe/stores/
kv.rs

1//! Key-value store abstraction for PXE local state.
2
3use std::collections::BTreeMap;
4use std::sync::RwLock;
5
6use async_trait::async_trait;
7use aztec_core::error::Error;
8
9/// Simple key-value store trait (async for future flexibility).
10#[async_trait]
11pub trait KvStore: Send + Sync {
12    /// Get a value by key.
13    async fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Error>;
14    /// Put a key-value pair.
15    async fn put(&self, key: &[u8], value: &[u8]) -> Result<(), Error>;
16    /// Delete a key.
17    async fn delete(&self, key: &[u8]) -> Result<(), Error>;
18    /// List all key-value pairs with a given prefix.
19    async fn list_prefix(&self, prefix: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>, Error>;
20}
21
22/// In-memory KV store for testing and ephemeral use.
23pub struct InMemoryKvStore {
24    data: RwLock<BTreeMap<Vec<u8>, Vec<u8>>>,
25}
26
27impl InMemoryKvStore {
28    pub fn new() -> Self {
29        Self {
30            data: RwLock::new(BTreeMap::new()),
31        }
32    }
33}
34
35impl Default for InMemoryKvStore {
36    fn default() -> Self {
37        Self::new()
38    }
39}
40
41#[async_trait]
42impl KvStore for InMemoryKvStore {
43    async fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Error> {
44        let data = self
45            .data
46            .read()
47            .map_err(|e| Error::InvalidData(format!("lock poisoned: {e}")))?;
48        Ok(data.get(key).cloned())
49    }
50
51    async fn put(&self, key: &[u8], value: &[u8]) -> Result<(), Error> {
52        let mut data = self
53            .data
54            .write()
55            .map_err(|e| Error::InvalidData(format!("lock poisoned: {e}")))?;
56        data.insert(key.to_vec(), value.to_vec());
57        Ok(())
58    }
59
60    async fn delete(&self, key: &[u8]) -> Result<(), Error> {
61        let mut data = self
62            .data
63            .write()
64            .map_err(|e| Error::InvalidData(format!("lock poisoned: {e}")))?;
65        data.remove(key);
66        Ok(())
67    }
68
69    async fn list_prefix(&self, prefix: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>, Error> {
70        let data = self
71            .data
72            .read()
73            .map_err(|e| Error::InvalidData(format!("lock poisoned: {e}")))?;
74        let results = data
75            .range(prefix.to_vec()..)
76            .take_while(|(k, _)| k.starts_with(prefix))
77            .map(|(k, v)| (k.clone(), v.clone()))
78            .collect();
79        Ok(results)
80    }
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86
87    #[tokio::test]
88    async fn put_get_delete() {
89        let store = InMemoryKvStore::new();
90        assert!(store.get(b"key1").await.unwrap().is_none());
91
92        store.put(b"key1", b"value1").await.unwrap();
93        assert_eq!(store.get(b"key1").await.unwrap().unwrap(), b"value1");
94
95        store.delete(b"key1").await.unwrap();
96        assert!(store.get(b"key1").await.unwrap().is_none());
97    }
98
99    #[tokio::test]
100    async fn list_prefix_filtering() {
101        let store = InMemoryKvStore::new();
102        store.put(b"contract:a", b"1").await.unwrap();
103        store.put(b"contract:b", b"2").await.unwrap();
104        store.put(b"key:x", b"3").await.unwrap();
105
106        let results = store.list_prefix(b"contract:").await.unwrap();
107        assert_eq!(results.len(), 2);
108        assert_eq!(results[0].0, b"contract:a");
109        assert_eq!(results[1].0, b"contract:b");
110    }
111}
aztec_pxe/stores/kv.rs

aztec_pxe/stores/
kv.rs
