Commit 895fb2bd6d

@@ -878,17 +878,17 @@ test "zig fmt: async function" {
     );
 }
 
-//test "zig fmt: whitespace fixes" {
-//    try testTransform("test \"\" {\r\n\tconst hi = x;\r\n}\n// zig fmt: off\ntest \"\"{\r\n\tconst a  = b;}\r\n",
-//        \\test "" {
-//        \\    const hi = x;
-//        \\}
-//        \\// zig fmt: off
-//        \\test ""{
-//        \\    const a  = b;}
-//        \\
-//    );
-//}
+test "zig fmt: whitespace fixes" {
+    try testTransform("test \"\" {\r\n\tconst hi = x;\r\n}\n// zig fmt: off\ntest \"\"{\r\n\tconst a  = b;}\r\n",
+        \\test "" {
+        \\    const hi = x;
+        \\}
+        \\// zig fmt: off
+        \\test ""{
+        \\    const a  = b;}
+        \\
+    );
+}
 
 test "zig fmt: while else err prong with no block" {
     try testCanonical(
@@ -1098,128 +1098,154 @@ test "zig fmt: aligned struct field" {
     );
 }
 
-//test "zig fmt: comment to disable/enable zig fmt first" {
-//    try testCanonical(
-//        \\// Test trailing comma syntax
-//        \\// zig fmt: off
-//        \\
-//        \\const struct_trailing_comma = struct { x: i32, y: i32, };
-//    );
-//}
-//
-//test "zig fmt: comment to disable/enable zig fmt" {
-//    try testTransform(
-//        \\const  a  =  b;
-//        \\// zig fmt: off
-//        \\const  c  =  d;
-//        \\// zig fmt: on
-//        \\const  e  =  f;
-//    ,
-//        \\const a = b;
-//        \\// zig fmt: off
-//        \\const  c  =  d;
-//        \\// zig fmt: on
-//        \\const e = f;
-//        \\
-//    );
-//}
-//
-//test "zig fmt: line comment following 'zig fmt: off'" {
-//    try testCanonical(
-//        \\// zig fmt: off
-//        \\// Test
-//        \\const  e  =  f;
-//    );
-//}
-//
-//test "zig fmt: doc comment following 'zig fmt: off'" {
-//    try testCanonical(
-//        \\// zig fmt: off
-//        \\/// test
-//        \\const  e  =  f;
-//    );
-//}
-//
-//test "zig fmt: line and doc comment following 'zig fmt: off'" {
-//    try testCanonical(
-//        \\// zig fmt: off
-//        \\// test 1
-//        \\/// test 2
-//        \\const  e  =  f;
-//    );
-//}
-//
-//test "zig fmt: doc and line comment following 'zig fmt: off'" {
-//    try testCanonical(
-//        \\// zig fmt: off
-//        \\/// test 1
-//        \\// test 2
-//        \\const  e  =  f;
-//    );
-//}
-//
-//test "zig fmt: alternating 'zig fmt: off' and 'zig fmt: on'" {
-//    try testCanonical(
-//        \\// zig fmt: off
-//        \\// zig fmt: on
-//        \\// zig fmt: off
-//        \\const  e  =  f;
-//        \\// zig fmt: off
-//        \\// zig fmt: on
-//        \\// zig fmt: off
-//        \\const  a  =  b;
-//        \\// zig fmt: on
-//        \\const c = d;
-//        \\// zig fmt: on
-//        \\
-//    );
-//}
-//
-//test "zig fmt: line comment following 'zig fmt: on'" {
-//    try testCanonical(
-//        \\// zig fmt: off
-//        \\const  e  =  f;
-//        \\// zig fmt: on
-//        \\// test
-//        \\const e = f;
-//        \\
-//    );
-//}
-//
-//test "zig fmt: doc comment following 'zig fmt: on'" {
-//    try testCanonical(
-//        \\// zig fmt: off
-//        \\const  e  =  f;
-//        \\// zig fmt: on
-//        \\/// test
-//        \\const e = f;
-//        \\
-//    );
-//}
-//
-//test "zig fmt: line and doc comment following 'zig fmt: on'" {
-//    try testCanonical(
-//        \\// zig fmt: off
-//        \\const  e  =  f;
-//        \\// zig fmt: on
-//        \\// test1
-//        \\/// test2
-//        \\const e = f;
-//        \\
-//    );
-//}
-//
-//test "zig fmt: doc and line comment following 'zig fmt: on'" {
-//    try testCanonical(
-//        \\// zig fmt: off
-//        \\const  e  =  f;
-//        \\// zig fmt: on
-//        \\/// test1
-//        \\// test2
-//        \\const e = f;
-//        \\
-//    );
-//}
+test "zig fmt: comment to disable/enable zig fmt first" {
+    try testCanonical(
+        \\// Test trailing comma syntax
+        \\// zig fmt: off
+        \\
+        \\const struct_trailing_comma = struct { x: i32, y: i32, };
+    );
+}
+
+test "zig fmt: comment to disable/enable zig fmt" {
+    try testTransform(
+        \\const  a  =  b;
+        \\// zig fmt: off
+        \\const  c  =  d;
+        \\// zig fmt: on
+        \\const  e  =  f;
+    ,
+        \\const a = b;
+        \\// zig fmt: off
+        \\const  c  =  d;
+        \\// zig fmt: on
+        \\const e = f;
+        \\
+    );
+}
+
+test "zig fmt: line comment following 'zig fmt: off'" {
+    try testCanonical(
+        \\// zig fmt: off
+        \\// Test
+        \\const  e  =  f;
+    );
+}
+
+test "zig fmt: doc comment following 'zig fmt: off'" {
+    try testCanonical(
+        \\// zig fmt: off
+        \\/// test
+        \\const  e  =  f;
+    );
+}
+
+test "zig fmt: line and doc comment following 'zig fmt: off'" {
+    try testCanonical(
+        \\// zig fmt: off
+        \\// test 1
+        \\/// test 2
+        \\const  e  =  f;
+    );
+}
+
+test "zig fmt: doc and line comment following 'zig fmt: off'" {
+    try testCanonical(
+        \\// zig fmt: off
+        \\/// test 1
+        \\// test 2
+        \\const  e  =  f;
+    );
+}
+
+test "zig fmt: alternating 'zig fmt: off' and 'zig fmt: on'" {
+    try testCanonical(
+        \\// zig fmt: off
+        \\// zig fmt: on
+        \\// zig fmt: off
+        \\const  e  =  f;
+        \\// zig fmt: off
+        \\// zig fmt: on
+        \\// zig fmt: off
+        \\const  a  =  b;
+        \\// zig fmt: on
+        \\const c = d;
+        \\// zig fmt: on
+        \\
+    );
+}
+
+test "zig fmt: line comment following 'zig fmt: on'" {
+    try testCanonical(
+        \\// zig fmt: off
+        \\const  e  =  f;
+        \\// zig fmt: on
+        \\// test
+        \\const e = f;
+        \\
+    );
+}
+
+test "zig fmt: doc comment following 'zig fmt: on'" {
+    try testCanonical(
+        \\// zig fmt: off
+        \\const  e  =  f;
+        \\// zig fmt: on
+        \\/// test
+        \\const e = f;
+        \\
+    );
+}
+
+test "zig fmt: line and doc comment following 'zig fmt: on'" {
+    try testCanonical(
+        \\// zig fmt: off
+        \\const  e  =  f;
+        \\// zig fmt: on
+        \\// test1
+        \\/// test2
+        \\const e = f;
+        \\
+    );
+}
+
+test "zig fmt: doc and line comment following 'zig fmt: on'" {
+    try testCanonical(
+        \\// zig fmt: off
+        \\const  e  =  f;
+        \\// zig fmt: on
+        \\/// test1
+        \\// test2
+        \\const e = f;
+        \\
+    );
+}
+
+test "zig fmt: 'zig fmt: (off|on)' works in the middle of code" {
+    try testTransform(
+        \\test "" {
+        \\    const x = 42;
+        \\
+        \\    if (foobar) |y| {
+        \\    // zig fmt: off
+        \\            }// zig fmt: on
+        \\
+        \\    const  z  = 420;
+        \\}
+        \\
+    ,
+        \\test "" {
+        \\    const x = 42;
+        \\
+        \\    if (foobar) |y| {
+        \\        // zig fmt: off
+        \\            }// zig fmt: on
+        \\
+        \\    const z = 420;
+        \\}
+        \\
+    );
+}
 
 test "zig fmt: pointer of unknown length" {
     try testCanonical(

@@ -30,6 +30,10 @@ pub fn renderTree(buffer: *std.ArrayList(u8), tree: ast.Tree) Error!void {
     _ = try renderComments(ais, tree, 0, comment_end_loc);
 
     try renderMembers(ais, tree, tree.rootDecls());
+
+    if (ais.disabled_offset) |disabled_offset| {
+        try writeFixingWhitespace(ais.underlying_writer, tree.source[disabled_offset..]);
+    }
 }
 
 /// Render all members in the given slice, keeping empty lines where appropriate
@@ -1971,6 +1975,7 @@ fn renderComments(ais: *Ais, tree: ast.Tree, start: usize, end: usize) Error!boo
         const comment_start = index + offset;
         const newline = comment_start +
             mem.indexOfScalar(u8, tree.source[comment_start..end], '\n').?;
+
         const untrimmed_comment = tree.source[comment_start..newline];
         const trimmed_comment = mem.trimRight(u8, untrimmed_comment, &std.ascii.spaces);
 
@@ -1993,6 +1998,17 @@ fn renderComments(ais: *Ais, tree: ast.Tree, start: usize, end: usize) Error!boo
 
         try ais.writer().print("{s}\n", .{trimmed_comment});
         index = newline + 1;
+
+        if (ais.disabled_offset) |disabled_offset| {
+            if (mem.eql(u8, trimmed_comment, "// zig fmt: on")) {
+                // write the source for which formatting was disabled directly
+                // to the underlying writer, fixing up invaild whitespace
+                try writeFixingWhitespace(ais.underlying_writer, tree.source[disabled_offset..index]);
+                ais.disabled_offset = null;
+            }
+        } else if (mem.eql(u8, trimmed_comment, "// zig fmt: off")) {
+            ais.disabled_offset = index;
+        }
     }
 
     if (index != start and mem.containsAtLeast(u8, tree.source[index - 1 .. end], 2, "\n")) {
@@ -2066,6 +2082,14 @@ fn tokenSliceForRender(tree: ast.Tree, token_index: ast.TokenIndex) []const u8 {
     return ret;
 }
 
+fn writeFixingWhitespace(writer: std.ArrayList(u8).Writer, slice: []const u8) Error!void {
+    for (slice) |byte| switch (byte) {
+        '\t' => try writer.writeAll(" " ** 4),
+        '\r' => {},
+        else => try writer.writeByte(byte),
+    };
+}
+
 fn nodeIsBlock(tag: ast.Node.Tag) bool {
     return switch (tag) {
         .block,
@@ -2145,6 +2169,14 @@ fn AutoIndentingStream(comptime UnderlyingWriter: type) type {
 
         underlying_writer: UnderlyingWriter,
 
+        /// Offset into the source at which formatting has been disabled with
+        /// a `zig fmt: off` comment.
+        ///
+        /// If non-null, the AutoIndentingStream will not write any bytes
+        /// to the underlying writer. It will however continue to track the
+        /// indentation level.
+        disabled_offset: ?usize = null,
+
         indent_count: usize = 0,
         indent_delta: usize,
         current_line_empty: bool = true,
@@ -2183,7 +2215,7 @@ fn AutoIndentingStream(comptime UnderlyingWriter: type) type {
             if (bytes.len == 0)
                 return @as(usize, 0);
 
-            try self.underlying_writer.writeAll(bytes);
+            if (self.disabled_offset == null) try self.underlying_writer.writeAll(bytes);
             if (bytes[bytes.len - 1] == '\n')
                 self.resetLine();
             return bytes.len;
@@ -2243,7 +2275,9 @@ fn AutoIndentingStream(comptime UnderlyingWriter: type) type {
         fn applyIndent(self: *Self) Error!void {
             const current_indent = self.currentIndent();
             if (self.current_line_empty and current_indent > 0) {
-                try self.underlying_writer.writeByteNTimes(' ', current_indent);
+                if (self.disabled_offset == null) {
+                    try self.underlying_writer.writeByteNTimes(' ', current_indent);
+                }
                 self.applied_indent = current_indent;
             }