zig fmt: support trailing comma on switch case items

ziglang · May 26, 2018 · 4405897 · 4405897
1 parent a630d3e
commit 4405897
Show file tree

Hide file tree

Showing 3 changed files with 79 additions and 25 deletions.
diff --git a/std/zig/parse.zig b/std/zig/parse.zig
@@ -1325,21 +1325,30 @@ pub fn parse(allocator: &mem.Allocator, source: []const u8) !ast.Tree {
                     continue;
                 } else {
                     prevToken(&tok_it, &tree);
-                    try stack.append(State{ .SwitchCaseItem = switch_case });
+                    stack.append(State{ .SwitchCaseItemCommaOrEnd = switch_case }) catch unreachable;
+                    try stack.append(State{ .RangeExpressionBegin = OptionalCtx{ .Required = try switch_case.items.addOne() } });
                     continue;
                 }
             },
-            State.SwitchCaseItem => |node| {
-                stack.append(State{ .SwitchCaseItemCommaOrEnd = node }) catch unreachable;
-                try stack.append(State{ .RangeExpressionBegin = OptionalCtx{ .Required = try node.items.addOne() } });
+            State.SwitchCaseItemOrEnd => |switch_case| {
+                const token = nextToken(&tok_it, &tree);
+                if (token.ptr.id == Token.Id.EqualAngleBracketRight) {
+                    switch_case.arrow_token = token.index;
+                    continue;
+                } else {
+                    prevToken(&tok_it, &tree);
+                    stack.append(State{ .SwitchCaseItemCommaOrEnd = switch_case }) catch unreachable;
+                    try stack.append(State{ .RangeExpressionBegin = OptionalCtx{ .Required = try switch_case.items.addOne() } });
+                    continue;
+                }
             },
-            State.SwitchCaseItemCommaOrEnd => |node| {
+            State.SwitchCaseItemCommaOrEnd => |switch_case| {
                 switch (expectCommaOrEnd(&tok_it, &tree, Token.Id.EqualAngleBracketRight)) {
                     ExpectCommaOrEndResult.end_token => |end_token| {
                         if (end_token) |t| {
-                            node.arrow_token = t;
+                            switch_case.arrow_token = t;
                         } else {
-                            stack.append(State{ .SwitchCaseItem = node }) catch unreachable;
+                            stack.append(State{ .SwitchCaseItemOrEnd = switch_case }) catch unreachable;
                         }
                         continue;
                     },
@@ -2828,8 +2837,8 @@ const State = union(enum) {
     SwitchCaseOrEnd: ListSave(ast.Node.Switch.CaseList),
     SwitchCaseCommaOrEnd: ListSave(ast.Node.Switch.CaseList),
     SwitchCaseFirstItem: &ast.Node.SwitchCase,
-    SwitchCaseItem: &ast.Node.SwitchCase,
     SwitchCaseItemCommaOrEnd: &ast.Node.SwitchCase,
+    SwitchCaseItemOrEnd: &ast.Node.SwitchCase,
 
     SuspendBody: &ast.Node.Suspend,
     AsyncAllocator: &ast.Node.AsyncAttribute,

diff --git a/std/zig/parser_test.zig b/std/zig/parser_test.zig
@@ -1,3 +1,27 @@
+test "zig fmt: switch cases trailing comma" {
+    try testTransform(
+        \\fn switch_cases(x: i32) void {
+        \\    switch (x) {
+        \\        1,2,3 => {},
+        \\        4,5, => {},
+        \\        6...8, => {},
+        \\        else => {},
+        \\    }
+        \\}
+    ,
+        \\fn switch_cases(x: i32) void {
+        \\    switch (x) {
+        \\        1, 2, 3 => {},
+        \\        4,
+        \\        5, => {},
+        \\        6 ... 8 => {},
+        \\        else => {},
+        \\    }
+        \\}
+        \\
+    );
+}
+
 test "zig fmt: slice align" {
     try testCanonical(
         \\const A = struct {
@@ -7,7 +31,7 @@ test "zig fmt: slice align" {
     );
 }
 
-test "zig fmt: first thing in file is line comment" {
+test "zig fmt: add trailing comma to array literal" {
     try testTransform(
         \\comptime {
         \\    return []u16{'m', 's', 'y', 's', '-' // hi
@@ -217,13 +241,11 @@ test "zig fmt: add comma on last switch prong" {
         \\test "aoeu" {
         \\    switch (self.init_arg_expr) {
         \\        InitArg.Type => |t| {},
-        \\        InitArg.None,
-        \\        InitArg.Enum => {},
+        \\        InitArg.None, InitArg.Enum => {},
         \\    }
         \\    switch (self.init_arg_expr) {
         \\        InitArg.Type => |t| {},
-        \\        InitArg.None,
-        \\        InitArg.Enum => {}, //line comment
+        \\        InitArg.None, InitArg.Enum => {}, //line comment
         \\    }
         \\}
         \\
@@ -1003,8 +1025,7 @@ test "zig fmt: switch" {
         \\    switch (0) {
         \\        0 => {},
         \\        1 => unreachable,
-        \\        2,
-        \\        3 => {},
+        \\        2, 3 => {},
         \\        4 ... 7 => {},
         \\        1 + 4 * 3 + 22 => {},
         \\        else => {

diff --git a/std/zig/render.zig b/std/zig/render.zig
@@ -939,17 +939,41 @@ fn renderExpression(allocator: &mem.Allocator, stream: var, tree: &ast.Tree, ind
         ast.Node.Id.SwitchCase => {
             const switch_case = @fieldParentPtr(ast.Node.SwitchCase, "base", base);
 
-            var it = switch_case.items.iterator(0);
-            while (it.next()) |node| {
-                if (it.peek()) |next_node| {
-                    try renderExpression(allocator, stream, tree, indent, node.*, Space.None);
+            assert(switch_case.items.len != 0);
+            const src_has_trailing_comma = blk: {
+                const last_node = switch_case.items.at(switch_case.items.len - 1).*;
+                const maybe_comma = tree.nextToken(last_node.lastToken());
+                break :blk tree.tokens.at(maybe_comma).id == Token.Id.Comma;
+            };
 
-                    const comma_token = tree.nextToken(node.*.lastToken());
-                    try renderToken(tree, stream, comma_token, indent, Space.Newline); // ,
-                    try renderExtraNewline(tree, stream, next_node.*);
-                    try stream.writeByteNTimes(' ', indent);
-                } else {
-                    try renderExpression(allocator, stream, tree, indent, node.*, Space.Space);
+            if (switch_case.items.len == 1 or !src_has_trailing_comma) {
+                var it = switch_case.items.iterator(0);
+                while (it.next()) |node| {
+                    if (it.peek()) |next_node| {
+                        try renderExpression(allocator, stream, tree, indent, node.*, Space.None);
+
+                        const comma_token = tree.nextToken(node.*.lastToken());
+                        try renderToken(tree, stream, comma_token, indent, Space.Space); // ,
+                        try renderExtraNewline(tree, stream, next_node.*);
+                    } else {
+                        try renderExpression(allocator, stream, tree, indent, node.*, Space.Space);
+                    }
+                }
+            } else {
+                var it = switch_case.items.iterator(0);
+                while (true) {
+                    const node = ??it.next();
+                    if (it.peek()) |next_node| {
+                        try renderExpression(allocator, stream, tree, indent, node.*, Space.None);
+
+                        const comma_token = tree.nextToken(node.*.lastToken());
+                        try renderToken(tree, stream, comma_token, indent, Space.Newline); // ,
+                        try renderExtraNewline(tree, stream, next_node.*);
+                        try stream.writeByteNTimes(' ', indent);
+                    } else {
+                        try renderTrailingComma(allocator, stream, tree, indent, node.*, Space.Space);
+                        break;
+                    }
                 }
             }