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

bark/persist/adaptor/
memory.rs
