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, TokenStream};
 use quote::{format_ident, quote};
 use scale_info::TypeDef;
 use subxt_metadata::{
@@ -61,147 +62,231 @@ 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)| {
+            let field_ty = primitive_type_alias(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
+        )*
     })
 }
+
+fn primitive_type_alias(type_path: &TypePath) -> TokenStream {
+    // Vec<T> is cast to [T]
+    if let Some(ty) = type_path.vec_type_param() {
+        return quote!([#ty]);
+    }
+    // String is cast to str
+    if type_path.is_string() {
+        return quote!(::core::primitive::str);
+    }
+    quote!(#type_path)
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::RuntimeGenerator;
+    use frame_metadata::v15;
+    use quote::{format_ident, quote};
+    use scale_info::{meta_type, MetaType};
+    use std::borrow::Cow;
+
+    use subxt_metadata::Metadata;
+
+    fn metadata_with_storage_entries(
+        storage_entries: impl IntoIterator<Item = (&'static str, MetaType)>,
+    ) -> Metadata {
+        let storage_entries: Vec<v15::StorageEntryMetadata> = storage_entries
+            .into_iter()
+            .map(|(name, key)| v15::StorageEntryMetadata {
+                name,
+                modifier: v15::StorageEntryModifier::Optional,
+                ty: v15::StorageEntryType::Map {
+                    hashers: vec![],
+                    key,
+                    value: meta_type::<bool>(),
+                },
+                default: vec![],
+                docs: vec![],
+            })
+            .collect();
+
+        let pallet_1 = v15::PalletMetadata {
+            name: "Pallet1",
+            storage: Some(v15::PalletStorageMetadata {
+                prefix: Default::default(),
+                entries: storage_entries,
+            }),
+            calls: None,
+            event: None,
+            constants: vec![],
+            error: None,
+            index: 0,
+            docs: vec![],
+        };
+
+        let extrinsic_metadata = v15::ExtrinsicMetadata {
+            version: 0,
+            signed_extensions: vec![],
+            address_ty: meta_type::<()>(),
+            call_ty: meta_type::<()>(),
+            signature_ty: meta_type::<()>(),
+            extra_ty: meta_type::<()>(),
+        };
+
+        let metadata: Metadata = v15::RuntimeMetadataV15::new(
+            vec![pallet_1],
+            extrinsic_metadata,
+            meta_type::<()>(),
+            vec![],
+            v15::OuterEnums {
+                call_enum_ty: meta_type::<()>(),
+                event_enum_ty: meta_type::<()>(),
+                error_enum_ty: meta_type::<()>(),
+            },
+            v15::CustomMetadata {
+                map: Default::default(),
+            },
+        )
+        .try_into()
+        .expect("can build valid metadata");
+        metadata
+    }
+
+    #[test]
+    fn borrow_type_replacements() {
+        let storage_entries = [
+            ("vector", meta_type::<Vec<u8>>()),
+            ("boxed", meta_type::<Box<u16>>()),
+            ("string", meta_type::<String>()),
+            ("static_string", meta_type::<&'static str>()),
+            ("cow_string", meta_type::<Cow<'_, str>>()),
+        ];
+
+        let expected_borrowed_types = [
+            quote!([::core::primitive::u8]),
+            quote!(::core::primitive::u16),
+            quote!(::core::primitive::str),
+            quote!(::core::primitive::str),
+            quote!(::core::primitive::str),
+        ];
+
+        let metadata = metadata_with_storage_entries(storage_entries);
+
+        let item_mod = syn::parse_quote!(
+            pub mod api {}
+        );
+        let generator = RuntimeGenerator::new(metadata);
+        let generated = generator
+            .generate_runtime(
+                item_mod,
+                Default::default(),
+                Default::default(),
+                "::subxt_path".into(),
+                false,
+            )
+            .expect("should be able to generate runtime");
+        let generated_str = generated.to_string();
+
+        for ((name, _), expected_type) in storage_entries
+            .into_iter()
+            .zip(expected_borrowed_types.into_iter())
+        {
+            let name_ident = format_ident!("{}", name);
+            let expected_storage_constructor = quote!(
+                fn #name_ident(
+                    &self,
+                    _0: impl ::std::borrow::Borrow<#expected_type>,
+                )
+            );
+            assert!(generated_str.contains(&expected_storage_constructor.to_string()));
+        }
+    }
+}

codegen/src/types/type_path.rs

@@ -61,6 +61,10 @@ impl TypePath {
         matches!(&self.0, TypePathInner::Type(ty) if ty.is_compact())
     }
 
+    pub(crate) fn is_string(&self) -> bool {
+        matches!(&self.0, TypePathInner::Type(ty) if ty.is_string())
+    }
+
     /// Returns the type parameters in a path which are inherited from the containing type.
     ///
     /// # Example
@@ -214,6 +218,15 @@ impl TypePathType {
         matches!(self, TypePathType::Compact { .. })
     }
 
+    pub(crate) fn is_string(&self) -> bool {
+        matches!(
+            self,
+            TypePathType::Primitive {
+                def: TypeDefPrimitive::Str
+            }
+        )
+    }
+
     fn to_syn_type(&self) -> syn::Type {
         match &self {
             TypePathType::Path { path, params } => {

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,9 @@ 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");
+    // With a dynamic query, we can just provide an empty Vec as the keys to iterate over all entries.
+    let keys = Vec::<()>::new();
+    let storage_query = subxt::dynamic::storage("System", "Account", keys);
 
     // Use that query to return an iterator over the results.
     let mut results = api