zig fmt: implement Tree.lastToken() for struct init

4 files changed, 98 insertions, 61 deletions

diff --git a/lib/std/zig/ast.zig b/lib/std/zig/ast.zig
index 1780cc4f8d..527c358103 100644
--- a/lib/std/zig/ast.zig
+++ b/lib/std/zig/ast.zig
@@ -259,6 +259,7 @@ pub const Tree = struct {
             .ArrayInitDotTwo,
             .ArrayInitDotTwoComma,
             .StructInitDot,
+            .StructInitDotComma,
             .StructInitDotTwo,
             .StructInitDotTwoComma,
             .EnumLiteral,
@@ -316,7 +317,9 @@ pub const Tree = struct {
             .ArrayInit,
             .ArrayInitComma,
             .StructInitOne,
+            .StructInitOneComma,
             .StructInit,
+            .StructInitComma,
             .CallOne,
             .CallOneComma,
             .Call,
@@ -607,13 +610,16 @@ pub const Tree = struct {
                 const extra = tree.extraData(datas[n].rhs, Node.Asm);
                 return extra.rparen + end_offset;
             },
-            .ArrayInit => {
+            .ArrayInit,
+            .StructInit,
+            => {
                 const elements = tree.extraData(datas[n].rhs, Node.SubRange);
                 assert(elements.end - elements.start > 0);
                 end_offset += 1; // for the rbrace
                 n = tree.extra_data[elements.end - 1]; // last element
             },
             .ArrayInitComma,
+            .StructInitComma,
             .ContainerDeclArgComma,
             .SwitchComma,
             => {
@@ -623,6 +629,7 @@ pub const Tree = struct {
                 n = tree.extra_data[members.end - 1]; // last parameter
             },
             .ArrayInitDot,
+            .StructInitDot,
             .Block,
             .ContainerDecl,
             .TaggedUnion,
@@ -633,6 +640,7 @@ pub const Tree = struct {
                 n = tree.extra_data[datas[n].rhs - 1]; // last statement
             },
             .ArrayInitDotComma,
+            .StructInitDotComma,
             .BlockSemicolon,
             .ContainerDeclComma,
             .TaggedUnionComma,
@@ -784,7 +792,9 @@ pub const Tree = struct {
                 }
             },
 
-            .ArrayInitOne => {
+            .ArrayInitOne,
+            .StructInitOne,
+            => {
                 end_offset += 1; // rbrace
                 n = datas[n].rhs;
                 assert(n != 0);
@@ -793,6 +803,7 @@ pub const Tree = struct {
             .CallOneComma,
             .AsyncCallOneComma,
             .ArrayInitOneComma,
+            .StructInitOneComma,
             => {
                 end_offset += 2; // ellipsis2 + rbracket, or comma + rparen
                 n = datas[n].rhs;
@@ -929,14 +940,6 @@ pub const Tree = struct {
                 n = extra.elem_type;
             },
 
-            // These are not supported by lastToken() because implementation would
-            // require recursion due to the optional comma followed by rbrace.
-            // TODO follow the pattern set by StructInitDotTwoComma which will allow
-            // lastToken to work for all of these.
-            .StructInit => unreachable, // TODO
-            .StructInitOne => unreachable, // TODO
-            .StructInitDot => unreachable, // TODO
-
             .TaggedUnionEnumTag => unreachable, // TODO
             .TaggedUnionEnumTagComma => unreachable, // TODO
             .SwitchRange => unreachable, // TODO
@@ -1118,7 +1121,8 @@ pub const Tree = struct {
     }
 
     pub fn structInitOne(tree: Tree, buffer: *[1]Node.Index, node: Node.Index) full.StructInit {
-        assert(tree.nodes.items(.tag)[node] == .StructInitOne);
+        assert(tree.nodes.items(.tag)[node] == .StructInitOne or
+            tree.nodes.items(.tag)[node] == .StructInitOneComma);
         const data = tree.nodes.items(.data)[node];
         buffer[0] = data.rhs;
         const fields = if (data.rhs == 0) buffer[0..0] else buffer[0..1];
@@ -1148,7 +1152,8 @@ pub const Tree = struct {
     }
 
     pub fn structInitDot(tree: Tree, node: Node.Index) full.StructInit {
-        assert(tree.nodes.items(.tag)[node] == .StructInitDot);
+        assert(tree.nodes.items(.tag)[node] == .StructInitDot or
+            tree.nodes.items(.tag)[node] == .StructInitDotComma);
         const data = tree.nodes.items(.data)[node];
         return tree.fullStructInit(.{
             .lbrace = tree.nodes.items(.main_token)[node],
@@ -1158,7 +1163,8 @@ pub const Tree = struct {
     }
 
     pub fn structInit(tree: Tree, node: Node.Index) full.StructInit {
-        assert(tree.nodes.items(.tag)[node] == .StructInit);
+        assert(tree.nodes.items(.tag)[node] == .StructInit or
+            tree.nodes.items(.tag)[node] == .StructInitComma);
         const data = tree.nodes.items(.data)[node];
         const fields_range = tree.extraData(data.rhs, Node.SubRange);
         return tree.fullStructInit(.{
@@ -2281,6 +2287,8 @@ pub const Node = struct {
         assert(@sizeOf(Tag) == 1);
     }
 
+    /// Note: The FooComma/FooSemicolon variants exist to ease the implementation of
+    /// Tree.lastToken()
     pub const Tag = enum {
         /// sub_list[lhs...rhs]
         Root,
@@ -2477,21 +2485,29 @@ pub const Node = struct {
         /// `lhs{.a = rhs}`. rhs can be omitted making it empty.
         /// main_token is the lbrace.
         StructInitOne,
+        /// `lhs{.a = rhs,}`. rhs can *not* be omitted.
+        /// main_token is the lbrace.
+        StructInitOneComma,
         /// `.{.a = lhs, .b = rhs}`. lhs and rhs can be omitted.
         /// main_token is the lbrace.
         /// No trailing comma before the rbrace.
         StructInitDotTwo,
         /// Same as `StructInitDotTwo` except there is known to be a trailing comma
-        /// before the final rbrace. This tag exists to facilitate lastToken() implemented
-        /// without recursion.
+        /// before the final rbrace.
         StructInitDotTwoComma,
         /// `.{.a = b, .c = d}`. `sub_list[lhs..rhs]`.
         /// main_token is the lbrace.
         StructInitDot,
+        /// Same as `StructInitDot` except there is known to be a trailing comma
+        /// before the final rbrace.
+        StructInitDotComma,
         /// `lhs{.a = b, .c = d}`. `sub_range_list[rhs]`.
         /// lhs can be omitted which means `.{.a = b, .c = d}`.
         /// main_token is the lbrace.
         StructInit,
+        /// Same as `StructInit` except there is known to be a trailing comma
+        /// before the final rbrace.
+        StructInitComma,
         /// `lhs(rhs)`. rhs can be omitted.
         CallOne,
         /// `lhs(rhs,)`. rhs can be omitted.
diff --git a/lib/std/zig/parse.zig b/lib/std/zig/parse.zig
index 9c07ad320a..21ebe43582 100644
--- a/lib/std/zig/parse.zig
+++ b/lib/std/zig/parse.zig
@@ -2147,7 +2147,7 @@ const Parser = struct {
             const comma_one = p.eatToken(.Comma);
             if (p.eatToken(.RBrace)) |_| {
                 return p.addNode(.{
-                    .tag = .StructInitOne,
+                    .tag = if (comma_one != null) .StructInitOneComma else .StructInitOne,
                     .main_token = lbrace,
                     .data = .{
                         .lhs = lhs,
@@ -2192,7 +2192,7 @@ const Parser = struct {
             }
             const span = try p.listToSpan(init_list.items);
             return p.addNode(.{
-                .tag = .StructInit,
+                .tag = if (p.token_tags[p.tok_i - 2] == .Comma) .StructInitComma else .StructInit,
                 .main_token = lbrace,
                 .data = .{
                     .lhs = lhs,
@@ -2709,12 +2709,8 @@ const Parser = struct {
                     if (field_init_one != 0) {
                         const comma_one = p.eatToken(.Comma);
                         if (p.eatToken(.RBrace)) |_| {
-                            const tag: Node.Tag = if (comma_one != null)
-                                .StructInitDotTwoComma
-                            else
-                                .StructInitDotTwo;
                             return p.addNode(.{
-                                .tag = tag,
+                                .tag = if (comma_one != null) .StructInitDotTwoComma else .StructInitDotTwo,
                                 .main_token = lbrace,
                                 .data = .{
                                     .lhs = field_init_one,
@@ -2730,12 +2726,8 @@ const Parser = struct {
                         const field_init_two = try p.expectFieldInit();
                         const comma_two = p.eatToken(.Comma);
                         if (p.eatToken(.RBrace)) |_| {
-                            const tag: Node.Tag = if (comma_two != null)
-                                .StructInitDotTwoComma
-                            else
-                                .StructInitDotTwo;
                             return p.addNode(.{
-                                .tag = tag,
+                                .tag = if (comma_two != null) .StructInitDotTwoComma else .StructInitDotTwo,
                                 .main_token = lbrace,
                                 .data = .{
                                     .lhs = field_init_one,
@@ -2784,8 +2776,9 @@ const Parser = struct {
                             }
                         }
                         const span = try p.listToSpan(init_list.items);
+                        const trailing_comma = p.token_tags[p.tok_i - 2] == .Comma;
                         return p.addNode(.{
-                            .tag = .StructInitDot,
+                            .tag = if (trailing_comma) .StructInitDotComma else .StructInitDot,
                             .main_token = lbrace,
                             .data = .{
                                 .lhs = span.start,
diff --git a/lib/std/zig/parser_test.zig b/lib/std/zig/parser_test.zig
index 3b0afff9a7..944722c296 100644
--- a/lib/std/zig/parser_test.zig
+++ b/lib/std/zig/parser_test.zig
@@ -466,102 +466,126 @@ test "zig fmt: anon literal in array" {
 
 test "zig fmt: anon struct literal 1 element" {
     try testCanonical(
-        \\const x = .{ .a = b };
+        \\test {
+        \\    const x = .{ .a = b };
+        \\}
         \\
     );
 }
 
 test "zig fmt: anon struct literal 1 element comma" {
     try testCanonical(
-        \\const x = .{
-        \\    .a = b,
-        \\};
+        \\test {
+        \\    const x = .{
+        \\        .a = b,
+        \\    };
+        \\}
         \\
     );
 }
 
 test "zig fmt: anon struct literal 2 element" {
     try testCanonical(
-        \\const x = .{ .a = b, .c = d };
+        \\test {
+        \\    const x = .{ .a = b, .c = d };
+        \\}
         \\
     );
 }
 
 test "zig fmt: anon struct literal 2 element comma" {
     try testCanonical(
-        \\const x = .{
-        \\    .a = b,
-        \\    .c = d,
-        \\};
+        \\test {
+        \\    const x = .{
+        \\        .a = b,
+        \\        .c = d,
+        \\    };
+        \\}
         \\
     );
 }
 
 test "zig fmt: anon struct literal 3 element" {
     try testCanonical(
-        \\const x = .{ .a = b, .c = d, .e = f };
+        \\test {
+        \\    const x = .{ .a = b, .c = d, .e = f };
+        \\}
         \\
     );
 }
 
 test "zig fmt: anon struct literal 3 element comma" {
     try testCanonical(
-        \\const x = .{
-        \\    .a = b,
-        \\    .c = d,
-        \\    .e = f,
-        \\};
+        \\test {
+        \\    const x = .{
+        \\        .a = b,
+        \\        .c = d,
+        \\        .e = f,
+        \\    };
+        \\}
         \\
     );
 }
 
 test "zig fmt: struct literal 1 element" {
     try testCanonical(
-        \\const x = X{ .a = b };
+        \\test {
+        \\    const x = X{ .a = b };
+        \\}
         \\
     );
 }
 
 test "zig fmt: struct literal 1 element comma" {
     try testCanonical(
-        \\const x = X{
-        \\    .a = b,
-        \\};
+        \\test {
+        \\    const x = X{
+        \\        .a = b,
+        \\    };
+        \\}
         \\
     );
 }
 
 test "zig fmt: struct literal 2 element" {
     try testCanonical(
-        \\const x = X{ .a = b, .c = d };
+        \\test {
+        \\    const x = X{ .a = b, .c = d };
+        \\}
         \\
     );
 }
 
 test "zig fmt: struct literal 2 element comma" {
     try testCanonical(
-        \\const x = X{
-        \\    .a = b,
-        \\    .c = d,
-        \\};
+        \\test {
+        \\    const x = X{
+        \\        .a = b,
+        \\        .c = d,
+        \\    };
+        \\}
         \\
     );
 }
 
 test "zig fmt: struct literal 3 element" {
     try testCanonical(
-        \\const x = X{ .a = b, .c = d, .e = f };
+        \\test {
+        \\    const x = X{ .a = b, .c = d, .e = f };
+        \\}
         \\
     );
 }
 
 test "zig fmt: struct literal 3 element comma" {
     try testCanonical(
-        \\const x = X{
-        \\    .a = b,
-        \\    .c = d,
-        \\    .e = f,
-        \\};
+        \\test {
+        \\    const x = X{
+        \\        .a = b,
+        \\        .c = d,
+        \\        .e = f,
+        \\    };
+        \\}
         \\
     );
 }
diff --git a/lib/std/zig/render.zig b/lib/std/zig/render.zig
index cc80b04d38..051b3f46b1 100644
--- a/lib/std/zig/render.zig
+++ b/lib/std/zig/render.zig
@@ -405,7 +405,7 @@ fn renderExpression(ais: *Ais, tree: ast.Tree, node: ast.Node.Index, space: Spac
         .ArrayInitComma,
         => return renderArrayInit(ais, tree, tree.arrayInit(node), space),
 
-        .StructInitOne => {
+        .StructInitOne, .StructInitOneComma => {
             var fields: [1]ast.Node.Index = undefined;
             return renderStructInit(ais, tree, tree.structInitOne(&fields, node), space);
         },
@@ -413,8 +413,12 @@ fn renderExpression(ais: *Ais, tree: ast.Tree, node: ast.Node.Index, space: Spac
             var fields: [2]ast.Node.Index = undefined;
             return renderStructInit(ais, tree, tree.structInitDotTwo(&fields, node), space);
         },
-        .StructInitDot => return renderStructInit(ais, tree, tree.structInitDot(node), space),
-        .StructInit => return renderStructInit(ais, tree, tree.structInit(node), space),
+        .StructInitDot,
+        .StructInitDotComma,
+        => return renderStructInit(ais, tree, tree.structInitDot(node), space),
+        .StructInit,
+        .StructInitComma,
+        => return renderStructInit(ais, tree, tree.structInit(node), space),
 
         .CallOne, .CallOneComma, .AsyncCallOne, .AsyncCallOneComma => {
             var params: [1]ast.Node.Index = undefined;