From 07bfec82f4163e3f204056b8cbf1d792a326571f Mon Sep 17 00:00:00 2001
From: Preston Evans <32944016+preston-evans98@users.noreply.github.com>
Date: Fri, 8 Sep 2023 02:18:41 -0700
Subject: [PATCH] Bug fix: StateMap::Keys are not consistent across platforms
 (#804)

* Bug fix: Introduce EncodeLike trait.

This PR removes our StateMap key encoding based on std::Hash, which was
not consistent across platforms. Instead, this PR introduces the
`EncodeLike<Ref, Target>` trait which marks that Ref can be encoded
like Target by the implementing codec.

This PR also removes the SingletonKey type, which required special
handling in codecs. Instead, of using this placeholder, this PR
implements auxiliary methods on the working set for dealing with
singletons

* Remove rollup config changes

* fix test: qualify conversion

* Add missing bounds for fuzzing

* fix docs

* clarify zsts in comment

* allow separate codecs for keys/values

* Split key and value codecs

* Fix fuzzing feature

* Introduce StateCodec trait to allow EncodeLike with SplitCodec

* add doc comments

* Fix fuzz and test targets

* Fix test

* fix docs
---
 examples/demo-prover/Cargo.lock               |  21 +++-
 examples/demo-prover/methods/guest/Cargo.lock |  17 ++-
 .../tests/custom_codec_must_be_used.rs        |  19 +++-
 .../sov-state/src/codec/bcs_codec.rs          |  24 +++-
 .../sov-state/src/codec/borsh_codec.rs        |  22 ++++
 .../sov-state/src/codec/json_codec.rs         |  23 ++++
 module-system/sov-state/src/codec/mod.rs      |  65 +++++++++++
 .../sov-state/src/codec/split_codec.rs        |  46 ++++++++
 .../sov-state/src/containers/accessory_map.rs |  86 ++++++++-------
 .../src/containers/accessory_value.rs         |  27 ++---
 .../sov-state/src/containers/accessory_vec.rs |  53 +++++----
 module-system/sov-state/src/containers/map.rs |  97 +++++++++-------
 .../sov-state/src/containers/value.rs         |  27 ++---
 module-system/sov-state/src/containers/vec.rs |  49 +++++----
 module-system/sov-state/src/scratchpad.rs     | 104 ++++++++++++++----
 module-system/sov-state/src/storage.rs        |  52 ++++-----
 .../src/state_machine/optimistic.rs           |   4 +-
 17 files changed, 520 insertions(+), 216 deletions(-)
 create mode 100644 module-system/sov-state/src/codec/split_codec.rs

diff --git a/examples/demo-prover/Cargo.lock b/examples/demo-prover/Cargo.lock
index 70c6937638..cd5af313f8 100644
--- a/examples/demo-prover/Cargo.lock
+++ b/examples/demo-prover/Cargo.lock
@@ -801,6 +801,7 @@ dependencies = [
  "sov-accounts",
  "sov-bank",
  "sov-blob-storage",
+ "sov-chain-state",
  "sov-cli",
  "sov-modules-api",
  "sov-modules-stf-template",
@@ -3671,14 +3672,26 @@ dependencies = [
  "schemars",
  "serde",
  "serde_json",
+ "sov-chain-state",
  "sov-modules-api",
- "sov-modules-macros",
- "sov-rollup-interface",
  "sov-sequencer-registry",
  "sov-state",
  "tracing",
 ]
 
+[[package]]
+name = "sov-chain-state"
+version = "0.1.0"
+dependencies = [
+ "anyhow",
+ "borsh",
+ "jsonrpsee 0.18.2",
+ "serde",
+ "serde_json",
+ "sov-modules-api",
+ "sov-state",
+]
+
 [[package]]
 name = "sov-cli"
 version = "0.1.0"
@@ -3767,6 +3780,7 @@ dependencies = [
  "risc0-adapter",
  "serde",
  "serde_json",
+ "sov-cli",
  "sov-db",
  "sov-modules-api",
  "sov-modules-stf-template",
@@ -3893,7 +3907,6 @@ dependencies = [
  "serde_json",
  "sov-bank",
  "sov-modules-api",
- "sov-rollup-interface",
  "sov-state",
 ]
 
@@ -3907,6 +3920,7 @@ dependencies = [
  "hex",
  "jmt",
  "serde",
+ "serde_json",
  "sha2 0.10.7",
  "sov-db",
  "sov-first-read-last-write-cache",
@@ -3948,7 +3962,6 @@ dependencies = [
  "serde",
  "serde_json",
  "sov-modules-api",
- "sov-rollup-interface",
  "sov-state",
  "thiserror",
 ]
diff --git a/examples/demo-prover/methods/guest/Cargo.lock b/examples/demo-prover/methods/guest/Cargo.lock
index 26e2576043..0ff1045690 100644
--- a/examples/demo-prover/methods/guest/Cargo.lock
+++ b/examples/demo-prover/methods/guest/Cargo.lock
@@ -454,6 +454,7 @@ dependencies = [
  "sov-accounts",
  "sov-bank",
  "sov-blob-storage",
+ "sov-chain-state",
  "sov-modules-api",
  "sov-modules-stf-template",
  "sov-rollup-interface",
@@ -2021,14 +2022,23 @@ dependencies = [
  "bincode",
  "borsh",
  "hex",
+ "sov-chain-state",
  "sov-modules-api",
- "sov-modules-macros",
- "sov-rollup-interface",
  "sov-sequencer-registry",
  "sov-state",
  "tracing",
 ]
 
+[[package]]
+name = "sov-chain-state"
+version = "0.1.0"
+dependencies = [
+ "anyhow",
+ "borsh",
+ "sov-modules-api",
+ "sov-state",
+]
+
 [[package]]
 name = "sov-demo-prover-guest"
 version = "0.1.0"
@@ -2139,7 +2149,6 @@ dependencies = [
  "risc0-zkvm-platform",
  "sov-bank",
  "sov-modules-api",
- "sov-rollup-interface",
  "sov-state",
  "zk-cycle-macros",
  "zk-cycle-utils",
@@ -2157,6 +2166,7 @@ dependencies = [
  "risc0-zkvm",
  "risc0-zkvm-platform",
  "serde",
+ "serde_json",
  "sha2 0.10.6",
  "sov-first-read-last-write-cache",
  "sov-rollup-interface",
@@ -2171,7 +2181,6 @@ dependencies = [
  "anyhow",
  "borsh",
  "sov-modules-api",
- "sov-rollup-interface",
  "sov-state",
  "thiserror",
 ]
diff --git a/module-system/sov-modules-macros/tests/custom_codec_must_be_used.rs b/module-system/sov-modules-macros/tests/custom_codec_must_be_used.rs
index 1dbfb88963..ef5b6c406d 100644
--- a/module-system/sov-modules-macros/tests/custom_codec_must_be_used.rs
+++ b/module-system/sov-modules-macros/tests/custom_codec_must_be_used.rs
@@ -2,7 +2,7 @@ use std::panic::catch_unwind;
 
 use sov_modules_api::default_context::DefaultContext;
 use sov_modules_api::{Context, ModuleInfo};
-use sov_state::codec::StateValueCodec;
+use sov_state::codec::{StateCodec, StateKeyCodec, StateValueCodec};
 use sov_state::{DefaultStorageSpec, ProverStorage, StateValue, WorkingSet};
 
 #[derive(ModuleInfo)]
@@ -26,6 +26,23 @@ impl CustomCodec {
     }
 }
 
+impl StateCodec for CustomCodec {
+    type KeyCodec = Self;
+    type ValueCodec = Self;
+    fn key_codec(&self) -> &Self::KeyCodec {
+        self
+    }
+    fn value_codec(&self) -> &Self::ValueCodec {
+        self
+    }
+}
+
+impl<K> StateKeyCodec<K> for CustomCodec {
+    fn encode_key(&self, _key: &K) -> Vec<u8> {
+        unimplemented!()
+    }
+}
+
 impl<V> StateValueCodec<V> for CustomCodec {
     type Error = String;
 
diff --git a/module-system/sov-state/src/codec/bcs_codec.rs b/module-system/sov-state/src/codec/bcs_codec.rs
index e7fa6e22ea..0635477a37 100644
--- a/module-system/sov-state/src/codec/bcs_codec.rs
+++ b/module-system/sov-state/src/codec/bcs_codec.rs
@@ -1,9 +1,19 @@
+use super::{StateCodec, StateKeyCodec};
 use crate::codec::StateValueCodec;
 
 /// A [`StateValueCodec`] that uses [`bcs`] for all keys and values.
 #[derive(Debug, Default, PartialEq, Eq, Clone, serde::Serialize, serde::Deserialize)]
 pub struct BcsCodec;
 
+impl<K> StateKeyCodec<K> for BcsCodec
+where
+    K: serde::Serialize,
+{
+    fn encode_key(&self, key: &K) -> Vec<u8> {
+        bcs::to_bytes(key).expect("Failed to serialize key")
+    }
+}
+
 impl<V> StateValueCodec<V> for BcsCodec
 where
     V: serde::Serialize + for<'a> serde::Deserialize<'a>,
@@ -11,10 +21,22 @@ where
     type Error = bcs::Error;
 
     fn encode_value(&self, value: &V) -> Vec<u8> {
-        bcs::to_bytes(value).expect("Failed to serialize key")
+        bcs::to_bytes(value).expect("Failed to serialize value")
     }
 
     fn try_decode_value(&self, bytes: &[u8]) -> Result<V, Self::Error> {
         bcs::from_bytes(bytes)
     }
 }
+
+impl StateCodec for BcsCodec {
+    type KeyCodec = Self;
+    type ValueCodec = Self;
+    fn key_codec(&self) -> &Self::KeyCodec {
+        self
+    }
+
+    fn value_codec(&self) -> &Self::ValueCodec {
+        self
+    }
+}
diff --git a/module-system/sov-state/src/codec/borsh_codec.rs b/module-system/sov-state/src/codec/borsh_codec.rs
index 4647755dea..92f59ca57a 100644
--- a/module-system/sov-state/src/codec/borsh_codec.rs
+++ b/module-system/sov-state/src/codec/borsh_codec.rs
@@ -1,9 +1,19 @@
+use super::{StateCodec, StateKeyCodec};
 use crate::codec::StateValueCodec;
 
 /// A [`StateValueCodec`] that uses [`borsh`] for all values.
 #[derive(Debug, Default, PartialEq, Eq, Clone, borsh::BorshDeserialize, borsh::BorshSerialize)]
 pub struct BorshCodec;
 
+impl<K> StateKeyCodec<K> for BorshCodec
+where
+    K: borsh::BorshSerialize + borsh::BorshDeserialize,
+{
+    fn encode_key(&self, value: &K) -> Vec<u8> {
+        value.try_to_vec().expect("Failed to serialize value")
+    }
+}
+
 impl<V> StateValueCodec<V> for BorshCodec
 where
     V: borsh::BorshSerialize + borsh::BorshDeserialize,
@@ -18,3 +28,15 @@ where
         V::try_from_slice(bytes)
     }
 }
+
+impl StateCodec for BorshCodec {
+    type KeyCodec = Self;
+    type ValueCodec = Self;
+    fn key_codec(&self) -> &Self::KeyCodec {
+        self
+    }
+
+    fn value_codec(&self) -> &Self::ValueCodec {
+        self
+    }
+}
diff --git a/module-system/sov-state/src/codec/json_codec.rs b/module-system/sov-state/src/codec/json_codec.rs
index 487b748179..c7d5ab35d8 100644
--- a/module-system/sov-state/src/codec/json_codec.rs
+++ b/module-system/sov-state/src/codec/json_codec.rs
@@ -1,11 +1,21 @@
 use serde_json;
 
+use super::{StateCodec, StateKeyCodec};
 use crate::codec::StateValueCodec;
 
 /// A [`StateValueCodec`] that uses [`serde_json`] for all values.
 #[derive(Debug, Default, PartialEq, Eq, Clone, serde::Serialize, serde::Deserialize)]
 pub struct JsonCodec;
 
+impl<K> StateKeyCodec<K> for JsonCodec
+where
+    K: serde::Serialize,
+{
+    fn encode_key(&self, key: &K) -> Vec<u8> {
+        serde_json::to_vec(key).expect("Failed to serialize value")
+    }
+}
+
 impl<V> StateValueCodec<V> for JsonCodec
 where
     V: serde::Serialize + for<'a> serde::Deserialize<'a>,
@@ -20,3 +30,16 @@ where
         serde_json::from_slice(bytes)
     }
 }
+
+impl StateCodec for JsonCodec {
+    type KeyCodec = Self;
+    type ValueCodec = Self;
+
+    fn key_codec(&self) -> &Self::KeyCodec {
+        self
+    }
+
+    fn value_codec(&self) -> &Self::ValueCodec {
+        self
+    }
+}
diff --git a/module-system/sov-state/src/codec/mod.rs b/module-system/sov-state/src/codec/mod.rs
index ad86a43a30..0a350b32e0 100644
--- a/module-system/sov-state/src/codec/mod.rs
+++ b/module-system/sov-state/src/codec/mod.rs
@@ -3,8 +3,10 @@
 mod bcs_codec;
 mod borsh_codec;
 mod json_codec;
+mod split_codec;
 
 pub use bcs_codec::BcsCodec;
+use borsh::BorshSerialize;
 pub use borsh_codec::BorshCodec;
 pub use json_codec::JsonCodec;
 
@@ -42,3 +44,66 @@ pub trait StateValueCodec<V> {
             .unwrap()
     }
 }
+
+/// A trait for types that can serialize keys for storage
+/// access.
+pub trait StateKeyCodec<K> {
+    fn encode_key(&self, key: &K) -> Vec<u8>;
+}
+
+/// A trait for types that can serialize keys and values, as well
+/// as deserializing values for storage access.
+pub trait StateCodec {
+    /// The codec used to serialize keys
+    type KeyCodec;
+    /// The codec used to serialize and deserialize values
+    type ValueCodec;
+
+    /// Returns a reference to the type's key codec
+    fn key_codec(&self) -> &Self::KeyCodec;
+    /// Returns a reference to the type's value codec
+    fn value_codec(&self) -> &Self::ValueCodec;
+}
+
+/// A trait for codecs which know how to serialize a type `Ref` as if it were
+/// some other type `Target`.
+///
+/// A good example of this is [`BorshCodec`], which knows how to serialize a
+/// `[T;N]` as if it were a `Vec<T>` even though the two types have different
+/// encodings by default.
+pub trait EncodeKeyLike<Ref: ?Sized, Target> {
+    /// Encodes a reference to `Ref` as if it were a reference to `Target`.
+    fn encode_key_like(&self, borrowed: &Ref) -> Vec<u8>;
+}
+
+// All items can be encoded like themselves by all codecs
+impl<C, T> EncodeKeyLike<T, T> for C
+where
+    C: StateKeyCodec<T>,
+{
+    fn encode_key_like(&self, borrowed: &T) -> Vec<u8> {
+        self.encode_key(borrowed)
+    }
+}
+
+// In borsh, a slice is encoded the same way as a vector except in edge case where
+// T is zero-sized, in which case Vec<T> is not borsh encodable.
+impl<T> EncodeKeyLike<[T], Vec<T>> for BorshCodec
+where
+    T: BorshSerialize,
+{
+    fn encode_key_like(&self, borrowed: &[T]) -> Vec<u8> {
+        borrowed.try_to_vec().unwrap()
+    }
+}
+
+#[test]
+fn test_borsh_slice_encode_alike() {
+    let codec = BorshCodec;
+    let slice = [1, 2, 3];
+    let vec = vec![1, 2, 3];
+    assert_eq!(
+        <BorshCodec as EncodeKeyLike<[i32], Vec<i32>>>::encode_key_like(&codec, &slice),
+        codec.encode_value(&vec)
+    );
+}
diff --git a/module-system/sov-state/src/codec/split_codec.rs b/module-system/sov-state/src/codec/split_codec.rs
new file mode 100644
index 0000000000..47d409c96a
--- /dev/null
+++ b/module-system/sov-state/src/codec/split_codec.rs
@@ -0,0 +1,46 @@
+//! This module defines a codec which delegates to one codec for keys and one codec for values.
+
+use super::{StateCodec, StateKeyCodec, StateValueCodec};
+
+/// A [`StateValueCodec`] that uses one pre-existing codec for keys and a different one values.
+#[derive(Debug, Default, PartialEq, Eq, Clone)]
+pub struct SplitCodec<KC, VC> {
+    pub key_codec: KC,
+    pub value_codec: VC,
+}
+
+impl<K, KC, VC> StateKeyCodec<K> for SplitCodec<KC, VC>
+where
+    KC: StateKeyCodec<K>,
+{
+    fn encode_key(&self, key: &K) -> Vec<u8> {
+        self.key_codec.encode_key(key)
+    }
+}
+
+impl<V, KC, VC> StateValueCodec<V> for SplitCodec<KC, VC>
+where
+    VC: StateValueCodec<V>,
+{
+    type Error = VC::Error;
+
+    fn encode_value(&self, value: &V) -> Vec<u8> {
+        self.value_codec.encode_value(value)
+    }
+
+    fn try_decode_value(&self, bytes: &[u8]) -> Result<V, Self::Error> {
+        self.value_codec.try_decode_value(bytes)
+    }
+}
+
+impl<KC, VC> StateCodec for SplitCodec<KC, VC> {
+    type KeyCodec = KC;
+    type ValueCodec = VC;
+    fn key_codec(&self) -> &Self::KeyCodec {
+        &self.key_codec
+    }
+
+    fn value_codec(&self) -> &Self::ValueCodec {
+        &self.value_codec
+    }
+}
diff --git a/module-system/sov-state/src/containers/accessory_map.rs b/module-system/sov-state/src/containers/accessory_map.rs
index 0f13fb023f..e148ce651e 100644
--- a/module-system/sov-state/src/containers/accessory_map.rs
+++ b/module-system/sov-state/src/containers/accessory_map.rs
@@ -1,9 +1,7 @@
-use std::borrow::Borrow;
-use std::hash::Hash;
 use std::marker::PhantomData;
 
 use super::StateMapError;
-use crate::codec::{BorshCodec, StateValueCodec};
+use crate::codec::{BorshCodec, EncodeKeyLike, StateCodec, StateKeyCodec, StateValueCodec};
 use crate::storage::StorageKey;
 use crate::{AccessoryWorkingSet, Prefix, StateReaderAndWriter, Storage};
 
@@ -14,7 +12,7 @@ use crate::{AccessoryWorkingSet, Prefix, StateReaderAndWriter, Storage};
 /// [`AccessoryStateMap`] is generic over:
 /// - a key type `K`;
 /// - a value type `V`;
-/// - a [`StateValueCodec`] `VC`.
+/// - a [`StateValueCodec`] `Codec`.
 #[derive(
     Debug,
     Clone,
@@ -24,9 +22,9 @@ use crate::{AccessoryWorkingSet, Prefix, StateReaderAndWriter, Storage};
     serde::Serialize,
     serde::Deserialize,
 )]
