From 9a7b1e08679ccf163f96553f3b5c2628d887928a Mon Sep 17 00:00:00 2001
From: trdthg <69898423+trdthg@users.noreply.github.com>
Date: Sat, 20 Jul 2024 07:52:48 +0800
Subject: [PATCH] Fmt: wrap exps into oneline if possible (#619)

* FIx sail-fmt block_comment indent

* fmt

* update test case

* add test case

* fmt: wrap exps into oneline if possibly
---
 src/lib/format_sail.ml                       | 70 +++++++++++++++
 test/format/comments.sail                    | 44 +++++++++
 test/format/default/comments.expect          | 39 ++++++++
 test/format/default/doc_comment.expect       |  9 ++
 test/format/default/function_quant.expect    | 20 ++---
 test/format/default/struct_update.expect     | 16 +---
 test/format/default/wrap_into_oneline.expect | 93 ++++++++++++++++++++
 test/format/doc_comment.sail                 | 15 +++-
 test/format/function_quant.sail              | 25 +++---
 test/format/wrap_into_oneline.sail           | 27 ++++++
 10 files changed, 313 insertions(+), 45 deletions(-)
 create mode 100644 test/format/default/wrap_into_oneline.expect
 create mode 100644 test/format/wrap_into_oneline.sail

diff --git a/src/lib/format_sail.ml b/src/lib/format_sail.ml
index 1f777a96b..5673b9fd2 100644
--- a/src/lib/format_sail.ml
+++ b/src/lib/format_sail.ml
@@ -218,8 +218,55 @@ module PPrintWrapper = struct
 
   let lines s = List.map string (Util.split_on_char '\n' s)
 
+  let count_indent line =
+    let rec loop i = if i < String.length line && line.[i] = ' ' then loop (i + 1) else i in
+    loop 0
+
+  let rtrim str =
+    let len = String.length str in
+    let rec find_end i =
+      if i < 0 then 0
+      else if str.[i] = ' ' || str.[i] = '\t' || str.[i] = '\n' || str.[i] = '\r' then find_end (i - 1)
+      else i + 1
+    in
+    let new_len = find_end (len - 1) in
+    String.sub str 0 new_len
+
+  let count_lines_min_indent lines =
+    let rec loop min_indent lines =
+      match lines with
+      | line :: rest_of_lines ->
+          (* Ignore empty line *)
+          if line = "" then loop min_indent rest_of_lines
+          else (
+            let indent = count_indent line in
+            let new_min_indent = min indent min_indent in
+            loop new_min_indent rest_of_lines
+          )
+      | [] -> min_indent
+    in
+    match lines with _ :: xs -> loop max_int xs | _ -> 0
+
+  let patch_comment_lines_indent col lines =
+    let min_indent = count_lines_min_indent lines in
+    let right_indent_count = col - min_indent in
+    let lines =
+      List.mapi
+        (fun i l ->
+          (* The first_line or empty_line remains unchanged *)
+          if i == 0 || l = "" then l
+          else if right_indent_count > 0 then String.make (abs right_indent_count) ' ' ^ l
+          else l
+        )
+        lines
+    in
+    lines
+
   let block_comment_lines col s =
     let lines = Util.split_on_char '\n' s in
+    (* Last line (before */) shouldn't be rtrimed *)
+    let lines = List.mapi (fun i l -> if i + 1 = List.length lines then l else rtrim l) lines in
+    let lines = patch_comment_lines_indent col lines in
     List.mapi
       (fun n line ->
         if n = 0 || col > String.length line then string line
@@ -393,6 +440,26 @@ module type CONFIG = sig
   val config : config
 end
 
+let rec can_chunks_list_wrap cqs =
+  match cqs with
+  | [] -> true
+  | [cq] -> (
+      match List.of_seq (Queue.to_seq cq) with
+      | [] -> true
+      | [c] -> (
+          match c with
+          (* Atom is ok *)
+          | Atom _ -> true
+          (* {{{ Atom }}} is ok *)
+          | Block (_, exps) -> can_chunks_list_wrap exps
+          | If_then_else (_, i, t, e) -> can_chunks_list_wrap [t; e]
+          | _ -> false
+        )
+      | c :: cq ->
+          can_chunks_list_wrap [Queue.of_seq (List.to_seq [c])] && can_chunks_list_wrap [Queue.of_seq (List.to_seq cq)]
+    )
+  | cq :: cqs -> can_chunks_list_wrap [cq] && can_chunks_list_wrap cqs
+
 module Make (Config : CONFIG) = struct
   let indent = Config.config.indent
   let preserve_structure = Config.config.preserve_structure
@@ -589,6 +656,9 @@ module Make (Config : CONFIG) = struct
              )
     | Pragma (pragma, arg) -> char '$' ^^ string pragma ^^ space ^^ string arg ^^ hardline
     | Block (always_hardline, exps) ->
+        let always_hardline =
+          match exps with [x] -> if can_chunks_list_wrap exps then false else always_hardline | _ -> always_hardline
+        in
         let exps =
           map_last
             (fun no_semi chunks -> doc_block_exp_chunks (opts |> nonatomic |> statement_like) no_semi chunks)
diff --git a/test/format/comments.sail b/test/format/comments.sail
index ede91e59d..a7d0993eb 100644
--- a/test/format/comments.sail
+++ b/test/format/comments.sail
@@ -1,4 +1,48 @@
+/*comment*/
+/*      comment          */
+/* first line comment
+                        */
+/* first line comment
+
+           */
+/* first line comment
+
+        last line comment         */
+
+/* first line comment
+             indent many
+           last line comment*/
+
+          /* first line comment
+      indent many
+           last line comment*/
+
+        /* first line comment
+             indent many
+      last line comment*/
 function a () -> int = {
+/*comment*/
+/*      comment          */
+/* first line comment
+                        */
+/* first line comment
+                    
+           */
+/* first line comment
+
+        last line comment         */
+
+/* first line comment
+             indent many
+           last line comment*/
+
+          /* first line comment
+      indent many
+           last line comment*/
+
+        /* first line comment
+             indent many
+      last line comment*/
     *R = baz; // comment
        // comment
      1// comment
diff --git a/test/format/default/comments.expect b/test/format/default/comments.expect
index ab2171cb2..b139155ee 100644
--- a/test/format/default/comments.expect
+++ b/test/format/default/comments.expect
@@ -1,4 +1,43 @@
+/*comment*/
+/*      comment          */
+/* first line comment
+                        */
+
+/* first line comment
+
+           */
+/* first line comment
+
+        last line comment         */
+/* first line comment
+             indent many
+           last line comment*/
+/* first line comment
+indent many
+     last line comment*/
+/* first line comment
+       indent many
+last line comment*/
 function a () -> int = {
+    /*comment*/
+    /*      comment          */
+    /* first line comment
+                            */
+    /* first line comment
+
+               */
+    /* first line comment
+
+            last line comment         */
+    /* first line comment
+                 indent many
+               last line comment*/
+    /* first line comment
+    indent many
+         last line comment*/
+    /* first line comment
+           indent many
+    last line comment*/
     *R = baz; // comment
 
     // comment
diff --git a/test/format/default/doc_comment.expect b/test/format/default/doc_comment.expect
index 92ec7f69b..e20061767 100644
--- a/test/format/default/doc_comment.expect
+++ b/test/format/default/doc_comment.expect
@@ -1,4 +1,13 @@
 default Order dec
 
+/*! first line comment
+             indent many
+           last line comment*/
+/*! first line comment
+      indent many
+           last line comment*/
+/*! first line comment
+             indent many
+      last line comment*/
 /*! A documentation comment */
 val main : unit -> unit
diff --git a/test/format/default/function_quant.expect b/test/format/default/function_quant.expect
index 5d0e0501f..b3bab1c1b 100644
--- a/test/format/default/function_quant.expect
+++ b/test/format/default/function_quant.expect
@@ -3,26 +3,16 @@ default Order dec
 $include <prelude.sail>
 
 function foo forall 'n 'm, 'n >= 0.
-(x : int('n), y : int('m)) -> unit = {
-    ()
-}
+(x : int('n), y : int('m)) -> unit = { () }
 
 function foo /* c */ forall 'n 'm, 'n >= 0.
-(x : int('n), y : int('m)) -> unit = {
-    ()
-}
+(x : int('n), y : int('m)) -> unit = { () }
 
 function foo forall /* c */ 'n 'm, 'n >= 0.
-(x : int('n), y : int('m)) -> unit = {
-    ()
-}
+(x : int('n), y : int('m)) -> unit = { () }
 
 function foo forall 'n 'm, /* c */ 'n >= 0.
-(x : int('n), y : int('m)) -> unit = {
-    ()
-}
+(x : int('n), y : int('m)) -> unit = { () }
 
 function foo forall 'very_long_identifier_that_will_cause_a_line_break 'm, 'n >= 0.
-(x : int('very_long_identifier_that_will_cause_a_line_break), y : int('m)) -> unit = {
-    ()
-}
+(x : int('very_long_identifier_that_will_cause_a_line_break), y : int('m)) -> unit = { () }
diff --git a/test/format/default/struct_update.expect b/test/format/default/struct_update.expect
index 88273bc44..8ecc92c38 100644
--- a/test/format/default/struct_update.expect
+++ b/test/format/default/struct_update.expect
@@ -44,22 +44,14 @@ enum E with f -> very_long_type_that_will_trigger_a_linebreak,
 }
 
 function has_loops () = {
-    foreach (n from 1 to 3) {
-        ()
-    };
-    foreach (n from 3 downto 1) {
-        ()
-    };
-    foreach (n from 0 to 4 by 2) {
-        ()
-    };
+    foreach (n from 1 to 3) { () };
+    foreach (n from 3 downto 1) { () };
+    foreach (n from 0 to 4 by 2) { () };
     foreach (n from 10000000000000000000000000000000 to 444444444444444444444444444444444444444444444444444444444444) {
         ()
     };
     while true do ();
-    while true do {
-        ()
-    };
+    while true do { () };
     repeat termination_measure { foo } () until true
 }
 
diff --git a/test/format/default/wrap_into_oneline.expect b/test/format/default/wrap_into_oneline.expect
new file mode 100644
index 000000000..c6fbcd8c1
--- /dev/null
+++ b/test/format/default/wrap_into_oneline.expect
@@ -0,0 +1,93 @@
+function f b = if b == bitone then { bitzero } else { bitone }
+function f b = if b == bitone then { bitzero } else { bitone }
+function f b = if b == bitone then {
+    let a = 1;
+    bitzero
+} else { bitone }
+
+function f b = if b == bitone then { { bitzero } } else { bitone }
+function f b = if b == bitone then { { { bitzero } } } else { bitone }
+function f b = if b == bitone then {
+    {
+        let a = 1;
+        bitzero
+    }
+} else { bitone }
+function f b = if b == bitone then {
+    {
+        {
+            let a = 1;
+            bitzero
+        }
+    }
+} else { bitone }
+function f b = if b == bitone then {
+    {
+        {
+            {
+                let a = 1;
+                bitzero
+            }
+        }
+    }
+} else {
+    {
+        {
+            {
+                let a = 1;
+                bitone
+            }
+        }
+    }
+}
+function f b = if b == bitone then {
+    {
+        {
+            {
+                let a = 1;
+                bitzero
+            }
+        }
+    }
+} else { { { bitone } } }
+function f b = if b == bitone then {
+    {
+        {
+            {
+                let a = 1;
+                bitzero
+            }
+        }
+    }
+} else {
+    {
+        let a = 1;
+        { bitone }
+    }
+}
+
+/* comment */
+function f /* comment */ b = if b == bitone then { bitzero } else { bitone }
+function f /* comment */ b = if b == bitone then { bitzero } else { bitone }
+function f b = /* comment */ if b == bitone then { bitzero } else { bitone }
+function f b = /* comment */ if b == bitone then { bitzero } else { bitone }
+
+// TODO function f b =/* comment */if b == bitone then bitzero else bitone
+function f b = if /* comment */ b == bitone then { bitzero } else { bitone }
+function f b = if b == /* comment */ bitone then { bitzero } else { bitone }
+function f b = if b == /* comment */ bitone then { bitzero } else { bitone }
+
+// TODO function f b = if b ==/* comment */bitone then bitzero else bitone
+function f b = if b == bitone /* comment */  then { bitzero } else { bitone }
+function f b = if b == bitone then {
+    /* comment */ bitzero
+} else { bitone }
+function f b = if b == bitone then {
+    bitzero /* comment */
+} else { bitone }
+function f b = if b == bitone then {
+    bitzero /* comment */
+} else { bitone }
+function f b = if b == bitone then {
+    bitzero /* comment */
+} else { bitone }
diff --git a/test/format/doc_comment.sail b/test/format/doc_comment.sail
index 7bbc77150..94f5671bb 100644
--- a/test/format/doc_comment.sail
+++ b/test/format/doc_comment.sail
@@ -1,8 +1,17 @@
 default Order dec
 
-/*! A documentation comment */
+/*! first line comment
+             indent many
+           last line comment*/
 
+          /*! first line comment
+      indent many
+           last line comment*/
 
+        /*! first line comment
+             indent many
+      last line comment*/
 
-val main :
-unit -> unit
+/*! A documentation comment */
+val main : 
+    unit -> unit
diff --git a/test/format/function_quant.sail b/test/format/function_quant.sail
index d30288de3..b3bab1c1b 100644
--- a/test/format/function_quant.sail
+++ b/test/format/function_quant.sail
@@ -2,22 +2,17 @@ default Order dec
 
 $include <prelude.sail>
 
-function foo forall 'n 'm, 'n >= 0. (x: int('n), y: int('m)) -> unit = {
-    ()
-}
+function foo forall 'n 'm, 'n >= 0.
+(x : int('n), y : int('m)) -> unit = { () }
 
-function foo /* c */ forall 'n 'm, 'n >= 0. (x: int('n), y: int('m)) -> unit = {
-    ()
-}
+function foo /* c */ forall 'n 'm, 'n >= 0.
+(x : int('n), y : int('m)) -> unit = { () }
 
-function foo forall/* c */'n 'm, 'n >= 0. (x: int('n), y: int('m)) -> unit = {
-    ()
-}
+function foo forall /* c */ 'n 'm, 'n >= 0.
+(x : int('n), y : int('m)) -> unit = { () }
 
