diff --git a/README.md b/README.md index c2ba96f..e0681ba 100644 --- a/README.md +++ b/README.md @@ -249,6 +249,13 @@ fn demo() -> StdResult<()> { } ``` +### Using custom types as a Map keys + +If you want to use a custom type as a key in a `Map`, you need to implement the `PrimaryKey` trait for it +as well as `KeyDeserialize` and `Prefixer` traits. + +Please take a look at the reference [example with Denom implementation](./tests/custom_types_serde.rs) for details. + ### Path Under the scenes, we create a `Path` from the `Map` when accessing a key. diff --git a/src/de.rs b/src/de.rs index 2377497..b356884 100644 --- a/src/de.rs +++ b/src/de.rs @@ -158,7 +158,7 @@ macro_rules! integer_de { integer_de!(for i8, u8, i16, u16, i32, u32, i64, u64, i128, u128, Uint64, Uint128, Int64, Int128); -fn parse_length(value: &[u8]) -> StdResult { +pub fn parse_length(value: &[u8]) -> StdResult { Ok(u16::from_be_bytes( value .try_into() @@ -170,7 +170,7 @@ fn parse_length(value: &[u8]) -> StdResult { /// Splits the first key from the value based on the provided number of key elements. /// The return value is ordered as (first_key, remainder). /// -fn split_first_key(key_elems: u16, value: &[u8]) -> StdResult<(Vec, &[u8])> { +pub fn split_first_key(key_elems: u16, value: &[u8]) -> StdResult<(Vec, &[u8])> { let mut index = 0; let mut first_key = Vec::new(); diff --git a/src/lib.rs b/src/lib.rs index 99f1718..e543073 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -33,6 +33,8 @@ mod snapshot; #[cfg(feature = "iterator")] pub use bound::{Bound, Bounder, PrefixBound, RawBound}; +pub use de::parse_length; +pub use de::split_first_key; pub use de::KeyDeserialize; pub use deque::Deque; pub use deque::DequeIter; diff --git a/tests/custom_types_serde.rs b/tests/custom_types_serde.rs new file mode 100644 index 0000000..5e584b4 --- /dev/null +++ b/tests/custom_types_serde.rs @@ -0,0 +1,105 @@ +use cosmwasm_std::{Addr, StdError, StdResult}; +use cw_storage_plus::{split_first_key, Key, KeyDeserialize, Prefixer, PrimaryKey}; +use std::u8; + +/// This file an example of the `PrimaryKey` and `KeyDeserialize` implementation for a custom type +/// `Denom` which can be either a native token or a cw20 token address. +/// +/// The idea is to store the Denom in the storage as a composite key with 2 elements: +/// 1. The prefix which is either `NATIVE_PREFIX` or `CW20_PREFIX` to differentiate between the +/// two types on a raw bytes level +/// 2. The value which is either the native token name or the cw20 token address + +/// Define a custom type which is Denom that can be either Native or Cw20 token address +#[derive(Clone, Debug, PartialEq, Eq)] +enum Denom { + Native(String), + Cw20(Addr), +} + +const NATIVE_PREFIX: u8 = 1; +const CW20_PREFIX: u8 = 2; + +impl PrimaryKey<'_> for Denom { + type Prefix = u8; + type SubPrefix = (); + type Suffix = String; + type SuperSuffix = Self; + + fn key(&self) -> Vec { + let (prefix, value) = match self { + Denom::Native(name) => (NATIVE_PREFIX, name.as_bytes()), + Denom::Cw20(addr) => (CW20_PREFIX, addr.as_bytes()), + }; + vec![Key::Val8([prefix]), Key::Ref(value)] + } +} + +impl Prefixer<'_> for Denom { + fn prefix(&self) -> Vec { + let (prefix, value) = match self { + Denom::Native(name) => (NATIVE_PREFIX.prefix(), name.prefix()), + Denom::Cw20(addr) => (CW20_PREFIX.prefix(), addr.prefix()), + }; + + let mut result: Vec = vec![]; + result.extend(prefix); + result.extend(value); + result + } +} + +impl KeyDeserialize for Denom { + type Output = Self; + + const KEY_ELEMS: u16 = 2; + + #[inline(always)] + fn from_vec(value: Vec) -> StdResult { + let (prefix, value) = split_first_key(Self::KEY_ELEMS, value.as_ref())?; + let value = value.to_vec(); + + match u8::from_vec(prefix)? { + NATIVE_PREFIX => Ok(Denom::Native(String::from_vec(value)?)), + CW20_PREFIX => Ok(Denom::Cw20(Addr::from_vec(value)?)), + _ => Err(StdError::generic_err("Invalid prefix")), + } + } +} + +#[cfg(test)] +mod test { + use crate::Denom; + use cosmwasm_std::testing::MockStorage; + use cosmwasm_std::{Addr, Uint64}; + use cw_storage_plus::Map; + + #[test] + fn round_trip_tests() { + let test_data = vec![ + Denom::Native("cosmos".to_string()), + Denom::Native("some_long_native_value_with_high_precise".to_string()), + Denom::Cw20(Addr::unchecked("contract1")), + Denom::Cw20(Addr::unchecked( + "cosmos1p7d8mnjttcszv34pk2a5yyug3474mhffasa7tg", + )), + ]; + + for denom in test_data { + verify_map_serde(denom); + } + } + + fn verify_map_serde(denom: Denom) { + let mut storage = MockStorage::new(); + let map: Map = Map::new("denom_map"); + let mock_value = Uint64::from(123u64); + + map.save(&mut storage, denom.clone(), &mock_value).unwrap(); + + assert!(map.has(&storage, denom.clone()), "key should exist"); + + let value = map.load(&storage, denom).unwrap(); + assert_eq!(value, mock_value, "value should match"); + } +}