Fixes replace_decls issue with nested and multiple types.

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)
     }
 }