perf(language-server): add reverse index for keyword/variable references

Optimize workspace-wide reference operations from O(D) to O(k) where
D = total documents and k = documents actually referencing the target.

Changes:
- Add reverse index data structures in DocumentsCacheHelper to track
  which documents reference each keyword/variable
- Use stable (source, name) tuple keys resilient to cache invalidation
- Implement diff-based updates to handle removed references after edits
- Add get_keyword_ref_users() and get_variable_ref_users() for O(1) lookup
- Update Find References to use reverse index with workspace scan fallback
- Update unused keyword/variable detection to use reverse index

Author: mikeleppane

packages/language_server/src/robotcode/language_server/robotframework/parts/diagnostics.py

@@ -18,8 +18,9 @@
     EnvironmentVariableDefinition,
     GlobalVariableDefinition,
     LibraryArgumentDefinition,
+    VariableDefinition,
 )
-from robotcode.robot.diagnostics.library_doc import LibraryDoc
+from robotcode.robot.diagnostics.library_doc import KeywordDoc, LibraryDoc
 from robotcode.robot.diagnostics.namespace import Namespace
 
 from ...common.parts.diagnostics import DiagnosticsCollectType, DiagnosticsResult
@@ -138,6 +139,40 @@ def collect_namespace_diagnostics(
                 ],
             )
 
