Always write a multiline struct literal if a field expr is multiline

2 files changed, 102 insertions, 11 deletions

diff --git a/std/zig/parser_test.zig b/std/zig/parser_test.zig
index 49a79f4ef7..85bb3d5f86 100644
--- a/std/zig/parser_test.zig
+++ b/std/zig/parser_test.zig
@@ -393,6 +393,50 @@ test "zig fmt: struct literal no trailing comma" {
     );
 }
 
+test "zig fmt: struct literal containing a multiline expression" {
+    try testTransform(
+        \\const a = A{ .x = if (f1()) 10 else 20 };
+        \\const a = A{ .x = if (f1()) 10 else 20, };
+        \\const a = A{ .x = if (f1())
+        \\    10 else 20 };
+        \\const a = A{ .x = if (f1()) 10 else 20, .y = f2() + 100 };
+        \\const a = A{ .x = if (f1()) 10 else 20, .y = f2() + 100, };
+        \\const a = A{ .x = if (f1())
+        \\    10 else 20};
+        \\const a = A{ .x = switch(g) {0 => "ok", else => "no"} };
+        \\
+    ,
+        \\const a = A{ .x = if (f1()) 10 else 20 };
+        \\const a = A{
+        \\    .x = if (f1()) 10 else 20,
+        \\};
+        \\const a = A{
+        \\    .x = if (f1())
+        \\        10
+        \\    else
+        \\        20,
+        \\};
+        \\const a = A{ .x = if (f1()) 10 else 20, .y = f2() + 100 };
+        \\const a = A{
+        \\    .x = if (f1()) 10 else 20,
+        \\    .y = f2() + 100,
+        \\};
+        \\const a = A{
+        \\    .x = if (f1())
+        \\        10
+        \\    else
+        \\        20,
+        \\};
+        \\const a = A{
+        \\    .x = switch (g) {
+        \\        0 => "ok",
+        \\        else => "no",
+        \\    },
+        \\};
+        \\
+    );
+}
+
 test "zig fmt: array literal with hint" {
     try testTransform(
         \\const a = []u8{
diff --git a/std/zig/render.zig b/std/zig/render.zig
index 80f188a198..245505b311 100644
--- a/std/zig/render.zig
+++ b/std/zig/render.zig
@@ -596,6 +596,28 @@ fn renderExpression(
                         return renderToken(tree, stream, suffix_op.rtoken, indent, start_col, space);
                     }
 
+                    const src_has_trailing_comma = blk: {
+                        const maybe_comma = tree.prevToken(suffix_op.rtoken);
+                        break :blk tree.tokens.at(maybe_comma).id == Token.Id.Comma;
+                    };
+
+                    const src_same_line = blk: {
+                        const loc = tree.tokenLocation(tree.tokens.at(lbrace).end, suffix_op.rtoken);
+                        break :blk loc.line == 0;
+                    };
+
+                    const expr_outputs_one_line = blk: {
+                        // render field expressions until a LF is found
+                        var it = field_inits.iterator(0);
+                        while (it.next()) |field_init| {
+                            var find_stream = FindByteOutStream.init('\n');
+                            var dummy_col: usize = 0;
+                            try renderExpression(allocator, &find_stream.stream, tree, 0, &dummy_col, field_init.*, Space.None);
+                            if (find_stream.byte_found) break :blk false;
+                        }
+                        break :blk true;
+                    };
+
                     if (field_inits.len == 1) blk: {
                         const field_init = field_inits.at(0).*.cast(ast.Node.FieldInitializer).?;
 
@@ -605,23 +627,18 @@ fn renderExpression(
                             }
                         }
 
+                        // if the expression outputs to multiline, make this struct multiline
+                        if (!expr_outputs_one_line or src_has_trailing_comma) {
+                            break :blk;
+                        }
+
                         try renderExpression(allocator, stream, tree, indent, start_col, suffix_op.lhs, Space.None);
                         try renderToken(tree, stream, lbrace, indent, start_col, Space.Space);
                         try renderExpression(allocator, stream, tree, indent, start_col, &field_init.base, Space.Space);
                         return renderToken(tree, stream, suffix_op.rtoken, indent, start_col, space);
                     }
 
-                    const src_has_trailing_comma = blk: {
-                        const maybe_comma = tree.prevToken(suffix_op.rtoken);
-                        break :blk tree.tokens.at(maybe_comma).id == Token.Id.Comma;
-                    };
-
-                    const src_same_line = blk: {
-                        const loc = tree.tokenLocation(tree.tokens.at(lbrace).end, suffix_op.rtoken);
-                        break :blk loc.line == 0;
-                    };
-
-                    if (!src_has_trailing_comma and src_same_line) {
+                    if (!src_has_trailing_comma and src_same_line and expr_outputs_one_line) {
                         // render all on one line, no trailing comma
                         try renderExpression(allocator, stream, tree, indent, start_col, suffix_op.lhs, Space.None);
                         try renderToken(tree, stream, lbrace, indent, start_col, Space.Space);
@@ -2092,3 +2109,33 @@ fn nodeCausesSliceOpSpace(base: *ast.Node) bool {
         else => true,
     };
 }
+
+// An OutStream that returns whether the given character has been written to it.
+// The contents are not written to anything.
+const FindByteOutStream = struct {
+    const Self = FindByteOutStream;
+    pub const Error = error{};
+    pub const Stream = std.io.OutStream(Error);
+
+    pub stream: Stream,
+    pub byte_found: bool,
+    byte: u8,
+
+    pub fn init(byte: u8) Self {
+        return Self{
+            .stream = Stream{ .writeFn = writeFn },
+            .byte = byte,
+            .byte_found = false,
+        };
+    }
+
+    fn writeFn(out_stream: *Stream, bytes: []const u8) Error!void {
+        const self = @fieldParentPtr(Self, "stream", out_stream);
+        if (self.byte_found) return;
+        self.byte_found = blk: {
+            for (bytes) |b|
+                if (b == self.byte) break :blk true;
+            break :blk false;
+        };
+    }
+};