Retain specific runtime APIs

cli/src/commands/metadata.rs

@@ -8,7 +8,7 @@ use codec::{Decode, Encode};
 use color_eyre::eyre;
 use frame_metadata::{RuntimeMetadata, RuntimeMetadataPrefixed};
 use std::io::{self, Write};
-use subxt_metadata::{metadata_v14_to_latest, retain_metadata_pallets};
+use subxt_metadata::{metadata_v14_to_latest, retain_metadata};
 
 /// Download metadata from a substrate node, for use with `subxt` codegen.
 #[derive(Debug, ClapParser)]
@@ -25,27 +25,39 @@ pub struct Opts {
     /// when using the option.
     #[clap(long, use_value_delimiter = true, value_parser)]
     pallets: Option<Vec<String>>,
+    /// Generate a subset of the metadata that contains only the
+    /// runtime APIs needed.
+    ///
+    /// The returned metadata is updated to the latest available version
+    /// when using the option.
+    #[clap(long, use_value_delimiter = true, value_parser)]
+    runtime_apis: Option<Vec<String>>,
 }
 
 pub async fn run(opts: Opts) -> color_eyre::Result<()> {
     let bytes = opts.file_or_url.fetch().await?;
     let mut metadata = <RuntimeMetadataPrefixed as Decode>::decode(&mut &bytes[..])?;
 
-    if let Some(pallets) = opts.pallets {
+    if opts.pallets.is_some() || opts.runtime_apis.is_some() {
         let mut metadata_v15 = match metadata.1 {
             RuntimeMetadata::V14(metadata_v14) => metadata_v14_to_latest(metadata_v14),
             RuntimeMetadata::V15(metadata_v15) => metadata_v15,
             _ => {
                 return Err(eyre::eyre!(
                     "Unsupported metadata version {:?}, expected V14.",
                     metadata.1
-                ))
+                ));
             }
         };
-
-        retain_metadata_pallets(&mut metadata_v15, |pallet_name| {
-            pallets.iter().any(|p| &**p == pallet_name)
-        });
+        let retain_pallets_fn: Box<dyn Fn(&str) -> bool> = match opts.pallets.as_ref() {
+            Some(pallets) => Box::new(|name| pallets.iter().any(|p| &**p == name)),
+            None => Box::new(|_| true),
+        };
+        let retain_runtime_apis_fn: Box<dyn Fn(&str) -> bool> = match opts.runtime_apis.as_ref() {
+            Some(apis) => Box::new(|name| apis.iter().any(|p| &**p == name)),
+            None => Box::new(|_| true),
+        };
+        retain_metadata(&mut metadata_v15, retain_pallets_fn, retain_runtime_apis_fn);
         metadata = metadata_v15.into();
     }
 

metadata/src/lib.rs

@@ -7,7 +7,7 @@ mod validation;
 
 use frame_metadata::{v14::RuntimeMetadataV14, v15::RuntimeMetadataV15};
 
