Skip to content

Commit 384cd77

Browse files
committed
perf: fix O(n²) sibling representative + batch graph edge insertion
sibling.py: replace linear scan (next over all fragments) with dict lookup for file representative. Was O(n²), now O(n). graph.py: use nx.DiGraph.add_weighted_edges_from() instead of per-edge add_edge() calls. Reduces 5M function calls to one bulk operation on large repos. Profiling on svelte (6K files): 82s → 33s.
1 parent 7852b22 commit 384cd77

2 files changed

Lines changed: 10 additions & 13 deletions

File tree

src/treemapper/diffctx/edges/structural/sibling.py

Lines changed: 7 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -30,19 +30,16 @@ def _group_files_by_dir(self, fragments: list[Fragment]) -> dict[Path, set[Path]
3030
return by_dir
3131

3232
def _build_file_representative_map(self, fragments: list[Fragment]) -> dict[Path, FragmentId]:
33+
frag_by_id: dict[FragmentId, Fragment] = {f.id: f for f in fragments}
3334
file_to_rep: dict[Path, FragmentId] = {}
3435
for f in fragments:
35-
self._update_file_representative(f, file_to_rep, fragments)
36-
return file_to_rep
37-
38-
def _update_file_representative(self, f: Fragment, file_to_rep: dict[Path, FragmentId], fragments: list[Fragment]) -> None:
39-
if f.path not in file_to_rep:
40-
file_to_rep[f.path] = f.id
41-
elif f.token_count > 0:
42-
existing = file_to_rep[f.path]
43-
existing_frag = next((fr for fr in fragments if fr.id == existing), None)
44-
if existing_frag and f.token_count > existing_frag.token_count:
36+
if f.path not in file_to_rep:
4537
file_to_rep[f.path] = f.id
38+
elif f.token_count > 0:
39+
existing = frag_by_id.get(file_to_rep[f.path])
40+
if existing and f.token_count > existing.token_count:
41+
file_to_rep[f.path] = f.id
42+
return file_to_rep
4643

4744
def _add_sibling_edges_for_dir(self, files: set[Path], file_to_rep: dict[Path, FragmentId], edges: EdgeDict) -> None:
4845
file_list = sorted(files)

src/treemapper/diffctx/graph.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -146,9 +146,9 @@ def build_graph(fragments: list[Fragment], repo_root: Path | None = None) -> Gra
146146

147147
all_edges = _apply_hub_suppression(all_edges, edge_categories)
148148

149-
for (src, dst), weight in all_edges.items():
150-
graph.add_edge(src, dst, weight)
151-
149+
graph._g.add_nodes_from(frag.id for frag in fragments)
150+
graph._g.add_weighted_edges_from((src, dst, w) for (src, dst), w in all_edges.items() if w > 0)
151+
graph._invalidate_cache()
152152
graph.edge_categories = edge_categories
153153

154154
return graph

0 commit comments

Comments
 (0)