From 406bea47b59a3efe4c27c56ac6cb1dba0c75847b Mon Sep 17 00:00:00 2001 From: Alex Ubillus Date: Fri, 22 May 2026 08:22:42 -0400 Subject: [PATCH] fix swift extension nodes duplicating across files (#969) tree-sitter-swift parses both `class Foo` and `extension Foo` as `class_declaration`, and node ids carry the file stem, so `extension Foo` in a sibling file produced a second `Foo` node instead of attaching to the original. Same-file extensions already dedupe via seen_ids; only the cross-file case leaked. Per-file extraction now tags `extension` class_declarations, and the corpus-level `extract()` runs a merge pass: when exactly one non-extension declaration shares the label, the extension nodes redirect onto it and their edges are rewritten (self-loops dropped, duplicates collapsed). Extensions of types outside the corpus and ambiguous label matches stay untouched. On a 25-file Swift project this collapses Parser from 6 split nodes (top of the god-node list, four entries) to one canonical node, and lets the generic cross-file call resolver attach previously ambiguous call edges to the right target. --- graphify/extract.py | 87 ++++++++++++++++++- tests/fixtures/swift_cross_file/Foo+Ext.swift | 3 + tests/fixtures/swift_cross_file/Foo.swift | 3 + tests/test_languages.py | 20 +++++ 4 files changed, 112 insertions(+), 1 deletion(-) create mode 100644 tests/fixtures/swift_cross_file/Foo+Ext.swift create mode 100644 tests/fixtures/swift_cross_file/Foo.swift diff --git a/graphify/extract.py b/graphify/extract.py index dbe8f0e4f..9f79f95f6 100644 --- a/graphify/extract.py +++ b/graphify/extract.py @@ -1241,6 +1241,11 @@ def _extract_generic(path: Path, config: LanguageConfig) -> dict: seen_ids: set[str] = set() function_bodies: list[tuple[str, object]] = [] pending_listen_edges: list[tuple[str, str, int]] = [] + # tree-sitter-swift parses both `class Foo` and `extension Foo` as + # `class_declaration`. Same-file pairs collapse via seen_ids, but cross-file + # extensions don't (file stem is part of the id), so they're collected here + # for a corpus-level merge after every file has been parsed. + swift_extensions: list[dict] = [] def add_node(nid: str, label: str, line: int) -> None: if nid not in seen_ids: @@ -1307,6 +1312,11 @@ def _extract_generic(path: Path, config: LanguageConfig) -> dict: add_node(class_nid, class_name, line) add_edge(file_nid, class_nid, "contains", line) + if config.ts_module == "tree_sitter_swift" and any( + c.type == "extension" for c in node.children + ): + swift_extensions.append({"nid": class_nid, "label": class_name}) + # Python-specific: inheritance if config.ts_module == "tree_sitter_python": args = node.child_by_field_name("superclasses") @@ -1953,7 +1963,10 @@ def _extract_generic(path: Path, config: LanguageConfig) -> dict: if src in valid_ids and (tgt in valid_ids or edge["relation"] in ("imports", "imports_from")): clean_edges.append(edge) - return {"nodes": nodes, "edges": clean_edges, "raw_calls": raw_calls} + result = {"nodes": nodes, "edges": clean_edges, "raw_calls": raw_calls} + if swift_extensions: + result["swift_extensions"] = swift_extensions + return result # ── Python rationale extraction ─────────────────────────────────────────────── @@ -4354,6 +4367,76 @@ def _resolve_cross_file_imports( return new_edges +def _merge_swift_extensions( + per_file: list[dict], + all_nodes: list[dict], + all_edges: list[dict], +) -> None: + """Collapse cross-file Swift `extension Foo` nodes into the canonical `Foo`. + + tree-sitter-swift reuses `class_declaration` for both `class Foo` and + `extension Foo`, and node ids carry the file stem, so each file that + extends `Foo` produces its own `Foo` node. The match is done by label: + when exactly one non-extension declaration shares the label, extension + nodes redirect onto it. Extensions of types outside the corpus (no match) + and ambiguous labels (more than one match) are left untouched — picking + arbitrarily would invent edges. + """ + extension_nids: set[str] = set() + extension_labels: dict[str, str] = {} + for result in per_file: + for ext in result.get("swift_extensions", []) or []: + extension_nids.add(ext["nid"]) + extension_labels[ext["nid"]] = ext["label"] + + if not extension_nids: + return + + label_to_canonical: dict[str, list[str]] = {} + for n in all_nodes: + if n.get("id") in extension_nids: + continue + label = n.get("label") + if not label: + continue + label_to_canonical.setdefault(label, []).append(n["id"]) + + remap: dict[str, str] = {} + for ext_nid in extension_nids: + candidates = label_to_canonical.get(extension_labels[ext_nid], []) + if len(candidates) != 1: + continue + canonical_nid = candidates[0] + if canonical_nid != ext_nid: + remap[ext_nid] = canonical_nid + + if not remap: + return + + all_nodes[:] = [n for n in all_nodes if n.get("id") not in remap] + + # Each extension file's `contains` edge ends up pointing at the canonical + # type — multiple files containing the same node is the intended shape: + # the type owns the methods, the files own their slice. Self-loops are + # dropped (e.g. an in-file extension method whose call already pointed at + # the canonical type). + rewritten: list[dict] = [] + seen_keys: set[tuple] = set() + for e in all_edges: + src = remap.get(e.get("source"), e.get("source")) + tgt = remap.get(e.get("target"), e.get("target")) + if src == tgt: + continue + e["source"] = src + e["target"] = tgt + key = (src, tgt, e.get("relation"), e.get("source_file"), e.get("source_location")) + if key in seen_keys: + continue + seen_keys.add(key) + rewritten.append(e) + all_edges[:] = rewritten + + def _resolve_cross_file_java_imports( per_file: list[dict], paths: list[Path], @@ -6470,6 +6553,8 @@ def extract( if e.get("target") in id_remap: e["target"] = id_remap[e["target"]] + _merge_swift_extensions(per_file, all_nodes, all_edges) + # Add cross-file class-level edges (Python only - uses Python parser internally) py_paths = [p for p in paths if p.suffix == ".py"] if py_paths: diff --git a/tests/fixtures/swift_cross_file/Foo+Ext.swift b/tests/fixtures/swift_cross_file/Foo+Ext.swift new file mode 100644 index 000000000..74fabeb39 --- /dev/null +++ b/tests/fixtures/swift_cross_file/Foo+Ext.swift @@ -0,0 +1,3 @@ +extension Foo { + func two() {} +} diff --git a/tests/fixtures/swift_cross_file/Foo.swift b/tests/fixtures/swift_cross_file/Foo.swift new file mode 100644 index 000000000..b71ab60cd --- /dev/null +++ b/tests/fixtures/swift_cross_file/Foo.swift @@ -0,0 +1,3 @@ +class Foo { + func one() {} +} diff --git a/tests/test_languages.py b/tests/test_languages.py index 3497f3f69..aa6500056 100644 --- a/tests/test_languages.py +++ b/tests/test_languages.py @@ -555,6 +555,26 @@ def test_swift_call_edges_have_call_context(): assert all(e.get("context") == "call" for e in call_edges) +def test_swift_extension_across_files_merges_into_canonical_type(): + """`extension Foo` in a separate file from `class Foo` must resolve to a + single Foo node. tree-sitter-swift parses both as `class_declaration` and + node ids carry the file stem, so without a corpus-level merge each file + would emit its own Foo.""" + from graphify.extract import extract + paths = sorted((FIXTURES / "swift_cross_file").glob("*.swift")) + r = extract(paths, cache_root=Path("/tmp/graphify-test-no-cache")) + foo_nodes = [n for n in r["nodes"] if n["label"] == "Foo"] + assert len(foo_nodes) == 1, f"Foo should appear once, got {len(foo_nodes)}: {[n['id'] for n in foo_nodes]}" + foo_id = foo_nodes[0]["id"] + method_targets = { + e["target"] for e in r["edges"] + if e["relation"] == "method" and e["source"] == foo_id + } + method_labels = {n["label"] for n in r["nodes"] if n["id"] in method_targets} + assert any("one" in l for l in method_labels), f"one() should attach to Foo, got {method_labels}" + assert any("two" in l for l in method_labels), f"extension method two() should attach to Foo, got {method_labels}" + + # ── Elixir ──────────────────────────────────────────────────────────────────── from graphify.extract import extract_elixir