-pub struct AccessoryStateMap<K, V, VC = BorshCodec> {
+pub struct AccessoryStateMap<K, V, Codec = BorshCodec> {
     _phantom: (PhantomData<K>, PhantomData<V>),
-    value_codec: VC,
+    codec: Codec,
     prefix: Prefix,
 }
 
@@ -38,12 +36,12 @@ impl<K, V> AccessoryStateMap<K, V> {
     }
 }
 
-impl<K, V, VC> AccessoryStateMap<K, V, VC> {
+impl<K, V, Codec> AccessoryStateMap<K, V, Codec> {
     /// Creates a new [`AccessoryStateMap`] with the given prefix and [`StateValueCodec`].
-    pub fn with_codec(prefix: Prefix, codec: VC) -> Self {
+    pub fn with_codec(prefix: Prefix, codec: Codec) -> Self {
         Self {
             _phantom: (PhantomData, PhantomData),
-            value_codec: codec,
+            codec,
             prefix,
         }
     }
@@ -54,10 +52,11 @@ impl<K, V, VC> AccessoryStateMap<K, V, VC> {
     }
 }
 
-impl<K, V, VC> AccessoryStateMap<K, V, VC>
+impl<K, V, Codec> AccessoryStateMap<K, V, Codec>
 where
-    K: Hash + Eq,
-    VC: StateValueCodec<V>,
+    Codec: StateCodec,
+    Codec::KeyCodec: StateKeyCodec<K>,
+    Codec::ValueCodec: StateValueCodec<V>,
 {
     /// Inserts a key-value pair into the map.
     ///
@@ -65,10 +64,10 @@ where
     /// map’s key type.
     pub fn set<Q, S: Storage>(&self, key: &Q, value: &V, working_set: &mut AccessoryWorkingSet<S>)
     where
-        K: Borrow<Q>,
-        Q: Hash + Eq + ?Sized,
+        Codec::KeyCodec: EncodeKeyLike<Q, K>,
+        Q: ?Sized,
     {
-        working_set.set_value(self.prefix(), key, value, &self.value_codec)
+        working_set.set_value(self.prefix(), key, value, &self.codec)
     }
 
     /// Returns the value corresponding to the key, or [`None`] if the map
@@ -76,9 +75,8 @@ where
     ///
     /// # Examples
     ///
-    /// The key may be any borrowed form of the map’s key type. Note that
-    /// [`Hash`] and [`Eq`] on the borrowed form must match those for the key
-    /// type.
+    /// The key may be any item that implements [`EncodeKeyLike`] the map's key type
+    /// using your chosen codec.
     ///
     /// ```
     /// use sov_state::{AccessoryStateMap, Storage, AccessoryWorkingSet};
@@ -93,7 +91,7 @@ where
     /// }
     /// ```
     ///
-    /// If the map's key type does not implement [`Borrow`] for your desired
+    /// If the map's key type does not implement [`EncodeKeyLike`] for your desired
     /// target type, you'll have to convert the key to something else. An
     /// example of this would be "slicing" an array to use in [`Vec`]-keyed
     /// maps:
@@ -110,10 +108,10 @@ where
     /// ```
     pub fn get<Q, S: Storage>(&self, key: &Q, working_set: &mut AccessoryWorkingSet<S>) -> Option<V>
     where
-        K: Borrow<Q>,
-        Q: Hash + Eq + ?Sized,
+        Codec::KeyCodec: EncodeKeyLike<Q, K>,
+        Q: ?Sized,
     {
-        working_set.get_value(self.prefix(), key, &self.value_codec)
+        working_set.get_value(self.prefix(), key, &self.codec)
     }
 
     /// Returns the value corresponding to the key or [`StateMapError`] if key is absent in
@@ -124,11 +122,14 @@ where
         working_set: &mut AccessoryWorkingSet<S>,
     ) -> Result<V, StateMapError>
     where
-        K: Borrow<Q>,
-        Q: Hash + Eq + ?Sized,
+        Codec::KeyCodec: EncodeKeyLike<Q, K>,
+        Q: ?Sized,
     {
         self.get(key, working_set).ok_or_else(|| {
-            StateMapError::MissingValue(self.prefix().clone(), StorageKey::new(self.prefix(), key))
+            StateMapError::MissingValue(
+                self.prefix().clone(),
+                StorageKey::new(self.prefix(), key, self.codec.key_codec()),
+            )
         })
     }
 
@@ -140,10 +141,10 @@ where
         working_set: &mut AccessoryWorkingSet<S>,
     ) -> Option<V>
     where
-        K: Borrow<Q>,
-        Q: Hash + Eq + ?Sized,
+        Codec::KeyCodec: EncodeKeyLike<Q, K>,
+        Q: ?Sized,
     {
-        working_set.remove_value(self.prefix(), key, &self.value_codec)
+        working_set.remove_value(self.prefix(), key, &self.codec)
     }
 
     /// Removes a key from the map, returning the corresponding value (or
@@ -156,11 +157,14 @@ where
         working_set: &mut AccessoryWorkingSet<S>,
     ) -> Result<V, StateMapError>
     where