+    def _is_keyword_used_anywhere(
+        self,
+        document: TextDocument,
+        kw: KeywordDoc,
+        namespace: Namespace,
+    ) -> bool:
+        """Check if keyword is used anywhere, using index with safe fallback."""
+        if self.parent.documents_cache.get_keyword_ref_users(kw):
+            return True
+
+        if namespace.get_keyword_references().get(kw):
+            return True
+
+        # Safe fallback: workspace scan if index might be incomplete
+        refs = self.parent.robot_references.find_keyword_references(document, kw, False, True)
+        return bool(refs)
+
+    def _is_variable_used_anywhere(
+        self,
+        document: TextDocument,
+        var: VariableDefinition,
+        namespace: Namespace,
+    ) -> bool:
+        """Check if variable is used anywhere, using index with safe fallback."""
+        if self.parent.documents_cache.get_variable_ref_users(var):
+            return True
+
+        if namespace.get_variable_references().get(var):
+            return True
+
+        # Safe fallback: workspace scan if index might be incomplete
+        refs = self.parent.robot_references.find_variable_references(document, var, False, True)
+        return bool(refs)
+
     @language_id("robotframework")
     @_logger.call
     def collect_unused_keyword_references(
@@ -166,10 +201,7 @@ def _collect_unused_keyword_references(
             for kw in (namespace.get_library_doc()).keywords.values():
                 check_current_task_canceled()
 
-                references = self.parent.robot_references.find_keyword_references(
-                    document, kw, False, True
-                )
-                if not references:
+                if not self._is_keyword_used_anywhere(document, kw, namespace):
                     result.append(
                         Diagnostic(
                             range=kw.name_range,
@@ -255,10 +287,7 @@ def _collect_unused_variable_references(
                 ):
                     continue
 
-                references = self.parent.robot_references.find_variable_references(
-                    document, var, False, True
-                )
-                if not references:
+                if not self._is_variable_used_anywhere(document, var, namespace):
                     result.append(
                         Diagnostic(
                             range=var.name_range,

packages/language_server/src/robotcode/language_server/robotframework/parts/references.py

@@ -224,23 +224,30 @@ def _find_variable_references(
         include_declaration: bool = True,
         stop_at_first: bool = False,
     ) -> List[Location]:
-        result = []
+        result: List[Location] = []
 
         if include_declaration and variable.source:
             result.append(Location(str(Uri.from_path(variable.source)), variable.name_range))
 
         if variable.type == VariableDefinitionType.LOCAL_VARIABLE:
             result.extend(self.find_variable_references_in_file(document, variable, False))
         else:
-            result.extend(
-                self._find_references_in_workspace(
-                    document,
-                    stop_at_first,
-                    self.find_variable_references_in_file,
-                    variable,
-                    False,
+            # Use reverse index for lookup instead of workspace scan
+            docs_to_search = self.parent.documents_cache.get_variable_ref_users(variable)
+            if docs_to_search:
+                for doc in docs_to_search:
+                    check_current_task_canceled()
+                    result.extend(self.find_variable_references_in_file(doc, variable, False))
+                    if result and stop_at_first:
+                        break
+            else:
+                # Fallback to workspace scan if index is empty
+                result.extend(
+                    self._find_references_in_workspace(
+                        document, stop_at_first, self.find_variable_references_in_file, variable, False
+                    )
                 )
-            )
+
         return result
 
     @_logger.call
@@ -317,20 +324,26 @@ def _find_keyword_references(
         include_declaration: bool = True,
         stop_at_first: bool = False,
     ) -> List[Location]:
-        result = []
+        result: List[Location] = []
 
         if include_declaration and kw_doc.source:
             result.append(Location(str(Uri.from_path(kw_doc.source)), kw_doc.range))
 
-        result.extend(
-            self._find_references_in_workspace(
-                document,
-                stop_at_first,
-                self.find_keyword_references_in_file,
-                kw_doc,
-                False,
+        # Use reverse index for lookup instead of workspace scan
+        docs_to_search = self.parent.documents_cache.get_keyword_ref_users(kw_doc)
+        if docs_to_search:
+            for doc in docs_to_search:
+                check_current_task_canceled()
+                result.extend(self.find_keyword_references_in_file(doc, kw_doc, False))
+                if result and stop_at_first:
+                    break
+        else:
+            # Fallback to workspace scan if index is empty
+            result.extend(
+                self._find_references_in_workspace(
+                    document, stop_at_first, self.find_keyword_references_in_file, kw_doc, False
+                )
             )
-        )
 
         return result
 

packages/robot/src/robotcode/robot/diagnostics/document_cache_helper.py

@@ -30,8 +30,9 @@
 from ..utils import get_robot_version
 from ..utils.stubs import Languages
 from .data_cache import CacheSection
+from .entities import VariableDefinition
 from .imports_manager import ImportsManager
-from .library_doc import LibraryDoc
+from .library_doc import KeywordDoc, LibraryDoc
 from .namespace import DocumentType, Namespace, NamespaceCacheData, NamespaceMetaData
 from .workspace_config import (
     AnalysisDiagnosticModifiersConfig,
@@ -99,6 +100,18 @@ def __init__(
         self._variables_users_lock = threading.RLock()
         self._variables_users: dict[str, weakref.WeakSet[TextDocument]] = {}
 
+        # Reference tracking for O(1) lookup of keyword/variable usages
+        # Uses (source, name) tuples as keys for stability across cache invalidation
+        self._ref_tracking_lock = threading.RLock()
+        self._keyword_ref_users: dict[tuple[str, str], weakref.WeakSet[TextDocument]] = {}
+        self._variable_ref_users: dict[tuple[str, str], weakref.WeakSet[TextDocument]] = {}
+        self._doc_keyword_refs: weakref.WeakKeyDictionary[
+            TextDocument, set[tuple[str, str]]
+        ] = weakref.WeakKeyDictionary()
+        self._doc_variable_refs: weakref.WeakKeyDictionary[
+            TextDocument, set[tuple[str, str]]
+        ] = weakref.WeakKeyDictionary()
+
         # Counter for periodic cleanup of stale dependency map entries
         self._track_count = 0
 
@@ -520,6 +533,8 @@ def __namespace_initialized(self, sender: Namespace) -> None:
     def __namespace_analysed(self, sender: Namespace) -> None:
         """Re-save namespace to cache after analysis to include diagnostics and analysis results."""
         if sender.document is not None:
+            self._track_references(sender.document, sender)
+
             imports_manager = self.get_imports_manager(sender.document)
             self._save_namespace_to_cache(sender, imports_manager)
 
@@ -585,6 +600,22 @@ def get_variables_users(self, variables_doc: LibraryDoc) -> list[TextDocument]:
                 return list(self._variables_users[var_key])
             return []
 
+    def get_keyword_ref_users(self, kw_doc: KeywordDoc) -> list[TextDocument]:
+        """Get documents that reference a keyword."""
+        with self._ref_tracking_lock:
+            key = (kw_doc.source or "", kw_doc.name)
+            if key in self._keyword_ref_users:
+                return list(self._keyword_ref_users[key])
+            return []
+
+    def get_variable_ref_users(self, var_def: VariableDefinition) -> list[TextDocument]:
+        """Get documents that reference a variable."""
+        with self._ref_tracking_lock:
+            key = (var_def.source or "", var_def.name)
+            if key in self._variable_ref_users:
+                return list(self._variable_ref_users[key])
+            return []
+
     def _cleanup_stale_dependency_maps(self) -> None:
         """Remove entries with empty WeakSets from dependency maps.
 
@@ -606,6 +637,59 @@ def _cleanup_stale_dependency_maps(self) -> None:
             for var_key in stale_var_keys:
                 del self._variables_users[var_key]
 
+        with self._ref_tracking_lock:
+            stale_kw_ref_keys = [k for k, v in self._keyword_ref_users.items() if len(v) == 0]
+            for kw_ref_key in stale_kw_ref_keys:
+                del self._keyword_ref_users[kw_ref_key]
+
+            stale_var_ref_keys = [k for k, v in self._variable_ref_users.items() if len(v) == 0]
+            for var_ref_key in stale_var_ref_keys:
+                del self._variable_ref_users[var_ref_key]
+
+    def _track_references(self, document: TextDocument, namespace: Namespace) -> None:
+        """Track keyword/variable references.
+
+        Uses diff-based updates: compares current references against previous
+        to handle documents that stop referencing items after edits.
+        """
+        with self._ref_tracking_lock:
+            self._update_keyword_refs(document, namespace)
+            self._update_variable_refs(document, namespace)
+
+    def _update_keyword_refs(self, document: TextDocument, namespace: Namespace) -> None:
+        """Update reverse index for keyword references."""
+        keyword_refs = namespace.get_keyword_references()
+        new_keys = {(kw.source or "", kw.name) for kw in keyword_refs}
+        old_keys = self._doc_keyword_refs.get(document, set())
+
+        for key in old_keys - new_keys:
+            if key in self._keyword_ref_users:
+                self._keyword_ref_users[key].discard(document)
+
+        for key in new_keys - old_keys:
+            if key not in self._keyword_ref_users:
+                self._keyword_ref_users[key] = weakref.WeakSet()
+            self._keyword_ref_users[key].add(document)
+
+        self._doc_keyword_refs[document] = new_keys
+
+    def _update_variable_refs(self, document: TextDocument, namespace: Namespace) -> None:
+        """Update reverse index for variable references."""
+        variable_refs = namespace.get_variable_references()
+        new_keys = {(var.source or "", var.name) for var in variable_refs}
+        old_keys = self._doc_variable_refs.get(document, set())
+
+        for key in old_keys - new_keys:
+            if key in self._variable_ref_users:
+                self._variable_ref_users[key].discard(document)
+
+        for key in new_keys - old_keys:
+            if key not in self._variable_ref_users:
+                self._variable_ref_users[key] = weakref.WeakSet()
+            self._variable_ref_users[key].add(document)
+
+        self._doc_variable_refs[document] = new_keys
+
     def get_initialized_namespace(self, document: TextDocument) -> Namespace:
         result: Namespace | None = document.get_data(self.INITIALIZED_NAMESPACE)
         if result is None: