Avoid silent failure for associated types. (#16)

tadeohepperle · niklasad1 · web-flow · commit 4ae5a1216fed · 2024-03-19T17:46:38.000+01:00
* dedup bug and error handling

* clippy and fmt

* add test for no skip case

* fix tricky situation with multiple params that share a type id

* Update typegen/src/utils.rs

Co-authored-by: Niklas Adolfsson &lt;niklasadolfsson1@gmail.com&gt;

---------

Co-authored-by: Niklas Adolfsson &lt;niklasadolfsson1@gmail.com&gt;
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/typegen/Cargo.toml b/typegen/Cargo.toml
@@ -14,6 +14,7 @@ quote = { workspace = true }
 scale-info = { workspace = true }
 syn = { workspace = true }
 thiserror = { workspace = true }
+smallvec = { workspace = true }
 
 [dev-dependencies]
 scale-bits = { workspace = true }
diff --git a/typegen/src/tests/mod.rs b/typegen/src/tests/mod.rs
@@ -15,7 +15,7 @@ use crate::{
         },
     },
     utils::ensure_unique_type_paths,
-    DerivesRegistry, TypeSubstitutes,
+    DerivesRegistry, TypeSubstitutes, TypegenError,
 };
 
 mod utils;
@@ -852,7 +852,7 @@ fn range_fields() {
 }
 
 #[test]