-        K: Borrow<Q>,
-        Q: Hash + Eq + ?Sized,
+        Codec::KeyCodec: EncodeKeyLike<Q, K>,
+        Q: ?Sized,
     {
         self.remove(key, working_set).ok_or_else(|| {
-            StateMapError::MissingValue(self.prefix().clone(), StorageKey::new(self.prefix(), key))
+            StateMapError::MissingValue(
+                self.prefix().clone(),
+                StorageKey::new(self.prefix(), key, self.codec.key_codec()),
+            )
         })
     }
 
@@ -170,19 +174,21 @@ where
     /// return the value beforing deletion.
     pub fn delete<Q, S: Storage>(&self, key: &Q, working_set: &mut AccessoryWorkingSet<S>)
     where
-        K: Borrow<Q>,
-        Q: Hash + Eq + ?Sized,
+        Codec::KeyCodec: EncodeKeyLike<Q, K>,
+        Q: ?Sized,
     {
-        working_set.delete_value(self.prefix(), key);
+        working_set.delete_value(self.prefix(), key, &self.codec);
     }
 }
 
 #[cfg(feature = "arbitrary")]
-impl<'a, K, V, VC> AccessoryStateMap<K, V, VC>
+impl<'a, K, V, Codec> AccessoryStateMap<K, V, Codec>
 where
