Handle traits when doing namespace item imports.

This reverts PR #6080 due to a pretty severe performance regression.

The reason this was split to a separate step was because it needs a
typed declaration due to the depedency on the type system and thus
cannot be used during collection.

However, the type system will be changed to work with parsed
declarations pretty soon, so revert this back to the old behavior and
temporarily check for a typed declaration instead.

Fixes #6182.

sway-core/src/semantic_analysis/ast_node/mod.rs

@@ -59,7 +59,6 @@ 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,
@@ -138,56 +137,6 @@ 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 decl = decl.expect_typed_ref();
-
-        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 collect_use_statement(
     handler: &Handler,
     engines: &Engines,

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

@@ -1,6 +1,8 @@
 use std::fmt;
 
-use super::{module::Module, namespace::Namespace, Ident, ResolvedTraitImplItem};
+use super::{
+    module::Module, namespace::Namespace, trait_map::TraitMap, Ident, ResolvedTraitImplItem,
+};
 use crate::{
     decl_engine::DeclRef,
     engine_threading::*,
@@ -87,6 +89,35 @@ 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.
@@ -179,6 +210,7 @@ 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) {
@@ -188,6 +220,41 @@ impl Root {
                     });
                 }
 
+                // We only handle trait imports when handling typed declarations,
+                // that is, when performing type-checking, and not when collecting.
+                // Update this once the type system is updated to refer to parsed
+                // declarations.
+                let handle_trait_import = match decl {
+                    ResolvedDeclaration::Parsed(_) => false,
+                    ResolvedDeclaration::Typed(_) => true,
+                };
+
+                if handle_trait_import {
+                    //  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| {
@@ -221,6 +288,12 @@ 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(())
     }