Fixing a generic types resolution bug (FuelLabs#736)

* Initial fix

* Some fixes + adding a test

* Initial fix

* Some fixes + adding a test

* fmt

* fixing more cases and imporving the test

* clippy

* Addressing review comment

sway-core/src/semantic_analysis/namespace.rs

@@ -118,6 +118,28 @@ impl Namespace {
         methods
     }
 
+    // Given a TypeInfo old_type with a set of methods available to it, make those same methods
+    // available to TypeInfo new_type. This is useful in situations where old_type is being
+    // monomorphized to new_type and and we want `get_methods_for_type()` to return the same set of
+    // methods for new_type as it does for old_type.
+    pub(crate) fn copy_methods_to_type(&mut self, old_type: TypeInfo, new_type: TypeInfo) {
+        // This map grabs all (trait name, vec of methods) from self.implemented_traits
+        // corresponding to `old_type`.
+        let mut methods: HashMap<TraitName, Vec<TypedFunctionDeclaration>> = HashMap::new();
+        for ((trait_name, type_info), l_methods) in &self.implemented_traits {
+            if *type_info == old_type {
+                methods.insert((*trait_name).clone(), l_methods.clone());
+            }
+        }
+
+        // Insert into `self.implemented_traits` the contents of the map above but with `new_type`
+        // as the `TypeInfo` key.
+        for (trait_name, l_methods) in &methods {
+            self.implemented_traits
+                .insert(((*trait_name).clone(), new_type.clone()), l_methods.clone());
+        }
+    }
+
     pub(crate) fn get_tuple_elems(
         &self,
         ty: TypeId,

sway-core/src/semantic_analysis/namespace/arena.rs

@@ -62,6 +62,7 @@ pub trait NamespaceWrapper {
     fn star_import(&self, from_module: Option<NamespaceRef>, path: Vec<Ident>)
         -> CompileResult<()>;
     fn get_methods_for_type(&self, r#type: TypeId) -> Vec<TypedFunctionDeclaration>;
+    fn copy_methods_to_type(&self, old_type: TypeInfo, new_type: TypeInfo);
     fn get_name_from_path(&self, path: &[Ident], name: &Ident) -> CompileResult<TypedDeclaration>;
     /// Used for calls that look like this:
     /// `foo::bar::function`
@@ -273,6 +274,13 @@ impl NamespaceWrapper for NamespaceRef {
     fn get_methods_for_type(&self, r#type: TypeId) -> Vec<TypedFunctionDeclaration> {
         read_module(|ns| ns.get_methods_for_type(r#type), *self)
     }
+    fn copy_methods_to_type(&self, old_type: TypeInfo, new_type: TypeInfo) {
+        write_module(
+            move |ns| ns.copy_methods_to_type(old_type.clone(), new_type),
+            *self,
+        )
+    }
+
     fn star_import(
         &self,
         from_module: Option<NamespaceRef>,
@@ -588,28 +596,54 @@ impl NamespaceWrapper for NamespaceRef {
         Ok(match ty {
             TypeInfo::Custom { ref name } => {
                 match self.get_symbol(name).ok(&mut warnings, &mut errors) {
-                    Some(TypedDeclaration::StructDeclaration(TypedStructDeclaration {
-                        name,
-                        fields,
-                        ..
-                    })) => crate::type_engine::insert_type(TypeInfo::Struct {
-                        name: name.as_str().to_string(),
-                        fields: fields
-                            .iter()
-                            .map(TypedStructField::as_owned_typed_struct_field)
-                            .collect::<Vec<_>>(),
-                    }),
-                    Some(TypedDeclaration::EnumDeclaration(TypedEnumDeclaration {
-                        name,
-                        variants,
-                        ..
-                    })) => crate::type_engine::insert_type(TypeInfo::Enum {
-                        name: name.as_str().to_string(),
-                        variant_types: variants
-                            .iter()
-                            .map(TypedEnumVariant::as_owned_typed_enum_variant)
-                            .collect(),
-                    }),
+                    Some(TypedDeclaration::StructDeclaration(decl)) => {
+                        let old_struct = TypeInfo::Struct {
+                            name: decl.name.as_str().to_string(),
+                            fields: decl
+                                .fields
+                                .iter()
+                                .map(TypedStructField::as_owned_typed_struct_field)
+                                .collect::<Vec<_>>(),
+                        };
+                        let mut new_struct = old_struct.clone();
+                        if !decl.type_parameters.is_empty() {
+                            let new_decl = decl.monomorphize();
+                            new_struct = TypeInfo::Struct {
+                                name: new_decl.name.as_str().to_string(),
+                                fields: new_decl
+                                    .fields
+                                    .iter()
+                                    .map(TypedStructField::as_owned_typed_struct_field)
+                                    .collect::<Vec<_>>(),
+                            };
+                            self.copy_methods_to_type(old_struct, new_struct.clone());
+                        }
+                        crate::type_engine::insert_type(new_struct)
+                    }
+                    Some(TypedDeclaration::EnumDeclaration(decl)) => {
+                        let old_enum = TypeInfo::Enum {
+                            name: decl.name.as_str().to_string(),
+                            variant_types: decl
+                                .variants
+                                .iter()
+                                .map(TypedEnumVariant::as_owned_typed_enum_variant)
+                                .collect(),
+                        };
+                        let mut new_enum = old_enum.clone();
+                        if !decl.type_parameters.is_empty() {
+                            let new_decl = decl.monomorphize();
+                            new_enum = TypeInfo::Enum {
+                                name: new_decl.name.as_str().to_string(),
+                                variant_types: new_decl
+                                    .variants
+                                    .iter()
+                                    .map(TypedEnumVariant::as_owned_typed_enum_variant)
+                                    .collect(),
+                            };
+                            self.copy_methods_to_type(old_enum, new_enum.clone());
+                        }
+                        crate::type_engine::insert_type(new_enum)
+                    }
                     Some(TypedDeclaration::GenericTypeForFunctionScope { name, .. }) => {
                         crate::type_engine::insert_type(TypeInfo::UnknownGeneric { name })
                     }

test/src/e2e_vm_tests/mod.rs

@@ -81,6 +81,7 @@ pub fn run(filter_regex: Option<regex::Regex>) {
         ("multi_item_import", ProgramState::Return(0)), // false
         ("tuple_indexing", ProgramState::Return(1)),
         ("tuple_access", ProgramState::Return(42)),
+        ("funcs_with_generic_types", ProgramState::Return(1)), // true
     ];
 
     let mut number_of_tests_run = positive_project_names.iter().fold(0, |acc, (name, res)| {

test/src/e2e_vm_tests/test_programs/funcs_with_generic_types/Forc.toml

@@ -0,0 +1,5 @@
+[project]
+author = "Fuel Labs <[email protected]>"
+license = "Apache-2.0"
+name = "funcs_with_generic_types"
+entry = "main.sw"

test/src/e2e_vm_tests/test_programs/funcs_with_generic_types/json_abi_oracle.json

@@ -0,0 +1 @@
+[]

test/src/e2e_vm_tests/test_programs/funcs_with_generic_types/src/main.sw

@@ -0,0 +1,103 @@
+script;
+
+/* ------------------*/
+
+struct Foo1 {
+    a: u64,
+    b: u64,
+}
+
+impl Foo1 {
+    fn trivial(self) -> bool {
+      false
+    }
+}
+
+fn func1() -> bool {
+    let f = Foo1 {a: 0, b: 0};
+    f.trivial()
+}
+
+/* ------------------*/
+
+enum Bar {
+    a: (),
+    b: (),
+}
+
+impl Bar {
+    fn trivial(self) -> bool {
+        false
+    }
+}
+
+fn func2(m: Bar) -> bool {
+    m.trivial()
+}
+
+/* ------------------*/
+
+struct Foo2<T> {
+    foo: T,
+}
+
+impl Foo2<T> {
+    fn trivial(self) -> bool {
+        false
+    }
+}
+
+fn func3(a: Foo2<u8>) -> Foo2<bool> {
+    if a.trivial() {
+        Foo2 {foo: false}
+    } else {
+        Foo2 {foo: true}
+    }
+}
+
+fn func4(b: Foo2<bool>) -> Foo2<u8> {
+    if b.trivial() {
+        Foo2 {foo: 0u8} 
+    } else {
+        Foo2 {foo: 1u8}
+    }
+}
+
+/* ------------------*/
+
+pub enum Rezult<T, E> {
+    Ok: T,
+    Err: E,
+}
+
+pub enum DumbError {
+    Error: (),
+}
+
+impl Rezult<T, E> {
+    fn trivial(self) -> bool {
+        false
+    }
+}
+
+pub fn func5(r: Rezult<u8, DumbError>) -> Rezult<u8, DumbError> {
+    if r.trivial() {
+        Rezult::Err(DumbError::Error)
+    } else {
+        Rezult::Ok(1u8)
+    }
+}
+
+pub fn func6(r: Rezult<bool, DumbError>) -> Rezult<bool, DumbError> {
+   if r.trivial() {
+        Rezult::Err(DumbError::Error)
+    } else {
+        Rezult::Ok(true)
+    }
+}
+
+/* ------------------*/
+
+fn main() -> bool {
+  true
+}