use serde::{de::DeserializeOwned, ser::Serialize};
use std::marker::PhantomData;
use cosmwasm_std::{to_vec, ReadonlyStorage, StdResult, Storage};
#[cfg(feature = "iterator")]
use cosmwasm_std::{Order, KV};
#[cfg(feature = "iterator")]
use crate::type_helpers::deserialize_kv;
use crate::type_helpers::{may_deserialize, must_deserialize};
pub fn typed<S, T>(storage: &mut S) -> TypedStorage<S, T>
where
S: Storage,
T: Serialize + DeserializeOwned,
{
TypedStorage::new(storage)
}
pub fn typed_read<S, T>(storage: &S) -> ReadonlyTypedStorage<S, T>
where
S: ReadonlyStorage,
T: Serialize + DeserializeOwned,
{
ReadonlyTypedStorage::new(storage)
}
pub struct TypedStorage<'a, S, T>
where
S: Storage,
T: Serialize + DeserializeOwned,
{
storage: &'a mut S,
data: PhantomData<T>,
}
impl<'a, S, T> TypedStorage<'a, S, T>
where
S: Storage,
T: Serialize + DeserializeOwned,
{
pub fn new(storage: &'a mut S) -> Self {
TypedStorage {
storage,
data: PhantomData,
}
}
pub fn save(&mut self, key: &[u8], data: &T) -> StdResult<()> {
self.storage.set(key, &to_vec(data)?);
Ok(())
}
pub fn load(&self, key: &[u8]) -> StdResult<T> {
let value = self.storage.get(key);
must_deserialize(&value)
}
pub fn may_load(&self, key: &[u8]) -> StdResult<Option<T>> {
let value = self.storage.get(key);
may_deserialize(&value)
}
#[cfg(feature = "iterator")]
pub fn range<'b>(
&'b self,
start: Option<&[u8]>,
end: Option<&[u8]>,
order: Order,
) -> Box<dyn Iterator<Item = StdResult<KV<T>>> + 'b> {
let mapped = self
.storage
.range(start, end, order)
.map(deserialize_kv::<T>);
Box::new(mapped)
}
pub fn update<A>(&mut self, key: &[u8], action: A) -> StdResult<T>
where
A: FnOnce(Option<T>) -> StdResult<T>,
{
let input = self.may_load(key)?;
let output = action(input)?;
self.save(key, &output)?;
Ok(output)
}
}
pub struct ReadonlyTypedStorage<'a, S, T>
where
S: ReadonlyStorage,
T: Serialize + DeserializeOwned,
{
storage: &'a S,
data: PhantomData<T>,
}
impl<'a, S, T> ReadonlyTypedStorage<'a, S, T>
where
S: ReadonlyStorage,
T: Serialize + DeserializeOwned,
{
pub fn new(storage: &'a S) -> Self {
ReadonlyTypedStorage {
storage,
data: PhantomData,
}
}
pub fn load(&self, key: &[u8]) -> StdResult<T> {
let value = self.storage.get(key);
must_deserialize(&value)
}
pub fn may_load(&self, key: &[u8]) -> StdResult<Option<T>> {
let value = self.storage.get(key);
may_deserialize(&value)
}
#[cfg(feature = "iterator")]
pub fn range<'b>(
&'b self,
start: Option<&[u8]>,
end: Option<&[u8]>,
order: Order,
) -> Box<dyn Iterator<Item = StdResult<KV<T>>> + 'b> {
let mapped = self
.storage
.range(start, end, order)
.map(deserialize_kv::<T>);
Box::new(mapped)
}
}
#[cfg(test)]
mod test {
use super::*;
use cosmwasm_std::testing::MockStorage;
use cosmwasm_std::StdError;
use serde::{Deserialize, Serialize};
use crate::prefixed;
#[derive(Serialize, Deserialize, PartialEq, Debug, Clone)]
struct Data {
pub name: String,
pub age: i32,
}
#[test]
fn store_and_load() {
let mut store = MockStorage::new();
let mut bucket = TypedStorage::<_, Data>::new(&mut store);
assert!(bucket.load(b"maria").is_err());
assert_eq!(bucket.may_load(b"maria").unwrap(), None);
let data = Data {
name: "Maria".to_string(),
age: 42,
};
bucket.save(b"maria", &data).unwrap();
let loaded = bucket.load(b"maria").unwrap();
assert_eq!(data, loaded);
}
#[test]
fn store_with_prefix() {
let mut store = MockStorage::new();
let mut space = prefixed(&mut store, b"data");
let mut bucket = typed::<_, Data>(&mut space);
let data = Data {
name: "Maria".to_string(),
age: 42,
};
bucket.save(b"maria", &data).unwrap();
let loaded = bucket.load(b"maria").unwrap();
assert_eq!(data, loaded);
}
#[test]
fn readonly_works() {
let mut store = MockStorage::new();
let mut bucket = typed::<_, Data>(&mut store);
let data = Data {
name: "Maria".to_string(),
age: 42,
};
bucket.save(b"maria", &data).unwrap();
let reader = typed_read::<_, Data>(&mut store);
assert!(reader.load(b"john").is_err());
assert_eq!(reader.may_load(b"john").unwrap(), None);
let loaded = reader.load(b"maria").unwrap();
assert_eq!(data, loaded);
}
#[test]
fn update_success() {
let mut store = MockStorage::new();
let mut bucket = typed::<_, Data>(&mut store);
let init = Data {
name: "Maria".to_string(),
age: 42,
};
bucket.save(b"maria", &init).unwrap();
let birthday = |mayd: Option<Data>| -> StdResult<Data> {
let mut d = mayd.ok_or(StdError::not_found("Data"))?;
d.age += 1;
Ok(d)
};
let output = bucket.update(b"maria", &birthday).unwrap();
let expected = Data {
name: "Maria".to_string(),
age: 43,
};
assert_eq!(output, expected);
let loaded = bucket.load(b"maria").unwrap();
assert_eq!(loaded, expected);
}
#[test]
fn update_fails_on_error() {
let mut store = MockStorage::new();
let mut bucket = typed::<_, Data>(&mut store);
let init = Data {
name: "Maria".to_string(),
age: 42,
};
bucket.save(b"maria", &init).unwrap();
let output = bucket.update(b"maria", |_d| {
Err(StdError::generic_err("cuz i feel like it"))
});
assert!(output.is_err());
let loaded = bucket.load(b"maria").unwrap();
assert_eq!(loaded, init);
}
#[test]
fn update_handles_on_no_data() {
let mut store = MockStorage::new();
let mut bucket = typed::<_, Data>(&mut store);
let init_value = Data {
name: "Maria".to_string(),
age: 42,
};
let output = bucket
.update(b"maria", |d| match d {
Some(_) => Err(StdError::generic_err("Ensure this was empty")),
None => Ok(init_value.clone()),
})
.unwrap();
assert_eq!(output, init_value);
let loaded = bucket.load(b"maria").unwrap();
assert_eq!(loaded, init_value);
}
#[test]
#[cfg(feature = "iterator")]
fn range_over_data() {
let mut store = MockStorage::new();
let mut bucket = typed::<_, Data>(&mut store);
let jose = Data {
name: "Jose".to_string(),
age: 42,
};
let maria = Data {
name: "Maria".to_string(),
age: 27,
};
bucket.save(b"maria", &maria).unwrap();
bucket.save(b"jose", &jose).unwrap();
let res_data: StdResult<Vec<KV<Data>>> =
bucket.range(None, None, Order::Ascending).collect();
let data = res_data.unwrap();
assert_eq!(data.len(), 2);
assert_eq!(data[0], (b"jose".to_vec(), jose.clone()));
assert_eq!(data[1], (b"maria".to_vec(), maria.clone()));
let read_bucket = typed_read::<_, Data>(&store);
let res_data: StdResult<Vec<KV<Data>>> =
read_bucket.range(None, None, Order::Ascending).collect();
let data = res_data.unwrap();
assert_eq!(data.len(), 2);
assert_eq!(data[0], (b"jose".to_vec(), jose));
assert_eq!(data[1], (b"maria".to_vec(), maria));
}
}