Fixes ReplaceDecls to no longer depend on filter hack. (#5541)

`ReplaceDecls` was using a hack to replace some function refs based on
the self function parameter, this didn't work for associated functions.

Now DeclMapping sources also contain a typeid that guarantees that the
replacement is done to the correct declaration in case we have nested
trait methods with the same name replaced.

Fixes #5504

## 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.

Co-authored-by: Joshua Batty <joshpbatty@gmail.com>

sway-core/src/decl_engine/associated_item_decl_id.rs

@@ -1,9 +1,11 @@
 use sway_error::error::CompileError;
-use sway_types::Span;
+use sway_types::{Named, Span, Spanned};
 
 use crate::{
     decl_engine::*,
+    engine_threading::DisplayWithEngines,
     language::ty::{self, TyFunctionDecl},
+    Engines,
 };
 
 #[derive(Debug, Eq, PartialEq, Hash, Clone)]
@@ -14,6 +16,17 @@ pub enum AssociatedItemDeclId {
     Type(DeclId<ty::TyTraitType>),
 }
 
+impl AssociatedItemDeclId {
+    pub fn span(&self, engines: &Engines) -> Span {
+        match self {
+            Self::TraitFn(decl_id) => engines.de().get(decl_id).span(),
+            Self::Function(decl_id) => engines.de().get(decl_id).span(),
+            Self::Constant(decl_id) => engines.de().get(decl_id).span(),
+            Self::Type(decl_id) => engines.de().get(decl_id).span(),
+        }
+    }
+}
+
 impl From<DeclId<ty::TyFunctionDecl>> for AssociatedItemDeclId {
     fn from(val: DeclId<ty::TyFunctionDecl>) -> Self {
         Self::Function(val)
@@ -97,6 +110,29 @@ impl std::fmt::Display for AssociatedItemDeclId {
     }
 }
 
+impl DisplayWithEngines for AssociatedItemDeclId {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>, engines: &Engines) -> std::fmt::Result {
+        match self {
+            Self::TraitFn(decl_id) => {
+                write!(
+                    f,
+                    "decl(trait function {})",
+                    engines.de().get(decl_id).name()
+                )
+            }
+            Self::Function(decl_id) => {
+                write!(f, "decl(function {})", engines.de().get(decl_id).name())
+            }
+            Self::Constant(decl_id) => {
+                write!(f, "decl(constant {})", engines.de().get(decl_id).name())
+            }
+            Self::Type(decl_id) => {
+                write!(f, "decl(type {})", engines.de().get(decl_id).name())
+            }
+        }
+    }
+}
+
 impl TryFrom<DeclRefMixedFunctional> for DeclRefFunction {
     type Error = CompileError;
     fn try_from(value: DeclRefMixedFunctional) -> Result<Self, Self::Error> {

sway-core/src/decl_engine/mapping.rs

@@ -1,17 +1,21 @@
-use std::fmt;
+use std::{collections::HashSet, fmt};
+
+use sway_error::handler::{ErrorEmitted, Handler};
 
 use crate::{
+    engine_threading::DebugWithEngines,
     language::ty::{TyTraitInterfaceItem, TyTraitItem},
     Engines, TypeId, UnifyCheck,
 };
 
-use super::{AssociatedItemDeclId, DeclEngineGet, InterfaceItemMap, ItemMap};
+use super::{AssociatedItemDeclId, InterfaceItemMap, ItemMap};
 
-type SourceDecl = AssociatedItemDeclId;
+type SourceDecl = (AssociatedItemDeclId, TypeId);
 type DestinationDecl = AssociatedItemDeclId;
 
 /// The [DeclMapping] is used to create a mapping between a [SourceDecl] (LHS)
 /// and a [DestinationDecl] (RHS).
+#[derive(Clone)]
 pub struct DeclMapping {
     mapping: Vec<(SourceDecl, DestinationDecl)>,
 }
@@ -26,7 +30,7 @@ impl fmt::Display for DeclMapping {
                 .map(|(source_type, dest_type)| {
                     format!(
                         "{} -> {}",
-                        source_type,
+                        source_type.0,
                         match dest_type {
                             AssociatedItemDeclId::TraitFn(decl_id) => decl_id.inner(),
                             AssociatedItemDeclId::Function(decl_id) => decl_id.inner(),
@@ -55,17 +59,27 @@ impl fmt::Debug for DeclMapping {
     }
 }
 
-impl DeclMapping {
-    pub(crate) fn new() -> Self {
-        Self {
-            mapping: Vec::new(),
-        }
-    }
-
-    pub(crate) fn insert(&mut self, k: SourceDecl, v: DestinationDecl) {
-        self.mapping.push((k, v))
+impl DebugWithEngines for DeclMapping {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>, engines: &Engines) -> fmt::Result {
+        write!(
+            f,
+            "DeclMapping {{ {} }}",
+            self.mapping
+                .iter()
+                .map(|(source_type, dest_type)| {
+                    format!(
+                        "{} -> {}",
+                        engines.help_out(source_type.0.clone()),
+                        engines.help_out(dest_type)
+                    )
+                })
+                .collect::<Vec<_>>()
+                .join(", ")
+        )
     }
+}
 
+impl DeclMapping {
     pub(crate) fn is_empty(&self) -> bool {
         self.mapping.is_empty()
     }
@@ -79,9 +93,15 @@ impl DeclMapping {
         for (interface_decl_name, interface_item) in interface_decl_refs.into_iter() {
             if let Some(new_item) = impld_decl_refs.get(&interface_decl_name) {
                 let interface_decl_ref = match interface_item {
-                    TyTraitInterfaceItem::TraitFn(decl_ref) => decl_ref.id().into(),
-                    TyTraitInterfaceItem::Constant(decl_ref) => decl_ref.id().into(),
-                    TyTraitInterfaceItem::Type(decl_ref) => decl_ref.id().into(),
+                    TyTraitInterfaceItem::TraitFn(decl_ref) => {
+                        (decl_ref.id().into(), interface_decl_name.1)
+                    }
+                    TyTraitInterfaceItem::Constant(decl_ref) => {
+                        (decl_ref.id().into(), interface_decl_name.1)
+                    }
+                    TyTraitInterfaceItem::Type(decl_ref) => {
+                        (decl_ref.id().into(), interface_decl_name.1)
+                    }
                 };
                 let new_decl_ref = match new_item {
                     TyTraitItem::Fn(decl_ref) => decl_ref.id().into(),
@@ -94,9 +114,9 @@ impl DeclMapping {
         for (decl_name, item) in item_decl_refs.into_iter() {
             if let Some(new_item) = impld_decl_refs.get(&decl_name) {
                 let interface_decl_ref = match item {
-                    TyTraitItem::Fn(decl_ref) => decl_ref.id().into(),
-                    TyTraitItem::Constant(decl_ref) => decl_ref.id().into(),
-                    TyTraitItem::Type(decl_ref) => decl_ref.id().into(),
+                    TyTraitItem::Fn(decl_ref) => (decl_ref.id().into(), decl_name.1),
+                    TyTraitItem::Constant(decl_ref) => (decl_ref.id().into(), decl_name.1),
+                    TyTraitItem::Type(decl_ref) => (decl_ref.id().into(), decl_name.1),
                 };
                 let new_decl_ref = match new_item {
                     TyTraitItem::Fn(decl_ref) => decl_ref.id().into(),
@@ -109,39 +129,41 @@ impl DeclMapping {
         DeclMapping { mapping }
     }
 
-    pub(crate) fn find_match(&self, decl_ref: SourceDecl) -> Option<DestinationDecl> {
-        for (source_decl_ref, dest_decl_ref) in self.mapping.iter() {
-            if *source_decl_ref == decl_ref {
-                return Some(dest_decl_ref.clone());
-            }
-        }
-        None
-    }
-
-    /// This method returns only associated item functions that have as self type the given type.
-    pub(crate) fn filter_functions_by_self_type(
+    pub(crate) fn find_match(
         &self,
-        self_type: TypeId,
+        _handler: &Handler,
         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);
+        decl_ref: AssociatedItemDeclId,
+        typeid: Option<TypeId>,
+        self_typeid: Option<TypeId>,
+    ) -> Result<Option<DestinationDecl>, ErrorEmitted> {
+        let mut dest_decl_refs = HashSet::<DestinationDecl>::new();
 
-                    let unify_check = UnifyCheck::non_dynamic_equality(engines);
-                    if let (left, Some(right)) = (self_type, func.parameters.first()) {
-                        if unify_check.check(left, right.type_argument.type_id) {
-                            mapping.push((source_decl_ref, dest_decl_ref));
-                        }
-                    }
+        if let Some(mut typeid) = typeid {
+            if engines.te().get(typeid).is_self_type() && self_typeid.is_some() {
+                // If typeid is `Self`, then we use the self_typeid instead.
+                typeid = self_typeid.unwrap();
+            }
+            for (source_decl_ref, dest_decl_ref) in self.mapping.iter() {
+                let unify_check = UnifyCheck::non_dynamic_equality(engines);
+                if source_decl_ref.0 == decl_ref && unify_check.check(source_decl_ref.1, typeid) {
+                    dest_decl_refs.insert(dest_decl_ref.clone());
                 }
-                AssociatedItemDeclId::Constant(_) => mapping.push((source_decl_ref, dest_decl_ref)),
-                AssociatedItemDeclId::Type(_) => mapping.push((source_decl_ref, dest_decl_ref)),
             }
         }
-        DeclMapping { mapping }
+
+        // At most one replacement should be found for decl_ref.
+        /* TODO uncomment this and close issue #5540
+        if dest_decl_refs.len() > 1 {
+            handler.emit_err(CompileError::InternalOwned(
+                format!(
+                    "Multiple replacements for decl {} implemented in {}",
+                    engines.help_out(decl_ref),
+                    engines.help_out(typeid),
+                ),
+                dest_decl_refs.iter().last().unwrap().span(engines),
+            ));
+        }*/
+        Ok(dest_decl_refs.iter().next().cloned())
     }
 }

sway-core/src/decl_engine/ref.rs

@@ -294,20 +294,35 @@ impl ReplaceDecls for DeclRefFunction {
     fn replace_decls_inner(
         &mut self,
         decl_mapping: &DeclMapping,
-        _handler: &Handler,
+        handler: &Handler,
         ctx: &mut 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()) {
+
+        let func = decl_engine.get(self);
+
+        if let Some(new_decl_ref) = decl_mapping.find_match(
+            handler,
+            ctx.engines(),
+            self.id.into(),
+            func.implementing_for_typeid,
+            ctx.self_type(),
+        )? {
             if let AssociatedItemDeclId::Function(new_decl_ref) = new_decl_ref {
                 self.id = new_decl_ref;
             }
             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 Some(new_decl_ref) = decl_mapping.find_match(
+                handler,
+                ctx.engines(),
+                parent.clone(),
+                func.implementing_for_typeid,
+                ctx.self_type(),
+            )? {
                 if let AssociatedItemDeclId::Function(new_decl_ref) = new_decl_ref {
                     self.id = new_decl_ref;
                 }

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

@@ -30,6 +30,7 @@ pub struct TyFunctionDecl {
     pub body: TyCodeBlock,
     pub parameters: Vec<TyFunctionParameter>,
     pub implementing_type: Option<TyDecl>,
+    pub implementing_for_typeid: Option<TypeId>,
     pub span: Span,
     pub call_path: CallPath,
     pub attributes: transform::AttributesMap,
@@ -159,6 +160,7 @@ impl HashWithEngines for TyFunctionDecl {
             span: _,
             attributes: _,
             implementing_type: _,
+            implementing_for_typeid: _,
             where_clause: _,
             is_trait_method_dummy: _,
         } = self;
@@ -183,6 +185,9 @@ impl SubstTypes for TyFunctionDecl {
             .for_each(|x| x.subst(type_mapping, engines));
         self.return_type.subst(type_mapping, engines);
         self.body.subst(type_mapping, engines);
+        if let Some(implementing_for) = self.implementing_for_typeid.as_mut() {
+            implementing_for.subst(type_mapping, engines);
+        }
     }
 }
 
@@ -193,7 +198,9 @@ impl ReplaceDecls for TyFunctionDecl {
         handler: &Handler,
         ctx: &mut TypeCheckContext,
     ) -> Result<(), ErrorEmitted> {
-        self.body.replace_decls(decl_mapping, handler, ctx)
+        let mut func_ctx = ctx.by_ref().with_self_type(self.implementing_for_typeid);
+        self.body
+            .replace_decls(decl_mapping, handler, &mut func_ctx)
     }
 }
 
@@ -296,6 +303,7 @@ impl TyFunctionDecl {
             name: name.clone(),
             body: TyCodeBlock::default(),
             implementing_type: None,
+            implementing_for_typeid: None,
             span: span.clone(),
             call_path: CallPath::from(Ident::dummy()),
             attributes: Default::default(),

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

@@ -8,8 +8,8 @@ use sway_types::{Ident, Named, Span, Spanned};
 
 use crate::{
     decl_engine::{
-        mapping::DeclMapping, DeclEngineReplace, DeclRefConstant, DeclRefFunction, DeclRefTraitFn,
-        DeclRefTraitType, ReplaceFunctionImplementingType,
+        DeclEngineReplace, DeclRefConstant, DeclRefFunction, DeclRefTraitFn, DeclRefTraitType,
+        ReplaceFunctionImplementingType,
     },
     engine_threading::*,
     language::{parsed, CallPath, Visibility},
@@ -44,6 +44,35 @@ pub enum TyTraitInterfaceItem {
     Type(DeclRefTraitType),
 }
 
+impl DisplayWithEngines for TyTraitInterfaceItem {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>, engines: &Engines) -> fmt::Result {
+        write!(f, "{:?}", engines.help_out(self))
+    }
+}
+
+impl DebugWithEngines for TyTraitInterfaceItem {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>, engines: &Engines) -> fmt::Result {
+        write!(
+            f,
+            "TyTraitItem {}",
+            match self {
+                TyTraitInterfaceItem::TraitFn(fn_ref) => format!(
+                    "fn {:?}",
+                    engines.help_out(&*engines.de().get_trait_fn(fn_ref))
+                ),
+                TyTraitInterfaceItem::Constant(const_ref) => format!(
+                    "const {:?}",
+                    engines.help_out(&*engines.de().get_constant(const_ref))
+                ),
+                TyTraitInterfaceItem::Type(type_ref) => format!(
+                    "type {:?}",
+                    engines.help_out(&*engines.de().get_type(type_ref))
+                ),
+            }
+        )
+    }
+}
+
 #[derive(Clone, Debug)]
 pub enum TyTraitItem {
     Fn(DeclRefFunction),
@@ -242,7 +271,6 @@ impl Spanned for TyTraitItem {
 
 impl SubstTypes for TyTraitDecl {
     fn subst_inner(&mut self, type_mapping: &TypeSubstMap, engines: &Engines) {
-        let mut decl_mapping = DeclMapping::new();
         self.type_parameters
             .iter_mut()
             .for_each(|x| x.subst(type_mapping, engines));
@@ -253,14 +281,12 @@ impl SubstTypes for TyTraitDecl {
                     let new_item_ref = item_ref
                         .clone()
                         .subst_types_and_insert_new_with_parent(type_mapping, engines);
-                    decl_mapping.insert(item_ref.id().into(), new_item_ref.id().into());
                     item_ref.replace_id(*new_item_ref.id());
                 }
                 TyTraitInterfaceItem::Constant(decl_ref) => {
                     let new_decl_ref = decl_ref
                         .clone()
                         .subst_types_and_insert_new(type_mapping, engines);
-                    decl_mapping.insert(decl_ref.id().into(), new_decl_ref.id().into());
                     decl_ref.replace_id(*new_decl_ref.id());
                 }
                 TyTraitInterfaceItem::Type(decl_ref) => {