improve printing algorithm performance for many nested callbacks

src/res_doc.ml

@@ -21,7 +21,7 @@ type t =
   | LineSuffix of t
   | LineBreak of lineStyle
   | Group of {mutable shouldBreak: bool; doc: t}
-  | CustomLayout of t list
+  | CustomLayout of t Lazy.t list
   | BreakParent
 
 let nil = Nil
@@ -102,13 +102,16 @@ let propagateForcedBreaks doc =
           let childForcesBreak = walk child in
           forceBreak || childForcesBreak)
         false children
-    | CustomLayout 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 _ = walk (Concat children) in
+      (* walking through all variation is a serious performance hit. Since we\
+       * don't intend to propagate break up from these variations, we might as
+       * well stop here for now, and later when a variation is chosen, call
+       * this function on that variation alone. *)
       false
   in
   let _ = walk doc in
@@ -119,7 +122,8 @@ let rec willBreak doc =
   match doc with
   | LineBreak (Hard | Literal) | BreakParent | Group {shouldBreak = true} ->
     true
-  | Group {doc} | Indent doc | CustomLayout (doc :: _) -> willBreak doc
+  | Group {doc} | Indent doc -> willBreak doc
+  | CustomLayout (doc :: _) -> willBreak (Lazy.force doc)
   | Concat docs -> List.exists willBreak docs
   | IfBreaks {yes; no} -> willBreak yes || willBreak no
   | _ -> false
@@ -157,7 +161,7 @@ let fits w stack =
     | _, Concat docs -> calculateConcat indent mode docs
     | _, CustomLayout (hd :: _) ->
       (* TODO: if we have nested custom layouts, what we should do here? *)
-      calculate indent mode hd
+      calculate indent mode (Lazy.force hd)
     | _, CustomLayout [] -> ()
   and calculateConcat indent mode docs =
     if result.contents == None then
@@ -238,9 +242,12 @@ let toString ~width doc =
         let rec findGroupThatFits groups =
           match groups with
           | [] -> Nil
-          | [lastGroup] -> lastGroup
+          | [lastGroup] -> Lazy.force lastGroup
           | doc :: docs ->
-            if fits (width - pos) ((ind, Flat, doc) :: rest) then doc
+            let computedDoc = Lazy.force doc in
+            propagateForcedBreaks computedDoc;
+            if fits (width - pos) ((ind, Flat, computedDoc) :: rest) then
+              computedDoc
             else findGroupThatFits docs
         in
         let doc = findGroupThatFits docs in
@@ -291,7 +298,9 @@ let debug t =
                (concat
                   [
                     line;
-                    join ~sep:(concat [text ","; line]) (List.map toDoc docs);
+                    join
+                      ~sep:(concat [text ","; line])
+                      (List.map (fun doc -> toDoc (Lazy.force doc)) docs);
                   ]);
              line;
              text ")";

src/res_doc.mli

@@ -16,7 +16,7 @@ val breakableGroup : forceBreak:bool -> t -> t
 (* `customLayout docs` will pick the layout that fits from `docs`.
  * This is a very expensive computation as every layout from the list
  * will be checked until one fits. *)
-val customLayout : t list -> t
+val customLayout : t Lazy.t list -> t
 val breakParent : t
 val join : sep:t -> t list -> t
 

src/res_printer.ml

@@ -1883,20 +1883,27 @@ and printValueBinding ~recFlag vb cmtTbl i =
     if ParsetreeViewer.isSinglePipeExpr vb.pvb_expr then
       Doc.customLayout
         [
-          Doc.group
-            (Doc.concat
-               [
-                 attrs; header; patternDoc; Doc.text " ="; Doc.space; printedExpr;
-               ]);
-          Doc.group
-            (Doc.concat
-               [
-                 attrs;
-                 header;
-                 patternDoc;
-                 Doc.text " =";
-                 Doc.indent (Doc.concat [Doc.line; printedExpr]);
-               ]);
+          lazy
+            (Doc.group
+               (Doc.concat
+                  [
+                    attrs;
+                    header;
+                    patternDoc;
+                    Doc.text " =";
+                    Doc.space;
+                    printedExpr;
+                  ]));
+          lazy
+            (Doc.group
+               (Doc.concat
+                  [
+                    attrs;
+                    header;
+                    patternDoc;
+                    Doc.text " =";
+                    Doc.indent (Doc.concat [Doc.line; printedExpr]);
+                  ]));
         ]
     else
       let shouldIndent =
