Refactor trait handling when doing namespace item imports. (#6080)

## Description

This splits off trait handling when doing namespace item imports to be
done as a separate step.
Instead of doing it as we build the import tree, we now do it at a later
stage.

This process requires a fully typed tree, which is not available when
doing imports for a parsed declaration tree.

## Checklist

- [ ] I have linked to any relevant issues.
- [ ] 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).
- [ ] If my change requires substantial documentation changes, I have
[requested support from the DevRel
team](https://github.com/FuelLabs/devrel-requests/issues/new/choose)
- [ ] I have added tests that prove my fix is effective or that my
feature works.
- [ ] 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/semantic_analysis/ast_node/mod.rs

@@ -7,7 +7,12 @@ pub(crate) use expression::*;
 pub(crate) use modes::*;
 
 use crate::{
-    language::{parsed::*, ty},
+    engine_threading::SpannedWithEngines,
+    language::{
+        parsed::*,
+        ty::{self, TyDecl},
+    },
+    namespace::TraitMap,
     semantic_analysis::*,
     type_system::*,
     Engines, Ident,
@@ -52,6 +57,7 @@ impl ty::TyAstNode {
             content: match node.content.clone() {
                 AstNodeContent::UseStatement(stmt) => {
                     handle_use_statement(&mut ctx, engines, &stmt, handler);
+                    handle_item_trait_imports(&mut ctx, engines, handler)?;
                     ty::TyAstNodeContent::SideEffect(ty::TySideEffect {
                         side_effect: ty::TySideEffectVariant::UseStatement(ty::TyUseStatement {
                             alias: stmt.alias,
@@ -130,6 +136,54 @@ impl ty::TyAstNode {
     }
 }
 
+fn handle_item_trait_imports(
+    ctx: &mut TypeCheckContext<'_>,
+    engines: &Engines,
+    handler: &Handler,
+) -> Result<(), ErrorEmitted> {
+    let mut impls_to_insert = TraitMap::default();
+
+    let root_mod = &ctx.namespace().root().module;
+    let dst_mod = ctx.namespace.module(engines);
+
+    for (_, (_, src, decl)) in dst_mod.current_items().use_item_synonyms.iter() {
+        let src_mod = root_mod.lookup_submodule(handler, engines, src)?;
+
+        //  if this is an enum or struct or function, import its implementations
+        if let Ok(type_id) = decl.return_type(&Handler::default(), engines) {
+            impls_to_insert.extend(
+                src_mod
+                    .current_items()
+                    .implemented_traits
+                    .filter_by_type_item_import(type_id, engines),
+                engines,
+            );
+        }
+        // if this is a trait, import its implementations
+        let decl_span = decl.span(engines);
+        if matches!(decl, TyDecl::TraitDecl(_)) {
+            // TODO: we only import local impls from the source namespace
+            // this is okay for now but we'll need to device some mechanism to collect all
+            // available trait impls
+            impls_to_insert.extend(
+                src_mod
+                    .current_items()
+                    .implemented_traits
+                    .filter_by_trait_decl_span(decl_span),
+                engines,
+            );
+        }
+    }
+
+    let dst_mod = ctx.namespace_mut().module_mut(engines);
+    dst_mod
+        .current_items_mut()
+        .implemented_traits
+        .extend(impls_to_insert, engines);
+
+    Ok(())
+}
+
 fn handle_use_statement(
     ctx: &mut TypeCheckContext<'_>,
     engines: &Engines,

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

@@ -1,8 +1,6 @@
 use std::fmt;
 
-use super::{
-    module::Module, namespace::Namespace, trait_map::TraitMap, Ident, ResolvedTraitImplItem,
-};
+use super::{module::Module, namespace::Namespace, Ident, ResolvedTraitImplItem};
 use crate::{
     decl_engine::DeclRef,
     engine_threading::*,
@@ -66,35 +64,6 @@ impl ResolvedDeclaration {
             ResolvedDeclaration::Typed(decl) => decl.visibility(engines.de()),
         }
     }
-
-    fn span(&self, engines: &Engines) -> sway_types::Span {
-        match self {
-            ResolvedDeclaration::Parsed(decl) => decl.span(engines),
-            ResolvedDeclaration::Typed(decl) => decl.span(engines),
-        }
-    }
-
-    pub(crate) fn return_type(
-        &self,
-        handler: &Handler,
-        engines: &Engines,
-    ) -> Result<TypeId, ErrorEmitted> {
-        match self {
-            ResolvedDeclaration::Parsed(_decl) => unreachable!(),
-            ResolvedDeclaration::Typed(decl) => decl.return_type(handler, engines),
-        }
-    }
-
-    fn is_trait(&self) -> bool {
-        match self {
-            ResolvedDeclaration::Parsed(decl) => {
-                matches!(decl, Declaration::TraitDeclaration(_))
-            }
-            ResolvedDeclaration::Typed(decl) => {
-                matches!(decl, TyDecl::TraitDecl(_))
-            }
-        }
-    }
 }
 
 /// The root module, from which all other modules can be accessed.
@@ -188,7 +157,6 @@ impl Root {
         self.check_module_privacy(handler, engines, src)?;
 
         let src_mod = self.module.lookup_submodule(handler, engines, src)?;
-        let mut impls_to_insert = TraitMap::default();
         match src_mod.current_items().symbols.get(item).cloned() {
             Some(decl) => {
                 if !decl.visibility(engines).is_public() && !is_ancestor(src, dst) {
@@ -198,29 +166,6 @@ impl Root {
                     });
                 }
 
-                //  if this is an enum or struct or function, import its implementations
-                if let Ok(type_id) = decl.return_type(&Handler::default(), engines) {
-                    impls_to_insert.extend(
-                        src_mod
-                            .current_items()
-                            .implemented_traits
-                            .filter_by_type_item_import(type_id, engines),
-                        engines,
-                    );
-                }
-                // if this is a trait, import its implementations
-                let decl_span = decl.span(engines);
-                if decl.is_trait() {
-                    // TODO: we only import local impls from the source namespace
-                    // this is okay for now but we'll need to device some mechanism to collect all available trait impls
-                    impls_to_insert.extend(
-                        src_mod
-                            .current_items()
-                            .implemented_traits
-                            .filter_by_trait_decl_span(decl_span),
-                        engines,
-                    );
-                }
                 // no matter what, import it this way though.
                 let dst_mod = self.module.lookup_submodule_mut(handler, engines, dst)?;
                 let check_name_clash = |name| {
@@ -254,12 +199,6 @@ impl Root {
             }
         };
 
-        let dst_mod = self.module.lookup_submodule_mut(handler, engines, dst)?;
-        dst_mod
-            .current_items_mut()
-            .implemented_traits
-            .extend(impls_to_insert, engines);
-
         Ok(())
     }