From 6930e4fa9048f17a548e287aad2b563a8d533e25 Mon Sep 17 00:00:00 2001
From: rajatnagda45 <215653880+rajatnagda45@users.noreply.github.com>
Date: Mon, 24 Aug 2026 13:15:23 +0100
Subject: [PATCH] fix(zig): extract methods declared on enums and unions

Zig lets enums and unions carry methods like structs, but only the struct branch recursed
into the container body, so enum/union methods and the calls in their bodies were dropped.
Walk enum/union children the same way, parenting methods to the container with the same id
scheme; variants and fields are not turned into nodes.
---
 graphify/extractors/zig.py |  6 ++++
 tests/test_languages.py    | 59 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+)

diff --git a/graphify/extractors/zig.py b/graphify/extractors/zig.py
index 9744c514..3fd9f762 100644
--- a/graphify/extractors/zig.py
+++ b/graphify/extractors/zig.py
@@ -124,6 +124,12 @@ def extract_zig(path: Path) -> dict:
                     type_nid = _make_id(stem, type_name)
                     add_node(type_nid, type_name, line)
                     add_edge(file_nid, type_nid, "contains", line)
+                    # Zig enums and tagged unions can declare methods just like
+                    # structs (`pub fn ...` inside the container). Recurse so those
+                    # methods — and the calls made from their bodies — are captured
+                    # rather than dropped along with the whole method layer.
+                    for child in value_node.children:
+                        walk(child, parent_struct_nid=type_nid)
                 return
 
             if value_node and value_node.type in ("builtin_function", "field_expression"):
diff --git a/tests/test_languages.py b/tests/test_languages.py
index 2c0ba620..f8784be8 100644
--- a/tests/test_languages.py
+++ b/tests/test_languages.py
@@ -3788,3 +3788,62 @@ def test_markdown_heading_id_is_stable_regardless_of_frontmatter():
     stem = _file_stem(Path(src_file))
     assert _make_id(stem, "Overview") in heading_ids
     assert _make_id(stem, "Details") in heading_ids
+
+
+# ── Zig ───────────────────────────────────────────────────────────────────────
+from graphify.extract import extract_zig
+
+_needs_zig = pytest.mark.skipif(
+    _ilu.find_spec("tree_sitter_zig") is None,
+    reason="tree-sitter-zig not installed",
+)
+
+
+@_needs_zig
+def test_zig_enum_and_union_methods_are_extracted(tmp_path):
+    """Methods declared inside a Zig enum or tagged union must be captured.
+
+    Only `struct` containers recursed into their members, so `pub fn` methods on
+    an `enum`/`union` — and every call made from those method bodies — were
+    dropped along with the whole method layer of the type.
+    """
+    src = (
+        "const Color = enum {\n"
+        "    red,\n"
+        "    green,\n"
+        "    pub fn isRed(self: Color) bool {\n"
+        "        return self == .red;\n"
+        "    }\n"
+        "};\n"
+        "\n"
+        "const Shape = union(enum) {\n"
+        "    circle: f32,\n"
+        "    pub fn area(self: Shape) f32 {\n"
+        "        return helper();\n"
+        "    }\n"
+        "};\n"
+        "\n"
+        "fn helper() f32 {\n"
+        "    return 1.0;\n"
+        "}\n"
+    )
+    f = tmp_path / "shapes.zig"
+    f.write_text(src)
+    r = extract_zig(f)
+    assert "error" not in r
+
+    method_targets = {
+        e["target"] for e in r["edges"] if e["relation"] == "method"
+    }
+    id_to_label = {n["id"]: n["label"] for n in r["nodes"]}
+    method_labels = {id_to_label[t] for t in method_targets}
+    assert ".isRed()" in method_labels, "enum method dropped"
+    assert ".area()" in method_labels, "union method dropped"
+
+    # A call made from an enum/union method body must resolve too.
+    calls = {
+        (id_to_label.get(e["source"], e["source"]),
+         id_to_label.get(e["target"], e["target"]))
+        for e in r["edges"] if e["relation"] == "calls"
+    }
+    assert (".area()", "helper()") in calls, "call from union method body dropped"