feat: flatten tagged text

Motivation: in JavaScript `JSON.stringify` does not like deeply nested
objects (it blows the stack). With a flatter structure we can show
longer outputs without running into this issue.

src/Lean/Widget/InteractiveDiagnostics.lean

@@ -75,22 +75,9 @@ private def msgToInteractive (msgData : MessageData) : IO (TaggedText MsgEmbed)
         { ef := subTt.map fun n => ⟨fun () => e.runMetaM (e.traverse n)⟩ }
         (TaggedText.text "")
 
-/-- Remove tags, leaving just the pretty-printed string. -/
-partial def TaggedText.stripTags (tt : TaggedText α) : String :=
-  go "" [tt]
-where go (acc : String) : List (TaggedText α) → String
-  | []               => acc
-  | text s :: ts     => go (acc ++ s) ts
-  | append a b :: ts => go acc (a :: b :: ts)
-  | tag _ a :: ts    => go acc (a :: ts)
-
-partial def TaggedText.stripTags₂ (tt : TaggedText (MsgEmbed)) : String :=
-  go "" [tt]
-where go (acc : String) : List (TaggedText (MsgEmbed)) → String
-  | []               => acc
-  | text s :: ts     => go (acc ++ s) ts
-  | append a b :: ts => go acc (a :: b :: ts)
-  | tag ⟨et⟩ _ :: ts => go acc (text et.stripTags :: ts)
+private partial def stripEmbeds (tt : TaggedText MsgEmbed) : String :=
+  let tt : TaggedText MsgEmbed := tt.rewrite fun ⟨et⟩ _ => TaggedText.text et.stripTags
+  tt.stripTags
 
 /-- Transform a Lean Message concerning the given text into an LSP Diagnostic. -/
 def msgToDiagnostic (text : FileMap) (m : Message) : ReaderT RpcSession IO Diagnostic := do
@@ -120,7 +107,7 @@ def msgToDiagnostic (text : FileMap) (m : Message) : ReaderT RpcSession IO Diagn
     fullRange := fullRange
     severity? := severity
     source? := source
-    message := tt.stripTags₂
+    message := stripEmbeds tt
     taggedMsg? := ttJson
   }
 

src/Lean/Widget/TaggedText.lean

@@ -15,57 +15,88 @@ Much like Lean 3 [`sf`](https://github.com/leanprover-community/mathlib/blob/bfa
 but with indentation already stringified. -/
 inductive TaggedText (α : Type u) where
   | text (s : String)
-  | append (a b : TaggedText α)
+  /- Invariant: non-empty and never directly nested. `append #[tag _ append #[]]` is okay. -/
+  | append (as : Array (TaggedText α))
   | tag (t : α) (a : TaggedText α)
-  deriving Inhabited, BEq, FromJson, ToJson
+  deriving Inhabited, BEq
 
 namespace TaggedText
 
--- TODO(WN): pick one depending on perf
-private partial def append' : TaggedText α → TaggedText α → TaggedText α
-  | text s₁,            text s₂            => text (s₁ ++ s₂)
-  | text s₁,            append (text s₂) b => append' (text <| s₁ ++ s₂) b
-  | append a (text s₁), text s₂            => append' a (text <| s₁ ++ s₂)
-  | append a (text s₁), append (text s₂) b => append' a (append' (text <| s₁ ++ s₂) b)
-  | a,                  b                  => append a b
-
-private partial def append'' : TaggedText α → TaggedText α → TaggedText α
-  | text s₁,            text s₂            => text (s₁ ++ s₂)
-  | text s₁,            append (text s₂) b => append (text <| s₁ ++ s₂) b
-  | append a (text s₁), text s₂            => append a (text <| s₁ ++ s₂)
-  | append a (text s₁), append (text s₂) b => append a (append (text <| s₁ ++ s₂) b)
-  | a,                  b                  => append a b
-
-instance : Append (TaggedText α) where
-  append := append'
-
-def map (f : α → β) : TaggedText α → TaggedText β :=
+partial def fromJson? [FromJson α] (j : Json) : Except String (TaggedText α) :=
+  (do
+    let j ← j.getObjVal? "text"
+    let s ← j.getObjValAs? String "s"
+    text s
+  ) <|>
+  (do
+    let _ : FromJson (TaggedText α) := ⟨fromJson?⟩
+    let j ← j.getObjVal? "append"
+    let as ← j.getObjValAs? (Array (TaggedText α)) "as"
+    append as
+  ) <|>
+  (do
+    let _ : FromJson (TaggedText α) := ⟨fromJson?⟩
+    let j ← j.getObjVal? "tag"
+    let t ← j.getObjValAs? α "t"
+    let a ← j.getObjValAs? (TaggedText α) "a"
+    tag t a
+  )
+
+instance [FromJson α] : FromJson (TaggedText α) :=
+  ⟨fromJson?⟩
+
+partial def toJson [ToJson α] : TaggedText α → Json
+  | text s    => Json.mkObj [("text", Json.mkObj [("s", s)])]
+  | append as => Json.mkObj [("append", Json.mkObj [("as", Json.arr <| as.map toJson)])]
+  | tag t a   => Json.mkObj [("tag", Json.mkObj [("t", ToJson.toJson t), ("a", toJson a)])]
+
+instance [ToJson α] : ToJson (TaggedText α) :=
+  ⟨toJson⟩
+
+def appendText (s₀ : String) : TaggedText α → TaggedText α
+  | text s    => text (s ++ s₀)
+  | append as => match as.back with
+    | text s => append <| as.set! (as.size - 1) <| text (s ++ s₀)
+    | a      => append <| as.push (text s₀)
+  | a         => append #[a, text s₀]
+
+def appendTag (acc : TaggedText α) (t₀ : α) (a₀ : TaggedText α) : TaggedText α :=
+  match acc with
+  | append as => append (as.push <| tag t₀ a₀)
+  | a         => append #[a, tag t₀ a₀]
+
+partial def toString [ToString α] : TaggedText α → String
+  | text s => s
+  | append as => " ++ ".intercalate (as.toList.map toString)
+  | tag t a => s!"([{t}] {toString a})"
+
+partial def map (f : α → β) : TaggedText α → TaggedText β :=
   go
 where go : TaggedText α → TaggedText β
   | text s => text s
