bark/persist/adaptor/
memory.rs

1//! In-memory implementation of the `StorageAdaptor` trait.
2//!
3//! This module provides a simple in-memory storage adaptor for testing
4//! and ephemeral use cases, along with a reusable test suite for validating
5//! any `StorageAdaptor` implementation.
6
7use std::collections::{BTreeMap, HashMap};
8
9use crate::persist::adaptor::{StorageAdaptor, Query, Record};
10
11/// In-memory storage adaptor for testing and simple use cases.
12///
13/// This implementation stores records in partition-keyed `HashMap`s.
14/// Each partition has its own map.
15///
16/// # Example
17///
18/// ```rust
19/// # use bitcoin::Network;
20/// # use bitcoin::bip32::Fingerprint;
21/// # use bark::WalletProperties;
22/// # use bark::persist::BarkPersister;
23/// # use bark::persist::adaptor::StorageAdaptorWrapper;
24/// # use bark::persist::adaptor::memory::MemoryStorageAdaptor;
25///
26/// # async fn example() -> anyhow::Result<()> {
27/// let storage = StorageAdaptorWrapper::new(MemoryStorageAdaptor::new());
28/// let properties = WalletProperties {
29///		network: Network::Testnet,
30///		fingerprint: Fingerprint::default(),
31///		server_pubkey: None,
32///	};
33///
34/// storage.init_wallet(&properties).await?;
35/// # Ok(())
36/// # }
37/// ```
38#[derive(Debug, Default)]
39pub struct MemoryStorageAdaptor {
40	/// Map from partition -> (pk -> record)
41	partitions: HashMap<u8, BTreeMap<Vec<u8>, Record>>,
42}
43
44impl MemoryStorageAdaptor {
45	/// Creates a new empty in-memory storage.
46	pub fn new() -> Self {
47		Self::default()
48	}
49
50	pub fn partitions(&self) -> &HashMap<u8, BTreeMap<Vec<u8>, Record>> {
51		&self.partitions
52	}
53}
54
55#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
56#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
57impl StorageAdaptor for MemoryStorageAdaptor {
58	async fn put(&mut self, record: Record) -> anyhow::Result<()> {
59		let partition = self.partitions.entry(record.partition).or_default();
60		partition.insert(record.pk.clone(), record);
61		Ok(())
62	}
63
64	async fn get(&self, partition: u8, pk: &[u8]) -> anyhow::Result<Option<Record>> {
65		if let Some(partition) = self.partitions.get(&partition) {
66			if let Some(record) = partition.get(pk) {
67				return Ok(Some(record.clone()));
68			}
69		}
70		Ok(None)
71	}
72
73	async fn delete(&mut self, partition: u8, pk: &[u8]) -> anyhow::Result<Option<Record>> {
74		if let Some(partition) = self.partitions.get_mut(&partition) {
75			return Ok(partition.remove(pk))
76		}
77
78		Ok(None)
79	}
80
81	async fn query(&self, query: Query) -> anyhow::Result<Vec<Record>> {
82		let Some(partition) = self.partitions.get(&query.partition) else {
83			return Ok(Vec::new());
84		};
85
86		let mut results: Vec<_> = partition
87			.values()
88			.filter(|r| {
89				let matches_start = if let Some(start) = &query.start {
90					match &r.sort_key {
91						Some(sort_key) => sort_key.cmp(start) >= std::cmp::Ordering::Equal,
92						None => false,
93					}
94				} else { true };
95
96				let matches_end = if let Some(end) = &query.end {
97					match &r.sort_key {
98						Some(sort_key) => sort_key.cmp(end) <= std::cmp::Ordering::Equal,
99						None => false,
100					}
101				} else { true };
102
103				matches_start && matches_end
104			})
105			.cloned()
106			.collect();
107
108		// Sort by sort key
109		results.sort_by(|a, b| {
110			match (&a.sort_key, &b.sort_key) {
111				(Some(ka), Some(kb)) => ka.cmp(kb),
112				(Some(_), None) => std::cmp::Ordering::Less,
113				(None, Some(_)) => std::cmp::Ordering::Greater,
114				(None, None) => std::cmp::Ordering::Equal,
115			}
116		});
117
118		// Apply limit
119		if let Some(limit) = query.limit {
120			results.truncate(limit);
121		}
122
123		Ok(results)
124	}
125}
126
127#[cfg(test)]
128mod tests {
129	use super::*;
130	use crate::persist::adaptor::test_suite;
131
132	/// Run the full test suite against MemoryStorageAdaptor.
133	#[tokio::test]
134	async fn memory_adaptor_full_test_suite() {
135		let mut storage = MemoryStorageAdaptor::new();
136		test_suite::run_all(&mut storage).await;
137	}
138}
bark/persist/adaptor/memory.rs

bark/persist/adaptor/
memory.rs
