From 998ba0d751a9a2dba6cbb63a11db612cad809c27 Mon Sep 17 00:00:00 2001
From: Sven Thiele <sthiele78@gmail.com>
Date: Mon, 3 Oct 2022 14:42:41 +0200
Subject: [PATCH] Support comments to skip formatting

The comment `% fmt: off` will turn formatting off
and `% fmt: on` turns it on again.

```
% fmt: off
This section will not be formatted
% fmt: on
```
---
 src/lib.rs   | 37 ++++++++++++++++++++++++++++++-------
 src/tests.rs | 42 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 72 insertions(+), 7 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 763d882..a9e0ec2 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -58,7 +58,8 @@ impl<'a> Formatter<'a> {
         }
         Ok(())
     }
-    fn process_comment(&mut self, buf: &[u8]) -> Result<()> {
+    /// Returns true if comment contains fmt: off statement
+    fn process_comment(&mut self, buf: &[u8]) -> Result<bool> {
         match self.state {
             No => {}
             SomeBlock | Block(StatementType::Other) => {
@@ -70,7 +71,10 @@ impl<'a> Formatter<'a> {
         write!(self.out, "{}", text.trim_end())?;
         self.state = SomeBlock;
 
-        Ok(())
+        if text.contains("fmt: off") {
+            return Ok(true);
+        }
+        Ok(false)
     }
     fn process_statement(&mut self, stmt_type: StatementType, buf: &[u8]) -> Result<()> {
         match (self.state, stmt_type) {
@@ -104,6 +108,8 @@ pub fn format_program(
     out: &mut dyn Write,
     debug: bool,
 ) -> Result<()> {
+    let mut skip_fmt = false;
+    let mut start_unformatted = 0;
     let mut formatter = Formatter {
         out,
         state: FormatterState::No,
@@ -153,16 +159,33 @@ pub fn format_program(
                 match node.kind() {
                     "statement" => {
                         let mut buf = Vec::new();
-                        let stmt_type =
-                            format_statement(&node, source_code, &mut buf, debug)?;
-
-                        formatter.process_statement(stmt_type, &buf)?;
+                        if !skip_fmt {
+                            let stmt_type = format_statement(&node, source_code, &mut buf, debug)?;
+                            formatter.process_statement(stmt_type, &buf)?;
+                        }
                         short_cut = true;
                     }
                     "single_comment" | "multi_comment" => {
                         let start_byte = node.start_byte();
                         let end_byte = node.end_byte();
-                        formatter.process_comment(&source_code[start_byte..end_byte])?;
+                        if skip_fmt {
+                            let last_text =
+                                std::str::from_utf8(&source_code[start_byte..end_byte]).unwrap();
+                            if last_text.contains("fmt: on") {
+                                let unformated_text = std::str::from_utf8(
+                                    &source_code[start_unformatted..start_byte],
+                                )
+                                .unwrap();
+                                write!(formatter.out, "{}", unformated_text)?;
+                                writeln!(formatter.out, "{}", last_text)?;
+                                skip_fmt = false;
+                            }
+                        } else if formatter.process_comment(&source_code[start_byte..end_byte])? {
+                            skip_fmt = true;
+                            start_unformatted = node.end_byte();
+                        } else {
+                            skip_fmt = false;
+                        }
                     }
                     _ => {}
                 }
diff --git a/src/tests.rs b/src/tests.rs
index 778f2ba..aca4b20 100644
--- a/src/tests.rs
+++ b/src/tests.rs
@@ -15,6 +15,48 @@ fn fmt_and_cmp(source_code: &str, res: &str) {
     assert_eq!(parse_res, res)
 }
 
+#[test]
+fn test_fmt_off_on() {
+    let source = r#"a(A):-b(A,V),c(A).
+% Comment1
+% fmt: off
+% Comment2 in fmt: off
+    index(A,I):-vary(A),I = #count{ B : vary(B),B <= A },not bounds(0,0).
+    counter(I,1)  :-index(A,I),bounds(L,U),L <= I,selected(A).counter(I,C+1):-index(A,I),bounds(L,U),C < U,selected(A),counter(I+1,C).
+    counter(I,C)  :-index(A,I),bounds(L,U),L < C+I,counter(I+1,C).
+% Comment3 in fmt: off
+% Comment4 in fmt: off
+:- bounds(L,U),0 < L,not counter(1,L).:- bounds(L,U),index(A,I),selected(A),counter(I+1,U).
+% Comment5 in fmt: off
+% fmt: on
+
+% Comment6
+exclude(M,A):-model(M),select(A,0).
+"#;
+    let result = r#"a(A) :-
+    b(A, V),
+    c(A).
+
+% Comment1
+% fmt: off
+% Comment2 in fmt: off
+    index(A,I):-vary(A),I = #count{ B : vary(B),B <= A },not bounds(0,0).
+    counter(I,1)  :-index(A,I),bounds(L,U),L <= I,selected(A).counter(I,C+1):-index(A,I),bounds(L,U),C < U,selected(A),counter(I+1,C).
+    counter(I,C)  :-index(A,I),bounds(L,U),L < C+I,counter(I+1,C).
+% Comment3 in fmt: off
+% Comment4 in fmt: off
+:- bounds(L,U),0 < L,not counter(1,L).:- bounds(L,U),index(A,I),selected(A),counter(I+1,U).
+% Comment5 in fmt: off
+% fmt: on
+
+% Comment6
+exclude(M, A) :-
+    model(M),
+    select(A, 0).
+"#;
+    fmt_and_cmp(source, result);
+}
+
 #[test]
 fn test_pass_new() {
     fmt_and_cmp(" \n \n ", "");