fix: tail-recursion in Format pretty-printer

src/Init/Data/Format/Basic.lean

@@ -77,6 +77,7 @@ private def spaceUptoLine : Format → Bool → Nat → SpaceResult
 private structure WorkItem where
   f : Format
   indent : Int
+  activeTags : Nat
 
 private structure WorkGroup where
   flatten : Bool
@@ -90,10 +91,13 @@ private partial def spaceUptoLine' : List WorkGroup → Nat → SpaceResult
 
 /-- A monad in which we can pretty-print `Format` objects. -/
 class MonadPrettyFormat (m : Type → Type) where
-  pushOutput (s : String) : m Unit
+  pushOutput (s : String)    : m Unit
   pushNewline (indent : Nat) : m Unit
-  currColumn : m Nat
-  withTag {α : Type} (t : Nat) : m α → m α
+  currColumn                 : m Nat
+  /-- Start a scope tagged with `n`. -/
+  startTag                   : Nat → m Unit
+  /-- Exit the scope of `n`-many opened tags. -/
+  endTags                    : Nat → m Unit
 open MonadPrettyFormat
 
 private def pushGroup (flb : FlattenBehavior) (items : List WorkItem) (gs : List WorkGroup) (w : Nat) [Monad m] [MonadPrettyFormat m] : m (List WorkGroup) := do
