fix: fix lsp memory leak (#1841)

* fix: fix lsp memery leak Signed-off-by: he1pa <18012015693@163.com> * ut Signed-off-by: he1pa <18012015693@163.com> --------- Signed-off-by: he1pa <18012015693@163.com>
kcl-lang · Feb 5, 2025 · 6affc15 · 6affc15
1 parent 98739db
commit 6affc15
Show file tree

Hide file tree

Showing 5 changed files with 88 additions and 2 deletions.
diff --git a/kclvm/sema/src/core/global_state.rs b/kclvm/sema/src/core/global_state.rs
@@ -23,6 +23,13 @@ pub struct GlobalState {
     pub(crate) sema_db: SemanticDB,
     // new and invalidate(changed and affected by changed) pkg from CachedScope::update()
     pub new_or_invalidate_pkgs: HashSet<String>,
+
+    pub ctx: GlobalStateContext,
+}
+
+#[derive(Default, Debug, Clone)]
+pub struct GlobalStateContext {
+    pub has_init_builtin: bool,
 }
 
 impl GlobalState {

diff --git a/kclvm/sema/src/core/symbol.rs b/kclvm/sema/src/core/symbol.rs
@@ -240,6 +240,7 @@ impl SymbolData {
         match id.get_kind() {
             SymbolKind::Schema => {
                 self.schemas.remove(id.get_id());
+                self.symbols_info.schema_builtin_symbols.remove(id);
             }
             SymbolKind::Attribute => {
                 self.attributes.remove(id.get_id());
@@ -887,11 +888,21 @@ impl SymbolData {
         &self.symbols_info.node_symbol_map
     }
 
+    #[inline]
+    pub fn get_node_symbol_map_mut(&mut self) -> &mut IndexMap<NodeKey, SymbolRef> {
+        &mut self.symbols_info.node_symbol_map
+    }
+
     #[inline]
     pub fn get_symbol_node_map(&self) -> &IndexMap<SymbolRef, NodeKey> {
         &self.symbols_info.symbol_node_map
     }
 
+    #[inline]
+    pub fn get_symbol_node_map_mut(&mut self) -> &mut IndexMap<SymbolRef, NodeKey> {
+        &mut self.symbols_info.symbol_node_map
+    }
+
     #[inline]
     pub fn get_fully_qualified_name_map(&self) -> &IndexMap<String, SymbolRef> {
         &self.symbols_info.fully_qualified_name_map
@@ -913,6 +924,9 @@ impl SymbolData {
         }
         for symbol in to_remove {
             self.remove_symbol(&symbol);
+            if let Some(node_id) = self.get_symbol_node_map_mut().remove(&symbol) {
+                self.get_node_symbol_map_mut().remove(&node_id);
+            }
         }
     }
 

diff --git a/kclvm/sema/src/namer/mod.rs b/kclvm/sema/src/namer/mod.rs
@@ -100,13 +100,18 @@ impl<'ctx> Namer<'ctx> {
 
     // serial namer pass
     pub fn find_symbols(program: &'ctx Program, gs: &'ctx mut GlobalState) {
+        let has_init_builtin = gs.ctx.has_init_builtin;
+        gs.ctx.has_init_builtin = true;
         let mut namer = Self::new(program, gs);
         namer.ctx.current_package_info = Some(PackageInfo::new(
             BUILTIN_SYMBOL_PKG_PATH.to_string(),
             "".to_string(),
             true,
         ));
-        namer.init_builtin_symbols();
+        if !has_init_builtin {
+            namer.init_builtin_symbols();
+        }
+
         namer
             .gs
             .get_packages_mut()

diff --git a/kclvm/sema/src/resolver/mod.rs b/kclvm/sema/src/resolver/mod.rs
@@ -208,14 +208,23 @@ pub fn resolve_program_with_opts(
             cached_scope.invalidate_pkgs.clear();
             cached_scope.update(program);
             resolver.scope_map = cached_scope.scope_map.clone();
-            resolver.node_ty_map = Rc::new(RefCell::new(cached_scope.node_ty_map.clone()));
             resolver.ctx.schema_mapping = cached_scope.schema_mapping.clone();
             cached_scope
                 .invalidate_pkgs
                 .insert(kclvm_ast::MAIN_PKG.to_string());
             for pkg in &cached_scope.invalidate_pkgs {
                 resolver.scope_map.remove(pkg);
             }
+            let mut nodes = vec![];
+            for node in cached_scope.node_ty_map.keys() {
+                if cached_scope.invalidate_pkgs.contains(&node.pkgpath) {
+                    nodes.push(node.clone());
+                }
+            }
+            for node in nodes {
+                cached_scope.node_ty_map.remove(&node);
+            }
+            resolver.node_ty_map = Rc::new(RefCell::new(cached_scope.node_ty_map.clone()));
         }
     }
     let scope = resolver.check_and_lint_all_pkgs();

diff --git a/kclvm/sema/src/resolver/tests.rs b/kclvm/sema/src/resolver/tests.rs
@@ -1098,3 +1098,54 @@ fn test_clear_cache_by_module() {
         assert_eq!(cached_scope.invalidate_pkgs, expect);
     };
 }
+
+#[test]
+fn clear_cache_test() {
+    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_cache = Arc::new(RwLock::new(CachedScope::default()));
+    let first_scope = resolve_program_with_opts(
+        &mut program,
+        Options {
+            merge_program: false,
+            type_erasure: false,
+            ..Default::default()
+        },
+        Some(scope_cache.clone()),
+    );
+
+    let mut program = load_program(
+        sess.clone(),
+        &["./src/resolver/test_data/cache/main.k"],
+        None,
+        None,
+    )
+    .unwrap()
+    .program;
+
+    let second_scope = resolve_program_with_opts(
+        &mut program,
+        Options {
+            merge_program: false,
+            type_erasure: false,
+            ..Default::default()
+        },
+        Some(scope_cache.clone()),
+    );
+    let first_node_ty_map_len = first_scope.node_ty_map.borrow().len();
+    let second_node_ty_map_len = second_scope.node_ty_map.borrow().len();
+    assert_eq!(first_node_ty_map_len, second_node_ty_map_len);
+    assert_eq!(first_scope.scope_map.len(), second_scope.scope_map.len());
+    assert_eq!(
+        first_scope.schema_mapping.len(),
+        second_scope.schema_mapping.len()
+    );
+}