From 0abeb0c4cd895c341d1ad32976d662a7c3cd9345 Mon Sep 17 00:00:00 2001
From: sofia <sofia@teascade.net>
Date: Sun, 3 Aug 2025 21:15:47 +0300
Subject: [PATCH] Flatten a pyramid with a helper function

---
 reid-lsp/src/analysis.rs | 11 +++++++++++
 reid-lsp/src/main.rs     | 40 ++++++++++++++--------------------------
 2 files changed, 25 insertions(+), 26 deletions(-)

diff --git a/reid-lsp/src/analysis.rs b/reid-lsp/src/analysis.rs
index 06e5dfa..bd2e9e4 100644
--- a/reid-lsp/src/analysis.rs
+++ b/reid-lsp/src/analysis.rs
@@ -44,6 +44,17 @@ pub struct StaticAnalysis {
     pub error: Option<ReidError>,
 }
 
+impl StaticAnalysis {
+    pub fn find_definition(&self, token_idx: usize) -> Option<&FullToken> {
+        let semantic_token = self.state.map.get(&token_idx)?;
+        let symbol_id = semantic_token.symbol?;
+        let definition_id = self.state.find_definition(&symbol_id);
+
+        let def_token_idx = self.state.symbol_to_token.get(&definition_id)?;
+        self.tokens.get(*def_token_idx)
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct SemanticToken {
     pub ty: Option<TypeKind>,
diff --git a/reid-lsp/src/main.rs b/reid-lsp/src/main.rs
index e0129d3..dcc757e 100644
--- a/reid-lsp/src/main.rs
+++ b/reid-lsp/src/main.rs
@@ -263,32 +263,20 @@ impl LanguageServer for Backend {
             });
 
             if let Some(token) = token {
-                dbg!(token);
-                if let Some(semantic_token) = analysis.state.map.get(&token.0) {
-                    dbg!(semantic_token);
-                    if let Some(symbol_id) = semantic_token.symbol {
-                        dbg!(symbol_id);
-                        let definition_id = analysis.state.find_definition(&symbol_id);
-                        if let Some(def_token_idx) = analysis.state.symbol_to_token.get(&definition_id) {
-                            dbg!(def_token_idx);
-                            if let Some(def_token) = analysis.tokens.get(*def_token_idx) {
-                                dbg!(def_token);
-                                return Ok(Some(GotoDefinitionResponse::Scalar(lsp_types::Location {
-                                    uri: params.text_document_position_params.text_document.uri,
-                                    range: Range {
-                                        start: lsp_types::Position {
-                                            line: def_token.position.1.max(1) - 1,
-                                            character: def_token.position.0.max(1) - 1,
-                                        },
-                                        end: lsp_types::Position {
-                                            line: def_token.position.1.max(1) - 1,
-                                            character: def_token.position.0.max(1) - 1 + def_token.token.len() as u32,
-                                        },
-                                    },
-                                })));
-                            }
-                        }
-                    }
+                if let Some(def_token) = analysis.find_definition(token.0) {
+                    return Ok(Some(GotoDefinitionResponse::Scalar(lsp_types::Location {
+                        uri: params.text_document_position_params.text_document.uri,
+                        range: Range {
+                            start: lsp_types::Position {
+                                line: def_token.position.1.max(1) - 1,
+                                character: def_token.position.0.max(1) - 1,
+                            },
+                            end: lsp_types::Position {
+                                line: def_token.position.1.max(1) - 1,
+                                character: def_token.position.0.max(1) - 1 + def_token.token.len() as u32,
+                            },
+                        },
+                    })));
                 }
             }
         };