feat(robot): add ProjectIndex for O(1) workspace-wide reference lookups

Add an incrementally maintained inverse reference index (ProjectIndex)
scoped per WorkspaceFolder. On each file build, references are updated
in-place instead of scanning all files. All six reference types are
supported: keywords, variables, namespace entries, keyword tags,
testcase tags, and metadata.

Author: d-biehl

hatch.toml

@@ -20,6 +20,7 @@ dependencies = [
   "pytest_asyncio>=0.23",
   "pytest-rerunfailures",
   "pytest-cov",
+  "pytest-mock",
   "mypy",
   "ruff",
   "debugpy",
@@ -92,7 +93,7 @@ matrix.rf.dependencies = [
 ]
 
 [envs.hatch-test]
-dependencies = ["pytest", "pytest-html", "pytest_asyncio>=0.23", "pyyaml"]
+dependencies = ["pytest", "pytest-html", "pytest-mock", "pytest_asyncio>=0.23", "pyyaml"]
 pre-install-commands = ["python ./scripts/install_packages.py"]
 
 [[envs.test.matrix]]

packages/analyze/src/robotcode/analyze/code/robot_framework_language_provider.py

@@ -57,8 +57,8 @@ def __init__(self, diagnostics_context: DiagnosticsContext) -> None:
         self.diagnostics_context.diagnostics.folder_analyzers.add(self.analyze_folder)
         self.diagnostics_context.diagnostics.document_analyzers.add(self.analyze_document)
         self.diagnostics_context.diagnostics.document_collectors.add(self.collect_diagnostics)
-        # self.diagnostics_context.diagnostics.document_collectors.add(self.collect_unused_keywords)
-        # self.diagnostics_context.diagnostics.document_collectors.add(self.collect_unused_variables)
+        self.diagnostics_context.diagnostics.document_collectors.add(self.collect_unused_keywords)
+        self.diagnostics_context.diagnostics.document_collectors.add(self.collect_unused_variables)
 
     def _update_python_path(self) -> None:
         root_path = (
@@ -126,22 +126,12 @@ def collect_diagnostics(self, sender: Any, document: TextDocument) -> Optional[L
     def collect_unused_keywords(self, sender: Any, document: TextDocument) -> Optional[List[Diagnostic]]:
         namespace = self._document_cache.get_namespace(document)
 
-        documents = (
-            [document]
-            if self._document_cache.get_document_type(document) != DocumentType.RESOURCE
-            else self.diagnostics_context.workspace.documents.documents
-        )
+        project_index = self._document_cache.get_project_index(document)
 
         result: List[Diagnostic] = []
 
         for kw in namespace.library_doc.keywords.values():
-            has_reference = False
-            for doc in documents:
-                refs = self._document_cache.get_namespace(doc).keyword_references
-                if refs.get(kw):
-                    has_reference = True
-                    break
-            if not has_reference:
+            if not project_index.find_keyword_references(kw):
                 result.append(
                     Diagnostic(
                         range=kw.name_range,
@@ -159,6 +149,7 @@ def collect_unused_variables(self, sender: Any, document: TextDocument) -> Optio
         result: List[Diagnostic] = []
 
         namespace = self._document_cache.get_namespace(document)
+        project_index = self._document_cache.get_project_index(document)
 
         for var, locations in namespace.variable_references.items():
             if var.type in (
@@ -182,16 +173,7 @@ def collect_unused_variables(self, sender: Any, document: TextDocument) -> Optio
                 )
                 and self._document_cache.get_document_type(document) == DocumentType.RESOURCE
             ):
-                if self.verbose_callback is not None:
-                    self.verbose_callback(f"Checking variable '{var.name}' {var.type} for usage. {document.uri}")
-                    self.verbose_callback(
-                        f"Searching references in {len(self.diagnostics_context.workspace.documents)} documents."
-                    )
-
-                has_reference = any(
-                    len(self._document_cache.get_namespace(doc).variable_references.get(var, set())) > 0
-                    for doc in self.diagnostics_context.workspace.documents.documents
-                )
+                has_reference = bool(project_index.find_variable_references(var))
 
             if not has_reference:
                 result.append(

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

@@ -232,15 +232,11 @@ def _find_variable_references(
         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,
-                )
-            )
+            refs = self.parent.documents_cache.get_project_index(document).find_variable_references(variable)
+            if stop_at_first and refs:
+                result.append(next(iter(refs)))
+            else:
+                result.extend(refs)
         return result
 
     @_logger.call
@@ -322,15 +318,11 @@ def _find_keyword_references(
         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,
-            )
-        )
+        refs = self.parent.documents_cache.get_project_index(document).find_keyword_references(kw_doc)
+        if stop_at_first and refs:
+            result.append(next(iter(refs)))
+        else:
+            result.extend(refs)
 
         return result
 
@@ -412,7 +404,7 @@ def references_LibraryImport(  # noqa: N802
         if entry is None:
             return None
 
-        result = self._find_references_in_workspace(document, False, self._find_library_alias_in_file, entry)
+        result = list(self.parent.documents_cache.get_project_index(document).find_namespace_references(entry))
 
         if context.include_declaration and entry.import_source:
             result.append(Location(str(Uri.from_path(entry.import_source)), entry.alias_range))

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

@@ -37,6 +37,7 @@
 from .imports_manager import ImportsManager
 from .library_doc import LibraryDoc
 from .namespace import DocumentType, Namespace, NamespaceBuilder
+from .project_index import ProjectIndex
 from .workspace_config import (
     AnalysisDiagnosticModifiersConfig,
     AnalysisRobotConfig,
@@ -81,6 +82,33 @@ def __init__(
         self._workspace_languages: weakref.WeakKeyDictionary[WorkspaceFolder, Optional[Languages]] = (
             weakref.WeakKeyDictionary()
         )
+        self._project_indexes_lock = threading.RLock()
+        self._project_indexes: weakref.WeakKeyDictionary[WorkspaceFolder, ProjectIndex] = weakref.WeakKeyDictionary()
+        self._default_project_index: Optional[ProjectIndex] = None
+
+    def get_project_index(self, document: TextDocument) -> ProjectIndex:
+        return self.get_project_index_for_uri(document.uri)
+
+    def get_project_index_for_uri(self, uri: Uri) -> ProjectIndex:
+        return self.get_project_index_for_workspace_folder(self.workspace.get_workspace_folder(uri))
+
+    def get_project_index_for_workspace_folder(self, folder: Optional[WorkspaceFolder]) -> ProjectIndex:
+        if folder is None:
+            if len(self.workspace.workspace_folders) == 1:
+                folder = self.workspace.workspace_folders[0]
+            else:
+                return self.default_project_index()
+
+        with self._project_indexes_lock:
+            if folder not in self._project_indexes:
+                self._project_indexes[folder] = ProjectIndex()
+            return self._project_indexes[folder]
+
+    def default_project_index(self) -> ProjectIndex:
+        with self._project_indexes_lock:
+            if self._default_project_index is None:
+                self._default_project_index = ProjectIndex()
+            return self._default_project_index
 
     def get_languages_for_document(self, document_or_uri: Union[TextDocument, Uri, str]) -> Optional[Languages]:
         if get_robot_version() < (6, 0):
@@ -352,6 +380,9 @@ def namespace_initialized(sender, namespace: Namespace) -> None: ...
     def namespace_invalidated(sender, namespace: Namespace) -> None: ...
 
     def _invalidate_namespace(self, sender: Namespace) -> None:
+        if sender.document is not None:
+            self.get_project_index(sender.document).remove_file(sender.source)
+
         document = sender.document
         if document is not None:
             document.remove_cache_entry(self.__get_general_namespace)
@@ -403,6 +434,9 @@ def __get_namespace_for_document_type(
 
         result = builder.build()
 
+        # Update the folder-scoped reference index
+        self.get_project_index(document).update_file(result.source, result)
+
         # When the namespace detects dependency changes, evict the
         # document cache entry so it gets rebuilt on next access.
         result.invalidated.add(self._invalidate_namespace)

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

@@ -0,0 +1,191 @@
+from __future__ import annotations
+
+import threading
+from collections import defaultdict
+from dataclasses import dataclass, field
+from typing import (
+    TYPE_CHECKING,
+    Dict,
+    Set,
+    TypeVar,
+)
+
+from robotcode.core.lsp.types import Location
+
+from .entities import LibraryEntry, VariableDefinition
+from .library_doc import KeywordDoc
+
+if TYPE_CHECKING:
+    from .namespace import Namespace
+
+
+_K = TypeVar("_K")
+
+
+@dataclass
+class _FileRefs:
+    """Tracks which references a single file contributed to the global index."""
+
+    keyword_references: Dict[KeywordDoc, Set[Location]] = field(default_factory=dict)
+    variable_references: Dict[VariableDefinition, Set[Location]] = field(default_factory=dict)
+    namespace_references: Dict[LibraryEntry, Set[Location]] = field(default_factory=dict)
+    keyword_tag_references: Dict[str, Set[Location]] = field(default_factory=dict)
+    testcase_tag_references: Dict[str, Set[Location]] = field(default_factory=dict)
+    metadata_references: Dict[str, Set[Location]] = field(default_factory=dict)
+
+
+class ProjectIndex:
+    """Workspace-wide inverse reference index.
+
+    Incrementally maintained: on file change only the affected file is
+    removed and re-inserted. All lookups are O(1).
+
+    Thread-safety: An RLock protects all mutation operations (update_file,
+    remove_file). Reads use the same lock — since writes are rare (only on
+    file changes) and short, they block reads minimally.
+    """
+
+    def __init__(self) -> None:
+        self._lock = threading.RLock()
+
+        self._keyword_references: Dict[KeywordDoc, Set[Location]] = defaultdict(set)
+        self._variable_references: Dict[VariableDefinition, Set[Location]] = defaultdict(set)
+        self._namespace_references: Dict[LibraryEntry, Set[Location]] = defaultdict(set)
+        self._keyword_tag_references: Dict[str, Set[Location]] = defaultdict(set)
+        self._testcase_tag_references: Dict[str, Set[Location]] = defaultdict(set)
+        self._metadata_references: Dict[str, Set[Location]] = defaultdict(set)
+
+        self._refs_by_file: Dict[str, _FileRefs] = {}
+
+    def update_file(self, source: str, namespace: Namespace) -> None:
+        """After NamespaceBuilder.build(): update this file's references."""
+        with self._lock:
+            self._remove_file_unlocked(source)
+
+            file_refs = _FileRefs()
+
+            self._merge_refs(
+                namespace.keyword_references,
+                self._keyword_references,
+                file_refs.keyword_references,
+            )
+            self._merge_refs(
+                namespace.variable_references,
+                self._variable_references,
+                file_refs.variable_references,
+            )
+            self._merge_refs(
+                namespace.namespace_references,
+                self._namespace_references,
+                file_refs.namespace_references,
+            )
+            self._merge_refs(
+                namespace.keyword_tag_references,
+                self._keyword_tag_references,
+                file_refs.keyword_tag_references,
+            )
+            self._merge_refs(
+                namespace.testcase_tag_references,
+                self._testcase_tag_references,
+                file_refs.testcase_tag_references,
+            )
+            self._merge_refs(
+                namespace.metadata_references,
+                self._metadata_references,
+                file_refs.metadata_references,
+            )
+
+            self._refs_by_file[source] = file_refs
+
+    def remove_file(self, source: str) -> None:
+        """File deleted or invalidated: remove all its references."""
+        with self._lock:
+            self._remove_file_unlocked(source)
+
+    def _remove_file_unlocked(self, source: str) -> None:
+        file_refs = self._refs_by_file.pop(source, None)
+        if file_refs is None:
+            return
+
+        self._subtract_refs(file_refs.keyword_references, self._keyword_references)
+        self._subtract_refs(file_refs.variable_references, self._variable_references)
+        self._subtract_refs(file_refs.namespace_references, self._namespace_references)
+        self._subtract_refs(file_refs.keyword_tag_references, self._keyword_tag_references)
+        self._subtract_refs(file_refs.testcase_tag_references, self._testcase_tag_references)
+        self._subtract_refs(file_refs.metadata_references, self._metadata_references)
+
+    @staticmethod
+    def _merge_refs(
+        source_refs: Dict[_K, Set[Location]],
+        global_refs: Dict[_K, Set[Location]],
+        file_refs: Dict[_K, Set[Location]],
+    ) -> None:
+        for key, locations in source_refs.items():
+            if locations:
+                copied = set(locations)
+                global_refs[key].update(copied)
+                file_refs[key] = copied
+
+    @staticmethod
+    def _subtract_refs(
+        file_refs: Dict[_K, Set[Location]],
+        global_refs: Dict[_K, Set[Location]],
+    ) -> None:
+        for key, locations in file_refs.items():
+            bucket = global_refs.get(key)
+            if bucket is not None:
+                bucket -= locations
+                if not bucket:
+                    del global_refs[key]
+
+    def find_keyword_references(self, kw: KeywordDoc) -> Set[Location]:
+        """O(1) lookup instead of O(N) workspace scan."""
+        with self._lock:
+            return set(self._keyword_references.get(kw, ()))
+
+    def find_variable_references(self, var: VariableDefinition) -> Set[Location]:
+        """O(1) lookup instead of O(N) workspace scan."""
+        with self._lock:
+            return set(self._variable_references.get(var, ()))
+
+    def find_namespace_references(self, entry: LibraryEntry) -> Set[Location]:
+        """O(1) lookup instead of O(N) workspace scan."""
+        with self._lock:
+            return set(self._namespace_references.get(entry, ()))
+
+    def find_keyword_tag_references(self, tag: str) -> Set[Location]:
+        with self._lock:
+            return set(self._keyword_tag_references.get(tag, ()))
+
+    def find_testcase_tag_references(self, tag: str) -> Set[Location]:
+        with self._lock:
+            return set(self._testcase_tag_references.get(tag, ()))
+
+    def find_metadata_references(self, key: str) -> Set[Location]:
+        with self._lock:
+            return set(self._metadata_references.get(key, ()))
+
+    @property
+    def keyword_references(self) -> Dict[KeywordDoc, Set[Location]]:
+        with self._lock:
+            return dict(self._keyword_references)
+
+    @property
+    def variable_references(self) -> Dict[VariableDefinition, Set[Location]]:
+        with self._lock:
+            return dict(self._variable_references)
+
+    @property
+    def namespace_references(self) -> Dict[LibraryEntry, Set[Location]]:
+        with self._lock:
+            return dict(self._namespace_references)
+
+    def clear(self) -> None:
+        with self._lock:
+            self._keyword_references.clear()
+            self._variable_references.clear()
+            self._namespace_references.clear()
+            self._keyword_tag_references.clear()
+            self._testcase_tag_references.clear()
+            self._metadata_references.clear()
+            self._refs_by_file.clear()