Fmt: optimize fmt for let_binder, exps, if_stmt

src/lib/chunk_ast.ml

@@ -354,7 +354,14 @@ let rec prerr_chunk indent = function
           Queue.iter (prerr_chunk (indent ^ "    ")) arg
         )
         [("vars", ex.vars); ("constr", ex.constr); ("typ", ex.typ)]
-  | Binder _ -> ()
+  | Binder (binder, x, y, z) ->
+      Printf.eprintf "%sBinder:%s\n" indent (binder_keyword binder);
+      List.iteri
+        (fun i arg ->
+          Printf.eprintf "%s  %d:\n" indent i;
+          Queue.iter (prerr_chunk (indent ^ "    ")) arg
+        )
+        [x; y; z]
   | Block_binder (binder, binding, exp) ->
       Printf.eprintf "%sBlock_binder:%s\n" indent (binder_keyword binder);
       List.iter
@@ -909,6 +916,7 @@ let rec chunk_exp comments chunks (E_aux (aux, l)) =
       in
       Queue.add (Block (true, block_chunks)) chunks
   | (E_let (LB_aux (LB_val (pat, exp), _), body) | E_internal_plet (pat, exp, body)) as binder ->
+      (* there need a way to find position of '=' *)
       let binder =
         match binder with
         | E_let _ -> Let_binder

src/lib/format_sail.ml

@@ -77,7 +77,25 @@ let rec map_last f = function
       let x = f false x in
       x :: map_last f xs
 
-let line_comment_opt = function Comment (Lexer.Comment_line, _, _, contents, _trailing) -> Some contents | _ -> None
+let is_comment chunk = match chunk with Comment (_, _, _, _, _) | Doc_comment _ -> true | _ -> false
+let is_block chunk = match chunk with Block (_, _) -> true | _ -> false
+let is_match chunk = match chunk with Match _ -> true | _ -> false
+let is_let chunk = match chunk with Binder (_, _, _, _) -> true | _ -> false
+let is_delta chunk = match chunk with Delim _ | Opt_delim _ -> true | _ -> false
+
+let rec count_chunks_by_rule ?(ignore_rule = is_delta) chunks target_rule =
+  let rec go chunks target_count other_count =
+    match chunks with
+    | [] -> (target_count, other_count)
+    | hd :: tl ->
+        let target_count, other_count =
+          if target_rule hd then (target_count + 1, other_count)
+          else if ignore_rule hd then (target_count, other_count)
+          else (target_count, other_count + 1)
+        in
+        go tl target_count other_count
+  in
+  go chunks 0 0
 
 (* Remove additional (> 1) trailing newlines at the end of a string *)
 let discard_extra_trailing_newlines s =
@@ -212,6 +230,15 @@ module PPrintWrapper = struct
 
   let surround n b opening contents closing = opening ^^ Nest (n, break b ^^ contents) ^^ break b ^^ closing
 
+  let surround_hardline h n b opening contents closing =
+    let b = if h then hardline else break b in
+    group (opening ^^ nest n (b ^^ contents) ^^ b ^^ closing)
+
+  let group_indent n doc = group (nest n doc)
+  let prefix_hardline h n b contents =
+    let b = if h then hardline else break b in
+    group (nest n (b ^^ contents))
+
   let repeat n doc =
     let rec go n acc = if n = 0 then acc else go (n - 1) (doc ^^ acc) in
     go n empty
@@ -378,10 +405,6 @@ let softline = break 0
 let prefix_parens n x y =
   x ^^ ifflat space (space ^^ char '(') ^^ nest n (softline ^^ y) ^^ softline ^^ ifflat empty (char ')')
 
-let surround_hardline h n b opening contents closing =
-  let b = if h then hardline else break b in
-  group (opening ^^ nest n (b ^^ contents) ^^ b ^^ closing)
-
 type config = { indent : int; preserve_structure : bool; line_width : int; ribbon_width : float }
 
 let default_config = { indent = 4; preserve_structure = false; line_width = 120; ribbon_width = 1. }
@@ -440,31 +463,67 @@ module type CONFIG = sig
   val config : config
 end
 