-  | append a b => append (go a) (go b)
+  | append as => append (as.map go)
   | tag t a => tag (f t) (go a)
 
-def mapM [Monad m] (f : α → m β) : TaggedText α → m (TaggedText β) :=
+partial def mapM [Monad m] (f : α → m β) : TaggedText α → m (TaggedText β) :=
   go
 where go : TaggedText α → m (TaggedText β)
   | text s => text s
-  | append a b => do append (← go a) (← go b)
+  | append as => do append (← as.mapM go)
   | tag t a => do tag (← f t) (← go a)
 
-def rewrite (f : α → TaggedText α → TaggedText β) : TaggedText α → TaggedText β :=
+partial def rewrite (f : α → TaggedText α → TaggedText β) : TaggedText α → TaggedText β :=
   go
 where go : TaggedText α → TaggedText β
   | text s => text s
-  | append a b => append (go a) (go b)
+  | append as => append (as.map go)
   | tag t a => f t a
 
 /-- Like `mapM` but allows rewriting the whole subtree at `tag` nodes. -/
-def rewriteM [Monad m] (f : α → TaggedText α → m (TaggedText β)) : TaggedText α → m (TaggedText β) :=
+partial def rewriteM [Monad m] (f : α → TaggedText α → m (TaggedText β)) : TaggedText α → m (TaggedText β) :=
   go
 where go : TaggedText α → m (TaggedText β)
   | text s => text s
-  | append a b => do append (← go a) (← go b)
+  | append as => do append (← as.mapM go)
   | tag t a => f t a
 
 instance [RpcEncoding α β] : RpcEncoding (TaggedText α) (TaggedText β) where
@@ -77,15 +108,25 @@ private structure TaggedState where
   column : Nat            := 0
 
 instance : Std.Format.MonadPrettyFormat (StateM TaggedState) where
-  pushOutput s       := modify fun ⟨out, col⟩ => ⟨out ++ text s, col + s.length⟩
-  pushNewline indent := modify fun ⟨out, col⟩ => ⟨out ++ text ("\n".pushn ' ' indent), indent⟩
+  pushOutput s       := modify fun ⟨out, col⟩ => ⟨out.appendText s, col + s.length⟩
+  pushNewline indent := modify fun ⟨out, col⟩ => ⟨out.appendText ("\n".pushn ' ' indent), indent⟩
   currColumn         := return (←get).column
   withTag t x        := modifyGet fun ⟨currOut, currCol⟩ =>
     let ⟨ret, out, col⟩ := x { column := currCol }
-    (ret, ⟨currOut ++ tag t out, col⟩)
+    (ret, ⟨currOut.appendTag t out, col⟩)
 
 def prettyTagged (f : Format) (w : Nat := Std.Format.defWidth) : TaggedText Nat :=
   (f.prettyM w : StateM TaggedState Unit) {} |>.snd.out
 
+/-- Remove tags, leaving just the pretty-printed string. -/
+partial def stripTags (tt : TaggedText α) : String :=
+  go "" #[tt]
+where go (acc : String) : Array (TaggedText α) → String
+  | #[] => acc
+  | ts  => match ts.back with
+    | text s    => go (acc ++ s) ts.pop
+    | append as => go acc (ts.pop ++ as)
+    | tag _ a   => go acc (ts.set! (ts.size - 1) a)
+
 end TaggedText
 end Lean.Widget