-function foo forall 'n   'm, /* c */ 'n >= 0. (x: int('n), y: int('m)) -> unit = {
-    ()
-}
+function foo forall 'n 'm, /* c */ 'n >= 0.
+(x : int('n), y : int('m)) -> unit = { () }
 
-function foo forall 'very_long_identifier_that_will_cause_a_line_break 'm, 'n >= 0. (x: int('very_long_identifier_that_will_cause_a_line_break), y: int('m)) -> unit = {
-    ()
-}
+function foo forall 'very_long_identifier_that_will_cause_a_line_break 'm, 'n >= 0.
+(x : int('very_long_identifier_that_will_cause_a_line_break), y : int('m)) -> unit = { () }
diff --git a/test/format/wrap_into_oneline.sail b/test/format/wrap_into_oneline.sail
new file mode 100644
index 000000000..0342c711b
--- /dev/null
+++ b/test/format/wrap_into_oneline.sail
@@ -0,0 +1,27 @@
+function f b = if b == bitone then bitzero else bitone
+function f b = if b == bitone then { bitzero } else { bitone }
+function f b = if b == bitone then { let a = 1;bitzero } else { bitone }
+
+function f b = if b == bitone then {  {bitzero } } else { bitone }
+function f b = if b == bitone then {  {{bitzero} } } else { bitone }
+function f b = if b == bitone then { { let a = 1; bitzero } } else { bitone }
+function f b = if b == bitone then { { {let a = 1; bitzero} } } else { bitone }
+function f b = if b == bitone then { { { {let a = 1; bitzero}} } } else { {{{let a= 1; bitone}}} }
+function f b = if b == bitone then { { { {let a = 1; bitzero}} } } else { {{bitone }}}
+function f b = if b == bitone then { { { {let a = 1; bitzero}} } } else { {let a = 1;{bitone }}}
+
+ /* comment */
+function/* comment */f b = if b == bitone then bitzero else bitone
+function f/* comment */b = if b == bitone then bitzero else bitone
+function f b/* comment */= if b == bitone then bitzero else bitone
+function f b = /* comment */if b == bitone then bitzero else bitone
+// TODO function f b =/* comment */if b == bitone then bitzero else bitone
+function f b = if/* comment */b == bitone then bitzero else bitone
+function f b = if b/* comment */== bitone then bitzero else bitone
+function f b = if b == /* comment */ bitone then bitzero else bitone
+// TODO function f b = if b ==/* comment */bitone then bitzero else bitone
+function f b = if b == bitone/* comment */then bitzero else bitone
+function f b = if b == bitone then/* comment */bitzero else bitone
+function f b = if b == bitone then bitzero/* comment */else bitone
+function f b = if b == bitone then bitzero else/* comment */bitone
+function f b = if b == bitone then bitzero else bitone/* comment */