From e327345ef3adc09663c69d78de5d04b549b88e3a Mon Sep 17 00:00:00 2001
From: Joao Matos <joao@tritao.eu>
Date: Sun, 7 Jul 2024 21:05:37 +0100
Subject: [PATCH] Collect nodes inside each function body and keep a reference
 to their lexical scope.

---
 .../src/language/parsed/declaration/function.rs  |  2 ++
 .../src/semantic_analysis/ast_node/code_block.rs | 16 ++++++++++++++++
 .../ast_node/declaration/declaration.rs          | 15 ++++++++++++---
 .../symbol_collection_context.rs                 |  8 +++++---
 .../to_parsed_lang/convert_parse_tree.rs         |  1 +
 5 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/sway-core/src/language/parsed/declaration/function.rs b/sway-core/src/language/parsed/declaration/function.rs
index 716f545ef0e..afd61ef54d3 100644
--- a/sway-core/src/language/parsed/declaration/function.rs
+++ b/sway-core/src/language/parsed/declaration/function.rs
@@ -1,6 +1,7 @@
 use crate::{
     engine_threading::*,
     language::{parsed::*, *},
+    namespace::LexicalScopeId,
     transform::{self, AttributeKind},
     type_system::*,
 };
@@ -28,6 +29,7 @@ pub struct FunctionDeclaration {
     pub where_clause: Vec<(Ident, Vec<TraitConstraint>)>,
     pub kind: FunctionDeclarationKind,
     pub implementing_type: Option<Declaration>,
+    pub lexical_scope: LexicalScopeId,
 }
 
 impl EqWithEngines for FunctionDeclaration {}
diff --git a/sway-core/src/semantic_analysis/ast_node/code_block.rs b/sway-core/src/semantic_analysis/ast_node/code_block.rs
index 485fe51fa4d..41fa19bb75a 100644
--- a/sway-core/src/semantic_analysis/ast_node/code_block.rs
+++ b/sway-core/src/semantic_analysis/ast_node/code_block.rs
@@ -5,6 +5,22 @@ use crate::language::{
 };
 
 impl ty::TyCodeBlock {
+    pub(crate) fn collect(
+        handler: &Handler,
+        engines: &Engines,
+        ctx: &mut SymbolCollectionContext,
+        code_block: &CodeBlock,
+    ) -> Result<(), ErrorEmitted> {
+        let _ = code_block
+            .contents
+            .iter()
+            .map(|node| ty::TyAstNode::collect(handler, engines, ctx, node))
+            .filter_map(|res| res.ok())
+            .collect::<Vec<_>>();
+
+        Ok(())
+    }
+
     pub(crate) fn type_check(
         handler: &Handler,
         ctx: TypeCheckContext,
diff --git a/sway-core/src/semantic_analysis/ast_node/declaration/declaration.rs b/sway-core/src/semantic_analysis/ast_node/declaration/declaration.rs
index 1945a8ba441..3f37865215a 100644
--- a/sway-core/src/semantic_analysis/ast_node/declaration/declaration.rs
+++ b/sway-core/src/semantic_analysis/ast_node/declaration/declaration.rs
@@ -2,10 +2,13 @@ use sway_error::handler::{ErrorEmitted, Handler};
 use sway_types::{Ident, Named, Spanned};
 
 use crate::{
-    decl_engine::{DeclEngineGet, DeclEngineInsert, DeclRef, ReplaceFunctionImplementingType},
+    decl_engine::{
+        parsed_engine::ParsedDeclEngineReplace, DeclEngineGet, DeclEngineInsert, DeclRef,
+        ReplaceFunctionImplementingType,
+    },
     language::{
         parsed::{self, StorageEntry},
-        ty::{self, FunctionDecl, TyDecl, TyStorageField},
+        ty::{self, FunctionDecl, TyCodeBlock, TyDecl, TyStorageField},
         CallPath,
     },
     namespace::{IsExtendingExistingImpl, IsImplSelf},
@@ -53,8 +56,14 @@ impl TyDecl {
             }
             parsed::Declaration::EnumVariantDeclaration(_decl) => {}
             parsed::Declaration::FunctionDeclaration(decl_id) => {
-                let fn_decl = engines.pe().get_function(decl_id).as_ref().clone();
+                let decl_id = *decl_id;
+                let mut fn_decl = engines.pe().get_function(&decl_id).as_ref().clone();
                 let _ = ctx.insert_parsed_symbol(handler, engines, fn_decl.name.clone(), decl);
+                let (_ret, lexical_scope_id) = ctx.scoped(engines, |scoped_ctx| {
+                    TyCodeBlock::collect(handler, engines, scoped_ctx, &fn_decl.body)
+                });
+                fn_decl.lexical_scope = lexical_scope_id;
+                engines.pe().replace(decl_id, fn_decl);
             }
             parsed::Declaration::TraitDeclaration(decl_id) => {
                 let trait_decl = engines.pe().get_trait(decl_id).as_ref().clone();
diff --git a/sway-core/src/semantic_analysis/symbol_collection_context.rs b/sway-core/src/semantic_analysis/symbol_collection_context.rs
index c072dce9cec..ca8d1c2b961 100644
--- a/sway-core/src/semantic_analysis/symbol_collection_context.rs
+++ b/sway-core/src/semantic_analysis/symbol_collection_context.rs
@@ -1,5 +1,6 @@
 use crate::{
     language::{parsed::Declaration, Visibility},
+    namespace::LexicalScopeId,
     namespace::ModulePath,
     semantic_analysis::Namespace,
     Engines,
@@ -40,15 +41,16 @@ impl SymbolCollectionContext {
         &mut self,
         engines: &Engines,
         with_scoped_ctx: impl FnOnce(&mut SymbolCollectionContext) -> Result<T, ErrorEmitted>,
-    ) -> Result<T, ErrorEmitted> {
-        self.namespace
+    ) -> (Result<T, ErrorEmitted>, LexicalScopeId) {
+        let lexical_scope_id: LexicalScopeId = self
+            .namespace
             .module_mut(engines)
             .write(engines, |m| m.push_new_lexical_scope());
         let ret = with_scoped_ctx(self);
         self.namespace
             .module_mut(engines)
             .write(engines, |m| m.pop_lexical_scope());
-        ret
+        (ret, lexical_scope_id)
     }
 
     /// Enter the lexical scope and produce a collection context ready for
diff --git a/sway-core/src/transform/to_parsed_lang/convert_parse_tree.rs b/sway-core/src/transform/to_parsed_lang/convert_parse_tree.rs
index f15e204b846..e96a20a1156 100644
--- a/sway-core/src/transform/to_parsed_lang/convert_parse_tree.rs
+++ b/sway-core/src/transform/to_parsed_lang/convert_parse_tree.rs
@@ -598,6 +598,7 @@ pub fn item_fn_to_function_declaration(
             .unwrap_or(vec![]),
         kind,
         implementing_type,
+        lexical_scope: 0,
     };
     let decl_id = engines.pe().insert(fn_decl);
     Ok(decl_id)