diff --git a/src/pages/storey/containers/map/key-impl.mdx b/src/pages/storey/containers/map/key-impl.mdx index e3fc2b54..7585a144 100644 --- a/src/pages/storey/containers/map/key-impl.mdx +++ b/src/pages/storey/containers/map/key-impl.mdx @@ -5,3 +5,265 @@ tags: ["storey", "containers"] import { Callout } from "nextra/components"; # Implementing key types + +In this section, we will implement a custom key type for use with [maps]. + +Let's say we have a `Denom` enum to represent different kinds of tokens: + +- native tokens +- [CW20] tokens + +```rust +enum Denom { + Native(String), + CW20(String), +} +``` + +## The Key trait + +We can implement the `Key` trait for this enum to make it usable as a key in a map: + +```rust template="storage" {1-2, 9-15} +use cw_storey::containers::CwKeySet; +use storey::containers::map::{key::DynamicKey, Key}; + +enum Denom { + Native(String), + CW20(String), +} + +impl Key for Denom { + type Kind = DynamicKey; + + fn encode(&self) -> Vec { + todo!() + } +} +``` + +The [`Kind`] associated type is used to signal to the framework whether the key is dynamically sized +or not. In this case, we use [`DynamicKey`] because the key size is not fixed. If it was always +exactly 8 bytes, we could use [`FixedSizeKey<8>`] instead. + + +Why does this matter? The framework uses this information to determine how to encode the key. If there are more keys following this one (e.g. in a multi-level map), the framework needs to know how to tell where this one ends during decoding. + +For a dynamically sized key, the framework will length-prefix it when necessary. For a fixed-size +key, it will use the static information you provide with `FixedSizeKey` to figure out how many +bytes to eat - a more performant solution. + + + +The [`encode`] method is used to serialize the key into a byte vector. Let's implement it now! + +```rust template="storage" {13-22} +use cw_storey::containers::CwKeySet; +use storey::containers::map::{key::DynamicKey, Key}; + +enum Denom { + Native(String), + CW20(String), +} + +impl Key for Denom { + type Kind = DynamicKey; + + fn encode(&self) -> Vec { + let (discriminant, data) = match self { + Denom::Native(data) => (0, data), + Denom::CW20(data) => (1, data), + }; + + let mut result = Vec::with_capacity(1 + data.len()); + result.push(discriminant); + result.extend_from_slice(data.as_bytes()); + + result + } +} +``` + +The code should be pretty self-explanatory. We use a simple encoding scheme where we write a single +byte discriminant followed by the actual data. The discriminant is how we tell a native token from a +CW20 token. + +One little improvement we can go for is to avoid hardcoding the discriminant. We'll want to reuse +these values in the decoding logic, so let's define them as constants: + +```rust template="storage" {9-12, 19-20} +use cw_storey::containers::CwKeySet; +use storey::containers::map::{key::DynamicKey, Key}; + +enum Denom { + Native(String), + CW20(String), +} + +impl Denom { + const NATIVE_DISCRIMINANT: u8 = 0; + const CW20_DISCRIMINANT: u8 = 1; +} + +impl Key for Denom { + type Kind = DynamicKey; + + fn encode(&self) -> Vec { + let (discriminant, data) = match self { + Denom::Native(data) => (Self::NATIVE_DISCRIMINANT, data), + Denom::CW20(data) => (Self::CW20_DISCRIMINANT, data), + }; + + let mut result = Vec::with_capacity(1 + data.len()); + result.push(discriminant); + result.extend_from_slice(data.as_bytes()); + + result + } +} +``` + +Alright. The `Key` trait allows us to access the data in a map using our custom key type. We still +need a way to decode the key back into the enum. This is used for example in iteration. + +## The OwnedKey trait + +Let's now implement the [`OwnedKey`] trait. + +```rust template="storage" {31-44} +use cw_storey::containers::CwKeySet; +use storey::containers::map::{key::DynamicKey, Key, OwnedKey}; + +enum Denom { + Native(String), + CW20(String), +} + +impl Denom { + const NATIVE_DISCRIMINANT: u8 = 0; + const CW20_DISCRIMINANT: u8 = 1; +} + +impl Key for Denom { + type Kind = DynamicKey; + + fn encode(&self) -> Vec { + let (discriminant, data) = match self { + Denom::Native(data) => (Self::NATIVE_DISCRIMINANT, data), + Denom::CW20(data) => (Self::CW20_DISCRIMINANT, data), + }; + + let mut result = Vec::with_capacity(1 + data.len()); + result.push(discriminant); + result.extend_from_slice(data.as_bytes()); + + result + } +} + +impl OwnedKey for Denom { + type Error = (); + + fn from_bytes(bytes: &[u8]) -> Result { + let discriminant = bytes[0]; + let data = String::from_utf8(bytes[1..].to_vec()).map_err(|_| ())?; + + match discriminant { + Self::NATIVE_DISCRIMINANT => Ok(Self::Native(data)), + Self::CW20_DISCRIMINANT => Ok(Self::CW20(data)), + _ => Err(()), + } + } +} +``` + +The [`from_bytes`] method should return an instance of the key type or an error if the data is +invalid. Here it does the following: + +- read the discriminant byte +- read the data bytes as a UTF-8 string, erroring out if it's not valid +- match the discriminant to the enum variant, or error if invalid +- return the deserialized key + + + What we have is functional. There's one last improvement you could make here - a proper error + type. We used `()` as a placeholder, but in production code it's good practice to define an enum. + In this case, the enum could hold variants like `InvalidDiscriminant` and `InvalidUTF8`. + + +## Using the thing + +Now that we have our key type implemented, we can use it in a map: + +```rust template="storage" showLineNumbers {1-3, 48-57} +use cw_storey::containers::{Item, Map, CwKeySet}; +use storey::containers::IterableAccessor; +use storey::containers::map::{key::DynamicKey, Key, OwnedKey}; + +#[derive(Debug, PartialEq)] +enum Denom { + Native(String), + CW20(String), +} + +impl Denom { + const NATIVE_DISCRIMINANT: u8 = 0; + const CW20_DISCRIMINANT: u8 = 1; +} + +impl Key for Denom { + type Kind = DynamicKey; + + fn encode(&self) -> Vec { + let (discriminant, data) = match self { + Denom::Native(data) => (Self::NATIVE_DISCRIMINANT, data), + Denom::CW20(data) => (Self::CW20_DISCRIMINANT, data), + }; + + let mut result = Vec::with_capacity(1 + data.len()); + result.push(discriminant); + result.extend_from_slice(data.as_bytes()); + + result + } +} + +impl OwnedKey for Denom { + type Error = (); + + fn from_bytes(bytes: &[u8]) -> Result { + let discriminant = bytes[0]; + let data = String::from_utf8(bytes[1..].to_vec()).map_err(|_| ())?; + + match discriminant { + Self::NATIVE_DISCRIMINANT => Ok(Self::Native(data)), + Self::CW20_DISCRIMINANT => Ok(Self::CW20(data)), + _ => Err(()), + } + } +} + +const MAP_IX: u8 = 1; + +let map: Map> = Map::new(MAP_IX); +let mut access = map.access(&mut storage); + +access.entry_mut(&Denom::Native("USDT".into())).set(&1000).unwrap(); +access.entry_mut(&Denom::CW20("some_addr_3824792".into())).set(&2000).unwrap(); + +assert_eq!(access.entry(&Denom::Native("USDT".into())).get().unwrap(), Some(1000)); +assert_eq!(access.entry(&Denom::CW20("some_addr_3824792".into())).get().unwrap(), Some(2000)); +``` + +Voilà! It works just like it would with any other key. + +[maps]: /docs/storey/containers/map +[CW20]: /docs/getting-started/cw20 +[`Kind`]: https://docs.rs/storey/latest/storey/containers/map/key/trait.Key.html#associatedtype.Kind +[`DynamicKey`]: https://docs.rs/storey/latest/storey/containers/map/key/struct.DynamicKey.html +[`FixedSizeKey<8>`]: + https://docs.rs/storey/latest/storey/containers/map/key/struct.FixedSizeKey.html +[`encode`]: https://docs.rs/storey/latest/storey/containers/map/key/trait.Key.html#tymethod.encode +[`OwnedKey`]: https://docs.rs/storey/latest/storey/containers/map/key/trait.OwnedKey.html +[`from_bytes`]: + https://docs.rs/storey/latest/storey/containers/map/key/trait.OwnedKey.html#tymethod.from_bytes