From 79b6802b8f06bb141f1cc5de2e83f476e9a31bae Mon Sep 17 00:00:00 2001
From: Sorawee Porncharoenwase <sorawee.pwase@gmail.com>
Date: Wed, 13 Mar 2024 07:47:19 +0700
Subject: [PATCH] fix: wrong renderer used due to memoization

When a document is printed multiple times, memoization is reused, as intended.
However, these printing could be done with different renderers.
Prior this commit, the memoization also remembers the renderer that was used.
Therefore, subsequent printing could use a stale, memoized renderer.
This commit fixes the issue by threading the renderer through the layout
function, so that it becomes "stateless" and can be memoized with no
stale state saving.

Fixes #2
---
 CHANGES.md     |  4 +++-
 lib/printer.ml | 16 ++++++++--------
 test/main.ml   | 21 ++++++++++++++++++++-
 3 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/CHANGES.md b/CHANGES.md
index ee26893..3c6e597 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -1,7 +1,9 @@
-## 0.5 (2024-02-16)
+## 0.5 (2024-03-13)
 
 * Improve performance in `two_columns` via the zipper data structure 
   and caching.
+* Fix a bug caused from a wrong renderer is used due to memoization.
+  Thanks to @zbyrn who reported the issue (#2)!
 
 ## 0.4 (2024-02-14)
 
diff --git a/lib/printer.ml b/lib/printer.ml
index 4dcaa99..a72f6c0 100644
--- a/lib/printer.ml
+++ b/lib/printer.ml
@@ -27,7 +27,7 @@ module Core (C : Signature.CostFactory) = struct
     global_id := id + 1;
     id
 
-  type measure = { last: int; cost: C.t; layout: unit -> unit }
+  type measure = { last: int; cost: C.t; layout: Signature.renderer -> unit }
 
   let (<==) (m1 : measure) (m2 : measure): bool =
     m1.last <= m2.last && C.le m1.cost m2.cost
@@ -236,9 +236,9 @@ module Core (C : Signature.CostFactory) = struct
   let (++) (m1 : measure) (m2 : measure): measure =
     { last = m2.last;
       cost = C.combine m1.cost m2.cost;
-      layout = fun () ->
-        m1.layout ();
-        m2.layout () }
+      layout = fun renderer ->
+        m1.layout renderer;
+        m2.layout renderer }
 
   let process_concat
       (process_left : measure -> measure_set)
@@ -395,11 +395,11 @@ module Core (C : Signature.CostFactory) = struct
         | Text (s, len_s) ->
           MeasureSet [{ last = c + len_s;
                         cost = C.text c len_s;
-                        layout = fun () -> render_tree renderer s }]
+                        layout = fun renderer -> render_tree renderer s }]
         | Newline _ ->
           MeasureSet [{ last = i;
                         cost = C.newline i;
-                        layout = fun () ->
+                        layout = fun renderer ->
                           renderer "\n";
                           renderer (String.make i ' ') }]
         | Concat (d1, d2) ->
@@ -421,7 +421,7 @@ module Core (C : Signature.CostFactory) = struct
         | Blank i ->
           MeasureSet [{ last = c + i;
                         cost = C.text 0 0;
-                        layout = fun () -> renderer (String.make i ' ') }]
+                        layout = fun renderer -> renderer (String.make i ' ') }]
         | Evaled ms -> ms
         | Fail -> failwith "fails to render"
       in
@@ -439,7 +439,7 @@ module Core (C : Signature.CostFactory) = struct
     (* so the memoization tables should be cleared. *)
     (* However, in OCaml, there is no need to do the same, *)
     (* since a doc is tied to a cost factory. *)
-    m.layout ();
+    m.layout renderer;
     { is_tainted ; cost = m.cost }
 end
 
diff --git a/test/main.ml b/test/main.ml
index 131351a..88181c9 100644
--- a/test/main.ml
+++ b/test/main.ml
@@ -371,6 +371,22 @@ let test_two_columns_performance () =
      |> ignore;
      true)
 
+let test_different_renderer () =
+  let w = 20 in
+  let cf = Printer.default_cost_factory ~page_width:w () in
+  let module P = Printer.Make (val cf) in
+  let open P in
+
+  let d = text "while (true) {" ^^
+          nest 4
+            (nl ^^ text "f();" ^^ nl ^^ text "if (done())" ^^
+             (let exit_d = text "exit();" in
+              (space ^^ exit_d) <|> nest 4 (nl ^^ exit_d))) ^^
+          nl ^^ text "}"
+  in
+  Alcotest.(check string) "same string" vert_layout (pretty_format d);
+  Alcotest.(check string) "same string" vert_layout (pretty_format d)
+
 let () =
   Alcotest.run "pretty expressive"
     [ "example doc",
@@ -391,4 +407,7 @@ let () =
         "regression phantom space", `Quick, test_two_columns_regression_phantom;
         "cost factory - overflow", `Quick, test_two_columns_factory_overflow;
         "cost factory - bias", `Quick, test_two_columns_factory_bias;
-        "performance", `Slow, test_two_columns_performance ] ]
+        "performance", `Slow, test_two_columns_performance ];
+
+      "regression test",
+      [ "different renderer (issue #2)", `Quick, test_different_renderer ] ]