From d21a1922eb5d76b9b0d0611eaeb42c91f83234ab Mon Sep 17 00:00:00 2001
From: Andrew Kelley <superjoe30@gmail.com>
Date: Mon, 4 Jun 2018 12:15:02 -0400
Subject: [PATCH] support `zig fmt: off` and `zig fmt: on` between top level
 decls

closes #1030
closes #1033
---
 std/special/compiler_rt/udivmoddi4_test.zig |  2 +
 std/special/compiler_rt/udivmodti4_test.zig |  2 +
 std/zig/parser_test.zig                     | 26 +++++++++++++
 std/zig/render.zig                          | 41 ++++++++++++++++++++-
 test/cases/syntax.zig                       |  1 +
 5 files changed, 70 insertions(+), 2 deletions(-)

diff --git a/std/special/compiler_rt/udivmoddi4_test.zig b/std/special/compiler_rt/udivmoddi4_test.zig
index 324626d3f9..34b9dda1ea 100644
--- a/std/special/compiler_rt/udivmoddi4_test.zig
+++ b/std/special/compiler_rt/udivmoddi4_test.zig
@@ -1,3 +1,5 @@
+// Disable formatting to avoid unnecessary source repository bloat.
+// zig fmt: off
 const __udivmoddi4 = @import("udivmoddi4.zig").__udivmoddi4;
 const assert = @import("std").debug.assert;
 
diff --git a/std/special/compiler_rt/udivmodti4_test.zig b/std/special/compiler_rt/udivmodti4_test.zig
index 48d65b43c6..f6b370c26e 100644
--- a/std/special/compiler_rt/udivmodti4_test.zig
+++ b/std/special/compiler_rt/udivmodti4_test.zig
@@ -1,3 +1,5 @@
+// Disable formatting to avoid unnecessary source repository bloat.
+// zig fmt: off
 const __udivmodti4 = @import("udivmodti4.zig").__udivmodti4;
 const assert = @import("std").debug.assert;
 
diff --git a/std/zig/parser_test.zig b/std/zig/parser_test.zig
index c28a70b770..91a56de827 100644
--- a/std/zig/parser_test.zig
+++ b/std/zig/parser_test.zig
@@ -1,3 +1,29 @@
+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: pointer of unknown length" {
     try testCanonical(
         \\fn foo(ptr: [*]u8) void {}
diff --git a/std/zig/render.zig b/std/zig/render.zig
index 147adc6221..7c9b53b77a 100644
--- a/std/zig/render.zig
+++ b/std/zig/render.zig
@@ -82,8 +82,45 @@ fn renderRoot(
 
     var start_col: usize = 0;
     var it = tree.root_node.decls.iterator(0);
-    while (it.next()) |decl| {
-        try renderTopLevelDecl(allocator, stream, tree, 0, &start_col, decl.*);
+    while (true) {
+        var decl = (it.next() ?? 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) {
+                Token.Id.LineComment => {},
+                Token.Id.DocComment => 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) {
+                        Token.Id.LineComment => {},
+                        Token.Id.Eof => {
+                            const start = tree.tokens.at(start_token_index + 1).start;
+                            try stream.write(tree.source[start..]);
+                            return;
+                        },
+                        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 stream.print("{}\n", tree.source[start..end_token.end]);
+                        while (tree.tokens.at(decl.firstToken()).start < end_token.end) {
+                            decl = (it.next() ?? return).*;
+                        }
+                        break :zig_fmt_loop;
+                    }
+                }
+            }
+        }
+
+        try renderTopLevelDecl(allocator, stream, tree, 0, &start_col, decl);
         if (it.peek()) |next_decl| {
             try renderExtraNewline(tree, stream, &start_col, next_decl.*);
         }
diff --git a/test/cases/syntax.zig b/test/cases/syntax.zig
index 20fb6cd203..b497b060c4 100644
--- a/test/cases/syntax.zig
+++ b/test/cases/syntax.zig
@@ -1,4 +1,5 @@
 // Test trailing comma syntax
+// zig fmt: off
 
 const struct_trailing_comma = struct { x: i32, y: i32, };
 const struct_no_comma = struct { x: i32, y: i32 };