From 7fdbaeca728312d22d877a711b46e157d3b01fe7 Mon Sep 17 00:00:00 2001
From: Andrew Kelley <superjoe30@gmail.com>
Date: Fri, 4 May 2018 18:35:43 -0400
Subject: [PATCH] zig fmt: same-line comment after non-block if expression

---
 std/zig/parser.zig      | 23 ++++++++++++++++++++---
 std/zig/parser_test.zig | 24 ++++++++++++++++++------
 2 files changed, 38 insertions(+), 9 deletions(-)

diff --git a/std/zig/parser.zig b/std/zig/parser.zig
index f303bccd8..74271f1aa 100644
--- a/std/zig/parser.zig
+++ b/std/zig/parser.zig
@@ -1891,12 +1891,13 @@ pub const Parser = struct {
                         }
                     );
 
-                    stack.append(State {
+                    stack.append(State {.LookForSameLineCommentDirect = &node.base }) catch unreachable;
+                    try stack.append(State {
                         .ExpectTokenSave = ExpectTokenSave {
                             .id = Token.Id.Pipe,
                             .ptr = &node.rpipe,
                         }
-                    }) catch unreachable;
+                    });
                     try stack.append(State { .Identifier = OptionalCtx { .Required = &node.value_symbol } });
                     try stack.append(State {
                         .OptionalTokenSave = OptionalTokenSave {
@@ -3122,6 +3123,7 @@ pub const Parser = struct {
                 stack.append(State { .Else = &node.@"else" }) catch unreachable;
                 try stack.append(State { .Expression = OptionalCtx { .Required = &node.body } });
                 try stack.append(State { .PointerPayload = OptionalCtx { .Optional = &node.payload } });
+                try stack.append(State { .LookForSameLineComment = &node.condition });
                 try stack.append(State { .ExpectToken = Token.Id.RParen });
                 try stack.append(State { .Expression = OptionalCtx { .Required = &node.condition } });
                 try stack.append(State { .ExpectToken = Token.Id.LParen });
@@ -3460,6 +3462,7 @@ pub const Parser = struct {
         PrintIndent,
         Indent: usize,
         PrintSameLineComment: ?&Token,
+        PrintLineComment: &Token,
     };
 
     pub fn renderSource(self: &Parser, stream: var, root_node: &ast.Node.Root) !void {
@@ -4517,7 +4520,18 @@ pub const Parser = struct {
                             }
                         }
 
-                        try stack.append(RenderState { .Expression = if_node.body });
+                        if (if_node.condition.same_line_comment) |comment| {
+                            try stack.append(RenderState { .Indent = indent });
+                            try stack.append(RenderState { .Expression = if_node.body });
+                            try stack.append(RenderState.PrintIndent);
+                            try stack.append(RenderState { .Indent = indent + indent_delta });
+                            try stack.append(RenderState { .Text = "\n" });
+                            try stack.append(RenderState { .PrintLineComment = comment });
+                        } else {
+                            try stack.append(RenderState { .Expression = if_node.body });
+                        }
+
+
                         try stack.append(RenderState { .Text = " " });
 
                         if (if_node.payload) |payload| {
@@ -4678,6 +4692,9 @@ pub const Parser = struct {
                     const comment_token = maybe_comment ?? break :blk;
                     try stream.print(" {}", self.tokenizer.getTokenSlice(comment_token));
                 },
+                RenderState.PrintLineComment => |comment_token| {
+                    try stream.write(self.tokenizer.getTokenSlice(comment_token));
+                },
             }
         }
     }
diff --git a/std/zig/parser_test.zig b/std/zig/parser_test.zig
index 3971f45af..e1d75d838 100644
--- a/std/zig/parser_test.zig
+++ b/std/zig/parser_test.zig
@@ -1,6 +1,14 @@
-// TODO
-//if (sr > n_uword_bits - 1)  // d > r
-//    return 0;
+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: switch with empty body" {
     try testCanonical(
@@ -1108,15 +1116,15 @@ fn testParse(source: []const u8, allocator: &mem.Allocator) ![]u8 {
     return buffer.toOwnedSlice();
 }
 
-fn testCanonical(source: []const u8) !void {
+fn testTransform(source: []const u8, expected_source: []const u8) !void {
     const needed_alloc_count = x: {
         // Try it once with unlimited memory, make sure it works
         var fixed_allocator = std.heap.FixedBufferAllocator.init(fixed_buffer_mem[0..]);
         var failing_allocator = std.debug.FailingAllocator.init(&fixed_allocator.allocator, @maxValue(usize));
         const result_source = try testParse(source, &failing_allocator.allocator);
-        if (!mem.eql(u8, result_source, source)) {
+        if (!mem.eql(u8, result_source, expected_source)) {
             warn("\n====== expected this output: =========\n");
-            warn("{}", source);
+            warn("{}", expected_source);
             warn("\n======== instead found this: =========\n");
             warn("{}", result_source);
             warn("\n======================================\n");
@@ -1147,3 +1155,7 @@ fn testCanonical(source: []const u8) !void {
     }
 }
 
+fn testCanonical(source: []const u8) !void {
+    return testTransform(source, source);
+}
+