From 3bb4c0c78978a80a3e441552591f0209f032bfc7 Mon Sep 17 00:00:00 2001
From: Curtis Tate Wilkinson <curtistatewilkinson@gmail.com>
Date: Tue, 15 Mar 2022 08:10:59 +1000
Subject: [PATCH] zig fmt: Resolve #11131 loss of comment on switch cases

Correct switch cases dropping comments in certain situations by
checking for the presence of the comment before collapsing to one line.
---
 lib/std/zig/parser_test.zig | 13 +++++++++++++
 lib/std/zig/render.zig      | 10 ++++++----
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/lib/std/zig/parser_test.zig b/lib/std/zig/parser_test.zig
index 0c79b3b187..c9c598b1d6 100644
--- a/lib/std/zig/parser_test.zig
+++ b/lib/std/zig/parser_test.zig
@@ -1791,6 +1791,19 @@ test "zig fmt: switch comment before prong" {
     );
 }
 
+test "zig fmt: switch comment after prong" {
+    try testCanonical(
+        \\comptime {
+        \\    switch (a) {
+        \\        0,
+        \\        // hi
+        \\        => {},
+        \\    }
+        \\}
+        \\
+    );
+}
+
 test "zig fmt: struct literal no trailing comma" {
     try testTransform(
         \\const a = foo{ .x = 1, .y = 2 };
diff --git a/lib/std/zig/render.zig b/lib/std/zig/render.zig
index 0f6fcac8b7..eaae725e9a 100644
--- a/lib/std/zig/render.zig
+++ b/lib/std/zig/render.zig
@@ -1505,16 +1505,18 @@ fn renderSwitchCase(
     const node_tags = tree.nodes.items(.tag);
     const token_tags = tree.tokens.items(.tag);
     const trailing_comma = token_tags[switch_case.ast.arrow_token - 1] == .comma;
+    const has_comment_before_arrow = blk: {
+        if (switch_case.ast.values.len == 0) break :blk false;
+        break :blk hasComment(tree, tree.firstToken(switch_case.ast.values[0]), switch_case.ast.arrow_token);
+    };
 
     // Render everything before the arrow
     if (switch_case.ast.values.len == 0) {
         try renderToken(ais, tree, switch_case.ast.arrow_token - 1, .space); // else keyword
-    } else if (switch_case.ast.values.len == 1) {
+    } else if (switch_case.ast.values.len == 1 and !has_comment_before_arrow) {
         // render on one line and drop the trailing comma if any
         try renderExpression(gpa, ais, tree, switch_case.ast.values[0], .space);
-    } else if (trailing_comma or
-        hasComment(tree, tree.firstToken(switch_case.ast.values[0]), switch_case.ast.arrow_token))
-    {
+    } else if (trailing_comma or has_comment_before_arrow) {
         // Render each value on a new line
         try renderExpressions(gpa, ais, tree, switch_case.ast.values, .comma);
     } else {