zig fmt: implement Tree.lastToken() for all nodes

2 files changed, 35 insertions, 6 deletions

diff --git a/lib/std/zig/ast.zig b/lib/std/zig/ast.zig
index 358c40b28a..783a578e01 100644
--- a/lib/std/zig/ast.zig
+++ b/lib/std/zig/ast.zig
@@ -538,6 +538,7 @@ pub const Tree = struct {
             .ArrayType,
             .SwitchCaseOne,
             .SwitchCase,
+            .SwitchRange,
             => n = datas[n].rhs,
 
             .FieldAccess,
@@ -580,8 +581,21 @@ pub const Tree = struct {
                 }
                 n = tree.extra_data[params.end - 1]; // last parameter
             },
-            .CallComma, .AsyncCallComma => {
-                end_offset += 2; // for the comma+rparen
+            .TaggedUnionEnumTag => {
+                const members = tree.extraData(datas[n].rhs, Node.SubRange);
+                if (members.end - members.start == 0) {
+                    end_offset += 4; // for the rparen + rparen + lbrace + rbrace
+                    n = datas[n].lhs;
+                } else {
+                    end_offset += 1; // for the rbrace
+                    n = tree.extra_data[members.end - 1]; // last parameter
+                }
+            },
+            .CallComma,
+            .AsyncCallComma,
+            .TaggedUnionEnumTagComma,
+            => {
+                end_offset += 2; // for the comma + rparen/rbrace
                 const params = tree.extraData(datas[n].rhs, Node.SubRange);
                 assert(params.end > params.start);
                 n = tree.extra_data[params.end - 1]; // last parameter
@@ -942,10 +956,6 @@ pub const Tree = struct {
                 const extra = tree.extraData(datas[n].rhs, Node.ArrayTypeSentinel);
                 n = extra.elem_type;
             },
-
-            .TaggedUnionEnumTag => unreachable, // TODO
-            .TaggedUnionEnumTagComma => unreachable, // TODO
-            .SwitchRange => unreachable, // TODO
         };
     }
 
diff --git a/lib/std/zig/parser_test.zig b/lib/std/zig/parser_test.zig
index bc1a36d7e3..c6ee4f6bd9 100644
--- a/lib/std/zig/parser_test.zig
+++ b/lib/std/zig/parser_test.zig
@@ -841,6 +841,25 @@ test "zig fmt: tagged union with enum values" {
     );
 }
 
+test "zig fmt: tagged union enum tag last token" {
+    try testCanonical(
+        \\test {
+        \\    const U = union(enum(u32)) {};
+        \\}
+        \\
+        \\test {
+        \\    const U = union(enum(u32)) { foo };
+        \\}
+        \\
+        \\test {
+        \\    const U = union(enum(u32)) {
+        \\        foo,
+        \\    };
+        \\}
+        \\
+    );
+}
+
 test "zig fmt: allowzero pointer" {
     try testCanonical(
         \\const T = [*]allowzero const u8;