Printer: make propagateForcedBreaks allocation-free (#325)

Instead of returning a copy of the doc with some `IfBreak`s resolved, make doc mutable to mark when an `IfBreak` has been resolved to the first branch.

This makes the function allocation-free, but with side effects on the doc.

Author: cristianoc

src/res_doc.ml

@@ -15,10 +15,10 @@ type t =
   | Text of string
   | Concat of t list
   | Indent of t
-  | IfBreaks of {yes: t; no: t}
+  | IfBreaks of {yes: t; no: t; mutable broken: bool} (* when broken is true, treat as the yes branch *)
   | LineSuffix of t
   | LineBreak of lineStyle
-  | Group of {shouldBreak: bool; doc: t}
+  | Group of {mutable shouldBreak: bool; doc: t}
   | CustomLayout of t list
   | BreakParent
 
@@ -43,7 +43,7 @@ let rec _concat acc l =
 let concat l = Concat(_concat [] l)
 
 let indent d = Indent d
-let ifBreaks t f = IfBreaks {yes = t; no = f}
+let ifBreaks t f = IfBreaks {yes = t; no = f; broken = false}
 let lineSuffix d = LineSuffix d
 let group d = Group {shouldBreak = false; doc = d}
 let breakableGroup ~forceBreak d = Group {shouldBreak = forceBreak; doc = d}
@@ -66,55 +66,52 @@ let rbracket = Text "]"
 let question = Text "?"
 let tilde = Text "~"
 let equal = Text "="
-let trailingComma = IfBreaks {yes = comma; no = nil}
+let trailingComma = ifBreaks comma nil
 let doubleQuote = Text "\""
 
 let propagateForcedBreaks doc =
   let rec walk doc = match doc with
   | Text _ | Nil | LineSuffix _ ->
-    (false, doc)
+    false
   | BreakParent ->
-    (true, Nil)
+    true
   | LineBreak (Hard | Literal) ->
-    (true, doc)
+    true
   | LineBreak (Classic | Soft) ->
-    (false, doc)
+    false
   | Indent children ->
-    let (childForcesBreak, newChildren) = walk children in
-    (childForcesBreak, Indent newChildren)
-  | IfBreaks {yes = trueDoc; no = falseDoc} ->
-    let (falseForceBreak, falseDoc) = walk falseDoc in
+    let childForcesBreak = walk children in
+    childForcesBreak
+  | IfBreaks ({yes = trueDoc; no = falseDoc} as ib) ->
+    let falseForceBreak = walk falseDoc in
     if falseForceBreak then
-      let (_, trueDoc) = walk trueDoc in
-      (true, trueDoc)
+      let _ = walk trueDoc in
+      ib.broken <- true;
+      true
     else
-      let forceBreak, trueDoc = walk trueDoc in
-      (forceBreak, IfBreaks {yes = trueDoc; no = falseDoc})
-  | Group {shouldBreak = forceBreak; doc = children} ->
-    let (childForcesBreak, newChildren) = walk children in
+      let forceBreak = walk trueDoc in
+      forceBreak
+  | Group ({shouldBreak = forceBreak; doc = children} as gr) ->
+    let childForcesBreak = walk children in
     let shouldBreak = forceBreak || childForcesBreak in
-    (shouldBreak, Group {shouldBreak; doc = newChildren})
+    gr.shouldBreak <- shouldBreak;
+    shouldBreak
   | Concat children ->
-    let (forceBreak, newChildren) = List.fold_left (fun (forceBreak, newChildren) child ->
-      let (childForcesBreak, newChild) = walk child in
-      (forceBreak || childForcesBreak, newChild::newChildren)
-    ) (false, []) children
-    in
-    (forceBreak, Concat (List.rev newChildren))
+    List.fold_left (fun forceBreak child ->
+      let childForcesBreak = walk child in
+      forceBreak || childForcesBreak
+    ) false children
   | CustomLayout children ->
     (* When using CustomLayout, we don't want to propagate forced breaks
      * from the children up. By definition it picks the first layout that fits
      * otherwise it takes the last of the list.
      * However we do want to propagate forced breaks in the sublayouts. They
      * might need to be broken. We just don't propagate them any higher here *)
-    let children = match walk (Concat children) with
-    | (_, Concat children) -> children
-    | _ -> assert false
-    in
-    (false, CustomLayout children)
+    let _ = walk (Concat children) in
+    false
   in
-  let (_, processedDoc) = walk doc in
-  processedDoc
+  let _ = walk doc in
+  ()
 
 (* See documentation in interface file *)
 let rec willBreak doc = match doc with
@@ -153,6 +150,7 @@ let fits w stack =
     | Break, LineBreak _ -> result := Some true
     | _, Group {shouldBreak = true; doc} -> calculate indent Break doc
     | _, Group {doc} -> calculate indent mode doc
+    | _, IfBreaks {yes = breakDoc; broken = true} -> calculate indent mode breakDoc
     | Break, IfBreaks {yes = breakDoc} -> calculate indent mode breakDoc
     | Flat, IfBreaks {no = flatDoc} -> calculate indent mode flatDoc
     | _, Concat docs -> calculateConcat indent mode docs
@@ -180,7 +178,7 @@ let fits w stack =
   calculateAll stack
 
 let toString ~width doc =
-  let doc = propagateForcedBreaks doc in
+  propagateForcedBreaks doc;
   let buffer = MiniBuffer.create 1000 in
 
   let rec process ~pos lineSuffices stack =
@@ -199,6 +197,8 @@ let toString ~width doc =
         process ~pos lineSuffices (List.append ops rest)
       | Indent doc ->
         process ~pos lineSuffices ((ind + 2, mode, doc)::rest)
+      | IfBreaks {yes = breakDoc; broken = true} ->
+        process ~pos lineSuffices ((ind, mode, breakDoc)::rest)
       | IfBreaks {yes = breakDoc; no = flatDoc} ->
         if mode = Break then
           process ~pos lineSuffices ((ind, mode, breakDoc)::rest)
@@ -309,6 +309,7 @@ let debug t =
           softLine;
           text ")";
         ]
+    | IfBreaks {yes = trueDoc; broken = true} -> toDoc trueDoc
     | IfBreaks {yes = trueDoc; no = falseDoc} ->
       group(
         concat [