@@ -111,21 +115,24 @@ private partial def be (w : Nat) [Monad m] [MonadPrettyFormat m] : List WorkGrou
   | g@{ items := i::is, .. }::gs => do
     let gs' (is' : List WorkItem) := { g with items := is' }::gs;
     match i.f with
-    | nil => be w (gs' is)
-    | tag t f =>
-      withTag t <| be w [{ g with items := [{ i with f }] }]
+    | nil =>
+      endTags i.activeTags
       be w (gs' is)
-    | append f₁ f₂ => be w (gs' ({ i with f := f₁ }::{ i with f := f₂ }::is))
+    | tag t f =>
+      startTag t
+      be w (gs' ({ i with f, activeTags := i.activeTags + 1 }::is))
+    | append f₁ f₂ => be w (gs' ({ i with f := f₁, activeTags := 0 }::{ i with f := f₂ }::is))
     | nest n f => be w (gs' ({ i with f, indent := i.indent + n }::is))
     | text s =>
       let p := s.posOf '\n'
       if p == s.bsize then
         pushOutput s
+        endTags i.activeTags
         be w (gs' is)
       else
         pushOutput (s.extract 0 p)
         pushNewline i.indent.toNat
-        let is := { i with f := s.extract (s.next p) s.bsize }::is
+        let is := { i with f := text (s.extract (s.next p) s.bsize) }::is
         -- after a hard line break, re-evaluate whether to flatten the remaining group
         pushGroup g.flb is gs w >>= be w
     | line =>
@@ -134,14 +141,17 @@ private partial def be (w : Nat) [Monad m] [MonadPrettyFormat m] : List WorkGrou
         if g.flatten then
           -- flatten line = text " "
           pushOutput " "
+          endTags i.activeTags
           be w (gs' is)
         else
           pushNewline i.indent.toNat
+          endTags i.activeTags
           be w (gs' is)
       | FlattenBehavior.fill =>
         let breakHere := do
           pushNewline i.indent.toNat
           -- make new `fill` group and recurse
+          endTags i.activeTags
           pushGroup FlattenBehavior.fill is gs w >>= be w
         -- if preceding fill item fit in a single line, try to fit next one too
         if g.flatten then
@@ -149,6 +159,7 @@ private partial def be (w : Nat) [Monad m] [MonadPrettyFormat m] : List WorkGrou
             | panic "unreachable"
           if g'.flatten then
             pushOutput " ";
+            endTags i.activeTags
             be w gs'  -- TODO: use `return`
           else
             breakHere
@@ -162,7 +173,7 @@ private partial def be (w : Nat) [Monad m] [MonadPrettyFormat m] : List WorkGrou
         pushGroup flb [{ i with f }] (gs' is) w >>= be w
 
 def prettyM (f : Format) (w : Nat) [Monad m] [MonadPrettyFormat m] : m Unit :=
-  be w [{ flb := FlattenBehavior.allOrNone, flatten := false, items := [{ f := f, indent := 0 }]}]
+  be w [{ flb := FlattenBehavior.allOrNone, flatten := false, items := [{ f := f, indent := 0, activeTags := 0 }]}]
 
 @[inline] def bracket (l : String) (f : Format) (r : String) : Format :=
   group (nest l.length $ l ++ f ++ r)
@@ -195,7 +206,8 @@ instance : MonadPrettyFormat (StateM State) where
   pushOutput s       := modify fun ⟨out, col⟩ => ⟨out ++ s, col + s.length⟩
   pushNewline indent := modify fun ⟨out, col⟩ => ⟨out ++ "\n".pushn ' ' indent, indent⟩
   currColumn         := return (←get).column
-  withTag _ x        := x
+  startTag n         := return ()
+  endTags n          := return ()
 
 /-- Pretty-print a `Format` object as a string with expected width `w`. -/
 @[export lean_format_pretty]

src/Lean/Widget/TaggedText.lean

@@ -18,7 +18,7 @@ inductive TaggedText (α : Type u) where
   /- Invariant: non-empty and never directly nested. `append #[tag _ append #[]]` is okay. -/
   | append (as : Array (TaggedText α))
   | tag (t : α) (a : TaggedText α)
-  deriving Inhabited, BEq
+  deriving Inhabited, BEq, Repr
 
 namespace TaggedText
 
@@ -63,6 +63,7 @@ def appendText (s₀ : String) : TaggedText α → TaggedText α
 def appendTag (acc : TaggedText α) (t₀ : α) (a₀ : TaggedText α) : TaggedText α :=
   match acc with
   | append as => append (as.push <| tag t₀ a₀)
+  | text ""   => tag t₀ a₀
   | a         => append #[a, tag t₀ a₀]
 
 partial def map (f : α → β) : TaggedText α → TaggedText β :=
@@ -99,16 +100,20 @@ instance [RpcEncoding α β] : RpcEncoding (TaggedText α) (TaggedText β) where
   rpcDecode a := a.mapM rpcDecode
 
 private structure TaggedState where
-  out    : TaggedText Nat := TaggedText.text ""
-  column : Nat            := 0
+  out      : TaggedText Nat              := TaggedText.text ""
+  tagStack : List (Nat × TaggedText Nat) := []
+  column   : Nat                         := 0
+  deriving Inhabited
 
 instance : Std.Format.MonadPrettyFormat (StateM TaggedState) where
-  pushOutput s       := modify fun ⟨out, col⟩ => ⟨out.appendText s, col + s.length⟩
-  pushNewline indent := modify fun ⟨out, col⟩ => ⟨out.appendText ("\n".pushn ' ' indent), indent⟩
+  pushOutput s       := modify fun ⟨out, ts, col⟩ => ⟨out.appendText s, ts, col + s.length⟩
+  pushNewline indent := modify fun ⟨out, ts, col⟩ => ⟨out.appendText ("\n".pushn ' ' indent), ts, indent⟩
   currColumn         := return (←get).column
-  withTag t x        := modifyGet fun ⟨currOut, currCol⟩ =>
-    let ⟨ret, out, col⟩ := x { column := currCol }
-    (ret, ⟨currOut.appendTag t out, col⟩)
+  startTag n         := modify fun ⟨out, ts, col⟩ => ⟨TaggedText.text "", (n, out) :: ts, col⟩
+  endTags n          := modify fun ⟨out, ts, col⟩ =>
+    let (ended, left) := (ts.take n, ts.drop n)
+    let out' := ended.foldl (init := out) fun acc (n, top) => top.appendTag n acc
+    ⟨out', left, col⟩
 
 def prettyTagged (f : Format) (w : Nat := Std.Format.defWidth) : TaggedText Nat :=
   (f.prettyM w : StateM TaggedState Unit) {} |>.snd.out