feat: Clear the cache according to the changed module.

kclvm/ast/src/ast.rs

@@ -397,6 +397,17 @@ impl Program {
         }
         None
     }
+
+    pub fn get_module(&self, module_name: &str) -> Option<&Module> {
+        for (_, modules) in &self.pkgs {
+            for module in modules {
+                if module.filename == module_name {
+                    return Some(module);
+                }
+            }
+        }
+        None
+    }
 }
 
 /// Module is an abstract syntax tree for a single KCL file.

kclvm/sema/src/resolver/mod.rs

@@ -203,7 +203,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;
+            cached_scope.invalidate_pkg_modules = None;
         }
     }
     if opts.type_erasure {

kclvm/sema/src/resolver/scope.rs

@@ -534,9 +534,9 @@ pub struct CachedScope {
     pub schema_mapping: IndexMap<String, Arc<RefCell<SchemaType>>>,
     pub node_ty_map: NodeTyMap,
     pub invalidate_pkgs: HashSet<String>,
-    /// Specify the invalid module in the main package, used for invalidate_module().
+    /// Specify the invalid module in the program packages, 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>>,
+    pub invalidate_pkg_modules: Option<HashSet<String>>,
     dependency_graph: DependencyGraph,
 }
 
@@ -560,7 +560,7 @@ impl DependencyGraph {
     pub fn update(
         &mut self,
         program: &ast::Program,
-        invalidate_main_pkg_modules: &Option<HashSet<String>>,
+        invalidate_pkg_modules: &Option<HashSet<String>>,
     ) -> Result<HashSet<String>, String> {
         let mut new_modules = HashMap::new();
         for (pkgpath, modules) in program.pkgs.iter() {
@@ -593,24 +593,26 @@ impl DependencyGraph {
         for new_module in new_modules.values() {
             self.add_new_module(new_module);
         }
+
         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);
-                        }
+
+        match invalidate_pkg_modules {
+            Some(modules) => {
+                for module_name in modules {
+                    let result = self.invalidate_module(module_name)?;
+                    for pkg in result {
+                        invalidated_set.insert(pkg);
+                    }
+                    self.remove_dependency_from_pkg(&module_name);
+                    if let Some(module) = program.get_module(module_name) {
+                        self.add_new_module(module);
                     }
                 }
-                None => {
+            }
+            None => {
+                if let Some(main_modules) = program.pkgs.get(kclvm_ast::MAIN_PKG) {
                     for module in main_modules {
-                        let result = self.invalidate_module(module)?;
+                        let result = self.invalidate_module(&module.filename)?;
                         for pkg in result {
                             invalidated_set.insert(pkg);
                         }
@@ -641,13 +643,9 @@ impl DependencyGraph {
         }
     }
 
-    fn invalidate_module(
-        &mut self,
-        changed_module: &ast::Module,
-    ) -> Result<HashSet<String>, String> {
-        let module_file = changed_module.filename.clone();
+    fn invalidate_module(&mut self, changed_module: &str) -> Result<HashSet<String>, String> {
         let mut invalidated_set = HashSet::new();
-        if let Some(pkgpaths) = self.module_map.get(&module_file).cloned() {
+        if let Some(pkgpaths) = self.module_map.get(changed_module).cloned() {
             let mut pkg_queue = VecDeque::new();
             for pkgpath in pkgpaths.iter() {
                 pkg_queue.push_back(self.node_map.get(pkgpath));
@@ -722,11 +720,11 @@ impl CachedScope {
             invalidate_pkgs: HashSet::default(),
             dependency_graph: DependencyGraph::default(),
             schema_mapping: scope.schema_mapping.clone(),
-            invalidate_main_pkg_modules: None,
+            invalidate_pkg_modules: None,
         };
         let invalidated_pkgs = cached_scope
             .dependency_graph
-            .update(program, &cached_scope.invalidate_main_pkg_modules);
+            .update(program, &cached_scope.invalidate_pkg_modules);
         cached_scope.invalidate_cache(invalidated_pkgs.as_ref());
         cached_scope
     }
@@ -736,7 +734,7 @@ impl CachedScope {
         self.node_ty_map.clear();
         self.dependency_graph.clear();
         self.invalidate_pkgs.clear();
-        self.invalidate_main_pkg_modules = None;
+        self.invalidate_pkg_modules = None;
     }
 
     pub fn invalidate_cache(&mut self, invalidated_pkgs: Result<&HashSet<String>, &String>) {
@@ -758,7 +756,7 @@ impl CachedScope {
         }
         let invalidated_pkgs = self
             .dependency_graph
-            .update(program, &self.invalidate_main_pkg_modules);
+            .update(program, &self.invalidate_pkg_modules);
         self.invalidate_cache(invalidated_pkgs.as_ref());
     }
 }

kclvm/sema/src/resolver/test_data/cache/b/b.k

@@ -0,0 +1 @@
+import ..c

kclvm/sema/src/resolver/test_data/cache/c/c.k

@@ -0,0 +1 @@
+a = 1 

kclvm/sema/src/resolver/test_data/cache/main.k

@@ -0,0 +1,5 @@
+
+# main -> a
+#      -> b -> c
+import .a
+import .b

kclvm/sema/src/resolver/tests.rs

@@ -9,6 +9,7 @@ use crate::ty::{Type, TypeKind};
 use anyhow::Result;
 use kclvm_ast::ast;
 use kclvm_ast::pos::ContainsPos;
+use kclvm_ast::MAIN_PKG;
 use kclvm_error::*;
 use kclvm_parser::load_program;
 use kclvm_parser::parse_file_force_errors;
@@ -17,6 +18,7 @@ use kclvm_parser::ParseSession;
 use kclvm_utils::path::PathPrefix;
 use parking_lot::lock_api::RwLock;
 use std::collections::HashMap;
+use std::collections::HashSet;
 use std::path::Path;
 use std::sync::Arc;
 
@@ -906,3 +908,180 @@ fn test_schema_index_signature_check() {
     let diags = scope.handler.diagnostics;
     assert!(diags.is_empty())
 }
+
+#[test]
+fn test_clear_cache_by_module() {
+    let sess = Arc::new(ParseSession::default());
+    let mut program = load_program(
+        sess.clone(),
+        &["./src/resolver/test_data/cache/main.k"],
+        None,
+        None,
+    )
+    .unwrap()
+    .program;
+
+    let scope = resolve_program_with_opts(
+        &mut program,
+        Options {
+            merge_program: false,
+            type_erasure: false,
+            ..Default::default()
+        },
+        None,
+    );
+    let cached_scope = Arc::new(RwLock::new(CachedScope::new(&scope, &program)));
+    // first compile
+    let _ = resolve_program_with_opts(
+        &mut program,
+        Options {
+            merge_program: false,
+            type_erasure: false,
+            ..Default::default()
+        },
+        Some(cached_scope.clone()),
+    );
+
+    // recompile and clear cache
+    let invalidate_module = std::fs::canonicalize(std::path::PathBuf::from(
+        "./src/resolver/test_data/cache/main.k",
+    ))
+    .unwrap()
+    .to_str()
+    .unwrap()
+    .to_string()
+    .adjust_canonicalization();
+
+    if let Some(mut cached_scope) = cached_scope.try_write() {
+        let mut invalidate_pkg_modules = HashSet::new();
+        invalidate_pkg_modules.insert(invalidate_module);
+        cached_scope.invalidate_pkg_modules = Some(invalidate_pkg_modules);
+    };
+
+    let _ = resolve_program_with_opts(
+        &mut program,
+        Options {
+            merge_program: false,
+            type_erasure: false,
+            ..Default::default()
+        },
+        Some(cached_scope.clone()),
+    );
+    if let Some(cached_scope) = cached_scope.try_write() {
+        // main - a
+        //      - b - c
+        // invalidate main, invalidate_pkgs main
+        let mut expect = HashSet::new();
+        expect.insert(MAIN_PKG.to_string());
+        assert_eq!(cached_scope.invalidate_pkgs, expect);
+    };
+
+    // recompile and clear cache
+    let invalidate_module = std::fs::canonicalize(std::path::PathBuf::from(
+        "./src/resolver/test_data/cache/a/a.k",
+    ))
+    .unwrap()
+    .to_str()
+    .unwrap()
+    .to_string()
+    .adjust_canonicalization();
+
+    if let Some(mut cached_scope) = cached_scope.try_write() {
+        let mut invalidate_pkg_modules = HashSet::new();
+        invalidate_pkg_modules.insert(invalidate_module);
+        cached_scope.invalidate_pkg_modules = Some(invalidate_pkg_modules);
+    };
+
+    let _ = resolve_program_with_opts(
+        &mut program,
+        Options {
+            merge_program: false,
+            type_erasure: false,
+            ..Default::default()
+        },
+        Some(cached_scope.clone()),
+    );
+
+    if let Some(cached_scope) = cached_scope.try_write() {
+        // main - a
+        //      - b - c
+        // invalidate a, invalidate_pkgs a, main
+        let mut expect = HashSet::new();
+        expect.insert(MAIN_PKG.to_string());
+        expect.insert("cache.a".to_string());
+        assert_eq!(cached_scope.invalidate_pkgs, expect);
+    };
+
+    // recompile and clear cache
+    let invalidate_module = std::fs::canonicalize(std::path::PathBuf::from(
+        "./src/resolver/test_data/cache/b/b.k",
+    ))
+    .unwrap()
+    .to_str()
+    .unwrap()
+    .to_string()
+    .adjust_canonicalization();
+
+    if let Some(mut cached_scope) = cached_scope.try_write() {
+        let mut invalidate_pkg_modules = HashSet::new();
+        invalidate_pkg_modules.insert(invalidate_module);
+        cached_scope.invalidate_pkg_modules = Some(invalidate_pkg_modules);
+    };
+
+    let _ = resolve_program_with_opts(
+        &mut program,
+        Options {
+            merge_program: false,
+            type_erasure: false,
+            ..Default::default()
+        },
+        Some(cached_scope.clone()),
+    );
+
+    if let Some(cached_scope) = cached_scope.try_write() {
+        // main - a
+        //      - b - c
+        // invalidate b, invalidate_pkgs b, main
+        let mut expect = HashSet::new();
+        expect.insert(MAIN_PKG.to_string());
+        expect.insert("cache.b".to_string());
+        assert_eq!(cached_scope.invalidate_pkgs, expect);
+    };
+
+    // recompile and clear cache
+    let invalidate_module = std::fs::canonicalize(std::path::PathBuf::from(
+        "./src/resolver/test_data/cache/c/c.k",
+    ))
+    .unwrap()
+    .to_str()
+    .unwrap()
+    .to_string()
+    .adjust_canonicalization();
+
+    if let Some(mut cached_scope) = cached_scope.try_write() {
+        let mut invalidate_pkg_modules = HashSet::new();
+        invalidate_pkg_modules.insert(invalidate_module);
+        cached_scope.invalidate_pkg_modules = Some(invalidate_pkg_modules);
+    };
+
+    let _ = resolve_program_with_opts(
+        &mut program,
+        Options {
+            merge_program: false,
+            type_erasure: false,
+            ..Default::default()
+        },
+        Some(cached_scope.clone()),
+    );
+
+    if let Some(cached_scope) = cached_scope.try_write() {
+        // main - a
+        //      - b - c
+        // invalidate c, invalidate_pkgs c, b, main
+        let mut expect = HashSet::new();
+        expect.insert(MAIN_PKG.to_string());
+        expect.insert("cache.b".to_string());
+        expect.insert("cache.c".to_string());
+        assert_eq!(cached_scope.invalidate_pkgs, expect);
+    };
+}

kclvm/tools/src/LSP/src/util.rs

@@ -208,9 +208,9 @@ pub(crate) fn compile_with_params(
     // Resolver
     if let Some(cached_scope) = params.scope_cache.as_ref() {
         if let Some(mut cached_scope) = cached_scope.try_write() {
-            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 mut invalidate_pkg_modules = HashSet::new();
+            invalidate_pkg_modules.insert(params.file);
+            cached_scope.invalidate_pkg_modules = Some(invalidate_pkg_modules);
         }
     }