use std::collections::BTreeMap;
#[cfg(feature = "iterator")]
use std::iter;
#[cfg(feature = "iterator")]
use std::ops::{Bound, RangeBounds};
use crate::errors::StdResult;
#[cfg(feature = "iterator")]
use crate::iterator::{Order, KV};
use crate::traits::{ReadonlyStorage, Storage};
#[derive(Default)]
pub struct MemoryStorage {
data: BTreeMap<Vec<u8>, Vec<u8>>,
}
impl MemoryStorage {
pub fn new() -> Self {
MemoryStorage::default()
}
}
impl ReadonlyStorage for MemoryStorage {
fn get(&self, key: &[u8]) -> StdResult<Option<Vec<u8>>> {
Ok(self.data.get(key).cloned())
}
#[cfg(feature = "iterator")]
fn range<'a>(
&'a self,
start: Option<&[u8]>,
end: Option<&[u8]>,
order: Order,
) -> StdResult<Box<dyn Iterator<Item = StdResult<KV>> + 'a>> {
let bounds = range_bounds(start, end);
match (bounds.start_bound(), bounds.end_bound()) {
(Bound::Included(start), Bound::Excluded(end)) if start > end => {
return Ok(Box::new(iter::empty()));
}
_ => {}
}
let iter = self.data.range(bounds);
Ok(match order {
Order::Ascending => Box::new(iter.map(clone_item).map(StdResult::Ok)),
Order::Descending => Box::new(iter.rev().map(clone_item).map(StdResult::Ok)),
})
}
}
#[cfg(feature = "iterator")]
fn range_bounds(start: Option<&[u8]>, end: Option<&[u8]>) -> impl RangeBounds<Vec<u8>> {
(
start.map_or(Bound::Unbounded, |x| Bound::Included(x.to_vec())),
end.map_or(Bound::Unbounded, |x| Bound::Excluded(x.to_vec())),
)
}
#[cfg(feature = "iterator")]
type BTreeMapPairRef<'a, T = Vec<u8>> = (&'a Vec<u8>, &'a T);
#[cfg(feature = "iterator")]
fn clone_item<T: Clone>(item_ref: BTreeMapPairRef<T>) -> KV<T> {
let (key, value) = item_ref;
(key.clone(), value.clone())
}
impl Storage for MemoryStorage {
fn set(&mut self, key: &[u8], value: &[u8]) -> StdResult<()> {
self.data.insert(key.to_vec(), value.to_vec());
Ok(())
}
fn remove(&mut self, key: &[u8]) -> StdResult<()> {
self.data.remove(key);
Ok(())
}
}
#[cfg(test)]
mod test {
use super::*;
#[cfg(feature = "iterator")]
fn iterator_test_suite<S: Storage>(store: &mut S) {
assert_eq!(store.get(b"foo").unwrap(), Some(b"bar".to_vec()));
assert_eq!(
store.range(None, None, Order::Ascending).unwrap().count(),
1
);
store.set(b"ant", b"hill").expect("error setting value");
store.set(b"ze", b"bra").expect("error setting value");
store.set(b"bye", b"bye").expect("error setting value");
store.remove(b"bye").expect("error removing key");
{
let iter = store.range(None, None, Order::Ascending).unwrap();
let elements: Vec<KV> = iter.filter_map(StdResult::ok).collect();
assert_eq!(
elements,
vec![
(b"ant".to_vec(), b"hill".to_vec()),
(b"foo".to_vec(), b"bar".to_vec()),
(b"ze".to_vec(), b"bra".to_vec()),
]
);
}
{
let iter = store.range(None, None, Order::Descending).unwrap();
let elements: Vec<KV> = iter.filter_map(StdResult::ok).collect();
assert_eq!(
elements,
vec![
(b"ze".to_vec(), b"bra".to_vec()),
(b"foo".to_vec(), b"bar".to_vec()),
(b"ant".to_vec(), b"hill".to_vec()),
]
);
}
{
let iter = store
.range(Some(b"f"), Some(b"n"), Order::Ascending)
.unwrap();
let elements: Vec<KV> = iter.filter_map(StdResult::ok).collect();
assert_eq!(elements, vec![(b"foo".to_vec(), b"bar".to_vec())]);
}
{
let iter = store
.range(Some(b"air"), Some(b"loop"), Order::Descending)
.unwrap();
let elements: Vec<KV> = iter.filter_map(StdResult::ok).collect();
assert_eq!(
elements,
vec![
(b"foo".to_vec(), b"bar".to_vec()),
(b"ant".to_vec(), b"hill".to_vec()),
]
);
}
{
let iter = store
.range(Some(b"foo"), Some(b"foo"), Order::Ascending)
.unwrap();
let elements: Vec<KV> = iter.filter_map(StdResult::ok).collect();
assert_eq!(elements, vec![]);
}
{
let iter = store
.range(Some(b"foo"), Some(b"foo"), Order::Descending)
.unwrap();
let elements: Vec<KV> = iter.filter_map(StdResult::ok).collect();
assert_eq!(elements, vec![]);
}
{
let iter = store
.range(Some(b"z"), Some(b"a"), Order::Ascending)
.unwrap();
let elements: Vec<KV> = iter.filter_map(StdResult::ok).collect();
assert_eq!(elements, vec![]);
}
{
let iter = store
.range(Some(b"z"), Some(b"a"), Order::Descending)
.unwrap();
let elements: Vec<KV> = iter.filter_map(StdResult::ok).collect();
assert_eq!(elements, vec![]);
}
{
let iter = store.range(Some(b"f"), None, Order::Ascending).unwrap();
let elements: Vec<KV> = iter.filter_map(StdResult::ok).collect();
assert_eq!(
elements,
vec![
(b"foo".to_vec(), b"bar".to_vec()),
(b"ze".to_vec(), b"bra".to_vec()),
]
);
}
{
let iter = store.range(Some(b"f"), None, Order::Descending).unwrap();
let elements: Vec<KV> = iter.filter_map(StdResult::ok).collect();
assert_eq!(
elements,
vec![
(b"ze".to_vec(), b"bra".to_vec()),
(b"foo".to_vec(), b"bar".to_vec()),
]
);
}
{
let iter = store.range(None, Some(b"f"), Order::Ascending).unwrap();
let elements: Vec<KV> = iter.filter_map(StdResult::ok).collect();
assert_eq!(elements, vec![(b"ant".to_vec(), b"hill".to_vec()),]);
}
{
let iter = store.range(None, Some(b"no"), Order::Descending).unwrap();
let elements: Vec<KV> = iter.filter_map(StdResult::ok).collect();
assert_eq!(
elements,
vec![
(b"foo".to_vec(), b"bar".to_vec()),
(b"ant".to_vec(), b"hill".to_vec()),
]
);
}
}
#[test]
fn get_and_set() {
let mut store = MemoryStorage::new();
assert_eq!(None, store.get(b"foo").unwrap());
store.set(b"foo", b"bar").unwrap();
assert_eq!(Some(b"bar".to_vec()), store.get(b"foo").unwrap());
assert_eq!(None, store.get(b"food").unwrap());
}
#[test]
fn delete() {
let mut store = MemoryStorage::new();
store.set(b"foo", b"bar").unwrap();
store.set(b"food", b"bank").unwrap();
store.remove(b"foo").unwrap();
assert_eq!(None, store.get(b"foo").unwrap());
assert_eq!(Some(b"bank".to_vec()), store.get(b"food").unwrap());
}
#[test]
#[cfg(feature = "iterator")]
fn iterator() {
let mut store = MemoryStorage::new();
store.set(b"foo", b"bar").expect("error setting value");
iterator_test_suite(&mut store);
}
}