Fixes replace_decls issue with nested and multiple types. (#4889)

## Description

When a single method uses multiple trait implementation errors are
thrown by the compiler similar issue occurs when a method requires a
trait implementation and then calls another method that requires another
implementation of the same trait.

This PR fixes both issues by filtering traits to be replaced by the
arguments type used and trait method self type, and by gathering decl
mapping of nested methods.

Fixes #4852

## Checklist

- [x] I have linked to any relevant issues.
- [x] I have commented my code, particularly in hard-to-understand
areas.
- [ ] I have updated the documentation where relevant (API docs, the
reference, and the Sway book).
- [x] I have added tests that prove my fix is effective or that my
feature works.
- [x] I have added (or requested a maintainer to add) the necessary
`Breaking*` or `New Feature` labels where relevant.
- [x] I have done my best to ensure that my PR adheres to [the Fuel Labs
Code Review
Standards](https://github.com/FuelLabs/rfcs/blob/master/text/code-standards/external-contributors.md).
- [x] I have requested a review from the relevant team or maintainers.

sway-core/src/decl_engine/mapping.rs

@@ -1,8 +1,11 @@
 use std::fmt;
 
-use crate::language::ty::{TyTraitInterfaceItem, TyTraitItem};
+use crate::{
+    language::ty::{TyTraitInterfaceItem, TyTraitItem},
+    Engines, TypeId, UnifyCheck,
+};
 
-use super::{AssociatedItemDeclId, InterfaceItemMap, ItemMap};
+use super::{AssociatedItemDeclId, DeclEngineGet, InterfaceItemMap, ItemMap};
 
 type SourceDecl = AssociatedItemDeclId;
 type DestinationDecl = AssociatedItemDeclId;
@@ -99,4 +102,30 @@ impl DeclMapping {
         }
         None
     }
+
+    /// This method returns only associated item functions that have as self type the given type.
+    pub(crate) fn filter_functions_by_self_type(
+        &self,
+        self_type: TypeId,
+        engines: &Engines,
+    ) -> DeclMapping {
+        let mut mapping: Vec<(SourceDecl, DestinationDecl)> = vec![];
+        for (source_decl_ref, dest_decl_ref) in self.mapping.iter().cloned() {
+            match dest_decl_ref {
+                AssociatedItemDeclId::TraitFn(_) => mapping.push((source_decl_ref, dest_decl_ref)),
+                AssociatedItemDeclId::Function(func_id) => {
+                    let func = engines.de().get(&func_id);
+
+                    let unify_check = UnifyCheck::non_dynamic_equality(engines);
+                    if let (left, Some(right)) = (self_type, func.parameters.get(0)) {
+                        if unify_check.check(left, right.type_argument.type_id) {
+                            mapping.push((source_decl_ref, dest_decl_ref));
+                        }
+                    }
+                }
+                AssociatedItemDeclId::Constant(_) => mapping.push((source_decl_ref, dest_decl_ref)),
+            }
+        }
+        DeclMapping { mapping }
+    }
 }

sway-core/src/decl_engine/mod.rs

@@ -30,7 +30,10 @@ pub(crate) use replace_decls::*;
 use sway_types::Ident;
 pub(crate) use template::*;
 
-use crate::language::ty::{TyTraitInterfaceItem, TyTraitItem};
+use crate::{
+    language::ty::{TyTraitInterfaceItem, TyTraitItem},
+    TypeId,
+};
 
-pub(crate) type InterfaceItemMap = BTreeMap<Ident, TyTraitInterfaceItem>;
-pub(crate) type ItemMap = BTreeMap<Ident, TyTraitItem>;
+pub(crate) type InterfaceItemMap = BTreeMap<(Ident, TypeId), TyTraitInterfaceItem>;
+pub(crate) type ItemMap = BTreeMap<(Ident, TypeId), TyTraitItem>;

sway-core/src/decl_engine/ref.rs

@@ -22,6 +22,7 @@
 
 use std::hash::{Hash, Hasher};
 
+use sway_error::handler::{ErrorEmitted, Handler};
 use sway_types::{Ident, Named, Span, Spanned};
 
 use crate::{
@@ -31,6 +32,7 @@ use crate::{
         self, TyAbiDecl, TyConstantDecl, TyEnumDecl, TyFunctionDecl, TyImplTrait, TyStorageDecl,
         TyStructDecl, TyTraitDecl, TyTraitFn,
     },
+    semantic_analysis::TypeCheckContext,
     type_system::*,
 };
 
