Commit c98f792ff8

@@ -210,6 +210,103 @@ test "zig fmt: comment to disable/enable zig fmt" {
     );
 }
 
+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: pointer of unknown length" {
     try testCanonical(
         \\fn foo(ptr: [*]u8) void {}

@@ -89,41 +89,98 @@ fn renderRoot(
     var it = tree.root_node.decls.iterator(0);
     while (true) {
         var decl = (it.next() orelse return).*;
-        // look for zig fmt: off comment
-        var start_token_index = decl.firstToken();
-        zig_fmt_loop: while (start_token_index != 0) {
-            start_token_index -= 1;
-            const start_token = tree.tokens.at(start_token_index);
-            switch (start_token.id) {
+
+        // This loop does the following:
+        //
+        //  - Iterates through line/doc comment tokens that precedes the current
+        //    decl.
+        //  - Figures out the first token index (`copy_start_token_index`) which
+        //    hasn't been copied to the output stream yet.
+        //  - Detects `zig fmt: (off|on)` in the line comment tokens, and
+        //    determines whether the current decl should be reformatted or not.
+        //
+        var token_index = decl.firstToken();
+        var fmt_active = true;
+        var found_fmt_directive = false;
+
+        var copy_start_token_index = token_index;
+
+        while (token_index != 0) {
+            token_index -= 1;
+            const token = tree.tokens.at(token_index);
+            switch (token.id) {
                 Token.Id.LineComment => {},
-                Token.Id.DocComment => continue,
+                Token.Id.DocComment => {
+                    copy_start_token_index = token_index;
+                    continue;
+                },
                 else => break,
             }
-            if (mem.eql(u8, mem.trim(u8, tree.tokenSlicePtr(start_token)[2..], " "), "zig fmt: off")) {
-                var end_token_index = start_token_index;
-                while (true) {
-                    end_token_index += 1;
-                    const end_token = tree.tokens.at(end_token_index);
-                    switch (end_token.id) {
+
+            if (mem.eql(u8, mem.trim(u8, tree.tokenSlicePtr(token)[2..], " "), "zig fmt: off")) {
+                if (!found_fmt_directive) {
+                    fmt_active = false;
+                    found_fmt_directive = true;
+                }
+            } else if (mem.eql(u8, mem.trim(u8, tree.tokenSlicePtr(token)[2..], " "), "zig fmt: on")) {
+                if (!found_fmt_directive) {
+                    fmt_active = true;
+                    found_fmt_directive = true;
+                }
+            }
+        }
+
+        if (!fmt_active) {
+            // Reformatting is disabled for the current decl and possibly some
+            // more decls that follow.
+            // Find the next `decl` for which reformatting is re-enabled.
+            token_index = decl.firstToken();
+
+            while (!fmt_active) {
+                decl = (it.next() orelse {
+                    // If there's no next reformatted `decl`, just copy the
+                    // remaining input tokens and bail out.
+                    const start = tree.tokens.at(copy_start_token_index).start;
+                    try copyFixingWhitespace(stream, tree.source[start..]);
+                    return;
+                }).*;
+                var decl_first_token_index = decl.firstToken();
+
+                while (token_index < decl_first_token_index) : (token_index += 1) {
+                    const token = tree.tokens.at(token_index);
+                    switch (token.id) {
                         Token.Id.LineComment => {},
-                        Token.Id.Eof => {
-                            const start = tree.tokens.at(start_token_index + 1).start;
-                            try copyFixingWhitespace(stream, tree.source[start..]);
-                            return;
-                        },
+                        Token.Id.Eof => unreachable,
                         else => continue,
                     }
-                    if (mem.eql(u8, mem.trim(u8, tree.tokenSlicePtr(end_token)[2..], " "), "zig fmt: on")) {
-                        const start = tree.tokens.at(start_token_index + 1).start;
-                        try copyFixingWhitespace(stream, tree.source[start..end_token.end]);
-                        try stream.writeByte('\n');
-                        while (tree.tokens.at(decl.firstToken()).start < end_token.end) {
-                            decl = (it.next() orelse return).*;
-                        }
-                        break :zig_fmt_loop;
+                    if (mem.eql(u8, mem.trim(u8, tree.tokenSlicePtr(token)[2..], " "), "zig fmt: on")) {
+                        fmt_active = true;
+                    } else if (mem.eql(u8, mem.trim(u8, tree.tokenSlicePtr(token)[2..], " "), "zig fmt: off")) {
+                        fmt_active = false;
                     }
                 }
             }
+
+            // Found the next `decl` for which reformatting is enabled. Copy
+            // the input tokens before the `decl` that haven't been copied yet.
+            var copy_end_token_index = decl.firstToken();
+            token_index = copy_end_token_index;
+            while (token_index != 0) {
+                token_index -= 1;
+                const token = tree.tokens.at(token_index);
+                switch (token.id) {
+                    Token.Id.LineComment => {},
+                    Token.Id.DocComment => {
+                        copy_end_token_index = token_index;
+                        continue;
+                    },
+                    else => break,
+                }
+            }
+
+            const start = tree.tokens.at(copy_start_token_index).start;
+            const end = tree.tokens.at(copy_end_token_index).start;
+            try copyFixingWhitespace(stream, tree.source[start..end]);
         }
 
         try renderTopLevelDecl(allocator, stream, tree, 0, &start_col, decl);