From ecf555c6931afdb32af4ad0f84311986e0fa3687 Mon Sep 17 00:00:00 2001
From: Lewis Gaul <lewis.gaul@gmail.com>
Date: Sat, 10 Apr 2021 15:39:26 +0100
Subject: [PATCH] zig fmt: render array init on one line if no trailing comma

Continue to insert a trailing comma if there is a comment or multiline
string literal present.
---
 lib/std/zig/parser_test.zig | 74 +++++++++++++++++++++++++++++++++----
 lib/std/zig/render.zig      | 20 +++++++++-
 2 files changed, 84 insertions(+), 10 deletions(-)

diff --git a/lib/std/zig/parser_test.zig b/lib/std/zig/parser_test.zig
index 321f37c422..2b2737a24c 100644
--- a/lib/std/zig/parser_test.zig
+++ b/lib/std/zig/parser_test.zig
@@ -1801,19 +1801,21 @@ test "zig fmt: array literal with hint" {
     );
 }
 
-test "zig fmt: array literal veritical column alignment" {
+test "zig fmt: array literal vertical column alignment" {
     try testTransform(
         \\const a = []u8{
         \\    1000, 200,
         \\    30, 4,
-        \\    50000, 60
+        \\    50000, 60,
         \\};
         \\const a = []u8{0,   1, 2, 3, 40,
         \\    4,5,600,7,
         \\           80,
-        \\    9, 10, 11, 0, 13, 14, 15};
+        \\    9, 10, 11, 0, 13, 14, 15,};
         \\const a = [12]u8{
         \\    31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31 };
+        \\const a = [12]u8{
+        \\    31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31, };
         \\
     ,
         \\const a = []u8{
@@ -1827,8 +1829,20 @@ test "zig fmt: array literal veritical column alignment" {
         \\    9,  10, 11,  0, 13,
         \\    14, 15,
         \\};
+        \\const a = [12]u8{ 31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31 };
         \\const a = [12]u8{
-        \\    31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31,
+        \\    31,
+        \\    28,
+        \\    31,
+        \\    30,
+        \\    31,
+        \\    30,
+        \\    31,
+        \\    31,
+        \\    30,
+        \\    31,
+        \\    30,
+        \\    31,
         \\};
         \\
     );
@@ -2031,10 +2045,7 @@ test "zig fmt: add trailing comma to array literal" {
         \\    return []u16{
         \\        'm', 's', 'y', 's', '-', // hi
         \\    };
-        \\    return []u16{
-        \\        'm', 's', 'y', 's',
-        \\        '-',
-        \\    };
+        \\    return []u16{ 'm', 's', 'y', 's', '-' };
         \\    return []u16{ 'm', 's', 'y', 's', '-' };
         \\}
         \\
@@ -4666,6 +4677,53 @@ test "zig fmt: insert trailing comma if there are comments between switch values
     );
 }
 
+test "zig fmt: make single-line if no trailing comma" {
+    try testTransform(
+        \\test "function call no trailing comma" {
+        \\    foo(
+        \\        1,
+        \\        2
+        \\    );
+        \\}
+        \\
+    ,
+        \\test "function call no trailing comma" {
+        \\    foo(1, 2);
+        \\}
+        \\
+    );
+
+    try testTransform(
+        \\test "struct no trailing comma" {
+        \\    const a = .{
+        \\        .foo = 1,
+        \\        .bar = 2
+        \\    };
+        \\}
+        \\
+    ,
+        \\test "struct no trailing comma" {
+        \\    const a = .{ .foo = 1, .bar = 2 };
+        \\}
+        \\
+    );
+
+    try testTransform(
+        \\test "array no trailing comma" {
+        \\    var stream = multiOutStream(.{
+        \\        fbs1.outStream(),
+        \\        fbs2.outStream()
+        \\    });
+        \\}
+        \\
+    ,
+        \\test "array no trailing comma" {
+        \\    var stream = multiOutStream(.{ fbs1.outStream(), fbs2.outStream() });
+        \\}
+        \\
+    );
+}
+
 test "zig fmt: error for invalid bit range" {
     try testError(
         \\var x: []align(0:0:0)u8 = bar;
diff --git a/lib/std/zig/render.zig b/lib/std/zig/render.zig
index 1aa2fe8af6..a3187d665a 100644
--- a/lib/std/zig/render.zig
+++ b/lib/std/zig/render.zig
@@ -1632,9 +1632,10 @@ fn renderArrayInit(
         }
     }
 
-    const contains_newlines = !tree.tokensOnSameLine(array_init.ast.lbrace, rbrace);
+    const contains_comment = hasComment(tree, array_init.ast.lbrace, rbrace);
+    const contains_multiline_string = hasMultilineString(tree, array_init.ast.lbrace, rbrace);
 
-    if (!trailing_comma and !contains_newlines) {
+    if (!trailing_comma and !contains_comment and !contains_multiline_string) {
         // Render all on one line, no trailing comma.
         if (array_init.ast.elements.len == 1) {
             // If there is only one element, we don't use spaces
@@ -2252,6 +2253,21 @@ fn hasComment(tree: ast.Tree, start_token: ast.TokenIndex, end_token: ast.TokenI
     return mem.indexOf(u8, tree.source[start..end], "//") != null;
 }
 
+/// Returns true if there exists a multiline string literal between the start
+/// of token `start_token` and the start of token `end_token`.
+fn hasMultilineString(tree: ast.Tree, start_token: ast.TokenIndex, end_token: ast.TokenIndex) bool {
+    const token_tags = tree.tokens.items(.tag);
+
+    for (token_tags[start_token..end_token]) |tag| {
+        switch (tag) {
+            .multiline_string_literal_line => return true,
+            else => continue,
+        }
+    }
+
+    return false;
+}
+
 /// Assumes that start is the first byte past the previous token and
 /// that end is the last byte before the next token.
 fn renderComments(ais: *Ais, tree: ast.Tree, start: usize, end: usize) Error!bool {