-fn generics() {
+fn generics_foo() {
     #[allow(unused)]
     #[derive(TypeInfo)]
     struct Foo<T> {
@@ -1584,3 +1584,115 @@ fn find_similar_type_paths() {
     ];
     assert_eq!(similar, expected);
 }
+
+#[test]
+fn assoc_types_skip_params() {
+    pub trait Config {
+        type Inner: TypeInfo;
+    }
+
+    pub struct A;
+    pub struct B;
+
+    impl Config for A {
+        type Inner = ();
+    }
+
+    impl Config for B {
+        type Inner = u32;
+    }
+
+    #[derive(TypeInfo)]
+    #[scale_info(skip_type_params(T))]
+    pub struct X<T: Config> {
+        pub inner: T::Inner,
+    }
+
+    let res_no_dedup = Testgen::new()
+        .with::<X<A>>()
+        .with::<X<B>>()
+        .try_gen_tests_mod(subxt_settings(), false);
+    assert!(matches!(
+        res_no_dedup,
+        Err(TypegenError::DuplicateTypePath(_))
+    ));
+
+    let dedup_code = Testgen::new()
+        .with::<X<A>>()
+        .with::<X<B>>()
+        .try_gen_tests_mod(Default::default(), true)
+        .unwrap();
+
+    let expected_code = quote! {
+        pub mod tests {
+            use super::types;
+            pub struct X1 {
+                pub inner: (),
+            }
+            pub struct X2 {
+                pub inner: ::core::primitive::u32,
+            }
+        }
+    };
+
+    assert_eq!(dedup_code.to_string(), expected_code.to_string());
+}
+
+#[test]
+fn assoc_types_no_skip_params() {
+    pub trait Config {
+        type Inner: TypeInfo;
+    }
+
+    #[derive(TypeInfo)]
+    pub struct A;
+
+    #[derive(TypeInfo)]
+    pub struct B;
+
+    impl Config for A {
+        type Inner = ();
+    }
+    impl Config for B {
+        type Inner = u32;
+    }
+
+    #[derive(TypeInfo)]
+    pub struct X<T: Config> {
+        pub inner: T::Inner,
+    }
+
+    let res_no_dedup = Testgen::new()
+        .with::<X<A>>()
+        .with::<X<B>>()
+        .try_gen_tests_mod(subxt_settings(), false);
+    assert!(matches!(
+        res_no_dedup,
+        Err(TypegenError::DuplicateTypePath(_))
+    ));
+
+    let dedup_code = Testgen::new()
+        .with::<X<A>>()
+        .with::<X<B>>()
+        .try_gen_tests_mod(Default::default(), true)
+        .unwrap();
+
+    // Unfortunately the structs A and B are still part of the type registry and X1 and X2 have generic parameters, linking to them.
+    let expected_code = quote! {
+        pub mod tests {
+            use super::types;
+            pub struct A;
+            pub struct B;
+            pub struct X1<_0> {
+                pub inner: (),
+                pub __ignore: ::core::marker::PhantomData<_0>
+            }
+            pub struct X2<_0> {
+                pub inner: ::core::primitive::u32,
+                pub __ignore: ::core::marker::PhantomData<_0>
+            }
+        }
+    };
+
+    assert_eq!(dedup_code.to_string(), expected_code.to_string());
+}
diff --git a/typegen/src/tests/utils.rs b/typegen/src/tests/utils.rs
@@ -7,6 +7,7 @@ use syn::parse_quote;
 
 use crate::typegen::ir::module_ir::ModuleIR;
 use crate::utils::ensure_unique_type_paths;
+use crate::TypegenError;
 use crate::{
     typegen::settings::substitutes::absolute_path, DerivesRegistry, TypeGenerator,
     TypeGeneratorSettings, TypeSubstitutes,
@@ -41,12 +42,25 @@ impl Testgen {
         module.to_token_stream()
     }
 
-    pub fn gen_tests_mod(self, settings: TypeGeneratorSettings) -> TokenStream {
-        let registry: PortableRegistry = self.registry.into();
+    pub fn try_gen_tests_mod(
+        self,
+        settings: TypeGeneratorSettings,
+        deduplicate: bool,
+    ) -> Result<TokenStream, TypegenError> {
+        let mut registry: PortableRegistry = self.registry.into();
+        if deduplicate {
+            ensure_unique_type_paths(&mut registry)
+        }
         let type_gen = TypeGenerator::new(&registry, &settings);
-        let module = type_gen.generate_types_mod().unwrap();
-        let module = get_mod(&module, TESTS_MOD_PATH).unwrap();
-        module.to_token_stream()
+        type_gen
+            .generate_types_mod()
+            .map(|module| get_mod(&module, TESTS_MOD_PATH).unwrap().to_token_stream())
+    }
+
+    pub fn gen_tests_mod(self, settings: TypeGeneratorSettings) -> TokenStream {
+        self.try_gen_tests_mod(settings, false)
+            .unwrap()
+            .to_token_stream()
     }
 }
 
@@ -160,7 +174,7 @@ pub(super) fn subxt_default_substitutes() -> TypeSubstitutes {
 
 const TESTS_MOD_PATH: &[&str] = &["scale_typegen", "tests"];
 
-fn get_mod<'a>(module: &'a ModuleIR, path_segs: &[&'static str]) -> Option<&'a ModuleIR> {
+fn get_mod<'a>(module: &'a ModuleIR, path_segs: &[&'static str]) -> Option<&'a ModuleIR<'a>> {
     let (mod_name, rest) = path_segs.split_first()?;
     let mod_ident: Ident = syn::parse_str(mod_name).unwrap();
     let module = module.children.get(&mod_ident)?;
diff --git a/typegen/src/typegen/error.rs b/typegen/src/typegen/error.rs
@@ -31,6 +31,10 @@ pub enum TypegenError {
     /// The settings do not fit the given type registry.
     #[error("Settings do not fit the given type registry: {0}")]
     SettingsValidation(SettingsValidationError),
+    /// There are two types with the the same type path but a different structure.
+    /// Use [`crate::utils::ensure_unique_type_paths`] on your [`scale_info::PortableRegistry`] to deduplicate type paths.
+    #[error("There are two types with the the same type path {0} but different structure. Use `scale_typegen::utils::ensure_unique_type_paths` on your `PortableRegistry` before, to avoid this error.")]
+    DuplicateTypePath(String),
 }
 
 /// Error attempting to do type substitution.
diff --git a/typegen/src/typegen/ir/module_ir.rs b/typegen/src/typegen/ir/module_ir.rs
@@ -9,23 +9,24 @@ use scale_info::form::PortableForm;
 
 /// Represents a Rust `mod`, containing generated types and child `mod`s.
 #[derive(Debug, Clone)]
-pub struct ModuleIR {
+pub struct ModuleIR<'a> {
     /// Name of this module.
     pub name: Ident,
     /// Root module identifier.
     pub root_mod: Ident,
     /// Submodules of this module.
-    pub children: BTreeMap<Ident, ModuleIR>,
+    pub children: BTreeMap<Ident, ModuleIR<'a>>,
     /// Types in this module.
-    pub types: BTreeMap<scale_info::Path<PortableForm>, TypeIR>,
+    pub types:
+        BTreeMap<scale_info::Path<PortableForm>, (&'a scale_info::Type<PortableForm>, TypeIR)>,
 }
 
-impl ToTokens for ModuleIR {
+impl<'a> ToTokens for ModuleIR<'a> {
     fn to_tokens(&self, tokens: &mut TokenStream) {
         let name = &self.name;
         let root_mod = &self.root_mod;
         let modules = self.children.values();
-        let types = self.types.values().clone();
+        let types = self.types.values().map(|e| &e.1).clone();
 
         tokens.extend(quote! {
             pub mod #name {
@@ -38,7 +39,7 @@ impl ToTokens for ModuleIR {
     }
 }
 
-impl ModuleIR {
+impl<'a> ModuleIR<'a> {
     /// Create a new [`Module`], with a reference to the root `mod` for resolving type paths.
     pub(crate) fn new(name: Ident, root_mod: Ident) -> Self {
         Self {
@@ -61,7 +62,7 @@ impl ModuleIR {
 
     /// Returns the generated types.
     pub fn types(&self) -> impl Iterator<Item = (&scale_info::Path<PortableForm>, &TypeIR)> {
-        self.types.iter()
+        self.types.iter().map(|(k, v)| (k, &v.1))
     }
 
     /// Returns the root `mod` used for resolving type paths.
@@ -71,7 +72,7 @@ impl ModuleIR {
 
     /// Recursively creates submodules for the given namespace and returns a mutable reference to the innermost module created this way.
     /// Returns itself, if the namespace is empty.
-    pub fn get_or_insert_submodule(&mut self, namespace: &[String]) -> &mut ModuleIR {
+    pub fn get_or_insert_submodule(&mut self, namespace: &[String]) -> &mut ModuleIR<'a> {
         if namespace.is_empty() {
             return self;
         }
diff --git a/typegen/src/typegen/mod.rs b/typegen/src/typegen/mod.rs
@@ -1,4 +1,6 @@
-use crate::TypegenError;
+use std::collections::btree_map::Entry;
+
+use crate::{utils::types_equal_extended_to_params, TypegenError};
 
 use self::{
     ir::module_ir::ModuleIR,
@@ -89,10 +91,24 @@ impl<'a> TypeGenerator<'a> {
             }
 
             // if the type is not a builtin type, insert it into the respective module
-            if let Some(type_ir) = self.create_type_ir(&ty.ty, &flat_derives_registry)? {
+            let ty = &ty.ty;
+            if let Some(type_ir) = self.create_type_ir(ty, &flat_derives_registry)? {
                 // Create the module this type should go into
                 let innermost_module = root_mod.get_or_insert_submodule(namespace);
-                innermost_module.types.insert(path.clone(), type_ir);
+                match innermost_module.types.entry(path.clone()) {
+                    Entry::Vacant(e) => {
+                        e.insert((ty, type_ir));
+                    }
+                    Entry::Occupied(e) => {
+                        // There is already a type with the same type path present.
+                        // We do not just want to override it, so we check if the two types are semantically similar (structure + generics).
+                        // If not, return an error, if yes, just keep the first one.
+                        let other_ty = &e.get().0;
+                        if !types_equal_extended_to_params(ty, other_ty) {
+                            return Err(TypegenError::DuplicateTypePath(ty.path.to_string()));
+                        }
+                    }
+                };
             }
         }
 
diff --git a/typegen/src/utils.rs b/typegen/src/utils.rs
@@ -1,5 +1,6 @@
 use scale_info::{form::PortableForm, Field, PortableRegistry, Type, TypeDef, TypeParameter};
-use std::collections::HashMap;
+use smallvec::{smallvec, SmallVec};
+use std::collections::{hash_map::Entry, HashMap};
 
 use crate::TypegenError;
 
@@ -81,25 +82,50 @@ pub fn ensure_unique_type_paths(types: &mut PortableRegistry) {
 /// all type ids mentioned in the TypeDef are either:
 /// - equal
 /// - or different, but map essentially to the same generic type parameter
-fn types_equal_extended_to_params(a: &Type<PortableForm>, b: &Type<PortableForm>) -> bool {
-    let collect_params = |type_params: &[TypeParameter<PortableForm>]| {
-        type_params
-            .iter()
-            .filter_map(|p| p.ty.map(|ty| (ty.id, p.name.clone())))
-            .collect::<HashMap<u32, String>>()
-    };
+pub(crate) fn types_equal_extended_to_params(
+    a: &Type<PortableForm>,
+    b: &Type<PortableForm>,
+) -> bool {
+    // We map each type ID to all type params if could refer to. Type IDs can refer to multiple parameters:
+    // E.g. Foo<A,B> can be parameterized as Foo<u8,u8>, so if 42 is the type id of u8, a field with id=42 could refer to either A or B.
+    fn collect_params(
+        type_params: &[TypeParameter<PortableForm>],
+    ) -> HashMap<u32, SmallVec<[&str; 2]>> {
+        let mut hm: HashMap<u32, SmallVec<[&str; 2]>> = HashMap::new();
+        for p in type_params {
+            if let Some(ty) = &p.ty {
+                match hm.entry(ty.id) {
+                    Entry::Occupied(mut e) => {
+                        e.get_mut().push(p.name.as_str());
+                    }
+                    Entry::Vacant(e) => {
+                        e.insert(smallvec![p.name.as_str()]);
+                    }
+                }
+            }
+        }
+        hm
+    }
 
     let type_params_a = collect_params(&a.type_params);
-    let type_params_b = collect_params(&a.type_params);
+    let type_params_b = collect_params(&b.type_params);
 
-    if type_params_a.len() != type_params_b.len() {
+    if a.type_params.len() != b.type_params.len() {
         return false;
     }
-
-    // returns true if the ids are the same OR if they point to the same generic parameter.
+    // Returns true if the ids are the same OR if they point to the same generic parameter.
     let ids_equal = |a: u32, b: u32| -> bool {
-        a == b
-            || matches!((type_params_a.get(&a), type_params_b.get(&b)), (Some(a_name), Some(b_name)) if a_name == b_name)
+        if a == b {
+            return true;
+        }
+        let Some(a_param_names) = type_params_a.get(&a) else {
+            return false;
+        };
+        let Some(b_param_names) = type_params_b.get(&b) else {
+            return false;
+        };
+        // Check if there is any intersection, meaning that both IDs map to the same generic type param:
+        a_param_names.iter().any(|a_p| b_param_names.contains(a_p))
     };
 
     let fields_equal = |a: &[Field<PortableForm>], b: &[Field<PortableForm>]| -> bool {
