A short guide to implementing Storey key types

src/pages/storey/containers/map/key-impl.mdx

@@ -5,3 +5,263 @@ 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, 8-14}
+use storey::containers::map::{key::DynamicKey, Key};
+
+enum Denom {
+    Native(String),
+    CW20(String),
+}
+
+impl Key for Denom {
+    type Kind = DynamicKey;
+
+    fn encode(&self) -> Vec<u8> {
+        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.
+
+<Callout>
+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<L>` to figure out how many
+bytes to eat - a more performant solution.
+
+</Callout>
+
+The [`encode`] method is used to serialize the key into a byte vector. Let's implement it now!
+
+```rust template="storage" {12-21}
+use storey::containers::map::{key::DynamicKey, Key};
+
+enum Denom {
+    Native(String),
+    CW20(String),
+}
+
+impl Key for Denom {
+    type Kind = DynamicKey;
+
+    fn encode(&self) -> Vec<u8> {
+        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" {8-11, 18-19}
+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<u8> {
+        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" {30-43}
+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<u8> {
+        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<Self, Self::Error> {
+        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
+
+<Callout>
+  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`.
+</Callout>
+
+## Using the thing
+
+Now that we have our key type implemented, we can use it in a map:
+
+```rust template="storage" {1-3, 49-64}
+use cw_storey::CwStorage;
+use cw_storey::containers::{Item, Map};
+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<u8> {
+        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<Self, Self::Error> {
+        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<Denom, Item<u64>> = Map::new(MAP_IX);
+let mut cw_storage = CwStorage(&mut storage);
+let mut access = map.access(&mut cw_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