chore: add ModuleEntry.ast

crates/els/diagnostics.rs

@@ -8,6 +8,7 @@ use std::time::Duration;
 
 use erg_common::consts::PYTHON_MODE;
 use erg_common::dict::Dict;
+use erg_common::shared::MappedRwLockReadGuard;
 use erg_common::spawn::{safe_yield, spawn_new_thread};
 use erg_common::style::*;
 use erg_common::traits::Stream;
@@ -49,7 +50,14 @@ impl<Checker: BuildRunnable, Parser: Parsable> Server<Checker, Parser> {
         } else {
             "exec"
         };
-        if let Some((old, new)) = self.analysis_result.get_ast(&uri).zip(self.get_ast(&uri)) {
+        let old = self.analysis_result.get_ast(&uri).or_else(|| {
+            let ent = self.get_shared()?.mod_cache.get(&path)?;
+            ent.ast.as_ref()?;
+            Some(MappedRwLockReadGuard::map(ent, |ent| {
+                ent.ast.as_ref().unwrap()
+            }))
+        });
+        if let Some((old, new)) = old.zip(self.get_ast(&uri)) {
             if ASTDiff::diff(old, &new).is_nop() {
                 crate::_log!(self, "no changes: {uri}");
                 return Ok(());
@@ -92,15 +100,18 @@ impl<Checker: BuildRunnable, Parser: Parsable> Server<Checker, Parser> {
             }
         };
         if let Some(shared) = self.get_shared() {
+            let ast = self.get_ast(&uri);
             if mode == "declare" {
                 shared.py_mod_cache.register(
                     path,
+                    ast,
                     artifact.object.clone(),
                     checker.get_context().unwrap().clone(),
                 );
             } else {
                 shared.mod_cache.register(
                     path,
+                    ast,
                     artifact.object.clone(),
                     checker.get_context().unwrap().clone(),
                 );

crates/els/tests/c.er

@@ -0,0 +1,3 @@
+{neighbor;} = import "b"
+
+print! neighbor

crates/els/tests/test.rs

@@ -6,6 +6,7 @@ use lsp_types::{
 };
 const FILE_A: &str = "tests/a.er";
 const FILE_B: &str = "tests/b.er";
+const FILE_C: &str = "tests/c.er";
 const FILE_IMPORTS: &str = "tests/imports.er";
 
 use els::{NormalizedUrl, Server};
@@ -128,6 +129,13 @@ fn test_references() -> Result<(), Box<dyn std::error::Error>> {
     let locations = client.request_references(uri.raw(), 1, 4)?.unwrap();
     assert_eq!(locations.len(), 1);
     assert_eq!(&locations[0].range, &oneline_range(1, 4, 5));
+    client.notify_open(FILE_C)?;
+    client.notify_open(FILE_B)?;
+    let uri_b = NormalizedUrl::from_file_path(Path::new(FILE_B).canonicalize()?)?;
+    let uri_c = NormalizedUrl::from_file_path(Path::new(FILE_C).canonicalize()?)?;
+    let locations = client.request_references(uri_b.raw(), 0, 2)?.unwrap();
+    assert_eq!(locations.len(), 1);
+    assert_eq!(NormalizedUrl::new(locations[0].uri.clone()), uri_c);
     Ok(())
 }
 

crates/erg_compiler/context/initialize/mod.rs

@@ -1132,7 +1132,7 @@ impl Context {
         let module = ModuleContext::new(ctx, dict! {});
         shared
             .mod_cache
-            .register(PathBuf::from("<builtins>"), None, module);
+            .register(PathBuf::from("<builtins>"), None, None, module);
     }
 
     pub fn new_module<S: Into<Str>>(

crates/erg_compiler/context/register.rs

@@ -8,7 +8,7 @@ use std::thread::sleep;
 use std::time::{Duration, SystemTime};
 
 use erg_common::config::ErgMode;
-use erg_common::consts::{ERG_MODE, PYTHON_MODE};
+use erg_common::consts::{ELS, ERG_MODE, PYTHON_MODE};
 use erg_common::dict::Dict;
 use erg_common::env::{is_pystd_main_module, is_std_decl_path};
 use erg_common::erg_util::BUILTIN_ERG_MODS;
@@ -27,6 +27,7 @@ use ast::{
     VarName,
 };
 use erg_parser::ast;
+use erg_parser::parse::SimpleParser;
 
 use crate::ty::constructors::{
     free_var, func, func0, func1, proc, ref_, ref_mut, tp_enum, unknown_len_array_t, v_enum,
@@ -2015,6 +2016,13 @@ impl Context {
         let mod_ctx = ModuleContext::new(self.clone(), dict! {});
         let mut builder = HIRBuilder::new_with_ctx(mod_ctx);
         let src = Input::file(path.to_path_buf()).read();
+        let ast = if ELS {
+            SimpleParser::parse(src.clone())
+                .ok()
+                .map(|artifact| artifact.ast)
+        } else {
+            None
+        };
         let mode = if path.to_string_lossy().ends_with("d.er") {
             "declare"
         } else {
@@ -2031,7 +2039,7 @@ impl Context {
         } else {
             &self.shared().mod_cache
         };
-        cache.register(path.to_path_buf(), hir, ctx);
+        cache.register(path.to_path_buf(), ast, hir, ctx);
     }
 
     /// If the path is like `foo/bar`, check if `bar` is a public module (the definition is in `foo/__init__.er`)
@@ -2105,6 +2113,13 @@ impl Context {
             .input
             .try_read()
             .map_err(|_| self.import_err(line!(), __name__, loc))?;
+        let ast = if ELS {
+            SimpleParser::parse(src.clone())
+                .ok()
+                .map(|artifact| artifact.ast)
+        } else {
+            None
+        };
         let name = __name__.clone();
         let _path = path.clone();
         let shared = self.shared.as_ref().unwrap().inherit(path.clone());
@@ -2114,6 +2129,7 @@ impl Context {
                 Ok(artifact) => {
                     shared.mod_cache.register(
                         _path.clone(),
+                        ast,
                         Some(artifact.object),
                         builder.pop_mod_ctx().unwrap(),
                     );
@@ -2121,9 +2137,12 @@ impl Context {
                 }
                 Err(artifact) => {
                     if let Some(hir) = artifact.object {
-                        shared
-                            .mod_cache
-                            .register(_path, Some(hir), builder.pop_mod_ctx().unwrap());
+                        shared.mod_cache.register(
+                            _path,
+                            ast,
+                            Some(hir),
+                            builder.pop_mod_ctx().unwrap(),
+                        );
                     }
                     shared.warns.extend(artifact.warns);
                     shared.errors.extend(artifact.errors);
@@ -2315,6 +2334,13 @@ impl Context {
             .input
             .try_read()
             .map_err(|_| self.import_err(line!(), __name__, loc))?;
+        let ast = if ELS {
+            SimpleParser::parse(src.clone())
+                .ok()
+                .map(|artifact| artifact.ast)
+        } else {
+            None
+        };
         let mut builder = HIRBuilder::new_with_cache(
             cfg,
             self.mod_name(&path),
@@ -2323,12 +2349,12 @@ impl Context {
         match builder.build(src, "declare") {
             Ok(artifact) => {
                 let ctx = builder.pop_mod_ctx().unwrap();
-                py_mod_cache.register(path.clone(), Some(artifact.object), ctx);
+                py_mod_cache.register(path.clone(), ast, Some(artifact.object), ctx);
                 Ok(path)
             }
             Err(artifact) => {
                 if let Some(hir) = artifact.object {
-                    py_mod_cache.register(path, Some(hir), builder.pop_mod_ctx().unwrap());
+                    py_mod_cache.register(path, ast, Some(hir), builder.pop_mod_ctx().unwrap());
                 }
                 Err(artifact.errors)
             }

crates/erg_compiler/lower.rs

@@ -731,7 +731,10 @@ impl ASTLowerer {
                     &self.cfg.input,
                     &self.module.context,
                 ) {
-                    Triple::Ok(vi) => vi,
+                    Triple::Ok(vi) => {
+                        self.inc_ref(attr.ident.inspect(), &vi, &attr.ident.name);
+                        vi
+                    }
                     Triple::Err(errs) => {
                         self.errs.push(errs);
                         VarInfo::ILLEGAL
@@ -757,7 +760,6 @@ impl ASTLowerer {
                         VarInfo::ILLEGAL
                     }
                 };
-                self.inc_ref(attr.ident.inspect(), &vi, &attr.ident.name);
                 let ident = hir::Identifier::new(attr.ident, None, vi);
                 let acc = hir::Accessor::Attr(hir::Attribute::new(obj, ident));
                 Ok(acc)

crates/erg_compiler/module/cache.rs

@@ -12,6 +12,7 @@ use erg_common::shared::{
     MappedRwLockReadGuard, MappedRwLockWriteGuard, RwLockReadGuard, RwLockWriteGuard, Shared,
 };
 use erg_common::Str;
+use erg_parser::ast::Module;
 
 use crate::context::ModuleContext;
 use crate::hir::HIR;
@@ -34,6 +35,8 @@ impl ModId {
 #[derive(Debug, Clone)]
 pub struct ModuleEntry {
     pub id: ModId, // builtin == 0, __main__ == 1
+    /// mainly for ELS
+    pub ast: Option<Module>,
     pub hir: Option<HIR>,
     pub module: Arc<ModuleContext>,
 }
@@ -49,9 +52,10 @@ impl fmt::Display for ModuleEntry {
 }
 
 impl ModuleEntry {
-    pub fn new(id: ModId, hir: Option<HIR>, ctx: ModuleContext) -> Self {
+    pub fn new(id: ModId, ast: Option<Module>, hir: Option<HIR>, ctx: ModuleContext) -> Self {
         Self {
             id,
+            ast,
             hir,
             module: Arc::new(ctx),
         }
@@ -60,6 +64,7 @@ impl ModuleEntry {
     pub fn builtin(ctx: ModuleContext) -> Self {
         Self {
             id: ModId::builtin(),
+            ast: None,
             hir: None,
             module: Arc::new(ctx),
         }
@@ -111,10 +116,16 @@ impl ModuleCache {
         self.cache.get_mut(path)
     }
 
-    pub fn register(&mut self, path: NormalizedPathBuf, hir: Option<HIR>, ctx: ModuleContext) {
+    pub fn register(
+        &mut self,
+        path: NormalizedPathBuf,
+        ast: Option<Module>,
+        hir: Option<HIR>,
+        ctx: ModuleContext,
+    ) {
         self.last_id += 1;
         let id = ModId::new(self.last_id);
-        let entry = ModuleEntry::new(id, hir, ctx);
+        let entry = ModuleEntry::new(id, ast, hir, ctx);
         self.cache.insert(path, entry);
     }
 
@@ -251,10 +262,11 @@ impl SharedModuleCache {
     pub fn register<P: Into<NormalizedPathBuf>>(
         &self,
         path: P,
+        ast: Option<Module>,
         hir: Option<HIR>,
         ctx: ModuleContext,
     ) {
-        self.0.borrow_mut().register(path.into(), hir, ctx);
+        self.0.borrow_mut().register(path.into(), ast, hir, ctx);
     }
 
     pub fn remove<Q: Eq + Hash + ?Sized>(&self, path: &Q) -> Option<ModuleEntry>
@@ -278,7 +290,12 @@ impl SharedModuleCache {
             return;
         };
         self.0.borrow_mut().clear();
-        self.register(builtin_path, None, Arc::try_unwrap(builtin.module).unwrap());
+        self.register(
+            builtin_path,
+            None,
+            None,
+            Arc::try_unwrap(builtin.module).unwrap(),
+        );
     }
 
     pub fn rename_path<P: Into<NormalizedPathBuf>>(&self, path: &Path, new: P) {

crates/erg_compiler/varinfo.rs

@@ -1,4 +1,5 @@
 use std::fmt;
+use std::path::Path;
 
 use erg_common::error::Location;
 use erg_common::pathutil::NormalizedPathBuf;
@@ -159,6 +160,14 @@ impl AbsLocation {
         Self::new(None, Location::Unknown)
     }
 
+    pub fn is_unknown(&self) -> bool {
+        self.module.is_none() && self.loc.is_unknown()
+    }
+
+    pub fn starts_with(&self, path: impl AsRef<Path>) -> bool {
+        self.module.as_ref().is_some_and(|p| p.starts_with(path))
+    }
+
     pub fn code(&self) -> Option<String> {
         use std::io::{BufRead, BufReader};
         self.module.as_ref().and_then(|module| {

crates/erg_parser/ast.rs

@@ -4966,7 +4966,7 @@ impl Module {
     }
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct AST {
     pub name: Str,
     pub module: Module,

crates/erg_parser/parse.rs

@@ -116,6 +116,12 @@ impl Parsable for SimpleParser {
     }
 }
 
+impl SimpleParser {
+    pub fn parse(code: String) -> Result<CompleteArtifact, IncompleteArtifact> {
+        <Self as Parsable>::parse(code)
+    }
+}
+
 enum ExprOrOp {
     Expr(Expr),
     Op(Token),