Typed Storage Keys

codegen/src/api/storage.rs

@@ -8,7 +8,7 @@ use quote::{format_ident, quote};
 use scale_info::TypeDef;
 use scale_typegen::{typegen::type_path::TypePath, TypeGenerator};
 use subxt_metadata::{
-    PalletMetadata, StorageEntryMetadata, StorageEntryModifier, StorageEntryType,
+    PalletMetadata, StorageEntryMetadata, StorageEntryModifier, StorageEntryType, StorageHasher,
 };
 
 use super::CodegenError;
@@ -75,43 +75,83 @@ fn generate_storage_entry_fns(
     let alias_module_name = format_ident!("{snake_case_name}");
     let alias_storage_path = quote!( types::#alias_module_name::#alias_name );
 
-    let storage_entry_map = |idx, id| {
-        let ident: Ident = format_ident!("_{}", idx);
+    struct MapEntryKey {
+        arg_name: Ident,
+        alias_type_def: TokenStream,
+        alias_type_path: TokenStream,
+        hasher: StorageHasher,
+    }
+
+    let map_entry_key = |idx, id, hasher| -> MapEntryKey {
+        let arg_name: Ident = format_ident!("_{}", idx);
         let ty_path = type_gen
             .resolve_type_path(id)
             .expect("type is in metadata; qed");
 
         let alias_name = format_ident!("Param{}", idx);
         let alias_type = primitive_type_alias(&ty_path);
 
-        let alias_type = quote!( pub type #alias_name = #alias_type; );
-        let path_to_alias = quote!( types::#alias_module_name::#alias_name );
+        let alias_type_def = quote!( pub type #alias_name = #alias_type; );
+        let alias_type_path = quote!( types::#alias_module_name::#alias_name );
 
-        (ident, alias_type, path_to_alias)
+        MapEntryKey {
+            arg_name,
+            alias_type_def,
+            alias_type_path,
+            hasher,
+        }
     };
 
-    let keys: Vec<(Ident, TokenStream, TokenStream)> = match storage_entry.entry_type() {
+    let keys: Vec<MapEntryKey> = match storage_entry.entry_type() {
         StorageEntryType::Plain(_) => vec![],
-        StorageEntryType::Map { key_ty, .. } => {
+        StorageEntryType::Map {
+            key_ty, hashers, ..
+        } => {
             match &type_gen
                 .resolve_type(*key_ty)
                 .expect("key type should be present")
                 .type_def
             {
                 // An N-map; return each of the keys separately.
-                TypeDef::Tuple(tuple) => tuple
-                    .fields
-                    .iter()
-                    .enumerate()
-                    .map(|(idx, f)| storage_entry_map(idx, f.id))
-                    .collect::<Vec<_>>(),
+                TypeDef::Tuple(tuple) => {
+                    let key_count = tuple.fields.len();
+                    let hasher_count = hashers.len();
+                    if hasher_count != 1 && hasher_count != key_count {
+                        return Err(CodegenError::InvalidStorageHasherCount {
+                            storage_entry_name: storage_entry.name().to_owned(),
+                            key_count,
+                            hasher_count,
+                        });
+                    }
+
+                    let mut map_entry_keys: Vec<MapEntryKey> = vec![];
+                    for (idx, field) in tuple.fields.iter().enumerate() {
+                        // Note: these are in bounds because of the checks above, qed;
+                        let hasher = if idx >= hasher_count {
+                            hashers[0]
+                        } else {
+                            hashers[idx]
+                        };
+                        map_entry_keys.push(map_entry_key(idx, field.id, hasher));
+                    }
+                    map_entry_keys
+                }
                 // A map with a single key; return the single key.
                 _ => {
-                    vec![storage_entry_map(0, *key_ty)]
+                    let Some(hasher) = hashers.first() else {
+                        return Err(CodegenError::InvalidStorageHasherCount {
+                            storage_entry_name: storage_entry.name().to_owned(),
+                            key_count: 1,
+                            hasher_count: 0,
+                        });
+                    };
+
+                    vec![map_entry_key(0, *key_ty, *hasher)]
                 }
             }
         }
     };
+
     let pallet_name = pallet.name();
     let storage_name = storage_entry.name();
     let Some(storage_hash) = pallet.storage_hash(storage_name) else {
@@ -133,6 +173,10 @@ fn generate_storage_entry_fns(
         StorageEntryModifier::Optional => quote!(()),
     };
 
+    // Note: putting `#crate_path::storage::address::StaticStorageKey` into this variable is necessary
+    // to get the line width below a certain limit. If not done, rustfmt will refuse to format the following big expression.
+    // for more information see [this post](https://users.rust-lang.org/t/rustfmt-silently-fails-to-work/75485/4).
+    let static_storage_key: TokenStream = quote!(#crate_path::storage::address::StaticStorageKey);
     let all_fns = (0..=keys.len()).map(|n_keys| {
         let keys_slice = &keys[..n_keys];
         let (fn_name, is_fetchable, is_iterable) = if n_keys == keys.len() {
@@ -146,20 +190,73 @@ fn generate_storage_entry_fns(
             };
             (fn_name, false, true)
         };
-        let is_fetchable_type = is_fetchable.then_some(quote!(#crate_path::storage::address::Yes)).unwrap_or(quote!(()));
-        let is_iterable_type = is_iterable.then_some(quote!(#crate_path::storage::address::Yes)).unwrap_or(quote!(()));
-        let key_impls = keys_slice.iter().map(|(field_name, _, _)| quote!( #crate_path::storage::address::make_static_storage_map_key(#field_name.borrow()) ));
-        let key_args = keys_slice.iter().map(|(field_name, _, path_to_alias )| {
-            quote!( #field_name: impl ::std::borrow::Borrow<#path_to_alias> )
-        });
+        let is_fetchable_type = is_fetchable
+            .then_some(quote!(#crate_path::storage::address::Yes))
+            .unwrap_or(quote!(()));
+        let is_iterable_type = is_iterable
+            .then_some(quote!(#crate_path::storage::address::Yes))
+            .unwrap_or(quote!(()));
+
+        let (keys, keys_type) = match keys_slice.len() {
+            0 => (quote!(()), quote!(())),
+            1 => {
+                let key = &keys_slice[0];
+                if key.hasher.ends_with_key() {
+                    let arg = &key.arg_name;
+                    let keys = quote!(#static_storage_key::new(#arg.borrow()));
+                    let path = &key.alias_type_path;
+                    let path = quote!(#static_storage_key<#path>);
+                    (keys, path)
+                } else {
+                    (quote!(()), quote!(()))
+                }
+            }
+            _ => {
+                let keys_iter = keys_slice.iter().map(
+                    |MapEntryKey {
+                         arg_name, hasher, ..
+                     }| {
+                        if hasher.ends_with_key() {
+                            quote!( #static_storage_key::new(#arg_name.borrow()) )
+                        } else {
+                            quote!(())
+                        }
+                    },
+                );
+                let keys = quote!( (#(#keys_iter,)*) );
+                let paths_iter = keys_slice.iter().map(
+                    |MapEntryKey {
+                         alias_type_path,
+                         hasher,
+                         ..
+                     }| {
+                        if hasher.ends_with_key() {
+                            quote!( #static_storage_key<#alias_type_path> )
+                        } else {
+                            quote!(())
+                        }
+                    },
+                );
+                let paths = quote!( (#(#paths_iter,)*) );
+                (keys, paths)
+            }
+        };
+
+        let key_args = keys_slice.iter().map(
+            |MapEntryKey {
+                 arg_name,
+                 alias_type_path,
+                 ..
+             }| quote!( #arg_name: impl ::std::borrow::Borrow<#alias_type_path> ),
+        );
 
         quote!(
             #docs
             pub fn #fn_name(
                 &self,
                 #(#key_args,)*
             ) -> #crate_path::storage::address::Address::<
-                #crate_path::storage::address::StaticStorageMapKey,
+                #keys_type,
                 #alias_storage_path,
                 #is_fetchable_type,
                 #is_defaultable_type,
@@ -168,14 +265,16 @@ fn generate_storage_entry_fns(
                 #crate_path::storage::address::Address::new_static(
                     #pallet_name,
                     #storage_name,
-                    vec![#(#key_impls,)*],
+                    #keys,
                     [#(#storage_hash,)*]
                 )
             }
         )
     });
 
-    let alias_types = keys.iter().map(|(_, alias_type, _)| alias_type);
+    let alias_types = keys
+        .iter()
+        .map(|MapEntryKey { alias_type_def, .. }| alias_type_def);
 
     let types_mod_ident = type_gen.types_mod_ident();
     // Generate type alias for the return type only, since
@@ -231,7 +330,7 @@ mod tests {
                 name,
                 modifier: v15::StorageEntryModifier::Optional,
                 ty: v15::StorageEntryType::Map {
-                    hashers: vec![],
+                    hashers: vec![v15::StorageHasher::Blake2_128Concat],
                     key,
                     value: meta_type::<bool>(),
                 },

codegen/src/error.rs

@@ -39,15 +39,21 @@ pub enum CodegenError {
     #[error("Call variant for type {0} must have all named fields. Make sure you are providing a valid substrate-based metadata")]
     InvalidCallVariant(u32),
     /// Type should be an variant/enum.
-    #[error(
-        "{0} type should be an variant/enum type. Make sure you are providing a valid substrate-based metadata"
-    )]
+    #[error("{0} type should be an variant/enum type. Make sure you are providing a valid substrate-based metadata")]
     InvalidType(String),
     /// Extrinsic call type could not be found.
-    #[error(
-        "Extrinsic call type could not be found. Make sure you are providing a valid substrate-based metadata"
-    )]
+    #[error("Extrinsic call type could not be found. Make sure you are providing a valid substrate-based metadata")]
     MissingCallType,
+    /// There are too many or too few hashers.
+    #[error("Could not Generate functions for storage entry {storage_entry_name}. There are {key_count} keys, but only {hasher_count} hashers. The number of hashers must equal the number of keys or be exactly 1.")]
+    InvalidStorageHasherCount {
+        /// The name of the storage entry
+        storage_entry_name: String,
+        /// Number of keys
+        key_count: usize,
+        /// Number of hashers
+        hasher_count: usize,
+    },
     /// Cannot generate types.
     #[error("Type Generation failed: {0}")]
     TypeGeneration(#[from] TypegenError),

metadata/src/lib.rs

@@ -475,6 +475,35 @@ pub enum StorageHasher {
     Identity,
 }
 
+impl StorageHasher {
+    /// The hash produced by a [`StorageHasher`] can have these two components, in order:
+    ///
+    /// 1. A fixed size hash. (not present for [`StorageHasher::Identity`]).
+    /// 2. The SCALE encoded key that was used as an input to the hasher (only present for
+    /// [`StorageHasher::Twox64Concat`], [`StorageHasher::Blake2_128Concat`] or [`StorageHasher::Identity`]).
+    ///
+    /// This function returns the number of bytes used to represent the first of these.
+    pub fn len_excluding_key(&self) -> usize {
+        match self {
+            StorageHasher::Blake2_128Concat => 16,
+            StorageHasher::Twox64Concat => 8,
+            StorageHasher::Blake2_128 => 16,
+            StorageHasher::Blake2_256 => 32,
+            StorageHasher::Twox128 => 16,
+            StorageHasher::Twox256 => 32,
+            StorageHasher::Identity => 0,
+        }
+    }
+
+    /// Returns true if the key used to produce the hash is appended to the hash itself.
+    pub fn ends_with_key(&self) -> bool {
+        matches!(
+            self,
+            StorageHasher::Blake2_128Concat | StorageHasher::Twox64Concat | StorageHasher::Identity
+        )
+    }
+}
+
 /// Is the storage entry optional, or does it have a default value.
 #[derive(Debug, Clone, Copy, Eq, PartialEq)]
 pub enum StorageEntryModifier {

subxt/examples/storage_iterating.rs

@@ -16,9 +16,10 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     // a time from the node, but we always iterate over one at a time).
     let mut results = api.storage().at_latest().await?.iter(storage_query).await?;
 
-    while let Some(Ok((key, value))) = results.next().await {
-        println!("Key: 0x{}", hex::encode(&key));
-        println!("Value: {:?}", value);
+    while let Some(Ok(kv)) = results.next().await {
+        println!("Keys decoded: {:?}", kv.keys);
+        println!("Key: 0x{}", hex::encode(&kv.key_bytes));
+        println!("Value: {:?}", kv.value);
     }
 
     Ok(())

subxt/examples/storage_iterating_dynamic.rs

@@ -7,16 +7,17 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let api = OnlineClient::<PolkadotConfig>::new().await?;
 
     // Build a dynamic storage query to iterate account information.
-    // With a dynamic query, we can just provide an empty Vec as the keys to iterate over all entries.
-    let keys = Vec::<()>::new();
+    // With a dynamic query, we can just provide an empty vector as the keys to iterate over all entries.
+    let keys: Vec<scale_value::Value> = vec![];
     let storage_query = subxt::dynamic::storage("System", "Account", keys);
 
     // Use that query to return an iterator over the results.
     let mut results = api.storage().at_latest().await?.iter(storage_query).await?;
 
-    while let Some(Ok((key, value))) = results.next().await {
-        println!("Key: 0x{}", hex::encode(&key));
-        println!("Value: {:?}", value.to_value()?);
+    while let Some(Ok(kv)) = results.next().await {
+        println!("Keys decoded: {:?}", kv.keys);
+        println!("Key: 0x{}", hex::encode(&kv.key_bytes));
+        println!("Value: {:?}", kv.value.to_value()?);
     }
 
     Ok(())

subxt/examples/storage_iterating_partial.rs

@@ -38,11 +38,11 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     // Get back an iterator of results.
     let mut results = api.storage().at_latest().await?.iter(storage_query).await?;
 
-    while let Some(Ok((key, value))) = results.next().await {
-        println!("Key: 0x{}", hex::encode(&key));
-        println!("Value: {:?}", value);
+    while let Some(Ok(kv)) = results.next().await {
+        println!("Keys decoded: {:?}", kv.keys);
+        println!("Key: 0x{}", hex::encode(&kv.key_bytes));
+        println!("Value: {:?}", kv.value);
     }
-
     Ok(())
 }
 

subxt/src/book/usage/storage.rs

@@ -49,7 +49,7 @@
 //! // A static query capable of iterating over accounts:
 //! let storage_query = polkadot::storage().system().account_iter();
 //! // A dynamic query to do the same:
-//! let storage_query = subxt::dynamic::storage("System", "Account", Vec::<u8>::new());
+//! let storage_query = subxt::dynamic::storage("System", "Account", ());
 //! ```
 //!
 //! Some storage entries are maps with multiple keys. As an example, we might end up with

subxt/src/error/mod.rs

@@ -14,6 +14,7 @@ crate::macros::cfg_unstable_light_client! {
 pub use dispatch_error::{
     ArithmeticError, DispatchError, ModuleError, TokenError, TransactionalError,
 };
+use subxt_metadata::StorageHasher;
 
 // Re-expose the errors we use from other crates here:
 pub use crate::config::ExtrinsicParamsError;
@@ -206,6 +207,23 @@ pub enum StorageAddressError {
         /// The number of fields in the metadata for this storage entry.
         fields: usize,
     },
+    /// We weren't given enough bytes to decode the storage address/key.
+    #[error("Not enough remaining bytes to decode the storage address/key")]
+    NotEnoughBytes,
+    /// We have leftover bytes after decoding the storage address.
+    #[error("We have leftover bytes after decoding the storage address")]
+    TooManyBytes,
+    /// The bytes of a storage address are not the expected address for decoding the storage keys of the address.
+    #[error("Storage address bytes are not the expected format. Addresses need to be at least 16 bytes (pallet ++ entry) and follow a structure given by the hashers defined in the metadata")]
+    UnexpectedAddressBytes,
+    /// An invalid hasher was used to reconstruct a value from a chunk of bytes that is part of a storage address. Hashers where the hash does not contain the original value are invalid for this purpose.
+    #[error("An invalid hasher was used to reconstruct a value with type ID {ty_id} from a hash formed by a {hasher:?} hasher. This is only possible for concat-style hashers or the identity hasher")]
+    HasherCannotReconstructKey {
+        /// Type id of the key's type.
+        ty_id: u32,
+        /// The invalid hasher that caused this error.
+        hasher: StorageHasher,
+    },
 }
 
 /// Something went wrong trying to access details in the metadata.