@@ -3936,15 +3943,16 @@ and printArgumentsWithCallbackInFirstPosition ~uncurried args cmtTbl =
    * }, longArgumet, veryLooooongArgument)
    *)
   let fitsOnOneLine =
-    Doc.concat
-      [
-        (if uncurried then Doc.text "(. " else Doc.lparen);
-        callback;
-        Doc.comma;
-        Doc.line;
-        printedArgs;
-        Doc.rparen;
-      ]
+    lazy
+      (Doc.concat
+         [
+           (if uncurried then Doc.text "(. " else Doc.lparen);
+           callback;
+           Doc.comma;
+           Doc.line;
+           printedArgs;
+           Doc.rparen;
+         ])
   in
 
   (* Thing.map(
@@ -3954,7 +3962,7 @@ and printArgumentsWithCallbackInFirstPosition ~uncurried args cmtTbl =
    *   arg3,
    * )
    *)
-  let breakAllArgs = printArguments ~uncurried args cmtTblCopy in
+  let breakAllArgs = lazy (printArguments ~uncurried args cmtTblCopy) in
 
   (* Sometimes one of the non-callback arguments will break.
    * There might be a single line comment in there, or a multiline string etc.
@@ -3971,7 +3979,7 @@ and printArgumentsWithCallbackInFirstPosition ~uncurried args cmtTbl =
    * In this case, we always want the arguments broken over multiple lines,
    * like a normal function call.
    *)
-  if Doc.willBreak printedArgs then breakAllArgs
+  if Doc.willBreak printedArgs then Lazy.force breakAllArgs
   else Doc.customLayout [fitsOnOneLine; breakAllArgs]
 
 and printArgumentsWithCallbackInLastPosition ~uncurried args cmtTbl =
@@ -4015,27 +4023,29 @@ and printArgumentsWithCallbackInLastPosition ~uncurried args cmtTbl =
 
   (* Thing.map(foo, (arg1, arg2) => MyModuleBlah.toList(argument)) *)
   let fitsOnOneLine =
-    Doc.concat
-      [
-        (if uncurried then Doc.text "(." else Doc.lparen);
-        printedArgs;
-        callback;
-        Doc.rparen;
-      ]
+    lazy
+      (Doc.concat
+         [
+           (if uncurried then Doc.text "(." else Doc.lparen);
+           printedArgs;
+           callback;
+           Doc.rparen;
+         ])
   in
 
   (* Thing.map(longArgumet, veryLooooongArgument, (arg1, arg2) =>
    *   MyModuleBlah.toList(argument)
    * )
    *)
   let arugmentsFitOnOneLine =
-    Doc.concat
-      [
-        (if uncurried then Doc.text "(." else Doc.lparen);
-        printedArgs;
-        Doc.breakableGroup ~forceBreak:true callback2;
-        Doc.rparen;
-      ]
+    lazy
+      (Doc.concat
+         [
+           (if uncurried then Doc.text "(." else Doc.lparen);
+           printedArgs;
+           Doc.breakableGroup ~forceBreak:true callback2;
+           Doc.rparen;
+         ])
   in
 
   (* Thing.map(
@@ -4045,7 +4055,7 @@ and printArgumentsWithCallbackInLastPosition ~uncurried args cmtTbl =
    *   (param1, parm2) => doStuff(param1, parm2)
    * )
    *)
-  let breakAllArgs = printArguments ~uncurried args cmtTblCopy2 in
+  let breakAllArgs = lazy (printArguments ~uncurried args cmtTblCopy2) in
 
   (* Sometimes one of the non-callback arguments will break.
    * There might be a single line comment in there, or a multiline string etc.
@@ -4062,7 +4072,7 @@ and printArgumentsWithCallbackInLastPosition ~uncurried args cmtTbl =
    * In this case, we always want the arguments broken over multiple lines,
    * like a normal function call.
    *)
-  if Doc.willBreak printedArgs then breakAllArgs
+  if Doc.willBreak printedArgs then Lazy.force breakAllArgs
   else Doc.customLayout [fitsOnOneLine; arugmentsFitOnOneLine; breakAllArgs]
 
 and printArguments ~uncurried