Associate lexical scopes with their respective declarations. (#6780)

This adds a mapping from a lexical scope to the declaration that
lexically corresponds to it.

This is necessary for some upcoming collection context PR.

## Checklist

- [x] I have linked to any relevant issues.
- [x] I have commented my code, particularly in hard-to-understand
areas.
- [x] 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)
- [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/semantic_analysis/ast_node/code_block.rs

@@ -11,15 +11,20 @@ impl ty::TyCodeBlock {
         ctx: &mut SymbolCollectionContext,
         code_block: &CodeBlock,
     ) -> Result<(), ErrorEmitted> {
-        let _ = ctx.scoped(engines, code_block.whole_block_span.clone(), |scoped_ctx| {
-            let _ = code_block
-                .contents
-                .iter()
-                .map(|node| ty::TyAstNode::collect(handler, engines, scoped_ctx, node))
-                .filter_map(|res| res.ok())
-                .collect::<Vec<_>>();
-            Ok(())
-        });
+        let _ = ctx.scoped(
+            engines,
+            code_block.whole_block_span.clone(),
+            None,
+            |scoped_ctx| {
+                let _ = code_block
+                    .contents
+                    .iter()
+                    .map(|node| ty::TyAstNode::collect(handler, engines, scoped_ctx, node))
+                    .filter_map(|res| res.ok())
+                    .collect::<Vec<_>>();
+                Ok(())
+            },
+        );
         Ok(())
     }
 

sway-core/src/semantic_analysis/ast_node/declaration/abi.rs

@@ -39,14 +39,10 @@ impl ty::TyAbiDecl {
         decl_id: &ParsedDeclId<AbiDeclaration>,
     ) -> Result<(), ErrorEmitted> {
         let abi_decl = engines.pe().get_abi(decl_id);
-        ctx.insert_parsed_symbol(
-            handler,
-            engines,
-            abi_decl.name.clone(),
-            Declaration::AbiDeclaration(*decl_id),
-        )?;
+        let decl = Declaration::AbiDeclaration(*decl_id);
+        ctx.insert_parsed_symbol(handler, engines, abi_decl.name.clone(), decl.clone())?;
 
-        let _ = ctx.scoped(engines, abi_decl.span.clone(), |scoped_ctx| {
+        let _ = ctx.scoped(engines, abi_decl.span.clone(), Some(decl), |scoped_ctx| {
             abi_decl.interface_surface.iter().for_each(|item| {
                 let _ = TyTraitItem::collect(handler, engines, scoped_ctx, item);
             });

sway-core/src/semantic_analysis/ast_node/declaration/enum.rs

@@ -16,15 +16,13 @@ impl ty::TyEnumDecl {
         decl_id: &ParsedDeclId<EnumDeclaration>,
     ) -> Result<(), ErrorEmitted> {
         let enum_decl = engines.pe().get_enum(decl_id);
-        ctx.insert_parsed_symbol(
-            handler,
-            engines,
-            enum_decl.name.clone(),
-            Declaration::EnumDeclaration(*decl_id),
-        )?;
+        let decl = Declaration::EnumDeclaration(*decl_id);
+        ctx.insert_parsed_symbol(handler, engines, enum_decl.name.clone(), decl.clone())?;
 
         // create a namespace for the decl, used to create a scope for generics
-        let _ = ctx.scoped(engines, enum_decl.span.clone(), |mut _ctx| Ok(()));
+        let _ = ctx.scoped(engines, enum_decl.span.clone(), Some(decl), |mut _ctx| {
+            Ok(())
+        });
         Ok(())
     }
 

sway-core/src/semantic_analysis/ast_node/declaration/function.rs

@@ -28,15 +28,11 @@ impl ty::TyFunctionDecl {
         decl_id: &ParsedDeclId<FunctionDeclaration>,
     ) -> Result<(), ErrorEmitted> {
         let fn_decl = engines.pe().get_function(decl_id);
-        let _ = ctx.insert_parsed_symbol(
-            handler,
-            engines,
-            fn_decl.name.clone(),
-            Declaration::FunctionDeclaration(*decl_id),
-        );
+        let decl = Declaration::FunctionDeclaration(*decl_id);
+        let _ = ctx.insert_parsed_symbol(handler, engines, fn_decl.name.clone(), decl.clone());
 
         // create a namespace for the function
-        let _ = ctx.scoped(engines, fn_decl.span.clone(), |scoped_ctx| {
+        let _ = ctx.scoped(engines, fn_decl.span.clone(), Some(decl), |scoped_ctx| {
             TyCodeBlock::collect(handler, engines, scoped_ctx, &fn_decl.body)
         });
         Ok(())

sway-core/src/semantic_analysis/ast_node/declaration/impl_trait.rs

@@ -39,27 +39,34 @@ impl TyImplSelfOrTrait {
         decl_id: &ParsedDeclId<ImplSelfOrTrait>,
     ) -> Result<(), ErrorEmitted> {
         let impl_trait = engines.pe().get_impl_self_or_trait(decl_id);
+
+        let decl = Declaration::ImplSelfOrTrait(*decl_id);
         ctx.insert_parsed_symbol(
             handler,
             engines,
             impl_trait.trait_name.suffix.clone(),
-            Declaration::ImplSelfOrTrait(*decl_id),
+            decl.clone(),
         )?;
 
-        let _ = ctx.scoped(engines, impl_trait.block_span.clone(), |scoped_ctx| {
-            impl_trait.items.iter().for_each(|item| match item {
-                ImplItem::Fn(decl_id) => {
-                    let _ = TyFunctionDecl::collect(handler, engines, scoped_ctx, decl_id);
-                }
-                ImplItem::Constant(decl_id) => {
-                    let _ = TyConstantDecl::collect(handler, engines, scoped_ctx, decl_id);
-                }
-                ImplItem::Type(decl_id) => {
-                    let _ = TyTraitType::collect(handler, engines, scoped_ctx, decl_id);
-                }
-            });
-            Ok(())
-        });
+        let _ = ctx.scoped(
+            engines,
+            impl_trait.block_span.clone(),
+            Some(decl),
+            |scoped_ctx| {
+                impl_trait.items.iter().for_each(|item| match item {
+                    ImplItem::Fn(decl_id) => {
+                        let _ = TyFunctionDecl::collect(handler, engines, scoped_ctx, decl_id);
+                    }
+                    ImplItem::Constant(decl_id) => {
+                        let _ = TyConstantDecl::collect(handler, engines, scoped_ctx, decl_id);
+                    }
+                    ImplItem::Type(decl_id) => {
+                        let _ = TyTraitType::collect(handler, engines, scoped_ctx, decl_id);
+                    }
+                });
+                Ok(())
+            },
+        );
         Ok(())
     }
 

sway-core/src/semantic_analysis/ast_node/declaration/struct.rs

@@ -16,15 +16,16 @@ impl ty::TyStructDecl {
         decl_id: &ParsedDeclId<StructDeclaration>,
     ) -> Result<(), ErrorEmitted> {
         let struct_decl = engines.pe().get_struct(decl_id);
-        ctx.insert_parsed_symbol(
-            handler,
-            engines,
-            struct_decl.name.clone(),
-            Declaration::StructDeclaration(*decl_id),
-        )?;
+        let decl = Declaration::StructDeclaration(*decl_id);
+        ctx.insert_parsed_symbol(handler, engines, struct_decl.name.clone(), decl.clone())?;
 
         // create a namespace for the decl, used to create a scope for generics
-        let _ = ctx.scoped(engines, struct_decl.span.clone(), |_scoped_ctx| Ok(()));
+        let _ = ctx.scoped(
+            engines,
+            struct_decl.span.clone(),
+            Some(decl),
+            |_scoped_ctx| Ok(()),
+        );
         Ok(())
     }
 

sway-core/src/semantic_analysis/ast_node/declaration/trait.rs

@@ -53,15 +53,11 @@ impl TyTraitDecl {
         decl_id: &ParsedDeclId<TraitDeclaration>,
     ) -> Result<(), ErrorEmitted> {
         let trait_decl = engines.pe().get_trait(decl_id);
-        ctx.insert_parsed_symbol(
-            handler,
-            engines,
-            trait_decl.name.clone(),
-            Declaration::TraitDeclaration(*decl_id),
-        )?;
+        let decl = Declaration::TraitDeclaration(*decl_id);
+        ctx.insert_parsed_symbol(handler, engines, trait_decl.name.clone(), decl.clone())?;
 
         // A temporary namespace for checking within the trait's scope.
-        let _ = ctx.scoped(engines, trait_decl.span.clone(), |scoped_ctx| {
+        let _ = ctx.scoped(engines, trait_decl.span.clone(), Some(decl), |scoped_ctx| {
             trait_decl.interface_surface.iter().for_each(|item| {
                 let _ = TyTraitItem::collect(handler, engines, scoped_ctx, item);
             });

sway-core/src/semantic_analysis/ast_node/declaration/trait_fn.rs

@@ -24,13 +24,11 @@ impl ty::TyTraitFn {
         decl_id: &ParsedDeclId<TraitFn>,
     ) -> Result<(), ErrorEmitted> {
         let trait_fn = engines.pe().get_trait_fn(decl_id);
-        ctx.insert_parsed_symbol(
-            handler,
-            engines,
-            trait_fn.name.clone(),
-            Declaration::TraitFnDeclaration(*decl_id),
-        )?;
-        let _ = ctx.scoped(engines, trait_fn.span.clone(), |_scoped_ctx| Ok(()));
+        let decl = Declaration::TraitFnDeclaration(*decl_id);
+        ctx.insert_parsed_symbol(handler, engines, trait_fn.name.clone(), decl.clone())?;
+        let _ = ctx.scoped(engines, trait_fn.span.clone(), Some(decl), |_scoped_ctx| {
+            Ok(())
+        });
         Ok(())
     }
 

sway-core/src/semantic_analysis/ast_node/expression/typed_expression.rs

@@ -216,7 +216,7 @@ impl ty::TyExpression {
                     .iter()
                     .map(|branch| {
                         // create a new namespace for this branch result
-                        ctx.scoped(engines, branch.span.clone(), |scoped_ctx| {
+                        ctx.scoped(engines, branch.span.clone(), None, |scoped_ctx| {
                             Self::collect(handler, engines, scoped_ctx, &branch.result)
                         })
                         .0

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

@@ -71,6 +71,9 @@ pub struct LexicalScope {
     pub children: Vec<LexicalScopeId>,
     /// The parent scope associated with this scope. Will be None for a root scope.
     pub parent: Option<LexicalScopeId>,
+    /// The declaration associated with this scope. This will initially be a [ParsedDeclId],
+    /// but can be replaced to be a [DeclId] once the declaration is type checked.
+    pub declaration: Option<ResolvedDeclaration>,
 }
 
 /// The set of items that exist within some lexical scope via declaration or importing.

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

@@ -277,11 +277,16 @@ impl Module {
     }
 
     /// Pushes a new scope to the module's lexical scope hierarchy.
-    pub fn push_new_lexical_scope(&mut self, span: Span) -> LexicalScopeId {
+    pub fn push_new_lexical_scope(
+        &mut self,
+        span: Span,
+        declaration: Option<ResolvedDeclaration>,
+    ) -> LexicalScopeId {
         let previous_scope_id = self.current_lexical_scope_id();
         let new_scoped_id = {
             self.lexical_scopes.push(LexicalScope {
                 parent: Some(previous_scope_id),
+                declaration,
                 ..Default::default()
             });
             self.lexical_scopes.len() - 1

sway-core/src/semantic_analysis/symbol_collection_context.rs

@@ -1,7 +1,6 @@
 use crate::{
     language::{parsed::Declaration, Visibility},
-    namespace::LexicalScopeId,
-    namespace::ModulePath,
+    namespace::{LexicalScopeId, ModulePath, ResolvedDeclaration},
     semantic_analysis::Namespace,
     Engines,
 };
@@ -41,12 +40,14 @@ impl SymbolCollectionContext {
         &mut self,
         engines: &Engines,
         span: Span,
+        decl: Option<Declaration>,
         with_scoped_ctx: impl FnOnce(&mut SymbolCollectionContext) -> Result<T, ErrorEmitted>,
     ) -> (Result<T, ErrorEmitted>, LexicalScopeId) {
-        let lexical_scope_id: LexicalScopeId = self
-            .namespace
-            .module_mut(engines)
-            .write(engines, |m| m.push_new_lexical_scope(span.clone()));
+        let decl = decl.map(ResolvedDeclaration::Parsed);
+        let lexical_scope_id: LexicalScopeId =
+            self.namespace.module_mut(engines).write(engines, |m| {
+                m.push_new_lexical_scope(span.clone(), decl.clone())
+            });
         let ret = with_scoped_ctx(self);
         self.namespace
             .module_mut(engines)