Retain Rust code items from mod decorated with subxt attribute (#721)

* refactor: Simplify collecting type substitutes

* Simplify ItemMod::from

* Preserve inner Rust items when expanding the subxt::subxt macro

* No named lifetimes

* Update codegen/src/api/mod.rs

Co-authored-by: Andrew Jones <ascjones@gmail.com>

* Move passing UI tests under `ui_tests`

Co-authored-by: Andrew Jones <ascjones@gmail.com>

codegen/src/api/mod.rs

@@ -204,9 +204,7 @@ impl RuntimeGenerator {
         })
         .collect::<HashMap<_, _>>();
 
-        for (path, substitute) in item_mod_ir.type_substitutes().iter() {
-            type_substitutes.insert(path.to_string(), substitute.clone());
-        }
+        type_substitutes.extend(item_mod_ir.type_substitutes().into_iter());
 
         let type_gen = TypeGenerator::new(
             &self.metadata.types,
@@ -302,7 +300,7 @@ impl RuntimeGenerator {
             }
         };
 
-        let mod_ident = item_mod_ir.ident;
+        let mod_ident = &item_mod_ir.ident;
         let pallets_with_constants: Vec<_> = pallets_with_mod_names
             .iter()
             .filter_map(|(pallet, pallet_mod_name)| {
@@ -324,9 +322,14 @@ impl RuntimeGenerator {
             })
             .collect();
 
+        let rust_items = item_mod_ir.rust_items();
+
         quote! {
             #[allow(dead_code, unused_imports, non_camel_case_types)]
             pub mod #mod_ident {
+                // Preserve any Rust items that were previously defined in the adorned module
+                #( #rust_items ) *
+
                 // Make it easy to access the root via `root_mod` at different levels:
                 use super::#mod_ident as root_mod;
                 // Identify the pallets composing the static metadata by name.

codegen/src/ir.rs

@@ -26,16 +26,13 @@ impl From<syn::ItemMod> for ItemMod {
                 abort!(module, "out-of-line subxt modules are not supported",)
             }
         };
-        let items = items
-            .into_iter()
-            .map(<Item as From<syn::Item>>::from)
-            .collect::<Vec<_>>();
+
         Self {
             vis: module.vis,
             mod_token: module.mod_token,
             ident: module.ident,
             brace,
-            items,
+            items: items.into_iter().map(From::from).collect(),
         }
     }
 }
@@ -57,6 +54,10 @@ impl ItemMod {
             })
             .collect()
     }
+
+    pub fn rust_items(&self) -> impl Iterator<Item = &syn::Item> {
+        self.items.iter().filter_map(Item::as_rust)
+    }
 }
 
 #[allow(clippy::large_enum_variant)]
@@ -66,6 +67,15 @@ pub enum Item {
     Subxt(SubxtItem),
 }
 
+impl Item {
+    pub fn as_rust(&self) -> Option<&syn::Item> {
+        match self {
+            Item::Rust(item) => Some(item),
+            _ => None,
+        }
+    }
+}
+
 impl From<syn::Item> for Item {
     fn from(item: syn::Item) -> Self {
         if let syn::Item::Use(ref use_) = item {

testing/ui-tests/src/correct/rust-items-preserved.rs

@@ -0,0 +1,30 @@
+#[subxt::subxt(runtime_metadata_path = "../../../artifacts/polkadot_metadata.scale")]
+pub mod node_runtime {
+    pub struct SomeStruct;
+    pub enum SomeEnum {
+        A,
+        B,
+    }
+    pub trait SomeTrait {
+        fn some_func(&self) -> u32;
+    }
+    impl SomeTrait for SomeStruct {
+        fn some_func(&self) -> u32 {
+            1
+        }
+    }
+    impl SomeTrait for SomeEnum {
+        fn some_func(&self) -> u32 {
+            2
+        }
+    }
+}
+
+fn main() {
+    use node_runtime::SomeTrait;
+
+    let unit = node_runtime::SomeStruct;
+    assert_eq!(unit.some_func(), 1);
+    let enumeration = node_runtime::SomeEnum::A;
+    assert_eq!(enumeration.some_func(), 2);
+}

testing/ui-tests/src/lib.rs

@@ -25,6 +25,8 @@ fn ui_tests() {
     let mut m = MetadataTestRunner::default();
     let t = trybuild::TestCases::new();
 
+    t.pass("src/correct/*.rs");
+
     // Check that storage maps with no keys are handled properly.
     t.pass(&m.path_to_ui_test_for_metadata(
         "storage_map_no_keys",