From 8d27a6e7bd186c55c1ae7620c8e2c15b86e9cbce Mon Sep 17 00:00:00 2001
From: sofia <sofia@teascade.net>
Date: Mon, 4 Aug 2025 00:40:52 +0300
Subject: [PATCH] Fix reference-linking for local functions

---
 reid-lsp/src/analysis.rs | 21 ++++++++++++++-------
 reid-lsp/src/main.rs     |  1 -
 2 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/reid-lsp/src/analysis.rs b/reid-lsp/src/analysis.rs
index ffc42f3..c0a45fa 100644
--- a/reid-lsp/src/analysis.rs
+++ b/reid-lsp/src/analysis.rs
@@ -602,14 +602,16 @@ pub fn analyze_context(
 
     for (ty, function) in &module.associated_functions {
         if let Some(source_id) = function.source {
-            if let Some(state) = map.get(&source_id) {
-                if let Some(symbol) = state.associated_functions.get(&(ty.clone(), function.name.clone())) {
-                    scope
-                        .associated_functions
-                        .insert((ty.clone(), function.name.clone()), (source_id, *symbol));
+            if source_id != module.module_id {
+                if let Some(state) = map.get(&source_id) {
+                    if let Some(symbol) = state.associated_functions.get(&(ty.clone(), function.name.clone())) {
+                        scope
+                            .associated_functions
+                            .insert((ty.clone(), function.name.clone()), (source_id, *symbol));
+                    }
                 }
+                continue;
             }
-            continue;
         }
 
         let idx = scope
@@ -621,6 +623,9 @@ pub fn analyze_context(
             .state
             .associated_functions
             .insert((ty.clone(), function.name.clone()), symbol);
+        scope
+            .associated_functions
+            .insert((ty.clone(), function.name.clone()), (module.module_id, symbol));
 
         for param in &function.parameters {
             scope.state.init_types(&param.meta, Some(param.ty.clone()));
@@ -647,7 +652,6 @@ pub fn analyze_context(
     for function in &module.functions {
         if let Some(source_id) = function.source {
             if source_id != module.module_id {
-                dbg!(source_id, &function.name);
                 if let Some(state) = map.get(&source_id) {
                     if let Some(symbol) = state.functions.get(&function.name) {
                         scope.functions.insert(function.name.clone(), (source_id, *symbol));
@@ -667,6 +671,9 @@ pub fn analyze_context(
         let function_symbol = scope.state.new_symbol(idx, SemanticKind::Function, module.module_id);
         scope.state.set_symbol(idx, function_symbol);
         scope.state.functions.insert(function.name.clone(), function_symbol);
+        scope
+            .functions
+            .insert(function.name.clone(), (module.module_id, function_symbol));
     }
 
     for function in &module.functions {
diff --git a/reid-lsp/src/main.rs b/reid-lsp/src/main.rs
index 89692c6..a877eae 100644
--- a/reid-lsp/src/main.rs
+++ b/reid-lsp/src/main.rs
@@ -307,7 +307,6 @@ impl LanguageServer for Backend {
             });
             if let Some(token) = token {
                 let reference_tokens = analysis.find_references(token.0, &self.state_map());
-                dbg!(&reference_tokens);
                 let mut locations = Vec::new();
                 if let Some(reference_tokens) = reference_tokens {
                     for (module_id, symbol_idx) in reference_tokens {