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::{Query, QueryRange, Record, StorageAdaptor, StorageAdaptorWrapper};
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///
25/// # async fn example() -> anyhow::Result<()> {
26/// let storage = StorageAdaptorWrapper::new_memory();
27/// let properties = WalletProperties {
28///		network: Network::Testnet,
29///		fingerprint: Fingerprint::default(),
30///		server_pubkey: None,
31///		server_mailbox_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
55impl StorageAdaptorWrapper<MemoryStorageAdaptor> {
56	pub fn new_memory() -> Self {
57		Self::new(MemoryStorageAdaptor::new())
58	}
59}
60
61#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
62#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
63impl StorageAdaptor for MemoryStorageAdaptor {
64	async fn put(&mut self, record: Record) -> anyhow::Result<()> {
65		let partition = self.partitions.entry(record.partition).or_default();
66		partition.insert(record.pk.clone(), record);
67		Ok(())
68	}
69
70	async fn get(&self, partition: u8, pk: &[u8]) -> anyhow::Result<Option<Record>> {
71		if let Some(partition) = self.partitions.get(&partition) {
72			if let Some(record) = partition.get(pk) {
73				return Ok(Some(record.clone()));
74			}
75		}
76		Ok(None)
77	}
78
79	async fn delete(&mut self, partition: u8, pk: &[u8]) -> anyhow::Result<Option<Record>> {
80		if let Some(partition) = self.partitions.get_mut(&partition) {
81			return Ok(partition.remove(pk))
82		}
83
84		Ok(None)
85	}
86
87	async fn query_sorted<R: QueryRange>(&self, query: Query<R>) -> anyhow::Result<Vec<Record>> {
88		let Some(partition) = self.partitions.get(&query.partition) else {
89			return Ok(Vec::new());
90		};
91
92		let mut results: Vec<_> = partition
93			.values()
94			.filter(|r| {
95				// Records without sort keys are excluded from query results
96				let Some(sort_key) = &r.sort_key else {
97					return false;
98				};
99
100				query.range.contains(sort_key)
101			})
102			.cloned()
103			.collect();
104
105		// Sort by sort key (all records have sort keys at this point)
106		results.sort_by(|a, b| {
107			match (&a.sort_key, &b.sort_key) {
108				(Some(ka), Some(kb)) => ka.cmp(kb),
109				_ => unreachable!("all records should have sort keys after filtering"),
110			}
111		});
112
113		// Apply limit
114		if let Some(limit) = query.limit {
115			results.truncate(limit);
116		}
117
118		Ok(results)
119	}
120
121	async fn get_all(&self, partition: u8) -> anyhow::Result<Vec<Record>> {
122		let Some(partition_map) = self.partitions.get(&partition) else {
123			return Ok(Vec::new());
124		};
125
126		Ok(partition_map.values().cloned().collect())
127	}
128}
129
130#[cfg(test)]
131mod tests {
132	use super::*;
133	use crate::persist::adaptor::test_suite;
134
135	/// Run the full test suite against MemoryStorageAdaptor.
136	#[cfg_attr(not(target_arch = "wasm32"), tokio::test)]
137	#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
138	async fn memory_adaptor_full_test_suite() {
139		let mut storage = MemoryStorageAdaptor::new();
140		test_suite::run_all(&mut storage).await;
141	}
142}
bark/persist/adaptor/memory.rs

bark/persist/adaptor/
memory.rs
