Refactor document data cleanup

Author: st0012

rust/rubydex/src/model/declaration.rs

@@ -357,10 +357,8 @@ impl Declaration {
         all_declarations!(self, it => it.definition_ids.is_empty())
     }
 
-    /// Returns true if this declaration is no longer anchored to anything
-    /// and can be removed. For most declarations, that means having no
-    /// definitions. Singleton classes are also kept alive by members and
-    /// inbound constant references.
+    /// Returns true if this declaration has no local anchors and can be removed.
+    /// Graph-level cleanup can still keep declarations alive based on external anchors.
     #[must_use]
     pub fn is_removable(&self) -> bool {
         if let Some(ns) = self.as_singleton_class() {

rust/rubydex/src/model/definitions.rs

@@ -179,6 +179,18 @@ impl Definition {
         }
     }
 
+    #[must_use]
+    pub fn has_extend_mixin(&self) -> bool {
+        let mixins = match self {
+            Definition::Class(definition) => definition.mixins(),
+            Definition::SingletonClass(definition) => definition.mixins(),
+            Definition::Module(definition) => definition.mixins(),
+            _ => return false,
+        };
+
+        mixins.iter().any(|mixin| matches!(mixin, Mixin::Extend(_)))
+    }
+
     #[must_use]
     pub fn is_deprecated(&self) -> bool {
         all_definitions!(self, it => it.flags().is_deprecated())

rust/rubydex/src/model/graph.rs

@@ -1017,69 +1017,115 @@ impl Graph {
     }
 
     /// Removes raw document data (refs, defs, names, strings) from maps.
-    /// Does not touch declarations or perform invalidation -- that is handled by `invalidate`.
+    ///
+    /// `invalidate` has already handled dependency invalidation and the normal
+    /// definition detachment path. This cleanup removes the raw map entries and
+    /// prunes declarations that become removable only after those entries are detached.
     fn remove_document_data(&mut self, document: &Document) {
+        self.remove_document_method_references(document);
+
+        let mut declarations_to_prune = self.remove_document_constant_references(document);
+        declarations_to_prune.extend(self.detach_unmapped_document_definitions(document));
+        self.prune_removable_declarations(declarations_to_prune);
+
+        self.remove_document_definitions(document);
+    }
+
+    fn remove_document_method_references(&mut self, document: &Document) {
         for ref_id in document.method_references() {
             if let Some(method_ref) = self.method_references.remove(ref_id) {
                 self.untrack_string(*method_ref.str());
             }
         }
+    }
 
-        let mut emptied_by_ref_removal: Vec<InvalidationItem> = Vec::new();
+    fn remove_document_constant_references(&mut self, document: &Document) -> Vec<DeclarationId> {
+        let mut declarations_to_prune = Vec::new();
         for ref_id in document.constant_references() {
             if let Some(constant_ref) = self.constant_references.remove(ref_id) {
                 // Detach from target declaration. References unresolved during invalidation
                 // were already detached; this catches the rest.
-                if let Some(NameRef::Resolved(resolved)) = self.names.get(constant_ref.name_id()) {
-                    let target_id = *resolved.declaration_id();
-                    if let Some(declaration) = self.declarations.get_mut(&target_id) {
-                        declaration.remove_constant_reference(ref_id);
-                        if declaration.is_removable() {
-                            emptied_by_ref_removal.push(InvalidationItem::Declaration(target_id));
-                        }
-                    }
+                if let Some(target_id) =
+                    self.detach_constant_reference_from_resolved_target(*constant_ref.name_id(), *ref_id)
+                {
+                    declarations_to_prune.push(target_id);
                 }
 
                 self.remove_name_dependent(*constant_ref.name_id(), NameDependent::Reference(*ref_id));
                 self.untrack_name(*constant_ref.name_id());
             }
         }
 
-        if !emptied_by_ref_removal.is_empty() {
-            self.invalidate_graph(emptied_by_ref_removal, IdentityHashMap::default());
+        declarations_to_prune
+    }
+
+    fn detach_constant_reference_from_resolved_target(
+        &mut self,
+        name_id: NameId,
+        ref_id: ConstantReferenceId,
+    ) -> Option<DeclarationId> {
+        let target_id = match self.names.get(&name_id) {
+            Some(NameRef::Resolved(resolved)) => *resolved.declaration_id(),
+            Some(NameRef::Unresolved(_)) | None => return None,
+        };
+
+        let declaration = self.declarations.get_mut(&target_id)?;
+        declaration.remove_constant_reference(&ref_id);
+        self.declaration_is_removable(target_id).then_some(target_id)
+    }
+
+    fn detach_unmapped_document_definitions(&mut self, document: &Document) -> Vec<DeclarationId> {
+        let unmapped_definition_ids = self.unmapped_document_definition_ids(document);
+
+        if unmapped_definition_ids.is_empty() {
+            return Vec::new();
         }
 
-        // Detach removed definitions from their declarations.
-        // Most definitions were already detached by invalidate_declaration via
-        // pending_detachments. Definitions not handled by pending_detachments are
-        // those where definition_to_declaration_id returns None, for example:
-        //   - methods inside `class << self` when <Foo> was unresolved by a prior deletion
-        //   - instance variables in class body (owned by singleton, but lookup resolves to class)
-        //   - definitions whose enclosing namespace name chain is broken
-        // Detach those by scanning declarations for the remainder.
-        let missed_def_ids: Vec<DefinitionId> = document
+        self.detach_definitions_from_any_declaration(&unmapped_definition_ids)
+    }
+
+    /// Definitions that cannot be mapped through `definition_id_to_declaration_id`
+    /// are skipped by the normal invalidation detachment path. That happens when
+    /// the definition's name chain is already broken, or when the resolved owner
+    /// differs from the lexical owner used by lookup, such as singleton-owned
+    /// methods and instance variables.
+    fn unmapped_document_definition_ids(&self, document: &Document) -> Vec<DefinitionId> {
+        document
             .definitions()
             .iter()
             .copied()
             .filter(|def_id| self.definition_id_to_declaration_id(*def_id).is_none())
-            .collect();
+            .collect()
+    }
 
-        if !missed_def_ids.is_empty() {
-            let mut emptied: Vec<InvalidationItem> = Vec::new();
-            for (decl_id, declaration) in &mut self.declarations {
-                let had_definitions = !declaration.definitions().is_empty();
-                for def_id in &missed_def_ids {
-                    declaration.remove_definition(def_id);
-                }
-                if had_definitions && declaration.is_removable() {
-                    emptied.push(InvalidationItem::Declaration(*decl_id));
-                }
+    fn detach_definitions_from_any_declaration(
+        &mut self,
+        document_definition_ids: &[DefinitionId],
+    ) -> Vec<DeclarationId> {
+        let mut touched_declarations = Vec::new();
+
+        for (decl_id, declaration) in &mut self.declarations {
+            let mut removed_document_definition = false;
+            for def_id in document_definition_ids {
+                removed_document_definition |= declaration.remove_definition(def_id);
             }
-            if !emptied.is_empty() {
-                self.invalidate_graph(emptied, IdentityHashMap::default());
+
+            // This is a fallback scan across every declaration. Only prune
+            // declarations that actually contained one of this document's
+            // definitions; otherwise unrelated already-removable declarations
+            // can be removed while pending resolution work still references them.
+            if removed_document_definition {
+                touched_declarations.push(*decl_id);
             }
         }
 
+        touched_declarations
+            .into_iter()
+            .filter(|decl_id| self.declaration_is_removable(*decl_id))
+            .collect()
+    }
+
+    fn remove_document_definitions(&mut self, document: &Document) {
         for def_id in document.definitions() {
             let definition = self.definitions.remove(def_id).unwrap();
 
@@ -1091,6 +1137,52 @@ impl Graph {
         }
     }
 
+    fn prune_removable_declarations(&mut self, candidate_declaration_ids: Vec<DeclarationId>) {
+        if candidate_declaration_ids.is_empty() {
+            return;
+        }
+
+        let prune_items: Vec<_> = candidate_declaration_ids
+            .into_iter()
+            .filter(|id| self.declaration_is_removable(*id))
+            .map(InvalidationItem::Declaration)
+            .collect();
+
+        if prune_items.is_empty() {
+            return;
+        }
+
+        self.invalidate_graph(prune_items, IdentityHashMap::default());
+    }
+
+    fn declaration_is_removable(&self, declaration_id: DeclarationId) -> bool {
+        let Some(declaration) = self.declarations.get(&declaration_id) else {
+            return false;
+        };
+
+        if let Some(singleton) = declaration.as_singleton_class()
+            && self.singleton_class_needed_for_extend(singleton)
+        {
+            return false;
+        }
+
+        declaration.is_removable()
+    }
+
+    fn singleton_class_needed_for_extend(&self, singleton: &Namespace) -> bool {
+        // `class Foo; extend Bar; end` records the extend on Foo's definition,
+        // but resolution installs Bar in Foo::<Foo>'s ancestor chain.
+        let Some(attached_object) = self.declarations.get(singleton.owner_id()) else {
+            return false;
+        };
+
+        attached_object
+            .definitions()
+            .iter()
+            .filter_map(|definition_id| self.definitions.get(definition_id))
+            .any(Definition::has_extend_mixin)
+    }
+
     /// Unified invalidation worklist. Processes declaration and name items in a single loop,
     /// where processing one item can push new items back onto the queue.
     fn invalidate_graph(
@@ -1160,7 +1252,7 @@ impl Graph {
         let Some(decl) = self.declarations.get(&decl_id) else {
             return;
         };
-        let should_remove = decl.is_removable() || !self.declarations.contains_key(decl.owner_id());
+        let should_remove = self.declaration_is_removable(decl_id) || !self.declarations.contains_key(decl.owner_id());
 
         if should_remove {
             // Queue members + singleton for removal
@@ -1213,7 +1305,7 @@ impl Graph {
                     }
                 }
 
-                if self.declarations.get(&owner_id).is_some_and(Declaration::is_removable) {
+                if self.declaration_is_removable(owner_id) {
                     queue.push(InvalidationItem::Declaration(owner_id));
                 }
             }
@@ -3894,6 +3986,22 @@ mod incremental_resolution_tests {
         );
     }
 
+    #[test]
+    fn deleting_singleton_method_preserves_singleton_class_needed_by_extend() {
+        let mut context = GraphTest::new();
+        context.index_uri("file:///bar.rb", "module Bar; end");
+        context.index_uri("file:///foo.rb", "class Foo; extend Bar; end");
+        context.index_uri("file:///method.rb", "def Foo.baz; end");
+        context.resolve();
+
+        assert_declaration_exists!(context, "Foo::<Foo>");
+
+        context.delete_uri("file:///method.rb");
+        context.resolve();
+
+        assert_declaration_exists!(context, "Foo::<Foo>");
+    }
+
     #[test]
     fn no_duplicate_definition_on_identical_file_delete_readd() {
         let source = "class Foo; def self.run; end; def run; end; end";