fix: find_ending_with too slow

codebase_rag/graph_updater.py

@@ -22,6 +22,7 @@ class FunctionRegistryTrie:
     def __init__(self) -> None:
         self.root: dict[str, Any] = {}
         self._entries: dict[str, str] = {}
+        self._suffix_index: dict[str, set[str]] = defaultdict(set)
 
     def insert(self, qualified_name: str, func_type: str) -> None:
         """Insert a function into the trie."""
@@ -38,6 +39,10 @@ def insert(self, qualified_name: str, func_type: str) -> None:
         current["__type__"] = func_type
         current["__qn__"] = qualified_name
 
+        if parts:
+            suffix = parts[-1]
+            self._suffix_index[suffix].add(qualified_name)
+
     def get(self, qualified_name: str, default: str | None = None) -> str | None:
         """Get function type by exact qualified name."""
         return self._entries.get(qualified_name, default)
@@ -68,6 +73,15 @@ def __delitem__(self, qualified_name: str) -> None:
         parts = qualified_name.split(".")
         self._cleanup_trie_path(parts, self.root)
 
+        # clean up suffix_index
+        if parts:
+            suffix = parts[-1]
+            if suffix in self._suffix_index:
+                self._suffix_index[suffix].discard(qualified_name)
+                if not self._suffix_index[suffix]:
+                    del self._suffix_index[suffix]
+
+
     def _cleanup_trie_path(self, parts: list[str], node: dict[str, Any]) -> bool:
         """Recursively clean up empty trie nodes.
 
@@ -134,7 +148,11 @@ def dfs(node: dict[str, Any]) -> None:
 
     def find_ending_with(self, suffix: str) -> list[str]:
         """Find all qualified names ending with the given suffix."""
-        return [qn for qn in self._entries.keys() if qn.endswith(f".{suffix}")]
+        last_part = suffix.split(".")[-1]
+        candidates = self._suffix_index.get(last_part, set())
+        if "." not in suffix:
+            return list(candidates)
+        return [qn for qn in candidates if qn.endswith(f".{suffix}")]
 
     def find_with_prefix(self, prefix: str) -> list[tuple[str, str]]:
         """Find all qualified names that start with the given prefix.