feat: clear resolver cache by main pkg module (#1390)

Signed-off-by: he1pa <[email protected]>
kcl-lang · Jun 4, 2024 · e392979 · e392979
1 parent 3e40c8f
commit e392979
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 3 deletions.
diff --git a/kclvm/sema/src/resolver/mod.rs b/kclvm/sema/src/resolver/mod.rs
@@ -190,6 +190,7 @@ pub fn resolve_program_with_opts(
             cached_scope
                 .invalidate_pkgs
                 .insert(kclvm_ast::MAIN_PKG.to_string());
+            cached_scope.invalidate_main_pkg_modules = None;
         }
     }
     let scope = resolver.check_and_lint(kclvm_ast::MAIN_PKG);

diff --git a/kclvm/sema/src/resolver/scope.rs b/kclvm/sema/src/resolver/scope.rs
@@ -514,6 +514,9 @@ pub struct CachedScope {
     pub schema_mapping: IndexMap<String, Arc<RefCell<SchemaType>>>,
     pub node_ty_map: Rc<RefCell<NodeTyMap>>,
     pub invalidate_pkgs: HashSet<String>,
+    /// Specify the invalid module in the main package, used for invalidate_module().
+    /// If it is None, all modules in the main package will be invalidated
+    pub invalidate_main_pkg_modules: Option<HashSet<String>>,
     dependency_graph: DependencyGraph,
 }
 
@@ -534,7 +537,11 @@ impl DependencyGraph {
         self.node_map.clear();
     }
 
-    pub fn update(&mut self, program: &ast::Program) -> Result<HashSet<String>, String> {
+    pub fn update(
+        &mut self,
+        program: &ast::Program,
+        invalidate_main_pkg_modules: &Option<HashSet<String>>,
+    ) -> Result<HashSet<String>, String> {
         let mut new_modules = HashMap::new();
         for (pkgpath, modules) in program.pkgs.iter() {
             if pkgpath == kclvm_ast::MAIN_PKG {
@@ -568,6 +575,30 @@ impl DependencyGraph {
         }
         let mut invalidated_set = HashSet::new();
         if let Some(main_modules) = program.pkgs.get(kclvm_ast::MAIN_PKG) {
+            match invalidate_main_pkg_modules {
+                Some(modules) => {
+                    for module in main_modules {
+                        if modules.contains(&module.filename) {
+                            let result = self.invalidate_module(module)?;
+                            for pkg in result {
+                                invalidated_set.insert(pkg);
+                            }
+                            self.remove_dependency_from_pkg(&module.filename);
+                            self.add_new_module(module);
+                        }
+                    }
+                }
+                None => {
+                    for module in main_modules {
+                        let result = self.invalidate_module(module)?;
+                        for pkg in result {
+                            invalidated_set.insert(pkg);
+                        }
+                        self.remove_dependency_from_pkg(&module.filename);
+                        self.add_new_module(module);
+                    }
+                }
+            }
             for module in main_modules {
                 let result = self.invalidate_module(module)?;
                 for pkg in result {
@@ -679,8 +710,11 @@ impl CachedScope {
             invalidate_pkgs: HashSet::default(),
             dependency_graph: DependencyGraph::default(),
             schema_mapping: scope.schema_mapping.clone(),
+            invalidate_main_pkg_modules: None,
         };
-        let invalidated_pkgs = cached_scope.dependency_graph.update(program);
+        let invalidated_pkgs = cached_scope
+            .dependency_graph
+            .update(program, &cached_scope.invalidate_main_pkg_modules);
         cached_scope.invalidate_cache(invalidated_pkgs.as_ref());
         cached_scope
     }
@@ -690,6 +724,7 @@ impl CachedScope {
         self.node_ty_map.borrow_mut().clear();
         self.dependency_graph.clear();
         self.invalidate_pkgs.clear();
+        self.invalidate_main_pkg_modules = None;
     }
 
     pub fn invalidate_cache(&mut self, invalidated_pkgs: Result<&HashSet<String>, &String>) {
@@ -709,7 +744,9 @@ impl CachedScope {
             self.clear();
             self.program_root = program.root.clone();
         }
-        let invalidated_pkgs = self.dependency_graph.update(program);
+        let invalidated_pkgs = self
+            .dependency_graph
+            .update(program, &self.invalidate_main_pkg_modules);
         self.invalidate_cache(invalidated_pkgs.as_ref());
     }
 }
diff --git a/kclvm/tools/src/LSP/src/util.rs b/kclvm/tools/src/LSP/src/util.rs
@@ -130,6 +130,14 @@ pub(crate) fn compile_with_params(
     diags.extend(sess.1.borrow().diagnostics.clone());
 
     // Resolver
+    if let Some(cached_scope) = params.scope_cache.as_ref() {
+        if let Ok(mut cached_scope) = cached_scope.try_lock() {
+            let mut invalidate_main_pkg_modules = HashSet::new();
+            invalidate_main_pkg_modules.insert(params.file);
+            cached_scope.invalidate_main_pkg_modules = Some(invalidate_main_pkg_modules);
+        }
+    }
+
     let prog_scope = resolve_program_with_opts(
         &mut program,
         kclvm_sema::resolver::Options {