feat: Add PrimaryKey implementation example for custom types

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.

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<usize> {
+pub fn parse_length(value: &[u8]) -> StdResult<usize> {
     Ok(u16::from_be_bytes(
         value
             .try_into()
@@ -170,7 +170,7 @@ fn parse_length(value: &[u8]) -> StdResult<usize> {
 /// 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>, &[u8])> {
+pub fn split_first_key(key_elems: u16, value: &[u8]) -> StdResult<(Vec<u8>, &[u8])> {
     let mut index = 0;
     let mut first_key = Vec::new();
 

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;

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<Key> {
+        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<Key> {
+        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<Key> = 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<u8>) -> StdResult<Self::Output> {
+        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<Denom, Uint64> = 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");
+    }
+}