From eb266b9be6440cf2d39b2976aaec19dfaee6cd82 Mon Sep 17 00:00:00 2001
From: Andrew Jones <ascjones@gmail.com>
Date: Fri, 28 Jan 2022 09:28:04 +0000
Subject: [PATCH] Only convert struct names to camel case for Call variant
 structs (#412)

* Only convert struct names to camel case for `Call` variant structs

* Use Cow for transform fn

* Use as_ref
---
 codegen/src/api/calls.rs           | 13 +++++++++---
 codegen/src/api/events.rs          |  8 ++++++--
 codegen/src/api/mod.rs             | 11 +++++++---
 codegen/src/types/composite_def.rs |  3 +--
 codegen/src/types/tests.rs         | 33 ++++++++++++++++++++++++++++++
 5 files changed, 58 insertions(+), 10 deletions(-)

diff --git a/codegen/src/api/calls.rs b/codegen/src/api/calls.rs
index 25633ff4dc..ea769a3d94 100644
--- a/codegen/src/api/calls.rs
+++ b/codegen/src/api/calls.rs
@@ -22,7 +22,10 @@ use frame_metadata::{
     PalletCallMetadata,
     PalletMetadata,
 };
-use heck::SnakeCase as _;
+use heck::{
+    CamelCase as _,
+    SnakeCase as _,
+};
 use proc_macro2::TokenStream as TokenStream2;
 use proc_macro_error::abort_call_site;
 use quote::{
@@ -37,8 +40,12 @@ pub fn generate_calls(
     call: &PalletCallMetadata<PortableForm>,
     types_mod_ident: &syn::Ident,
 ) -> TokenStream2 {
-    let struct_defs =
-        super::generate_structs_from_variants(type_gen, call.ty.id(), "Call");
+    let struct_defs = super::generate_structs_from_variants(
+        type_gen,
+        call.ty.id(),
+        |name| name.to_camel_case().into(),
+        "Call",
+    );
     let (call_structs, call_fns): (Vec<_>, Vec<_>) = struct_defs
         .iter()
         .map(|struct_def| {
diff --git a/codegen/src/api/events.rs b/codegen/src/api/events.rs
index 0d0f51f6fe..6c2389b52b 100644
--- a/codegen/src/api/events.rs
+++ b/codegen/src/api/events.rs
@@ -29,8 +29,12 @@ pub fn generate_events(
     event: &PalletEventMetadata<PortableForm>,
     types_mod_ident: &syn::Ident,
 ) -> TokenStream2 {
-    let struct_defs =
-        super::generate_structs_from_variants(type_gen, event.ty.id(), "Event");
+    let struct_defs = super::generate_structs_from_variants(
+        type_gen,
+        event.ty.id(),
+        |name| name.into(),
+        "Event",
+    );
     let event_structs = struct_defs.iter().map(|struct_def| {
         let pallet_name = &pallet.name;
         let event_struct = &struct_def.name;
diff --git a/codegen/src/api/mod.rs b/codegen/src/api/mod.rs
index fa1ab9e879..2d41abfadf 100644
--- a/codegen/src/api/mod.rs
+++ b/codegen/src/api/mod.rs
@@ -352,19 +352,24 @@ impl RuntimeGenerator {
     }
 }
 
-pub fn generate_structs_from_variants<'a>(
+pub fn generate_structs_from_variants<'a, F>(
     type_gen: &'a TypeGenerator,
     type_id: u32,
+    variant_to_struct_name: F,
     error_message_type_name: &str,
-) -> Vec<CompositeDef> {
+) -> Vec<CompositeDef>
+where
+    F: Fn(&str) -> std::borrow::Cow<str>,
+{
     let ty = type_gen.resolve_type(type_id);
     if let scale_info::TypeDef::Variant(variant) = ty.type_def() {
         variant
             .variants()
             .iter()
             .map(|var| {
+                let struct_name = variant_to_struct_name(var.name());
                 let fields = CompositeDefFields::from_scale_info_fields(
-                    var.name(),
+                    struct_name.as_ref(),
                     var.fields(),
                     &[],
                     type_gen,
diff --git a/codegen/src/types/composite_def.rs b/codegen/src/types/composite_def.rs
index b6e5548f06..86d1aa31a4 100644
--- a/codegen/src/types/composite_def.rs
+++ b/codegen/src/types/composite_def.rs
@@ -22,7 +22,6 @@ use super::{
     TypeParameter,
     TypePath,
 };
-use heck::CamelCase as _;
 use proc_macro2::TokenStream;
 use proc_macro_error::abort_call_site;
 use quote::{
@@ -85,7 +84,7 @@ impl CompositeDef {
             }
         }
 
-        let name = format_ident!("{}", ident.to_camel_case());
+        let name = format_ident!("{}", ident);
 
         Self {
             name,
diff --git a/codegen/src/types/tests.rs b/codegen/src/types/tests.rs
index a1f2fe094a..639ce2a949 100644
--- a/codegen/src/types/tests.rs
+++ b/codegen/src/types/tests.rs
@@ -853,3 +853,36 @@ fn modules() {
         .to_string()
     )
 }
+
+#[test]
+fn dont_force_struct_names_camel_case() {
+    #[allow(unused)]
+    #[derive(TypeInfo)]
+    struct AB;
+
+    let mut registry = Registry::new();
+    registry.register_type(&meta_type::<AB>());
+    let portable_types: PortableRegistry = registry.into();
+
+    let type_gen = TypeGenerator::new(
+        &portable_types,
+        "root",
+        Default::default(),
+        Default::default(),
+    );
+    let types = type_gen.generate_types_mod();
+    let tests_mod = get_mod(&types, MOD_PATH).unwrap();
+
+    assert_eq!(
+        tests_mod.into_token_stream().to_string(),
+        quote! {
+            pub mod tests {
+                use super::root;
+
+                #[derive(::subxt::codec::Encode, ::subxt::codec::Decode, Debug)]
+                pub struct AB;
+            }
+        }
+        .to_string()
+    )
+}