-let rec can_chunks_list_wrap cqs =
+type chunk_wrap_opt = { nest : int; in_braces : bool }
+let nest_add opt = { opt with nest = opt.nest + 1 }
+let in_braces opt = { opt with in_braces = true }
+
+let rec can_chunks_list_wrap ?(opt = { nest = 0; in_braces = false }) cqs =
   match cqs with
   | [] -> true
   | [cq] -> (
+      let count = ref 0 in
       match List.of_seq (Queue.to_seq cq) with
       | [] -> true
-      | [c] -> (
+      | [c] ->
+          (* if deep > 0, then not wrap
+           * i.e.
+           * - foo(a, bar(a, b))
+           *)
+          let can_wrap_nested =
+            match c with
+            | Atom _ | String_literal _ | Delim _ | Opt_delim _ -> true
+            | Block (_, exps) -> can_chunks_list_wrap ~opt exps
+            | _ -> false
+          in
+          if not can_wrap_nested then count := !count + 1;
+          if opt.nest > 0 && !count > 1 && not can_wrap_nested then false
+          else (
+            match c with
+            (* Atom *)
+            | Atom _ | String_literal _ -> true
+            (* unit => (), *)
+            | Delim _ | Opt_delim _ -> true (* Xs[unsigned(r)], *)
+            | App (id, chunks) -> can_chunks_list_wrap ~opt:(opt |> nest_add) chunks
+            (* (v128 >> shift)[64..0] *)
+            | Binary (x, op, y) -> can_chunks_list_wrap ~opt [x] && can_chunks_list_wrap ~opt [y]
+            (* if regval[31..16] then *)
+            | Index (exp, ix) -> can_chunks_list_wrap ~opt:(opt |> in_braces) [exp] && can_chunks_list_wrap ~opt [ix]
+            (* if (a > 1) then {1} else {2} *)
+            | If_then_else (_, i, t, e) ->
+                can_chunks_list_wrap ~opt [i] && can_chunks_list_wrap ~opt [t] && can_chunks_list_wrap ~opt [e]
+            (* {{{ Atom }}} *)
+            | Block (_, exps) -> can_chunks_list_wrap ~opt:(opt |> in_braces) exps
+            (* case => (), // comment *)
+            (* then 2 // comment *)
+            | Comment (_, _, _, _, trailing) -> (not opt.in_braces) && trailing
+            | _ -> false
+          )
+      | c :: cq -> (
           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
+          | Comment (_, _, _, _, _) -> false
+          | _ ->
+              can_chunks_list_wrap ~opt [Queue.of_seq (List.to_seq [c])]
+              && can_chunks_list_wrap ~opt [Queue.of_seq (List.to_seq cq)]
         )
-      | 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
+  | cq :: cqs -> can_chunks_list_wrap ~opt [cq] && can_chunks_list_wrap ~opt cqs
 
 module Make (Config : CONFIG) = struct
   let indent = Config.config.indent
   let preserve_structure = Config.config.preserve_structure
 
-  let rec doc_chunk ?(ungroup_tuple = false) ?(toplevel = false) opts = function
+  let rec doc_chunk ?(ungroup_tuple = false) ?(toplevel = false) ?(wrap = true)
+      ?(force_block_comment_end_with_hardline = false) opts = function
     | Atom s -> string s
     | Chunks chunks -> doc_chunks opts chunks
     | Delim s -> string s ^^ space
@@ -509,7 +568,7 @@ module Make (Config : CONFIG) = struct
         let doc =
           infix indent spacing (string op)
             (doc_chunks (opts |> lhs_prec |> expression_like) lhs)
-            (doc_chunks (opts |> rhs_prec |> expression_like) rhs)
+            (doc_chunks ~force_block_comment_end_with_hardline:true (opts |> rhs_prec |> expression_like) rhs)
         in
         if outer_prec > opts.precedence then parens doc else doc
     | Ternary (x, op1, y, op2, z) ->
@@ -526,16 +585,53 @@ module Make (Config : CONFIG) = struct
         if outer_prec > opts.precedence then parens doc else doc
     | If_then_else (bracing, i, t, e) ->
         let insert_braces = opts.statement || bracing.then_brace || bracing.else_brace in
-        let i = doc_chunks (opts |> nonatomic |> expression_like) i in
-        let t =
-          if insert_braces && (not preserve_structure) && not bracing.then_brace then doc_chunk opts (Block (true, [t]))
-          else doc_chunks (opts |> nonatomic |> expression_like) t
+        let i_doc, i_wrap = doc_chunks_exps (opts |> nonatomic |> expression_like) i in
+        let t_wrap = can_chunks_list_wrap [t] in
+        let e_wrap = can_chunks_list_wrap [e] in
+        (* if any one can't wrap, then all should expand
+           i.e.
+
+              { 3 }
+
+              will expand to
+
+              {
+                3
+              }
+        *)
+        let wrap = i_wrap && t_wrap && e_wrap in
+        let t_doc, t_wrap =
+          if insert_braces && (not preserve_structure) && not bracing.then_brace then
+            (doc_chunk opts (Block (true, [t])), false)
+          else doc_chunks_exps ~wrap (opts |> nonatomic |> expression_like) t
         in
-        let e =
-          if insert_braces && (not preserve_structure) && not bracing.else_brace then doc_chunk opts (Block (true, [e]))
-          else doc_chunks (opts |> nonatomic |> expression_like) e
+        let e_doc, e_wrap =
+          if insert_braces && (not preserve_structure) && not bracing.else_brace then
+            (doc_chunk opts (Block (true, [e])), false)
+          else doc_chunks_exps ~wrap (opts |> nonatomic |> expression_like) e
         in
-        separate space [string "if"; i; string "then"; t; string "else"; e] |> atomic_parens opts
+        let doc_i = separate space [string "if"; i_doc] in
+        let doc_t = separate space [string "then"; t_doc] in
+        let doc_e = separate space [string "else"; e_doc] in
+        let sep is_wrap = if is_wrap then space else hardline in
+        (* add hardline or space
+           * - if_stmt:
+                if
+                  let a = 1 in
+                  a
+                then {
+                  // comment
+                  1
+                } else {
+                  2
+                }
+
+                - i has no braces, use hardline
+                - t has braces, use space
+                - e force to expand
+        *)
+        if wrap then doc_i ^^ sep i_wrap ^^ doc_t ^^ space ^^ doc_e |> atomic_parens opts
+        else doc_i ^^ sep i_wrap ^^ doc_t ^^ sep t_wrap ^^ doc_e |> atomic_parens opts
     | If_then (bracing, i, t) ->
         let i = doc_chunks (opts |> nonatomic |> expression_like) i in
         let t =
@@ -552,9 +648,10 @@ module Make (Config : CONFIG) = struct
           (char ']')
         |> atomic_parens opts
     | Index (exp, ix) ->
-        let exp_doc = doc_chunks (opts |> atomic |> expression_like) exp in
-        let ix_doc = doc_chunks (opts |> nonatomic |> expression_like) ix in
-        exp_doc ^^ surround indent 0 (char '[') ix_doc (char ']') |> subatomic_parens opts
+        let exp_doc, exp_wrap = doc_chunks_exps (opts |> atomic |> expression_like) exp in
+        let exp_doc = group (empty ^^ nest indent exp_doc ^^ empty) in
+        let ix_doc = doc_chunks_wrapped (opts |> nonatomic |> expression_like) ix in
+        exp_doc ^^ char '[' ^^ ix_doc ^^ char ']' |> subatomic_parens opts
     | Exists ex ->
         let ex_doc =
           doc_chunks (atomic opts) ex.vars
@@ -599,7 +696,9 @@ module Make (Config : CONFIG) = struct
                by forcing exp on a newline if the comment contains linebreaks
             *)
             match block_comment_lines col contents with
-            | [l] -> blank n ^^ string "/*" ^^ l ^^ string "*/" ^^ space
+            | [l] ->
+                let end_of_comment = if force_block_comment_end_with_hardline then require_hardline else space in
+                blank n ^^ string "/*" ^^ l ^^ string "*/" ^^ end_of_comment
             | ls -> blank n ^^ group (align (string "/*" ^^ separate hardline ls ^^ string "*/")) ^^ require_hardline
           )
       end
@@ -657,7 +756,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
+          match exps with
+          | [x] -> if not wrap then true else if can_chunks_list_wrap exps then false else always_hardline
+          | _ -> always_hardline
         in
         let exps =
           map_last
@@ -673,19 +774,17 @@ module Make (Config : CONFIG) = struct
             [string (binder_keyword binder); doc_chunks (atomic opts) x; char '='; doc_chunks (nonatomic opts) y]
         else
           separate space [string (binder_keyword binder); doc_chunks (atomic opts) x; char '=']
-          ^^ nest 4 (hardline ^^ doc_chunks (nonatomic opts) y)
+          ^^ nest indent (hardline ^^ doc_chunks (nonatomic opts) y)
     | Binder (binder, x, y, z) ->
-        prefix indent 1
-          (separate space
-             [
-               string (binder_keyword binder);
-               doc_chunks (atomic opts) x;
-               char '=';
-               doc_chunks (nonatomic opts) y;
-               string "in";
-             ]
+        let x, _ = doc_chunks_exps (atomic opts) x in
+        let y, _ = doc_chunks_exps (atomic opts) y in
+        let doc = separate space [string (binder_keyword binder); x; char '='; y; string "in"; hardline] in
+        Queue.fold
+          (fun acc chunk ->
+            let doc = doc_chunk ~force_block_comment_end_with_hardline:true opts chunk in
+            acc ^^ doc
           )
-          (doc_chunks (nonatomic opts) z)
+          doc z
     | Match m ->
         let kw1, kw2 = match_keywords m.kind in
         string kw1 ^^ space
@@ -732,9 +831,50 @@ module Make (Config : CONFIG) = struct
     | Field (exp, id) -> doc_chunks (subatomic opts) exp ^^ char '.' ^^ doc_id id
     | Raw str -> separate hardline (lines str)
 
+  (* Format chunks in [chunks], add prefix if can't wrap  *)
+  (* 1. wrapped
+      [if cond then x else y]
+  *)
+  (* 2. expanded
+      [
+        let a = 1 in
+        a
+      ]
+  *)
+  and doc_chunks_wrapped ?(ungroup_tuple = false) ?(force_hardline = false) opts chunks =
+    let doc = Queue.fold (fun doc chunk -> doc ^^ doc_chunk ~ungroup_tuple opts chunk) empty chunks in
+    let wrapped = can_chunks_list_wrap [chunks] in
+    if wrapped then doc else surround_hardline true indent 1 empty doc empty
+
+  (* Format chunks, add prefix if can't wrap  *)
+  and doc_chunks_exps ?(ungroup_tuple = false) ?(wrap = true) opts chunks =
+    let doc = Queue.fold (fun doc chunk -> doc ^^ doc_chunk ~wrap ~ungroup_tuple opts chunk) empty chunks in
+    let wrapped = can_chunks_list_wrap [chunks] in
+
+    let rec count_chunks_block_like chunks =
+      count_chunks_by_rule chunks (fun chunk -> is_block chunk || is_match chunk)
+    in
+    (* i.e.
+        no_prefix:
+          let a = match {
+            ...
+          }
+
+        prefix:
+          let f =
+            let a = 1 in
+            b
+    *)
+    let no_prefix = count_chunks_block_like (List.of_seq (Queue.to_seq chunks)) = (1, 0) in
+    let doc =
+      (* add indent for multi line exps which can't be wrapped *)
+      if wrapped || no_prefix then doc else prefix_hardline true indent 0 doc
+    in
+    (doc, wrapped)
+
   and doc_pexp_chunks_pair opts pexp =
     let pat = doc_chunks opts pexp.pat in
-    let body = doc_chunks opts pexp.body in
+    let body, _ = doc_chunks_exps opts pexp.body in
     match pexp.guard with
     | None -> (pat, body)
     | Some guard -> (separate space [pat; string "if"; doc_chunks opts guard], body)
@@ -753,9 +893,11 @@ module Make (Config : CONFIG) = struct
     ^^
     match pexp.guard with
     | None ->
+        let body_doc, _ = doc_chunks_exps opts pexp.body in
+
         (if pexp.funcl_space then space else empty)
         ^^ group (doc_chunks ~ungroup_tuple:true opts pexp.pat ^^ return_typ)
-        ^^ string "=" ^^ space ^^ doc_chunks opts pexp.body
+        ^^ string "=" ^^ space ^^ body_doc
     | Some guard ->
         parens (separate space [doc_chunks opts pexp.pat; string "if"; doc_chunks opts guard])
         ^^ return_typ ^^ string "=" ^^ space ^^ doc_chunks opts pexp.body
@@ -796,8 +938,10 @@ module Make (Config : CONFIG) = struct
     let doc = splice_into_doc chunks empty in
     (group doc, !requires_hardline)
 
-  and doc_chunks ?(ungroup_tuple = false) opts chunks =
-    Queue.fold (fun doc chunk -> doc ^^ doc_chunk ~ungroup_tuple opts chunk) empty chunks
+  and doc_chunks ?(ungroup_tuple = false) ?(force_block_comment_end_with_hardline = false) opts chunks =
+    Queue.fold
+      (fun doc chunk -> doc ^^ doc_chunk ~force_block_comment_end_with_hardline ~ungroup_tuple opts chunk)
+      empty chunks
 
   let to_string doc =
     let b = Buffer.create 1024 in

test/format/default/demo.sail

@@ -0,0 +1,6 @@
+
+function rX r = match r {
+    _ =>
+        //
+        Xs[unsigned(r)],
+}