-pub use retain::retain_metadata_pallets;
+pub use retain::retain_metadata;
 pub use validation::{
     get_call_hash, get_constant_hash, get_pallet_hash, get_runtime_api_hash, get_storage_hash,
     MetadataHasher, NotFound,

metadata/src/retain.rs

@@ -106,17 +106,12 @@ fn update_extrinsic_types(
 }
 
 /// Collect all type IDs needed to represent the runtime APIs.
-fn collect_runtime_api_types(
-    apis: &[RuntimeApiMetadata<PortableForm>],
-    type_ids: &mut HashSet<u32>,
-) {
-    for api in apis {
-        for method in &api.methods {
-            for input in &method.inputs {
-                type_ids.insert(input.ty.id);
-            }
-            type_ids.insert(method.output.id);
+fn collect_runtime_api_types(api: &RuntimeApiMetadata<PortableForm>, type_ids: &mut HashSet<u32>) {
+    for method in &api.methods {
+        for input in &method.inputs {
+            type_ids.insert(input.ty.id);
         }
+        type_ids.insert(method.output.id);
     }
 }
 
@@ -165,7 +160,7 @@ where
     let Some(call_ty) = extrinsic_ty.ty.type_params
         .iter_mut()
         .find(|ty| ty.name == "Call")
-        .and_then(|ty| ty.ty) else { return };
+        .and_then(|ty| ty.ty) else { return; };
 
     let call_ty = metadata
         .types
@@ -182,7 +177,7 @@ where
 }
 
 /// Generate a subset of the metadata that contains only the
-/// types needed to represent the provided pallets.
+/// types needed to represent the provided pallets and runtime APIs.
 ///
 /// # Note
 ///
@@ -193,26 +188,29 @@ where
 ///
 /// Panics if the [`scale_info::PortableRegistry`] did not retain all needed types,
 /// or the metadata does not contain the "sp_runtime::DispatchError" type.
-pub fn retain_metadata_pallets<F>(metadata: &mut RuntimeMetadataV15, mut filter: F)
-where
+pub fn retain_metadata<F, G>(
+    metadata: &mut RuntimeMetadataV15,
+    mut pallets_filter: F,
+    mut runtime_apis_filter: G,
+) where
     F: FnMut(&str) -> bool,
+    G: FnMut(&str) -> bool,
 {
     let mut type_ids = HashSet::new();
 
     // There is a special RuntimeCall type which points to all pallets and call types by default.
     // This brings in a significant chunk of types. We trim this down to only include variants
     // for the pallets we're retaining, to avoid this.
-    retain_pallets_in_runtime_call_type(metadata, &mut filter);
+    retain_pallets_in_runtime_call_type(metadata, &mut pallets_filter);
 
     // Filter our pallet list to only those pallets we want to keep. Keep hold of all
-    //type IDs in the pallets we're keeping.
+    // type IDs in the pallets we're keeping. Retain all, if no filter specified.
     metadata.pallets.retain(|pallet| {
-        if filter(&pallet.name) {
+        let should_retain = pallets_filter(&pallet.name);
+        if should_retain {
             collect_pallet_types(pallet, &mut type_ids);
-            true
-        } else {
-            false
         }
+        should_retain
     });
 
     // Keep the extrinsic stuff referenced in our metadata.
@@ -221,8 +219,15 @@ where
     // Keep the "runtime" type ID, since it's referenced in our metadata.
     type_ids.insert(metadata.ty.id);
 
-    // Keep the runtime APIs types.
-    collect_runtime_api_types(&metadata.apis, &mut type_ids);
+    // Keep only the runtime API types that the filter allows for. Keep hold of all
+    // type IDs in the runtime apis we're keeping. Retain all, if no filter specified.
+    metadata.apis.retain(|api| {
+        let should_retain = runtime_apis_filter(&api.name);
+        if should_retain {
+            collect_runtime_api_types(api, &mut type_ids);
+        }
+        should_retain
+    });
 
     // Additionally, subxt depends on the `DispatchError` type existing; we use the same
     // logic here that is used when building our `Metadata`.
@@ -275,10 +280,32 @@ mod tests {
         // Retain one pallet at a time ensuring the test does not panic.
         for pallet in &metadata_cache.pallets {
             let mut metadata = metadata_cache.clone();
-            retain_metadata_pallets(&mut metadata, |pallet_name| pallet_name == pallet.name);
+            retain_metadata(
+                &mut metadata,
+                |pallet_name| pallet_name == pallet.name,
+                |_| true,
+            );
 
             assert_eq!(metadata.pallets.len(), 1);
             assert_eq!(metadata.pallets.get(0).unwrap().name, pallet.name);
         }
     }
+
+    #[test]
+    fn retain_one_runtime_api() {
+        let metadata_cache = load_metadata();
+
+        // Retain one runtime API at a time ensuring the test does not panic.
+        for runtime_api in &metadata_cache.apis {
+            let mut metadata = metadata_cache.clone();
+            retain_metadata(
+                &mut metadata,
+                |_| true,
+                |runtime_api_name| runtime_api_name == runtime_api.name,
+            );
+
+            assert_eq!(metadata.apis.len(), 1);
+            assert_eq!(metadata.apis.get(0).unwrap().name, runtime_api.name);
+        }
+    }
 }

testing/ui-tests/src/utils/pallet_metadata_test_runner.rs

@@ -5,7 +5,7 @@
 use codec::{Decode, Encode};
 use frame_metadata::{v15::RuntimeMetadataV15, RuntimeMetadataPrefixed};
 use std::io::Read;
-use subxt_metadata::{metadata_v14_to_latest, retain_metadata_pallets};
+use subxt_metadata::{metadata_v14_to_latest, retain_metadata};
 
 static TEST_DIR_PREFIX: &str = "subxt_generated_pallets_ui_tests_";
 static METADATA_FILE: &str = "../../artifacts/polkadot_metadata_full.scale";
@@ -59,7 +59,11 @@ impl PalletMetadataTestRunner {
 
         // Build custom metadata containing only this pallet.
         let mut metadata = self.metadata.clone();
-        retain_metadata_pallets(&mut metadata, |pallet_filter| pallet_filter == pallet.name);
+        retain_metadata(
+            &mut metadata,
+            |pallet_filter| pallet_filter == pallet.name,
+            |_| true,
+        );
 
         let mut tmp_dir = std::env::temp_dir();
         tmp_dir.push(format!("{TEST_DIR_PREFIX}{index}"));