zig fmt: Allow one-line for loops

2 files changed, 49 insertions, 18 deletions

diff --git a/std/zig/parser_test.zig b/std/zig/parser_test.zig
index 0b521f5eaa..f37dbf10f1 100644
--- a/std/zig/parser_test.zig
+++ b/std/zig/parser_test.zig
@@ -1737,6 +1737,10 @@ test "zig fmt: switch" {
 test "zig fmt: while" {
     try testCanonical(
         \\test "while" {
+        \\    while (10 < 1) unreachable;
+        \\
+        \\    while (10 < 1) unreachable else unreachable;
+        \\
         \\    while (10 < 1) {
         \\        unreachable;
         \\    }
@@ -1803,10 +1807,27 @@ test "zig fmt: while" {
 test "zig fmt: for" {
     try testCanonical(
         \\test "for" {
+        \\    for (a) continue;
+        \\
+        \\    for (a)
+        \\        continue;
+        \\
+        \\    for (a) {
+        \\        continue;
+        \\    }
+        \\
         \\    for (a) |v| {
         \\        continue;
         \\    }
         \\
+        \\    for (a) |v| continue;
+        \\
+        \\    for (a) |v| continue else return;
+        \\
+        \\    for (a) |v| {
+        \\        continue;
+        \\    } else return;
+        \\
         \\    for (a) |v|
         \\        continue;
         \\
@@ -1820,6 +1841,17 @@ test "zig fmt: for" {
         \\    for (a) |v, i|
         \\        continue;
         \\
+        \\    for (a) |b| switch (b) {
+        \\        c => {},
+        \\        d => {},
+        \\    };
+        \\
+        \\    for (a) |b|
+        \\        switch (b) {
+        \\            c => {},
+        \\            d => {},
+        \\        };
+        \\
         \\    const res = for (a) |v, i| {
         \\        break v;
         \\    } else {
diff --git a/std/zig/render.zig b/std/zig/render.zig
index fb4e372a05..d4d00f6a54 100644
--- a/std/zig/render.zig
+++ b/std/zig/render.zig
@@ -1444,18 +1444,24 @@ fn renderExpression(
             try renderExpression(allocator, stream, tree, indent, start_col, for_node.array_expr, Space.None);
 
             const rparen = tree.nextToken(for_node.array_expr.lastToken());
-            const rparen_space = if (for_node.payload != null or
-                for_node.body.id == ast.Node.Id.Block) Space.Space else Space.Newline;
-            try renderToken(tree, stream, rparen, indent, start_col, rparen_space); // )
+
+            const has_payload = for_node.payload != null;
+            const body_is_block = for_node.body.id == ast.Node.Id.Block;
+            const src_one_line_to_body = !body_is_block and tree.tokensOnSameLine(rparen, for_node.body.firstToken());
+            const body_indent = if (!body_is_block and !src_one_line_to_body) indent + indent_delta else indent;
+
+            const rparen_space = if (has_payload or body_is_block or src_one_line_to_body) Space.Space else Space.Newline;
+            try renderToken(tree, stream, rparen, body_indent, start_col, rparen_space); // )
 
             if (for_node.payload) |payload| {
-                const payload_space = if (for_node.body.id == ast.Node.Id.Block) Space.Space else Space.Newline;
-                try renderExpression(allocator, stream, tree, indent, start_col, payload, payload_space);
+                const payload_space = if (body_is_block or src_one_line_to_body) Space.Space else Space.Newline;
+                try renderExpression(allocator, stream, tree, body_indent, start_col, payload, payload_space); // |x|
             }
 
             const body_space = blk: {
-                if (for_node.@"else" != null) {
-                    if (for_node.body.id == ast.Node.Id.Block) {
+                if (for_node.@"else") |@"else"| {
+                    const src_one_line_to_else = tree.tokensOnSameLine(for_node.body.lastToken(), @"else".firstToken());
+                    if (body_is_block or src_one_line_to_else) {
                         break :blk Space.Space;
                     } else {
                         break :blk Space.Newline;
@@ -1464,19 +1470,12 @@ fn renderExpression(
                     break :blk space;
                 }
             };
-            if (for_node.body.id == ast.Node.Id.Block) {
-                try renderExpression(allocator, stream, tree, indent, start_col, for_node.body, body_space);
-            } else {
-                try stream.writeByteNTimes(' ', indent + indent_delta);
-                try renderExpression(allocator, stream, tree, indent, start_col, for_node.body, body_space);
-            }
 
-            if (for_node.@"else") |@"else"| {
-                if (for_node.body.id != ast.Node.Id.Block) {
-                    try stream.writeByteNTimes(' ', indent);
-                }
+            if (!body_is_block and !src_one_line_to_body) try stream.writeByteNTimes(' ', body_indent);
+            try renderExpression(allocator, stream, tree, body_indent, start_col, for_node.body, body_space);
 
-                return renderExpression(allocator, stream, tree, indent, start_col, &@"else".base, space);
+            if (for_node.@"else") |@"else"| {
+                return renderExpression(allocator, stream, tree, body_indent, start_col, &@"else".base, space);
             }
         },