all tests passing again

2 files changed, 106 insertions, 76 deletions

diff --git a/std/zig/parser_test.zig b/std/zig/parser_test.zig
index 7b5358d238..29b231a4db 100644
--- a/std/zig/parser_test.zig
+++ b/std/zig/parser_test.zig
@@ -1,66 +1,66 @@
-//test "zig fmt: same-line comment after a statement" {
-//    try testCanonical(
-//        \\test "" {
-//        \\    a = b;
-//        \\    debug.assert(H.digest_size <= H.block_size); // HMAC makes this assumption
-//        \\    a = b;
-//        \\}
-//        \\
-//    );
-//}
-//
-//test "zig fmt: same-line comment after var decl in struct" {
-//    try testCanonical(
-//        \\pub const vfs_cap_data = extern struct {
-//        \\    const Data = struct {}; // when on disk.
-//        \\};
-//        \\
-//    );
-//}
-//
-//test "zig fmt: same-line comment after field decl" {
-//    try testCanonical(
-//        \\pub const dirent = extern struct {
-//        \\    d_name: u8,
-//        \\    d_name: u8, // comment 1
-//        \\    d_name: u8,
-//        \\    d_name: u8, // comment 2
-//        \\    d_name: u8,
-//        \\};
-//        \\
-//    );
-//}
-//
-//test "zig fmt: same-line comment after switch prong" {
-//    try testCanonical(
-//        \\test "" {
-//        \\    switch (err) {
-//        \\        error.PathAlreadyExists => {}, // comment 2
-//        \\        else => return err, // comment 1
-//        \\    }
-//        \\}
-//        \\
-//    );
-//}
-//
-//test "zig fmt: same-line comment after non-block if expression" {
-//    try testCanonical(
-//        \\comptime {
-//        \\    if (sr > n_uword_bits - 1) // d > r
-//        \\        return 0;
-//        \\}
-//        \\
-//    );
-//}
-//
-//test "zig fmt: same-line comment on comptime expression" {
-//    try testCanonical(
-//        \\test "" {
-//        \\    comptime assert(@typeId(T) == builtin.TypeId.Int); // must pass an integer to absInt
-//        \\}
-//        \\
-//    );
-//}
+test "zig fmt: same-line comment after a statement" {
+    try testCanonical(
+        \\test "" {
+        \\    a = b;
+        \\    debug.assert(H.digest_size <= H.block_size); // HMAC makes this assumption
+        \\    a = b;
+        \\}
+        \\
+    );
+}
+
+test "zig fmt: same-line comment after var decl in struct" {
+    try testCanonical(
+        \\pub const vfs_cap_data = extern struct {
+        \\    const Data = struct {}; // when on disk.
+        \\};
+        \\
+    );
+}
+
+test "zig fmt: same-line comment after field decl" {
+    try testCanonical(
+        \\pub const dirent = extern struct {
+        \\    d_name: u8,
+        \\    d_name: u8, // comment 1
+        \\    d_name: u8,
+        \\    d_name: u8, // comment 2
+        \\    d_name: u8,
+        \\};
+        \\
+    );
+}
+
+test "zig fmt: same-line comment after switch prong" {
+    try testCanonical(
+        \\test "" {
+        \\    switch (err) {
+        \\        error.PathAlreadyExists => {}, // comment 2
+        \\        else => return err, // comment 1
+        \\    }
+        \\}
+        \\
+    );
+}
+
+test "zig fmt: same-line comment after non-block if expression" {
+    try testCanonical(
+        \\comptime {
+        \\    if (sr > n_uword_bits - 1) // d > r
+        \\        return 0;
+        \\}
+        \\
+    );
+}
+
+test "zig fmt: same-line comment on comptime expression" {
+    try testCanonical(
+        \\test "" {
+        \\    comptime assert(@typeId(T) == builtin.TypeId.Int); // must pass an integer to absInt
+        \\}
+        \\
+    );
+}
 
 test "zig fmt: switch with empty body" {
     try testCanonical(
diff --git a/std/zig/render.zig b/std/zig/render.zig
index 3fa7c4c171..00a5613765 100644
--- a/std/zig/render.zig
+++ b/std/zig/render.zig
@@ -14,8 +14,13 @@ const RenderState = union(enum) {
     Statement: &ast.Node,
     PrintIndent,
     Indent: usize,
+    MaybeSemiColon: &ast.Node,
+    Token: ast.TokenIndex,
+    NonBreakToken: ast.TokenIndex,
 };
 
+const indent_delta = 4;
+
 pub fn render(allocator: &mem.Allocator, stream: var, tree: &ast.Tree) !void {
     var stack = SegmentedList(RenderState, 32).init(allocator);
     defer stack.deinit();
@@ -44,7 +49,6 @@ pub fn render(allocator: &mem.Allocator, stream: var, tree: &ast.Tree) !void {
         }
     }
 
-    const indent_delta = 4;
     var indent: usize = 0;
     while (stack.pop()) |state| {
         switch (state) {
@@ -92,7 +96,7 @@ pub fn render(allocator: &mem.Allocator, stream: var, tree: &ast.Tree) !void {
                             try stream.print("{} ", tree.tokenSlice(visib_token));
                         }
                         try stream.print("{}: ", tree.tokenSlice(field.name_token));
-                        try stack.push(RenderState { .Text = "," });
+                        try stack.push(RenderState { .Token = field.lastToken() + 1 });
                         try stack.push(RenderState { .Expression = field.type_expr});
                     },
                     ast.Node.Id.UnionTag => {
@@ -129,9 +133,7 @@ pub fn render(allocator: &mem.Allocator, stream: var, tree: &ast.Tree) !void {
                         try stream.print("{}", tree.tokenSlice(tag.name_token));
                     },
                     ast.Node.Id.Comptime => {
-                        if (decl.requireSemiColon()) {
-                            try stack.push(RenderState { .Text = ";" });
-                        }
+                        try stack.push(RenderState { .MaybeSemiColon = decl });
                         try stack.push(RenderState { .Expression = decl });
                     },
                     ast.Node.Id.LineComment => {
@@ -143,7 +145,7 @@ pub fn render(allocator: &mem.Allocator, stream: var, tree: &ast.Tree) !void {
             },
 
             RenderState.VarDecl => |var_decl| {
-                try stack.push(RenderState { .Text = ";" });
+                try stack.push(RenderState { .Token = var_decl.semicolon_token });
                 if (var_decl.init_node) |init_node| {
                     try stack.push(RenderState { .Expression = init_node });
                     const text = if (init_node.id == ast.Node.Id.MultilineStringLiteral) " =" else " = ";
@@ -895,7 +897,7 @@ pub fn render(allocator: &mem.Allocator, stream: var, tree: &ast.Tree) !void {
                 ast.Node.Id.SwitchCase => {
                     const switch_case = @fieldParentPtr(ast.Node.SwitchCase, "base", base);
 
-                    try stack.push(RenderState { .Text = "," });
+                    try stack.push(RenderState { .Token = switch_case.lastToken() + 1 });
                     try stack.push(RenderState { .Expression = switch_case.expr });
                     if (switch_case.payload) |payload| {
                         try stack.push(RenderState { .Text = " " });
@@ -1072,14 +1074,13 @@ pub fn render(allocator: &mem.Allocator, stream: var, tree: &ast.Tree) !void {
                     }
 
                     try stack.push(RenderState { .Expression = if_node.body });
-                    try stack.push(RenderState { .Text = " " });
 
                     if (if_node.payload) |payload| {
-                        try stack.push(RenderState { .Expression = payload });
                         try stack.push(RenderState { .Text = " " });
+                        try stack.push(RenderState { .Expression = payload });
                     }
 
-                    try stack.push(RenderState { .Text = ")" });
+                    try stack.push(RenderState { .NonBreakToken = if_node.condition.lastToken() + 1 });
                     try stack.push(RenderState { .Expression = if_node.condition });
                     try stack.push(RenderState { .Text = "(" });
                 },
@@ -1217,17 +1218,46 @@ pub fn render(allocator: &mem.Allocator, stream: var, tree: &ast.Tree) !void {
                         try stack.push(RenderState { .VarDecl = var_decl});
                     },
                     else => {
-                        if (base.requireSemiColon()) {
-                            try stack.push(RenderState { .Text = ";" });
-                        }
+                        try stack.push(RenderState { .MaybeSemiColon = base });
                         try stack.push(RenderState { .Expression = base });
                     },
                 }
             },
             RenderState.Indent => |new_indent| indent = new_indent,
             RenderState.PrintIndent => try stream.writeByteNTimes(' ', indent),
+            RenderState.Token => |token_index| try renderToken(tree, stream, token_index, indent, true),
+            RenderState.NonBreakToken => |token_index| try renderToken(tree, stream, token_index, indent, false),
+            RenderState.MaybeSemiColon => |base| {
+                if (base.requireSemiColon()) {
+                    const semicolon_index = base.lastToken() + 1;
+                    assert(tree.tokens.at(semicolon_index).id == Token.Id.Semicolon);
+                    try renderToken(tree, stream, semicolon_index, indent, true);
+                }
+            },
+        }
+    }
+}
+
+fn renderToken(tree: &ast.Tree, stream: var, token_index: ast.TokenIndex, indent: usize, line_break: bool) !void {
+    const token = tree.tokens.at(token_index);
+    try stream.write(tree.tokenSlicePtr(token));
+
+    const next_token = tree.tokens.at(token_index + 1);
+    if (next_token.id == Token.Id.LineComment) {
+        const loc = tree.tokenLocationPtr(token.end, next_token);
+        if (loc.line == 0) {
+            try stream.print(" {}", tree.tokenSlicePtr(next_token));
+            if (!line_break) {
+                try stream.write("\n");
+                try stream.writeByteNTimes(' ', indent + indent_delta);
+                return;
+            }
         }
     }
+
+    if (!line_break) {
+        try stream.writeByte(' ');
+    }
 }
 
 fn renderComments(tree: &ast.Tree, stream: var, node: var, indent: usize) !void {