-    K: arbitrary::Arbitrary<'a> + Hash + Eq,
-    V: arbitrary::Arbitrary<'a> + Hash + Eq,
-    VC: StateValueCodec<V> + Default,
+    K: arbitrary::Arbitrary<'a>,
+    V: arbitrary::Arbitrary<'a>,
+    Codec: StateCodec + Default,
+    Codec::KeyCodec: StateKeyCodec<K>,
+    Codec::ValueCodec: StateValueCodec<V>,
 {
     pub fn arbitrary_workset<S>(
         u: &mut arbitrary::Unstructured<'a>,
@@ -195,8 +201,8 @@ where
 
         let prefix = Prefix::arbitrary(u)?;
         let len = u.arbitrary_len::<(K, V)>()?;
-        let codec = VC::default();
-        let map = AccessoryStateMap::with_codec(prefix, codec);
+        let codec = Codec::default();
+        let map = Self::with_codec(prefix, codec);
 
         (0..len).try_fold(map, |map, _| {
             let key = K::arbitrary(u)?;
diff --git a/module-system/sov-state/src/containers/accessory_value.rs b/module-system/sov-state/src/containers/accessory_value.rs
index 2c95a82007..dda3641e8c 100644
--- a/module-system/sov-state/src/containers/accessory_value.rs
+++ b/module-system/sov-state/src/containers/accessory_value.rs
@@ -3,7 +3,7 @@ use std::marker::PhantomData;
 use borsh::{BorshDeserialize, BorshSerialize};
 use thiserror::Error;
 
-use crate::codec::{BorshCodec, StateValueCodec};
+use crate::codec::{BorshCodec, StateCodec, StateValueCodec};
 use crate::{AccessoryWorkingSet, Prefix, StateReaderAndWriter, Storage};
 
 /// Container for a single value stored as "accessory" state, outside of the
@@ -18,9 +18,9 @@ use crate::{AccessoryWorkingSet, Prefix, StateReaderAndWriter, Storage};
     serde::Serialize,
     serde::Deserialize,
 )]
-pub struct AccessoryStateValue<V, VC = BorshCodec> {
+pub struct AccessoryStateValue<V, Codec = BorshCodec> {
     _phantom: PhantomData<V>,
-    codec: VC,
+    codec: Codec,
     prefix: Prefix,
 }
 
@@ -39,9 +39,9 @@ impl<V> AccessoryStateValue<V> {
     }
 }
 
-impl<V, VC> AccessoryStateValue<V, VC> {
+impl<V, Codec> AccessoryStateValue<V, Codec> {
     /// Creates a new [`AccessoryStateValue`] with the given prefix and codec.
-    pub fn with_codec(prefix: Prefix, codec: VC) -> Self {
+    pub fn with_codec(prefix: Prefix, codec: Codec) -> Self {
         Self {
             _phantom: PhantomData,
             codec,
@@ -55,18 +55,19 @@ impl<V, VC> AccessoryStateValue<V, VC> {
     }
 }
 
-impl<V, VC> AccessoryStateValue<V, VC>
+impl<V, Codec> AccessoryStateValue<V, Codec>
 where
-    VC: StateValueCodec<V>,
+    Codec: StateCodec,
+    Codec::ValueCodec: StateValueCodec<V>,
 {
     /// Sets a value in the AccessoryStateValue.
     pub fn set<S: Storage>(&self, value: &V, working_set: &mut AccessoryWorkingSet<S>) {
-        working_set.set_value(self.prefix(), &SingletonKey, value, &self.codec)
+        working_set.set_singleton(self.prefix(), value, &self.codec)
     }
 
     /// Gets a value from the AccessoryStateValue or None if the value is absent.
     pub fn get<S: Storage>(&self, working_set: &mut AccessoryWorkingSet<S>) -> Option<V> {
-        working_set.get_value(self.prefix(), &SingletonKey, &self.codec)
+        working_set.get_singleton(self.prefix(), &self.codec)
     }
 
     /// Gets a value from the AccessoryStateValue or Error if the value is absent.
@@ -80,7 +81,7 @@ where
 
     /// Removes a value from the AccessoryStateValue, returning the value (or None if the key is absent).
     pub fn remove<S: Storage>(&self, working_set: &mut AccessoryWorkingSet<S>) -> Option<V> {
-        working_set.remove_value(self.prefix(), &SingletonKey, &self.codec)
+        working_set.remove_singleton(self.prefix(), &self.codec)
     }
 
     /// Removes a value and from the AccessoryStateValue, returning the value (or Error if the key is absent).
@@ -94,10 +95,6 @@ where
 
     /// Deletes a value from the AccessoryStateValue.
     pub fn delete<S: Storage>(&self, working_set: &mut AccessoryWorkingSet<S>) {
-        working_set.delete_value(self.prefix(), &SingletonKey);
+        working_set.delete_singleton(self.prefix());
     }
 }
-
-// SingletonKey is very similar to the unit type `()` i.e. it has only one value.
-#[derive(Debug, PartialEq, Eq, Hash)]
-struct SingletonKey;
diff --git a/module-system/sov-state/src/containers/accessory_vec.rs b/module-system/sov-state/src/containers/accessory_vec.rs
index b7ff94fd3f..3028d83b73 100644
--- a/module-system/sov-state/src/containers/accessory_vec.rs
+++ b/module-system/sov-state/src/containers/accessory_vec.rs
@@ -1,7 +1,7 @@
 use std::iter::FusedIterator;
 use std::marker::PhantomData;
 
-use crate::codec::{BorshCodec, StateValueCodec};
+use crate::codec::{BorshCodec, StateCodec, StateKeyCodec, StateValueCodec};
 use crate::{
     AccessoryStateMap, AccessoryStateValue, AccessoryWorkingSet, Prefix, StateVecError, Storage,
 };
@@ -15,19 +15,18 @@ use crate::{
     serde::Serialize,
     serde::Deserialize,
 )]
-pub struct AccessoryStateVec<V, VC = BorshCodec>
-where
-    VC: StateValueCodec<V>,
-{
+pub struct AccessoryStateVec<V, Codec = BorshCodec> {
     _phantom: PhantomData<V>,
     prefix: Prefix,
-    len_value: AccessoryStateValue<usize, VC>,
-    elems: AccessoryStateMap<usize, V, VC>,
+    len_value: AccessoryStateValue<usize, Codec>,
+    elems: AccessoryStateMap<usize, V, Codec>,
 }
 
 impl<V> AccessoryStateVec<V>
 where
-    BorshCodec: StateValueCodec<V>,
+    BorshCodec: StateCodec + Clone,
+    <BorshCodec as StateCodec>::ValueCodec: StateValueCodec<V> + StateValueCodec<usize>,
+    <BorshCodec as StateCodec>::KeyCodec: StateKeyCodec<usize>,
 {
     /// Crates a new [`AccessoryStateVec`] with the given prefix and the default
     /// [`StateValueCodec`] (i.e. [`BorshCodec`]).
@@ -36,18 +35,20 @@ where
     }
 }
 
-impl<V, VC> AccessoryStateVec<V, VC>
+impl<V, Codec> AccessoryStateVec<V, Codec>
 where
-    VC: StateValueCodec<V> + StateValueCodec<usize> + Clone,
+    Codec: StateCodec + Clone,
+    Codec::ValueCodec: StateValueCodec<V> + StateValueCodec<usize>,
+    Codec::KeyCodec: StateKeyCodec<usize>,
 {
     /// Creates a new [`AccessoryStateVec`] with the given prefix and codec.
-    pub fn with_codec(prefix: Prefix, codec: VC) -> Self {
+    pub fn with_codec(prefix: Prefix, codec: Codec) -> Self {
         // Differentiating the prefixes for the length and the elements
         // shouldn't be necessary, but it's best not to rely on implementation
         // details of `StateValue` and `StateMap` as they both have the right to
         // reserve the whole key space for themselves.
         let len_value =
-            AccessoryStateValue::<usize, VC>::with_codec(prefix.extended(b"l"), codec.clone());
+            AccessoryStateValue::<usize, Codec>::with_codec(prefix.extended(b"l"), codec.clone());
         let elems = AccessoryStateMap::with_codec(prefix.extended(b"e"), codec);
         Self {
             _phantom: PhantomData,
@@ -169,7 +170,7 @@ where
     pub fn iter<'a, 'ws, S: Storage>(
         &'a self,
         working_set: &'ws mut AccessoryWorkingSet<'ws, S>,
-    ) -> AccessoryStateVecIter<'a, 'ws, V, VC, S> {
+    ) -> AccessoryStateVecIter<'a, 'ws, V, Codec, S> {
         let len = self.len(working_set);
         AccessoryStateVecIter {
             state_vec: self,
@@ -183,20 +184,24 @@ where
 /// An [`Iterator`] over a [`AccessoryStateVec`]
 ///
 /// See [`AccessoryStateVec::iter`] for more details.
-pub struct AccessoryStateVecIter<'a, 'ws, V, VC, S>
+pub struct AccessoryStateVecIter<'a, 'ws, V, Codec, S>
 where
-    VC: StateValueCodec<V>,
+    Codec: StateCodec + Clone,
+    Codec::ValueCodec: StateValueCodec<V> + StateValueCodec<usize>,
+    Codec::KeyCodec: StateKeyCodec<usize>,
     S: Storage,
 {
-    state_vec: &'a AccessoryStateVec<V, VC>,
+    state_vec: &'a AccessoryStateVec<V, Codec>,
     ws: &'ws mut AccessoryWorkingSet<'ws, S>,
     len: usize,
     next_i: usize,
 }
 
-impl<'a, 'ws, V, VC, S> Iterator for AccessoryStateVecIter<'a, 'ws, V, VC, S>
+impl<'a, 'ws, V, Codec, S> Iterator for AccessoryStateVecIter<'a, 'ws, V, Codec, S>
 where
-    VC: StateValueCodec<V> + StateValueCodec<usize> + Clone,
+    Codec: StateCodec + Clone,
+    Codec::ValueCodec: StateValueCodec<V> + StateValueCodec<usize>,
+    Codec::KeyCodec: StateKeyCodec<usize>,
     S: Storage,
 {
     type Item = V;
@@ -211,9 +216,11 @@ where
     }
 }
 
-impl<'a, 'ws, V, VC, S> ExactSizeIterator for AccessoryStateVecIter<'a, 'ws, V, VC, S>
+impl<'a, 'ws, V, Codec, S> ExactSizeIterator for AccessoryStateVecIter<'a, 'ws, V, Codec, S>
 where
-    VC: StateValueCodec<V> + StateValueCodec<usize> + Clone,
+    Codec: StateCodec + Clone,
+    Codec::ValueCodec: StateValueCodec<V> + StateValueCodec<usize>,
+    Codec::KeyCodec: StateKeyCodec<usize>,
     S: Storage,
 {
     fn len(&self) -> usize {
@@ -221,9 +228,11 @@ where
     }
 }
 
-impl<'a, 'ws, V, VC, S> FusedIterator for AccessoryStateVecIter<'a, 'ws, V, VC, S>
+impl<'a, 'ws, V, Codec, S> FusedIterator for AccessoryStateVecIter<'a, 'ws, V, Codec, S>
 where
-    VC: StateValueCodec<V> + StateValueCodec<usize> + Clone,
+    Codec: StateCodec + Clone,
+    Codec::ValueCodec: StateValueCodec<V> + StateValueCodec<usize>,
+    Codec::KeyCodec: StateKeyCodec<usize>,
     S: Storage,
 {
 }
diff --git a/module-system/sov-state/src/containers/map.rs b/module-system/sov-state/src/containers/map.rs
index f596b293d1..9eccfdbb40 100644
--- a/module-system/sov-state/src/containers/map.rs
+++ b/module-system/sov-state/src/containers/map.rs
@@ -1,10 +1,8 @@
-use std::borrow::Borrow;
-use std::hash::Hash;
 use std::marker::PhantomData;
 
 use thiserror::Error;
 
-use crate::codec::{BorshCodec, StateValueCodec};
+use crate::codec::{BorshCodec, EncodeKeyLike, StateCodec, StateKeyCodec, StateValueCodec};
 use crate::storage::StorageKey;
 use crate::{Prefix, StateReaderAndWriter, Storage, WorkingSet};
 
@@ -14,7 +12,7 @@ use crate::{Prefix, StateReaderAndWriter, Storage, WorkingSet};
 /// [`StateMap`] is generic over:
 /// - a key type `K`;
 /// - a value type `V`;
-/// - a [`StateValueCodec`] `VC`.
+/// - a [`StateValueCodec`] `Codec`.
 #[derive(
     Debug,
     Clone,
@@ -24,9 +22,9 @@ use crate::{Prefix, StateReaderAndWriter, Storage, WorkingSet};
     serde::Serialize,
     serde::Deserialize,
 )]
-pub struct StateMap<K, V, VC = BorshCodec> {
+pub struct StateMap<K, V, Codec = BorshCodec> {
     _phantom: (PhantomData<K>, PhantomData<V>),
-    value_codec: VC,
+    codec: Codec,
     prefix: Prefix,
 }
 
@@ -45,26 +43,31 @@ impl<K, V> StateMap<K, V> {
     }
 }
 
-impl<K, V, VC> StateMap<K, V, VC> {
+impl<K, V, Codec> StateMap<K, V, Codec> {
     /// Creates a new [`StateMap`] with the given prefix and [`StateValueCodec`].
-    pub fn with_codec(prefix: Prefix, codec: VC) -> Self {
+    pub fn with_codec(prefix: Prefix, codec: Codec) -> Self {
         Self {
             _phantom: (PhantomData, PhantomData),
-            value_codec: codec,
+            codec,
             prefix,
         }
     }
 
+    pub fn codec(&self) -> &Codec {
+        &self.codec
+    }
+
     /// Returns the prefix used when this [`StateMap`] was created.
     pub fn prefix(&self) -> &Prefix {
         &self.prefix
     }
 }
 
-impl<K, V, VC> StateMap<K, V, VC>
+impl<K, V, Codec> StateMap<K, V, Codec>
 where
-    K: Hash + Eq,
-    VC: StateValueCodec<V>,
+    Codec: StateCodec,
+    Codec::KeyCodec: StateKeyCodec<K>,
+    Codec::ValueCodec: StateValueCodec<V>,
 {
     /// Inserts a key-value pair into the map.
     ///
@@ -72,10 +75,10 @@ where
     /// map’s key type.
     pub fn set<Q, S: Storage>(&self, key: &Q, value: &V, working_set: &mut WorkingSet<S>)
     where
-        K: Borrow<Q>,
-        Q: Hash + Eq + ?Sized,
+        Codec::KeyCodec: EncodeKeyLike<Q, K>,
+        Q: ?Sized,
     {
-        working_set.set_value(self.prefix(), key, value, &self.value_codec)
+        working_set.set_value(self.prefix(), key, value, &self.codec)
     }
 
     /// Returns the value corresponding to the key, or [`None`] if the map
@@ -83,9 +86,8 @@ where
     ///
     /// # Examples
     ///
-    /// The key may be any borrowed form of the map’s key type. Note that
-    /// [`Hash`] and [`Eq`] on the borrowed form must match those for the key
-    /// type.
+    /// The key may be any item that implements [`EncodeKeyLike`] the map's key type
+    /// using your chosen codec.
     ///
     /// ```
     /// use sov_state::{StateMap, Storage, WorkingSet};
@@ -100,7 +102,7 @@ where
     /// }
     /// ```
     ///
-    /// If the map's key type does not implement [`Borrow`] for your desired
+    /// If the map's key type does not implement [`EncodeKeyLike`] for your desired
     /// target type, you'll have to convert the key to something else. An
     /// example of this would be "slicing" an array to use in [`Vec`]-keyed
     /// maps:
@@ -117,10 +119,12 @@ where
     /// ```
     pub fn get<Q, S: Storage>(&self, key: &Q, working_set: &mut WorkingSet<S>) -> Option<V>
     where
-        K: Borrow<Q>,
-        Q: Hash + Eq + ?Sized,
+        Codec: StateCodec,
+        Codec::KeyCodec: EncodeKeyLike<Q, K>,
+        Codec::ValueCodec: StateValueCodec<V>,
+        Q: ?Sized,
     {
-        working_set.get_value(self.prefix(), key, &self.value_codec)
+        working_set.get_value(self.prefix(), key, &self.codec)
     }
 
     /// Returns the value corresponding to the key or [`StateMapError`] if key is absent in
@@ -131,11 +135,16 @@ where
         working_set: &mut WorkingSet<S>,
     ) -> Result<V, StateMapError>
     where
-        K: Borrow<Q>,
-        Q: Hash + Eq + ?Sized,
+        Codec: StateCodec,
+        Codec::KeyCodec: EncodeKeyLike<Q, K>,
+        Codec::ValueCodec: StateValueCodec<V>,
+        Q: ?Sized,
     {
         self.get(key, working_set).ok_or_else(|| {
-            StateMapError::MissingValue(self.prefix().clone(), StorageKey::new(self.prefix(), key))
+            StateMapError::MissingValue(
+                self.prefix().clone(),
+                StorageKey::new(self.prefix(), key, self.codec.key_codec()),
+            )
         })
     }
 
@@ -143,10 +152,12 @@ where
     /// [`None`] if the key is absent).
     pub fn remove<Q, S: Storage>(&self, key: &Q, working_set: &mut WorkingSet<S>) -> Option<V>
     where
-        K: Borrow<Q>,
-        Q: Hash + Eq + ?Sized,
+        Codec: StateCodec,
+        Codec::KeyCodec: EncodeKeyLike<Q, K>,
+        Codec::ValueCodec: StateValueCodec<V>,
+        Q: ?Sized,
     {
-        working_set.remove_value(self.prefix(), key, &self.value_codec)
+        working_set.remove_value(self.prefix(), key, &self.codec)
     }
 
     /// Removes a key from the map, returning the corresponding value (or
@@ -159,11 +170,16 @@ where
         working_set: &mut WorkingSet<S>,
     ) -> Result<V, StateMapError>
     where
-        K: Borrow<Q>,
-        Q: Hash + Eq + ?Sized,
+        Codec: StateCodec,
+        Codec::KeyCodec: EncodeKeyLike<Q, K>,
+        Codec::ValueCodec: StateValueCodec<V>,
+        Q: ?Sized,
     {
         self.remove(key, working_set).ok_or_else(|| {
-            StateMapError::MissingValue(self.prefix().clone(), StorageKey::new(self.prefix(), key))
+            StateMapError::MissingValue(
+                self.prefix().clone(),
+                StorageKey::new(self.prefix(), key, self.codec.key_codec()),
+            )
         })
     }
 
@@ -173,19 +189,22 @@ where
     /// return the value beforing deletion.
     pub fn delete<Q, S: Storage>(&self, key: &Q, working_set: &mut WorkingSet<S>)
     where
-        K: Borrow<Q>,
-        Q: Hash + Eq + ?Sized,
+        Codec: StateCodec,
+        Codec::KeyCodec: EncodeKeyLike<Q, K>,
+        Q: ?Sized,
     {
-        working_set.delete_value(self.prefix(), key);
+        working_set.delete_value(self.prefix(), key, &self.codec);
     }
 }
 
 #[cfg(feature = "arbitrary")]
-impl<'a, K, V, VC> StateMap<K, V, VC>
+impl<'a, K, V, Codec> StateMap<K, V, Codec>
 where
-    K: arbitrary::Arbitrary<'a> + Hash + Eq,
-    V: arbitrary::Arbitrary<'a> + Hash + Eq,
-    VC: StateValueCodec<V> + Default,
+    K: arbitrary::Arbitrary<'a>,
+    V: arbitrary::Arbitrary<'a>,
+    Codec: StateCodec + Default,
+    Codec::KeyCodec: StateKeyCodec<K>,
+    Codec::ValueCodec: StateValueCodec<V>,
 {
     pub fn arbitrary_workset<S>(
         u: &mut arbitrary::Unstructured<'a>,
@@ -198,7 +217,7 @@ where
 
         let prefix = Prefix::arbitrary(u)?;
         let len = u.arbitrary_len::<(K, V)>()?;
-        let codec = VC::default();
+        let codec = Codec::default();
         let map = StateMap::with_codec(prefix, codec);
 
         (0..len).try_fold(map, |map, _| {
diff --git a/module-system/sov-state/src/containers/value.rs b/module-system/sov-state/src/containers/value.rs
index c4d0c26abd..388c4e7dc7 100644
--- a/module-system/sov-state/src/containers/value.rs
+++ b/module-system/sov-state/src/containers/value.rs
@@ -2,7 +2,7 @@ use std::marker::PhantomData;
 
 use thiserror::Error;
 
-use crate::codec::{BorshCodec, StateValueCodec};
+use crate::codec::{BorshCodec, StateCodec, StateValueCodec};
 use crate::{Prefix, StateReaderAndWriter, Storage, WorkingSet};
 
 /// Container for a single value.
@@ -15,9 +15,9 @@ use crate::{Prefix, StateReaderAndWriter, Storage, WorkingSet};
     serde::Serialize,
     serde::Deserialize,
 )]
-pub struct StateValue<V, VC = BorshCodec> {
+pub struct StateValue<V, Codec = BorshCodec> {
     _phantom: PhantomData<V>,
-    codec: VC,
+    codec: Codec,
     prefix: Prefix,
 }
 
@@ -36,9 +36,9 @@ impl<V> StateValue<V> {
     }
 }
 
-impl<V, VC> StateValue<V, VC> {
+impl<V, Codec> StateValue<V, Codec> {
     /// Creates a new [`StateValue`] with the given prefix and codec.
-    pub fn with_codec(prefix: Prefix, codec: VC) -> Self {
+    pub fn with_codec(prefix: Prefix, codec: Codec) -> Self {
         Self {
             _phantom: PhantomData,
             codec,
@@ -52,18 +52,19 @@ impl<V, VC> StateValue<V, VC> {
     }
 }
 
-impl<V, VC> StateValue<V, VC>
+impl<V, Codec> StateValue<V, Codec>
 where
-    VC: StateValueCodec<V>,
+    Codec: StateCodec,
+    Codec::ValueCodec: StateValueCodec<V>,
 {
     /// Sets a value in the StateValue.
     pub fn set<S: Storage>(&self, value: &V, working_set: &mut WorkingSet<S>) {
-        working_set.set_value(self.prefix(), &SingletonKey, value, &self.codec)
+        working_set.set_singleton(self.prefix(), value, &self.codec)
     }
 
     /// Gets a value from the StateValue or None if the value is absent.
     pub fn get<S: Storage>(&self, working_set: &mut WorkingSet<S>) -> Option<V> {
-        working_set.get_value(self.prefix(), &SingletonKey, &self.codec)
+        working_set.get_singleton(self.prefix(), &self.codec)
     }
 
     /// Gets a value from the StateValue or Error if the value is absent.
@@ -74,7 +75,7 @@ where
 
     /// Removes a value from the StateValue, returning the value (or None if the key is absent).
     pub fn remove<S: Storage>(&self, working_set: &mut WorkingSet<S>) -> Option<V> {
-        working_set.remove_value(self.prefix(), &SingletonKey, &self.codec)
+        working_set.remove_singleton(self.prefix(), &self.codec)
     }
 
     /// Removes a value and from the StateValue, returning the value (or Error if the key is absent).
@@ -85,10 +86,6 @@ where
 
     /// Deletes a value from the StateValue.
     pub fn delete<S: Storage>(&self, working_set: &mut WorkingSet<S>) {
-        working_set.delete_value(self.prefix(), &SingletonKey);
+        working_set.delete_singleton(self.prefix());
     }
 }
-
-// SingletonKey is very similar to the unit type `()` i.e. it has only one value.
-#[derive(Debug, PartialEq, Eq, Hash)]
-struct SingletonKey;
diff --git a/module-system/sov-state/src/containers/vec.rs b/module-system/sov-state/src/containers/vec.rs
index 1d0b3b8c92..1873064e3e 100644
--- a/module-system/sov-state/src/containers/vec.rs
+++ b/module-system/sov-state/src/containers/vec.rs
@@ -3,7 +3,7 @@ use std::marker::PhantomData;
 
 use thiserror::Error;
 
-use crate::codec::{BorshCodec, StateValueCodec};
+use crate::codec::{BorshCodec, StateCodec, StateKeyCodec, StateValueCodec};
 use crate::{Prefix, StateMap, StateValue, Storage, WorkingSet};
 
 #[derive(
@@ -15,14 +15,11 @@ use crate::{Prefix, StateMap, StateValue, Storage, WorkingSet};
     serde::Serialize,
     serde::Deserialize,
 )]
-pub struct StateVec<V, VC = BorshCodec>
-where
-    VC: StateValueCodec<V>,
-{
+pub struct StateVec<V, Codec = BorshCodec> {
     _phantom: PhantomData<V>,
     prefix: Prefix,
-    len_value: StateValue<usize, VC>,
-    elems: StateMap<usize, V, VC>,
+    len_value: StateValue<usize, Codec>,
+    elems: StateMap<usize, V, Codec>,
 }
 
 /// Error type for `StateVec` get method.
@@ -45,12 +42,14 @@ where
     }
 }
 
-impl<V, VC> StateVec<V, VC>
+impl<V, Codec> StateVec<V, Codec>
 where
-    VC: StateValueCodec<V> + StateValueCodec<usize> + Clone,
+    Codec: StateCodec + Clone,
+    Codec::ValueCodec: StateValueCodec<V> + StateValueCodec<usize>,
+    Codec::KeyCodec: StateKeyCodec<usize>,
 {
     /// Creates a new [`StateVec`] with the given prefix and codec.
-    pub fn with_codec(prefix: Prefix, codec: VC) -> Self {
+    pub fn with_codec(prefix: Prefix, codec: Codec) -> Self {
         // Differentiating the prefixes for the length and the elements
         // shouldn't be necessary, but it's best not to rely on implementation
         // details of `StateValue` and `StateMap` as they both have the right to
@@ -173,7 +172,7 @@ where
     pub fn iter<'a, 'ws, S: Storage>(
         &'a self,
         working_set: &'ws mut WorkingSet<S>,
-    ) -> StateVecIter<'a, 'ws, V, VC, S> {
+    ) -> StateVecIter<'a, 'ws, V, Codec, S> {
         let len = self.len(working_set);
         StateVecIter {
             state_vec: self,
@@ -187,20 +186,24 @@ where
 /// An [`Iterator`] over a [`StateVec`]
 ///
 /// See [`StateVec::iter`] for more details.
-pub struct StateVecIter<'a, 'ws, V, VC, S>
+pub struct StateVecIter<'a, 'ws, V, Codec, S>
 where
-    VC: StateValueCodec<V>,
+    Codec: StateCodec + Clone,
+    Codec::ValueCodec: StateValueCodec<V> + StateValueCodec<usize>,
+    Codec::KeyCodec: StateKeyCodec<usize>,
     S: Storage,
 {
-    state_vec: &'a StateVec<V, VC>,
+    state_vec: &'a StateVec<V, Codec>,
     ws: &'ws mut WorkingSet<S>,
     len: usize,
     next_i: usize,
 }
 
-impl<'a, 'ws, V, VC, S> Iterator for StateVecIter<'a, 'ws, V, VC, S>
+impl<'a, 'ws, V, Codec, S> Iterator for StateVecIter<'a, 'ws, V, Codec, S>
 where
-    VC: StateValueCodec<V> + StateValueCodec<usize> + Clone,
+    Codec: StateCodec + Clone,
+    Codec::ValueCodec: StateValueCodec<V> + StateValueCodec<usize>,
+    Codec::KeyCodec: StateKeyCodec<usize>,
     S: Storage,
 {
     type Item = V;
@@ -215,9 +218,11 @@ where
     }
 }
 
-impl<'a, 'ws, V, VC, S> ExactSizeIterator for StateVecIter<'a, 'ws, V, VC, S>
+impl<'a, 'ws, V, Codec, S> ExactSizeIterator for StateVecIter<'a, 'ws, V, Codec, S>
 where
-    VC: StateValueCodec<V> + StateValueCodec<usize> + Clone,
+    Codec: StateCodec + Clone,
+    Codec::ValueCodec: StateValueCodec<V> + StateValueCodec<usize>,
+    Codec::KeyCodec: StateKeyCodec<usize>,
     S: Storage,
 {
     fn len(&self) -> usize {
@@ -225,9 +230,11 @@ where
     }
 }
 
-impl<'a, 'ws, V, VC, S> FusedIterator for StateVecIter<'a, 'ws, V, VC, S>
+impl<'a, 'ws, V, Codec, S> FusedIterator for StateVecIter<'a, 'ws, V, Codec, S>
 where
-    VC: StateValueCodec<V> + StateValueCodec<usize> + Clone,
+    Codec: StateCodec + Clone,
+    Codec::ValueCodec: StateValueCodec<V> + StateValueCodec<usize>,
+    Codec::KeyCodec: StateKeyCodec<usize>,
     S: Storage,
 {
 }
@@ -302,7 +309,7 @@ mod test {
         ws: &mut WorkingSet<S>,
     ) where
         S: Storage,
-        BorshCodec: StateValueCodec<T> + StateValueCodec<usize>,
+        BorshCodec: StateValueCodec<T>,
         T: Eq + Debug,
     {
         match action {
diff --git a/module-system/sov-state/src/scratchpad.rs b/module-system/sov-state/src/scratchpad.rs
index f0174f59c6..558b025a92 100644
--- a/module-system/sov-state/src/scratchpad.rs
+++ b/module-system/sov-state/src/scratchpad.rs
@@ -1,11 +1,10 @@
 use std::collections::HashMap;
 use std::fmt::Debug;
-use std::hash::Hash;
 
 use sov_first_read_last_write_cache::{CacheKey, CacheValue};
 use sov_rollup_interface::stf::Event;
 
-use crate::codec::StateValueCodec;
+use crate::codec::{EncodeKeyLike, StateCodec, StateValueCodec};
 use crate::internal_cache::{OrderedReadsAndWrites, StorageInternalCache};
 use crate::storage::{StorageKey, StorageValue};
 use crate::{Prefix, Storage};
@@ -310,50 +309,113 @@ pub(crate) trait StateReaderAndWriter {
 
     fn delete(&mut self, key: &StorageKey);
 
-    fn set_value<K, V, VC>(&mut self, prefix: &Prefix, storage_key: &K, value: &V, codec: &VC)
+    fn set_value<Q, K, V, Codec>(
+        &mut self,
+        prefix: &Prefix,
+        storage_key: &Q,
+        value: &V,
+        codec: &Codec,
+    ) where
+        Q: ?Sized,
+        Codec: StateCodec,
+        Codec::KeyCodec: EncodeKeyLike<Q, K>,
+        Codec::ValueCodec: StateValueCodec<V>,
+    {
+        let storage_key = StorageKey::new(prefix, storage_key, codec.key_codec());
+        let storage_value = StorageValue::new(value, codec.value_codec());
+        self.set(&storage_key, storage_value);
+    }
+
+    fn set_singleton<V, Codec>(&mut self, prefix: &Prefix, value: &V, codec: &Codec)
     where
-        K: Hash + Eq + ?Sized,
-        VC: StateValueCodec<V>,
+        Codec: StateCodec,
+        Codec::ValueCodec: StateValueCodec<V>,
     {
-        let storage_key = StorageKey::new(prefix, storage_key);
-        let storage_value = StorageValue::new(value, codec);
+        let storage_key = StorageKey::singleton(prefix);
+        let storage_value = StorageValue::new(value, codec.value_codec());
         self.set(&storage_key, storage_value);
     }
 
-    fn get_decoded<V, VC>(&mut self, storage_key: &StorageKey, codec: &VC) -> Option<V>
+    fn get_decoded<V, Codec>(&mut self, storage_key: &StorageKey, codec: &Codec) -> Option<V>
     where
-        VC: StateValueCodec<V>,
+        Codec: StateCodec,
+        Codec::ValueCodec: StateValueCodec<V>,
     {
         let storage_value = self.get(storage_key)?;
 
-        Some(codec.decode_value_unwrap(storage_value.value()))
+        Some(
+            codec
+                .value_codec()
+                .decode_value_unwrap(storage_value.value()),
+        )
+    }
+
+    fn get_value<Q, K, V, Codec>(
+        &mut self,
+        prefix: &Prefix,
+        storage_key: &Q,
+        codec: &Codec,
+    ) -> Option<V>
+    where
+        Q: ?Sized,
+        Codec: StateCodec,
+        Codec::KeyCodec: EncodeKeyLike<Q, K>,
+        Codec::ValueCodec: StateValueCodec<V>,
+    {
+        let storage_key = StorageKey::new(prefix, storage_key, codec.key_codec());
+        self.get_decoded(&storage_key, codec)
     }
 
-    fn get_value<K, V, VC>(&mut self, prefix: &Prefix, storage_key: &K, codec: &VC) -> Option<V>
+    fn get_singleton<V, Codec>(&mut self, prefix: &Prefix, codec: &Codec) -> Option<V>
     where
-        K: Hash + Eq + ?Sized,
-        VC: StateValueCodec<V>,
+        Codec: StateCodec,
+        Codec::ValueCodec: StateValueCodec<V>,
     {
-        let storage_key = StorageKey::new(prefix, storage_key);
+        let storage_key = StorageKey::singleton(prefix);
         self.get_decoded(&storage_key, codec)
     }
 
-    fn remove_value<K, V, VC>(&mut self, prefix: &Prefix, storage_key: &K, codec: &VC) -> Option<V>
+    fn remove_value<Q, K, V, Codec>(
+        &mut self,
+        prefix: &Prefix,
+        storage_key: &Q,
+        codec: &Codec,
+    ) -> Option<V>
     where
-        K: Hash + Eq + ?Sized,
-        VC: StateValueCodec<V>,
+        Q: ?Sized,
+        Codec: StateCodec,
+        Codec::KeyCodec: EncodeKeyLike<Q, K>,
+        Codec::ValueCodec: StateValueCodec<V>,
     {
-        let storage_key = StorageKey::new(prefix, storage_key);
+        let storage_key = StorageKey::new(prefix, storage_key, codec.key_codec());
         let storage_value = self.get_decoded(&storage_key, codec)?;
         self.delete(&storage_key);
         Some(storage_value)
     }
 
-    fn delete_value<K>(&mut self, prefix: &Prefix, storage_key: &K)
+    fn remove_singleton<V, Codec>(&mut self, prefix: &Prefix, codec: &Codec) -> Option<V>
     where
-        K: Hash + Eq + ?Sized,
+        Codec: StateCodec,
+        Codec::ValueCodec: StateValueCodec<V>,
     {
-        let storage_key = StorageKey::new(prefix, storage_key);
+        let storage_key = StorageKey::singleton(prefix);
+        let storage_value = self.get_decoded(&storage_key, codec)?;
+        self.delete(&storage_key);
+        Some(storage_value)
+    }
+
+    fn delete_value<Q, K, Codec>(&mut self, prefix: &Prefix, storage_key: &Q, codec: &Codec)
+    where
+        Q: ?Sized,
+        Codec: StateCodec,
+        Codec::KeyCodec: EncodeKeyLike<Q, K>,
+    {
+        let storage_key = StorageKey::new(prefix, storage_key, codec.key_codec());
+        self.delete(&storage_key);
+    }
+
+    fn delete_singleton(&mut self, prefix: &Prefix) {
+        let storage_key = StorageKey::singleton(prefix);
         self.delete(&storage_key);
     }
 }
diff --git a/module-system/sov-state/src/storage.rs b/module-system/sov-state/src/storage.rs
index 92dcfbd809..d3c1445905 100644
--- a/module-system/sov-state/src/storage.rs
+++ b/module-system/sov-state/src/storage.rs
@@ -1,5 +1,4 @@
 use std::fmt::Display;
-use std::hash::Hash;
 use std::sync::Arc;
 
 use anyhow::ensure;
@@ -9,7 +8,7 @@ use serde::de::DeserializeOwned;
 use serde::{Deserialize, Serialize};
 use sov_first_read_last_write_cache::{CacheKey, CacheValue};
 
-use crate::codec::StateValueCodec;
+use crate::codec::{EncodeKeyLike, StateKeyCodec, StateValueCodec};
 use crate::internal_cache::OrderedReadsAndWrites;
 use crate::utils::AlignedVec;
 use crate::witness::Witness;
@@ -57,11 +56,12 @@ impl Display for StorageKey {
 
 impl StorageKey {
     /// Creates a new StorageKey that combines a prefix and a key.
-    pub fn new<K>(prefix: &Prefix, key: &K) -> Self
+    pub fn new<K, Q, KC>(prefix: &Prefix, key: &Q, codec: &KC) -> Self
     where
-        K: Hash + ?Sized,
+        KC: EncodeKeyLike<Q, K>,
+        Q: ?Sized,
     {
-        let encoded_key = nohash_serialize(key);
+        let encoded_key = codec.encode_key_like(key);
         let encoded_key = AlignedVec::new(encoded_key);
 
         let full_key = Vec::<u8>::with_capacity(prefix.len() + encoded_key.len());
@@ -73,26 +73,13 @@ impl StorageKey {
             key: Arc::new(full_key.into_inner()),
         }
     }
-}
-
-// Serializes a value into a `Vec<u8>` using `std::hash::Hasher`
-// writer methods, but without actually ever hashing anything.
-fn nohash_serialize<T: Hash>(item: T) -> Vec<u8> {
-    struct NoHasher(Vec<u8>);
 
-    impl std::hash::Hasher for NoHasher {
-        fn finish(&self) -> u64 {
-            0
-        }
-
-        fn write(&mut self, bytes: &[u8]) {
-            self.0.extend_from_slice(bytes);
+    /// Creates a new StorageKey that combines a prefix and a key.
+    pub fn singleton(prefix: &Prefix) -> Self {
+        Self {
+            key: Arc::new(prefix.as_aligned_vec().clone().into_inner()),
         }
     }
-
-    let mut hasher = NoHasher(vec![]);
-    item.hash(&mut hasher);
-    hasher.0
 }
 
 /// A serialized value suitable for storing. Internally uses an [`Arc<Vec<u8>>`] for cheap cloning.
@@ -235,21 +222,21 @@ pub trait Storage: Clone {
         proof: StorageProof<Self::Proof>,
     ) -> Result<(StorageKey, Option<StorageValue>), anyhow::Error>;
 
-    fn verify_proof<K, V>(
+    fn verify_proof<K, V, Codec>(
         &self,
         state_root: [u8; 32],
         proof: StorageProof<Self::Proof>,
         expected_key: &K,
-        storage_map: &StateMap<K, V>,
+        storage_map: &StateMap<K, V, Codec>,
     ) -> Result<Option<StorageValue>, anyhow::Error>
     where
-        K: Hash + Eq,
+        Codec: StateKeyCodec<K>,
     {
         let (storage_key, storage_value) = self.open_proof(state_root, proof)?;
 
         // We have to check that the storage key is the same as the external key
         ensure!(
-            storage_key == StorageKey::new(storage_map.prefix(), expected_key),
+            storage_key == StorageKey::new(storage_map.prefix(), expected_key, storage_map.codec()),
             "The storage key from the proof doesn't match the expected storage key."
         );
 
@@ -288,15 +275,18 @@ pub trait NativeStorage: Storage {
     fn get_with_proof(&self, key: StorageKey, witness: &Self::Witness)
         -> StorageProof<Self::Proof>;
 
-    fn get_with_proof_from_state_map<K, V>(
+    fn get_with_proof_from_state_map<Q, K, V, Codec>(
         &self,
-        key: &K,
-        state_map: &StateMap<K, V>,
+        key: &Q,
+        state_map: &StateMap<K, V, Codec>,
         witness: &Self::Witness,
     ) -> StorageProof<Self::Proof>
     where
-        K: Hash + Eq,
+        Codec: EncodeKeyLike<Q, K>,
     {
-        self.get_with_proof(StorageKey::new(state_map.prefix(), key), witness)
+        self.get_with_proof(
+            StorageKey::new(state_map.prefix(), key, state_map.codec()),
+            witness,
+        )
     }
 }
diff --git a/rollup-interface/src/state_machine/optimistic.rs b/rollup-interface/src/state_machine/optimistic.rs
index 84668f7dfe..e33af6599e 100644
--- a/rollup-interface/src/state_machine/optimistic.rs
+++ b/rollup-interface/src/state_machine/optimistic.rs
@@ -34,11 +34,11 @@ pub struct Attestation<StateProof: BorshSerialize> {
 /// The contents of a challenge to an attestation, which are contained as a public output of the proof
 /// Generic over an address type and a validity condition
 #[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize, Serialize, Deserialize)]
-pub struct ChallengeContents<Address, VC> {
+pub struct ChallengeContents<Address, Codec> {
     /// The rollup address of the originator of this challenge
     pub challenger_address: Address,
     /// The state transition that was proven
-    pub state_transition: StateTransition<VC, Address>,
+    pub state_transition: StateTransition<Codec, Address>,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, Serialize, Deserialize)]