Storage: Support iterating over NMaps with partial keys

codegen/src/api/storage.rs

@@ -2,9 +2,10 @@
 // This file is dual-licensed as Apache-2.0 or GPL-3.0.
 // see LICENSE for license details.
 
+use crate::types::TypePath;
 use crate::{types::TypeGenerator, CratePath};
 use heck::ToSnakeCase as _;
-use proc_macro2::TokenStream as TokenStream2;
+use proc_macro2::{Ident, TokenStream as TokenStream2};
 use quote::{format_ident, quote};
 use scale_info::TypeDef;
 use subxt_metadata::{
@@ -61,147 +62,105 @@ fn generate_storage_entry_fns(
     crate_path: &CratePath,
     should_gen_docs: bool,
 ) -> Result<TokenStream2, CodegenError> {
-    let (fields, key_impl) = match storage_entry.entry_type() {
-        StorageEntryType::Plain(_) => (vec![], quote!(vec![])),
+    let keys: Vec<(Ident, TypePath)> = match storage_entry.entry_type() {
+        StorageEntryType::Plain(_) => vec![],
         StorageEntryType::Map { key_ty, .. } => {
             match &type_gen.resolve_type(*key_ty).type_def {
                 // An N-map; return each of the keys separately.
-                TypeDef::Tuple(tuple) => {
-                    let fields = tuple
-                        .fields
-                        .iter()
-                        .enumerate()
-                        .map(|(i, f)| {
-                            let field_name = format_ident!("_{}", syn::Index::from(i));
-                            let field_type = type_gen.resolve_type_path(f.id);
-                            (field_name, field_type)
-                        })
-                        .collect::<Vec<_>>();
-
-                    let keys = fields
-                        .iter()
-                        .map(|(field_name, _)| {
-                            quote!( #crate_path::storage::address::make_static_storage_map_key(#field_name.borrow()) )
-                        });
-                    let key_impl = quote! {
-                        vec![ #( #keys ),* ]
-                    };
-
-                    (fields, key_impl)
-                }
+                TypeDef::Tuple(tuple) => tuple
+                    .fields
+                    .iter()
+                    .enumerate()
+                    .map(|(i, f)| {
+                        let ident: Ident = format_ident!("_{}", syn::Index::from(i));
+                        let ty_path = type_gen.resolve_type_path(f.id);
+                        (ident, ty_path)
+                    })
+                    .collect::<Vec<_>>(),
                 // A map with a single key; return the single key.
                 _ => {
+                    let ident = format_ident!("_0");
                     let ty_path = type_gen.resolve_type_path(*key_ty);
-                    let fields = vec![(format_ident!("_0"), ty_path)];
-                    let key_impl = quote! {
-                        vec![ #crate_path::storage::address::make_static_storage_map_key(_0.borrow()) ]
-                    };
-                    (fields, key_impl)
+                    vec![(ident, ty_path)]
                 }
             }
         }
     };
-
     let pallet_name = pallet.name();
     let storage_name = storage_entry.name();
     let Some(storage_hash) = pallet.storage_hash(storage_name) else {
         return Err(CodegenError::MissingStorageMetadata(pallet_name.into(), storage_name.into()));
     };
 
-    let fn_name = format_ident!("{}", storage_entry.name().to_snake_case());
+    let snake_case_name = storage_entry.name().to_snake_case();
     let storage_entry_ty = match storage_entry.entry_type() {
         StorageEntryType::Plain(ty) => *ty,
         StorageEntryType::Map { value_ty, .. } => *value_ty,
     };
     let storage_entry_value_ty = type_gen.resolve_type_path(storage_entry_ty);
-
     let docs = storage_entry.docs();
     let docs = should_gen_docs
         .then_some(quote! { #( #[doc = #docs ] )* })
         .unwrap_or_default();
 
-    let key_args = fields.iter().map(|(field_name, field_type)| {
-        // The field type is translated from `std::vec::Vec<T>` to `[T]`. We apply
-        // Borrow to all types, so this just makes it a little more ergonomic.
-        //
-        // TODO [jsdw]: Support mappings like `String -> str` too for better borrow
-        // ergonomics.
-        let field_ty = match field_type.vec_type_param() {
-            Some(ty) => quote!([#ty]),
-            _ => quote!(#field_type),
-        };
-        quote!( #field_name: impl ::std::borrow::Borrow<#field_ty> )
-    });
-
-    let is_map_type = matches!(storage_entry.entry_type(), StorageEntryType::Map { .. });
-
-    // Is the entry iterable?
-    let is_iterable_type = if is_map_type {
-        quote!(#crate_path::storage::address::Yes)
-    } else {
-        quote!(())
-    };
-
-    let has_default_value = match storage_entry.modifier() {
-        StorageEntryModifier::Default => true,
-        StorageEntryModifier::Optional => false,
+    let is_defaultable_type = match storage_entry.modifier() {
+        StorageEntryModifier::Default => quote!(#crate_path::storage::address::Yes),
+        StorageEntryModifier::Optional => quote!(()),
     };
 
-    // Does the entry have a default value?
-    let is_defaultable_type = if has_default_value {
-        quote!(#crate_path::storage::address::Yes)
-    } else {
-        quote!(())
-    };
+    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() {
+            let fn_name = format_ident!("{snake_case_name}");
+            (fn_name, true, false)
+        } else {
+            let fn_name = if n_keys == 0 {
+                format_ident!("{snake_case_name}_iter")
+            } else {
+                format_ident!("{snake_case_name}_iter{}", n_keys)
+            };
+            (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, field_type)| {
+            // The field type is translated from `std::vec::Vec<T>` to `[T]`. We apply
+            // Borrow to all types, so this just makes it a little more ergonomic.
+            // TODO [jsdw]: Support mappings like `String -> str` too for better borrow
+            // ergonomics.
+            let field_ty = match field_type.vec_type_param() {
+                Some(ty) => quote!([#ty]),
+                _ => quote!(#field_type),
+            };
+            quote!( #field_name: impl ::std::borrow::Borrow<#field_ty> )
+        });
 
-    // If the item is a map, we want a way to access the root entry to do things like iterate over it,
-    // so expose a function to create this entry, too:
-    let root_entry_fn = if is_map_type {
-        let fn_name_root = format_ident!("{}_root", fn_name);
         quote!(
             #docs
-            pub fn #fn_name_root(
+            pub fn #fn_name(
                 &self,
+                #(#key_args,)*
             ) -> #crate_path::storage::address::Address::<
                 #crate_path::storage::address::StaticStorageMapKey,
                 #storage_entry_value_ty,
-                (),
+                #is_fetchable_type,
                 #is_defaultable_type,
                 #is_iterable_type
             > {
                 #crate_path::storage::address::Address::new_static(
                     #pallet_name,
                     #storage_name,
-                    Vec::new(),
+                    vec![#(#key_impls,)*],
                     [#(#storage_hash,)*]
                 )
             }
         )
-    } else {
-        quote!()
-    };
+    });
 
     Ok(quote! {
-        // Access a specific value from a storage entry
-        #docs
-        pub fn #fn_name(
-            &self,
-            #( #key_args, )*
-        ) -> #crate_path::storage::address::Address::<
-            #crate_path::storage::address::StaticStorageMapKey,
-            #storage_entry_value_ty,
-            #crate_path::storage::address::Yes,
-            #is_defaultable_type,
-            #is_iterable_type
-        > {
-            #crate_path::storage::address::Address::new_static(
-                #pallet_name,
-                #storage_name,
-                #key_impl,
-                [#(#storage_hash,)*]
-            )
-        }
+        #( #all_fns
 
-        #root_entry_fn
+        )*
     })
 }

subxt/examples/storage_iterating.rs

@@ -9,7 +9,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let api = OnlineClient::<PolkadotConfig>::new().await?;
 
     // Build a storage query to iterate over account information.
-    let storage_query = polkadot::storage().system().account_root();
+    let storage_query = polkadot::storage().system().account_iter();
 
     // Get back an iterator of results (here, we are fetching 10 items at
     // a time from the node, but we always iterate over one at a time).

subxt/examples/storage_iterating_dynamic.rs

@@ -6,7 +6,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let api = OnlineClient::<PolkadotConfig>::new().await?;
 
     // Build a dynamic storage query to iterate account information.
-    let storage_query = subxt::dynamic::storage_root("System", "Account");
+    let storage_query = subxt::dynamic::storage_iter("System", "Account");
 
     // Use that query to return an iterator over the results.
     let mut results = api

subxt/src/book/usage/storage.rs

@@ -39,19 +39,26 @@
 //! ```
 //!
 //! As well as accessing specific entries, some storage locations can also be iterated over (such as
-//! the map of account information). To do this, suffix `_root` onto the query constructor (this
+//! the map of account information). To do this, suffix `_iter` onto the query constructor (this
 //! will only be available on static constructors when iteration is actually possible):
 //!
 //! ```rust,no_run
 //! #[subxt::subxt(runtime_metadata_path = "../artifacts/polkadot_metadata_small.scale")]
 //! pub mod polkadot {}
 //!
 //! // A static query capable of iterating over accounts:
-//! let storage_query = polkadot::storage().system().account_root();
+//! let storage_query = polkadot::storage().system().account_iter();
 //! // A dynamic query to do the same:
-//! let storage_query = subxt::dynamic::storage_root("System", "Account");
+//! let storage_query = subxt::dynamic::storage_iter("System", "Account");
 //! ```
 //!
+//! Some storage entries require multiple keys to be accessed, for example
+//! `polkadot::storage().preimage().preimage_for(_, _)`, which takes a `H256` and a `u32` as keys.
+//! For these multi-key storage entries, the generated static interface also offers _partial iteration_.
+//! This means iterating over a subset of storage entries, where just a subset of the keys match.
+//! These generated constructors have a suffix `_iter1` / `_iter2` / ... added, to iterate over all
+//! storage entries given the first 1 / 2 / ... keys.
+//!
 //! All valid storage queries implement [`crate::storage::StorageAddress`]. As well as describing
 //! how to build a valid storage query, this trait also has some associated types that determine the
 //! shape of the result you'll get back, and determine what you can do with it (ie, can you iterate

subxt/src/dynamic.rs

@@ -26,7 +26,7 @@ pub use crate::tx::dynamic as tx;
 pub use crate::constants::dynamic as constant;
 
 // Lookup storage values dynamically.
-pub use crate::storage::{dynamic as storage, dynamic_root as storage_root};
+pub use crate::storage::{dynamic as storage, dynamic_iter as storage_iter};
 
 // Execute runtime API function call dynamically.
 pub use crate::runtime_api::dynamic as runtime_api_call;

subxt/src/storage/mod.rs

@@ -21,11 +21,11 @@ pub use crate::rpc::types::StorageKey;
 /// entry lives and how to properly decode it.
 pub mod address {
     pub use super::storage_address::{
-        dynamic, dynamic_root, make_static_storage_map_key, Address, DynamicAddress,
+        dynamic, dynamic_iter, make_static_storage_map_key, Address, DynamicAddress,
         StaticStorageMapKey, StorageAddress, Yes,
     };
 }
 
 // For consistency with other modules, also expose
 // the basic address stuff at the root of the module.
-pub use storage_address::{dynamic, dynamic_root, Address, DynamicAddress, StorageAddress};
+pub use storage_address::{dynamic, dynamic_iter, Address, DynamicAddress, StorageAddress};

subxt/src/storage/storage_address.rs

@@ -50,7 +50,7 @@ pub trait StorageAddress {
 pub struct Yes;
 
 /// A concrete storage address. This can be created from static values (ie those generated
-/// via the `subxt` macro) or dynamic values via [`dynamic`] and [`dynamic_root`].
+/// via the `subxt` macro) or dynamic values via [`dynamic`] and [`dynamic_iter`].
 pub struct Address<StorageKey, ReturnTy, Fetchable, Defaultable, Iterable> {
     pallet_name: Cow<'static, str>,
     entry_name: Cow<'static, str>,
@@ -229,7 +229,7 @@ pub fn make_static_storage_map_key<T: codec::Encode>(t: T) -> StaticStorageMapKe
 }
 
 /// Construct a new dynamic storage lookup to the root of some entry.
-pub fn dynamic_root(
+pub fn dynamic_iter(
     pallet_name: impl Into<String>,
     entry_name: impl Into<String>,
 ) -> DynamicAddress<Value> {

subxt/src/storage/storage_type.rs

@@ -177,7 +177,7 @@ where
     /// let api = OnlineClient::<PolkadotConfig>::new().await.unwrap();
     ///
     /// // Address to the root of a storage entry that we'd like to iterate over.
-    /// let address = polkadot::storage().xcm_pallet().version_notifiers_root();
+    /// let address = polkadot::storage().xcm_pallet().version_notifiers_iter();
     ///
     /// // Iterate over keys and values at that address.
     /// let mut iter = api

testing/integration-tests/src/full_client/client/mod.rs

@@ -107,7 +107,7 @@ async fn fetch_keys() {
     let ctx = test_context().await;
     let api = ctx.client();
 
-    let addr = node_runtime::storage().system().account_root();
+    let addr = node_runtime::storage().system().account_iter();
     let keys = api
         .storage()
         .at_latest()
@@ -124,7 +124,7 @@ async fn test_iter() {
     let ctx = test_context().await;
     let api = ctx.client();
 
-    let addr = node_runtime::storage().system().account_root();
+    let addr = node_runtime::storage().system().account_iter();
     let mut iter = api
         .storage()
         .at_latest()
@@ -409,7 +409,7 @@ async fn unsigned_extrinsic_is_same_shape_as_polkadotjs() {
     let expected_tx_bytes = hex::decode(
         "b004060700d43593c715fdd31c61141abd04a99fd6822c8558854ccde39a5684e7a56da27d0f0090c04bb6db2b"
     )
-    .unwrap();
+        .unwrap();
 
     // Make sure our encoding is the same as the encoding polkadot UI created.
     assert_eq!(actual_tx_bytes, expected_tx_bytes);

testing/integration-tests/src/full_client/codegen/mod.rs

@@ -8,7 +8,7 @@
 /// Generate by running this at the root of the repository:
 ///
 /// ```
-/// cargo run --bin subxt -- codegen --file artifacts/polkadot_metadata_full.scale | rustfmt > testing/integration-tests/src/codegen/polkadot.rs
+/// cargo run --bin subxt -- codegen --file artifacts/polkadot_metadata_full.scale | rustfmt > testing/integration-tests/src/full_client/codegen/polkadot.rs
 /// ```
 #[rustfmt::skip]
 #[allow(clippy::all)]