@@ -195,14 +197,15 @@ where
     pub(crate) fn replace_decls_and_insert_new_with_parent(
         &self,
         decl_mapping: &DeclMapping,
-        engines: &Engines,
-    ) -> Self {
-        let decl_engine = engines.de();
+        handler: &Handler,
+        ctx: &TypeCheckContext,
+    ) -> Result<Self, ErrorEmitted> {
+        let decl_engine = ctx.engines().de();
         let mut decl = decl_engine.get(&self.id);
-        decl.replace_decls(decl_mapping, engines);
-        decl_engine
+        decl.replace_decls(decl_mapping, handler, ctx)?;
+        Ok(decl_engine
             .insert(decl)
-            .with_parent(decl_engine, self.id.into())
+            .with_parent(decl_engine, self.id.into()))
     }
 }
 
@@ -334,23 +337,30 @@ where
 }
 
 impl ReplaceDecls for DeclRefFunction {
-    fn replace_decls_inner(&mut self, decl_mapping: &DeclMapping, engines: &Engines) {
+    fn replace_decls_inner(
+        &mut self,
+        decl_mapping: &DeclMapping,
+        _handler: &Handler,
+        ctx: &TypeCheckContext,
+    ) -> Result<(), ErrorEmitted> {
+        let engines = ctx.engines();
         let decl_engine = engines.de();
         if let Some(new_decl_ref) = decl_mapping.find_match(self.id.into()) {
             if let AssociatedItemDeclId::Function(new_decl_ref) = new_decl_ref {
                 self.id = new_decl_ref;
             }
-            return;
+            return Ok(());
         }
         let all_parents = decl_engine.find_all_parents(engines, &self.id);
         for parent in all_parents.iter() {
             if let Some(new_decl_ref) = decl_mapping.find_match(parent.clone()) {
                 if let AssociatedItemDeclId::Function(new_decl_ref) = new_decl_ref {
                     self.id = new_decl_ref;
                 }
-                return;
+                return Ok(());
             }
         }
+        Ok(())
     }
 }
 

sway-core/src/decl_engine/replace_decls.rs

@@ -1,17 +1,32 @@
+use sway_error::handler::{ErrorEmitted, Handler};
+
 use crate::{
     engine_threading::Engines,
     language::ty::{self, TyDecl},
+    semantic_analysis::TypeCheckContext,
 };
 
 use super::DeclMapping;
 
 pub trait ReplaceDecls {
-    fn replace_decls_inner(&mut self, decl_mapping: &DeclMapping, engines: &Engines);
+    fn replace_decls_inner(
+        &mut self,
+        decl_mapping: &DeclMapping,
+        handler: &Handler,
+        ctx: &TypeCheckContext,
+    ) -> Result<(), ErrorEmitted>;
 
-    fn replace_decls(&mut self, decl_mapping: &DeclMapping, engines: &Engines) {
+    fn replace_decls(
+        &mut self,
+        decl_mapping: &DeclMapping,
+        handler: &Handler,
+        ctx: &TypeCheckContext,
+    ) -> Result<(), ErrorEmitted> {
         if !decl_mapping.is_empty() {
-            self.replace_decls_inner(decl_mapping, engines);
+            self.replace_decls_inner(decl_mapping, handler, ctx)?;
         }
+
+        Ok(())
     }
 }
 

sway-core/src/language/ty/ast_node.rs

@@ -10,6 +10,7 @@ use crate::{
     decl_engine::*,
     engine_threading::*,
     language::{parsed::TreeType, ty::*, Visibility},
+    semantic_analysis::TypeCheckContext,
     transform::AttributeKind,
     type_system::*,
     types::*,
@@ -87,17 +88,24 @@ impl ReplaceSelfType for TyAstNode {
 }
 
 impl ReplaceDecls for TyAstNode {
-    fn replace_decls_inner(&mut self, decl_mapping: &DeclMapping, engines: &Engines) {
+    fn replace_decls_inner(
+        &mut self,
+        decl_mapping: &DeclMapping,
+        handler: &Handler,
+        ctx: &TypeCheckContext,
+    ) -> Result<(), ErrorEmitted> {
         match self.content {
             TyAstNodeContent::ImplicitReturnExpression(ref mut exp) => {
-                exp.replace_decls(decl_mapping, engines)
+                exp.replace_decls(decl_mapping, handler, ctx)
             }
             TyAstNodeContent::Declaration(TyDecl::VariableDecl(ref mut decl)) => {
-                decl.body.replace_decls(decl_mapping, engines);
+                decl.body.replace_decls(decl_mapping, handler, ctx)
             }
-            TyAstNodeContent::Declaration(_) => {}
-            TyAstNodeContent::Expression(ref mut expr) => expr.replace_decls(decl_mapping, engines),
-            TyAstNodeContent::SideEffect(_) => (),
+            TyAstNodeContent::Declaration(_) => Ok(()),
+            TyAstNodeContent::Expression(ref mut expr) => {
+                expr.replace_decls(decl_mapping, handler, ctx)
+            }
+            TyAstNodeContent::SideEffect(_) => Ok(()),
         }
     }
 }

sway-core/src/language/ty/code_block.rs

@@ -1,8 +1,10 @@
 use std::hash::Hasher;
 
+use sway_error::handler::{ErrorEmitted, Handler};
+
 use crate::{
-    decl_engine::*, engine_threading::*, language::ty::*, type_system::*,
-    types::DeterministicallyAborts,
+    decl_engine::*, engine_threading::*, language::ty::*, semantic_analysis::TypeCheckContext,
+    type_system::*, types::DeterministicallyAborts,
 };
 
 #[derive(Clone, Debug)]
@@ -41,10 +43,24 @@ impl ReplaceSelfType for TyCodeBlock {
 }
 
 impl ReplaceDecls for TyCodeBlock {
-    fn replace_decls_inner(&mut self, decl_mapping: &DeclMapping, engines: &Engines) {
-        self.contents
-            .iter_mut()
-            .for_each(|x| x.replace_decls(decl_mapping, engines));
+    fn replace_decls_inner(
+        &mut self,
+        decl_mapping: &DeclMapping,
+        handler: &Handler,
+        ctx: &TypeCheckContext,
+    ) -> Result<(), ErrorEmitted> {
+        handler.scope(|handler| {
+            for x in self.contents.iter_mut() {
+                match x.replace_decls(decl_mapping, handler, ctx) {
+                    Ok(res) => res,
+                    Err(_) => {
+                        continue;
+                    }
+                };
+            }
+
+            Ok(())
+        })
     }
 }
 

sway-core/src/language/ty/declaration/constant.rs

@@ -1,11 +1,13 @@
 use std::hash::{Hash, Hasher};
 
+use sway_error::handler::{ErrorEmitted, Handler};
 use sway_types::{Ident, Named, Span, Spanned};
 
 use crate::{
     decl_engine::{DeclMapping, ReplaceDecls},
     engine_threading::*,
     language::{ty::*, CallPath, Visibility},
+    semantic_analysis::TypeCheckContext,
     transform,
     type_system::*,
 };
@@ -103,9 +105,16 @@ impl ReplaceSelfType for TyConstantDecl {
 }
 
 impl ReplaceDecls for TyConstantDecl {
-    fn replace_decls_inner(&mut self, decl_mapping: &DeclMapping, engines: &Engines) {
+    fn replace_decls_inner(
+        &mut self,
+        decl_mapping: &DeclMapping,
+        handler: &Handler,
+        ctx: &TypeCheckContext,
+    ) -> Result<(), ErrorEmitted> {
         if let Some(expr) = &mut self.value {
-            expr.replace_decls(decl_mapping, engines);
+            expr.replace_decls(decl_mapping, handler, ctx)
+        } else {
+            Ok(())
         }
     }
 }

sway-core/src/language/ty/declaration/function.rs

@@ -10,6 +10,7 @@ use crate::{
     decl_engine::*,
     engine_threading::*,
     language::{parsed, ty::*, Inline, Purity, Visibility},
+    semantic_analysis::TypeCheckContext,
     transform,
     type_system::*,
     types::*,
@@ -123,8 +124,13 @@ impl ReplaceSelfType for TyFunctionDecl {
 }
 
 impl ReplaceDecls for TyFunctionDecl {
-    fn replace_decls_inner(&mut self, decl_mapping: &DeclMapping, engines: &Engines) {
-        self.body.replace_decls(decl_mapping, engines);
+    fn replace_decls_inner(
+        &mut self,
+        decl_mapping: &DeclMapping,
+        handler: &Handler,
+        ctx: &TypeCheckContext,
+    ) -> Result<(), ErrorEmitted> {
+        self.body.replace_decls(decl_mapping, handler, ctx)
     }
 }
 

sway-core/src/language/ty/expression/expression.rs

@@ -7,6 +7,7 @@ use crate::{
     decl_engine::*,
     engine_threading::*,
     language::{ty::*, Literal},
+    semantic_analysis::TypeCheckContext,
     type_system::*,
     types::*,
 };
@@ -59,8 +60,13 @@ impl ReplaceSelfType for TyExpression {
 }
 
 impl ReplaceDecls for TyExpression {
-    fn replace_decls_inner(&mut self, decl_mapping: &DeclMapping, engines: &Engines) {
-        self.expression.replace_decls(decl_mapping, engines);
+    fn replace_decls_inner(
+        &mut self,
+        decl_mapping: &DeclMapping,
+        handler: &Handler,
+        ctx: &TypeCheckContext,
+    ) -> Result<(), ErrorEmitted> {
+        self.expression.replace_decls(decl_mapping, handler, ctx)
     }
 }