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

@@ -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:

tests/fixtures/swift_cross_file/Foo+Ext.swift

@@ -0,0 +1,3 @@
+extension Foo {
+    func two() {}
+}

tests/fixtures/swift_cross_file/Foo.swift

@@ -0,0 +1,3 @@
+class Foo {
+    func one() {}
+}

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