chore(els): improve goto definition/references

crates/els/definition.rs

@@ -36,7 +36,7 @@ impl<Checker: BuildRunnable, Parser: Parsable> Server<Checker, Parser> {
     ) -> ELSResult<GotoDefinitionResponse> {
         let uri = NormalizedUrl::new(params.text_document_position_params.text_document.uri);
         let pos = params.text_document_position_params.position;
-        if let Some(token) = self.file_cache.get_token(&uri, pos) {
+        if let Some(token) = self.file_cache.get_symbol(&uri, pos) {
             if let Some(vi) = self.get_definition(&uri, &token)? {
                 // If the target variable is an imported one, jump to the definition file.
                 // Else if the target variable is an alias, jump to the definition of it.

crates/els/file_cache.rs

@@ -13,7 +13,7 @@ use erg_common::dict::Dict;
 use erg_common::shared::Shared;
 use erg_common::traits::DequeStream;
 use erg_compiler::erg_parser::lex::Lexer;
-use erg_compiler::erg_parser::token::{Token, TokenStream};
+use erg_compiler::erg_parser::token::{Token, TokenCategory, TokenStream};
 
 use crate::_log;
 use crate::server::ELSResult;
@@ -139,9 +139,23 @@ impl FileCache {
                 return Some(tok.clone());
             }
         }
+        for tok in tokens.iter() {
+            if util::roughly_pos_in_loc(tok, pos) {
+                return Some(tok.clone());
+            }
+        }
         None
     }
 
+    /// a{pos}\n -> \n -> a
+    pub fn get_symbol(&self, uri: &NormalizedUrl, pos: Position) -> Option<Token> {
+        let mut token = self.get_token(uri, pos)?;
+        while !matches!(token.category(), TokenCategory::Symbol) {
+            token = self.get_token_relatively(uri, pos, -1)?;
+        }
+        Some(token)
+    }
+
     pub fn get_token_relatively(
         &self,
         uri: &NormalizedUrl,

crates/els/references.rs

@@ -4,6 +4,7 @@ use erg_compiler::varinfo::AbsLocation;
 
 use lsp_types::{Location, Position, ReferenceParams, Url};
 
+use crate::_log;
 use crate::server::{ELSResult, Server};
 use crate::util::{self, NormalizedUrl};
 
@@ -12,15 +13,15 @@ impl<Checker: BuildRunnable, Parser: Parsable> Server<Checker, Parser> {
         &mut self,
         params: ReferenceParams,
     ) -> ELSResult<Option<Vec<Location>>> {
+        _log!("references: {params:?}");
         let uri = NormalizedUrl::new(params.text_document_position.text_document.uri);
         let pos = params.text_document_position.position;
         let result = self.show_refs_inner(&uri, pos);
         Ok(Some(result))
     }
 
     fn show_refs_inner(&self, uri: &NormalizedUrl, pos: Position) -> Vec<lsp_types::Location> {
-        if let Some(tok) = self.file_cache.get_token(uri, pos) {
-            // send_log(format!("token: {tok}"))?;
+        if let Some(tok) = self.file_cache.get_symbol(uri, pos) {
             if let Some(visitor) = self.get_visitor(uri) {
                 if let Some(vi) = visitor.get_info(&tok) {
                     return self.get_refs_from_abs_loc(&vi.def_loc);

crates/els/util.rs

@@ -117,6 +117,18 @@ pub(crate) fn pos_in_loc<L: Locational>(loc: &L, pos: Position) -> bool {
     }
 }
 
+pub(crate) fn roughly_pos_in_loc<L: Locational>(loc: &L, pos: Position) -> bool {
+    let ln_begin = loc.ln_begin().unwrap_or(0);
+    let ln_end = loc.ln_end().unwrap_or(0);
+    let in_lines = (ln_begin..=ln_end).contains(&(pos.line + 1));
+    if ln_begin == ln_end {
+        in_lines
+            && (loc.col_begin().unwrap_or(0)..=loc.col_end().unwrap_or(0)).contains(&pos.character)
+    } else {
+        in_lines
+    }
+}
+
 pub(crate) fn pos_to_byte_index(src: &str, pos: Position) -> usize {
     if src.is_empty() {
         return 0;