fix: hover and highlight info from builtin assert docstring roles (#14039)

This PR fixes a bug where the builting docstring roles for asserting
equalities did not properly highlight their contents for downstream
consumers of rich docstring info, and exposes a structure that was
mistakenly made private.

I used Claude while making this PR.

Author: david-christiansen

src/Lean/Elab/DocString/Builtin.lean

@@ -115,13 +115,22 @@ structure Data.ModuleName where
 deriving TypeName, Repr
 
 
-private def onlyCode [Monad m] [MonadError m] (xs : TSyntaxArray `inline) : m StrLit := do
-  if h : xs.size = 1 then
-    match xs[0] with
-    | `(inline|code($s)) => return s
+
+private def onlyCodes [Monad m] [MonadError m] (stxs : TSyntaxArray `inline) : m (Array StrLit) := do
+  let mut codes := #[]
+  for stx in stxs do
+    match stx with
+    | `(inline|code($s)) => codes := codes.push s
+    | `(inline|$s:str) =>
+      unless s.getString.all (·.isWhitespace) do
+        throwErrorAt stx "Expected code"
     | other => throwErrorAt other "Expected code"
-  else
-    throwError "Expected precisely 1 code argument"
+  return codes
+
+private def onlyCode [Monad m] [MonadError m] (stxs : TSyntaxArray `inline) : m StrLit := do
+  let codes ← onlyCodes stxs
+  if h : codes.size = 1 then return codes[0]
+  else throwError "Expected precisely 1 code argument"
 
 private def strLitRange [Monad m] [MonadFileMap m] (s : StrLit) : m Lean.Syntax.Range := do
   let pos := (s.raw.getPos? (canonicalOnly := true)).get!
@@ -1171,10 +1180,21 @@ def option (xs : TSyntaxArray `inline) : DocM (Inline ElabInline) := do
         let decl ← getOptionDecl optionName
         pushInfoLeaf <| .ofOptionInfo { stx := id, optionName, declName := decl.declName }
         validateOptionValue optionName decl val
+        let valStr :=
+          match val with
+          | .ofString v => v.quote
+          | v => toString v
+        let valHl : DocHighlight ←
+          match val with
+          | .ofBool b =>
+            let constName := if b then ``Bool.true else ``Bool.false
+            let sig ← PrettyPrinter.ppSignature constName
+            pure <| .const constName sig.fmt
+          | _ => pure <| .literal stx[3].getKind none
         let code := #[
           ("set_option", some .keyword), (" ", none),
-          (toString stx[1][0].getId, some <| .option optionName decl.declName), (" ", none),
-          (toString stx[2].getAtomVal, some <| .literal stx[2].getKind none)
+          (toString optionName, some <| .option optionName decl.declName), (" ", none),
+          (valStr, some valHl)
         ]
         return .other {name := ``Data.SetOption, val := .mk <| Data.SetOption.mk ⟨code⟩} #[
           .code s.getString
@@ -1205,35 +1225,51 @@ def option (xs : TSyntaxArray `inline) : DocM (Inline ElabInline) := do
           return .code s.getString
 
 
-private def assertContents : ParserFn :=
-  whitespace >>
-  nodeFn nullKind
-    (termParser.fn >>
-     chFn '=' (trailingWs := true) >>
-     termParser.fn >>
-     optionalFn (symbolFn ":" >> termParser.fn))
+private def assertTermContents : ParserFn :=
+  whitespace >> termParser.fn
 
 /--
-Asserts that an equality holds.
+Asserts that two terms are definitionally equal.
+
+The terms are provided as two separate code elements, optionally followed by a third code element
+that contains the type at which the terms are elaborated e.g.
+`` {assert'}[`Nat.zero` `Nat.zero` `Nat`] ``.
 
-This doesn't use the equality type because it is needed in the prelude, before the = notation is
-introduced.
+Unlike `assert`, this role doesn't use the equality type, because it is needed in the prelude,
+before the `=` notation is introduced.
 -/
 @[builtin_doc_role]
 def assert' (xs : TSyntaxArray `inline) : DocM (Inline ElabInline) := do
-  let s ← onlyCode xs
-  let stx ← parseStrLit assertContents s
-  let ty? ←
-    withoutErrToSorry <|
-    if stx[3][1].isMissing then -- no colon
-      pure none
-    else -- type after colon
-      some <$> elabType stx[3][1]
-  let lhs ← elabExtraTerm stx[0] ty?
-  let rhs ← elabExtraTerm stx[2] ty?
-  unless ← Meta.withTransparency .all <| Meta.isDefEq lhs rhs do
-    throwErrorAt stx m!"Expected {lhs} = {rhs}, which is {← Meta.whnf lhs} = {← Meta.whnf rhs}, reducing to {← Meta.reduceAll lhs} = {← Meta.reduceAll rhs} but they are not equal."
-  pure (.code s.getString)
+  let codes ← onlyCodes xs
+  let (lhsCode, rhsCode, tyCode?) ←
+    match h : codes.size with
+    | 2 => pure (codes[0], codes[1], none)
+    | 3 => pure (codes[0], codes[1], some codes[2])
+    | _ => throwError "Expected two or three code arguments: the two sides of the equality, \
+        optionally followed by their type, but got {codes.size} arguments."
+  let lhsStx ← parseStrLit assertTermContents lhsCode
+  let rhsStx ← parseStrLit assertTermContents rhsCode
+  let tyStx? ← tyCode?.mapM (parseStrLit assertTermContents)
+  withSaveInfoContext do
+    let ty? ← withoutErrToSorry <|
+      match tyStx? with
+      | some tyStx => some <$> elabType tyStx
+      | none => pure none
+    let lhs ← elabExtraTerm lhsStx ty?
+    let rhs ← elabExtraTerm rhsStx ty?
+    unless ← Meta.withTransparency .all <| Meta.isDefEq lhs rhs do
+      throwError m!"Expected {lhs} = {rhs}, which is {← Meta.whnf lhs} = {← Meta.whnf rhs}, reducing to {← Meta.reduceAll lhs} = {← Meta.reduceAll rhs} but they are not equal."
+  let str := lhsCode.getString ++ " = " ++ rhsCode.getString ++
+    (tyCode?.map (" : " ++ ·.getString) |>.getD "")
+  let trees ← getInfoTrees
+  if trees.size > 0 then
+    let mut code := (← highlightSyntax trees lhsStx) ++ " = " ++ (← highlightSyntax trees rhsStx)
+    if let some tyStx := tyStx? then
+      code := code ++ " : " ++ (← highlightSyntax trees tyStx)
+    return .other {name := ``Data.LeanTerm, val := .mk <| Data.LeanTerm.mk code} #[.code str]
+  else
+    -- No info
+    return .code str
 
 /--
 Asserts that an equality holds.
@@ -1242,14 +1278,20 @@ Asserts that an equality holds.
 def assert (xs : TSyntaxArray `inline) : DocM (Inline ElabInline) := do
   let s ← onlyCode xs
   let stx ← parseStrLit termParser.fn s
-  let ty ← withoutErrToSorry <| elabType stx
-  match_expr (← Meta.whnf ty) with
-  | Eq _ lhs rhs =>
-    unless (← Meta.isDefEq lhs rhs) do
-      throwErrorAt stx m!"Expected {lhs} = {rhs}, but they are not definitionally equal"
-  | _ => throwErrorAt stx m!"Expected equality type"
-
-  pure (.code s.getString)
+  withSaveInfoContext do
+    let ty ← withoutErrToSorry <| elabType stx
+    match_expr (← Meta.whnf ty) with
+    | Eq _ lhs rhs =>
+      unless (← Meta.isDefEq lhs rhs) do
+        throwErrorAt stx m!"Expected {lhs} = {rhs}, but they are not definitionally equal"
+    | _ => throwErrorAt stx m!"Expected equality type"
+  let trees ← getInfoTrees
+  if trees.size > 0 then
+    let tm := Data.LeanTerm.mk (← highlightSyntax trees stx)
+    return .other {name := ``Data.LeanTerm, val := .mk tm} #[.code s.getString]
+  else
+    -- No info
+    return .code s.getString
 
 /--
 Opens a namespace in the remainder of the documentation comment.

src/Lean/Elab/DocString/Builtin/Keywords.lean

@@ -23,7 +23,7 @@ open scoped Lean.Doc.Syntax
 set_option linter.missingDocs true
 
 /-- The code represents an atom drawn from some syntax. -/
-structure Data.Atom where
+public structure Data.Atom where
   /-- The syntax kind's name. -/
   name : Name
   /-- The syntax category -/

tests/elab/versoDocs.lean

@@ -754,3 +754,147 @@ info:  <kw>(</kw> <lit kind="num" type="Nat">2</lit>  <kw>*</kw>  <kw>(</kw>
   doc.text.flatMap (findInBlock ``Data.LeanTerm) |>.map docCodeStr |>.forM (IO.println ·)
 
 end HygieneInfoTests
+
+/-!
+Test that the {lit}`{option}` role, when given full {lit}`set_option` syntax, stores the actual
+option name and value in its {lit}`Data.SetOption` display code.
+-/
+
+section SetOptionRoleTests
+open Doc Elab
+
+private partial def findSetOptionInInline : Inline ElabInline → Array DocCode
+  | .other container _ =>
+    if let some (so : Data.SetOption) := container.val.get? Data.SetOption then
+      #[so.term]
+    else #[]
+  | .emph xs | .bold xs | .concat xs | .link xs _ | .footnote _ xs =>
+    xs.flatMap findSetOptionInInline
+  | .text .. | .code .. | .math .. | .linebreak .. | .image .. => #[]
+
+private partial def findSetOptionInBlock : Block ElabInline ElabBlock → Array DocCode
+  | .para inlines => inlines.flatMap findSetOptionInInline
+  | .concat blocks | .blockquote blocks => blocks.flatMap findSetOptionInBlock
+  | .dl items => items.flatMap fun ⟨x, y⟩ =>
+    x.flatMap findSetOptionInInline ++ y.flatMap findSetOptionInBlock
+  | .ol _ xs | .ul xs => xs.flatMap fun ⟨x⟩ => x.flatMap findSetOptionInBlock
+  | .other .. | .code .. => #[]
+
+/--
+Setting options:
+ * {option}`set_option pp.all true`
+ * {option}`set_option maxHeartbeats 1000`
+ * {option}`set_option trace.profiler.output "out.json"`
+-/
+def setOptionDisplay := ()
+
+/--
+info:  <kw>set_option</kw> ⏎
+<option name="pp.all" decl="Lean.pp.all">pp.all</option> ⏎
+<const name="Bool.true" sig="Bool.true : Bool">true</const>
+ <kw>set_option</kw> ⏎
+<option name="maxHeartbeats" decl="Lean.maxHeartbeats">maxHeartbeats</option> ⏎
+<lit kind="num" type=none>1000</lit>
+ <kw>set_option</kw> ⏎
+<option name="trace.profiler.output" decl="Lean.trace.profiler.output">trace.profiler.output</option> ⏎
+<lit kind="str" type=none>"out.json"</lit>
+-/
+#guard_msgs in
+#eval show TermElabM Unit from do
+  let some (.inr doc) ← findInternalDocString? (← getEnv) ``setOptionDisplay
+    | throwError "expected verso doc"
+  doc.text.flatMap findSetOptionInBlock |>.map docCodeStr |>.forM (IO.println ·)
+
+end SetOptionRoleTests
+
+/-!
+Test that the {lit}`assert` and {lit}`assert'` roles save elaboration info and store highlighted
+{lit}`Data.LeanTerm` payloads (previously they returned bare code with no hover information), and
+that {lit}`assert'` takes the two sides of the equality (and optionally the type at which they are
+compared) as separate code elements, so that it can be used without the {lit}`=` notation.
+-/
+
+section AssertRoleTests
+open Doc Elab
+
+/--
+{assert}`Nat.zero = Nat.zero`
+
+{assert'}[`Nat.zero` `Nat.zero`]
+
+{assert'}[`Nat.zero` `Nat.zero` `Nat`]
+-/
+def assertDisplay := ()
+
+/--
+info:  <const name="Nat.zero" sig="Nat.zero : Nat">Nat.zero</const> ⏎
+<kw>=</kw> ⏎
+<const name="Nat.zero" sig="Nat.zero : Nat">Nat.zero</const>
+ <const name="Nat.zero" sig="Nat.zero : Nat">Nat.zero</const> = ⏎
+<const name="Nat.zero" sig="Nat.zero : Nat">Nat.zero</const>
+ <const name="Nat.zero" sig="Nat.zero : Nat">Nat.zero</const> = ⏎
+<const name="Nat.zero" sig="Nat.zero : Nat">Nat.zero</const> : ⏎
+<const name="Nat" sig="Nat : Type">Nat</const>
+-/
+#guard_msgs in
+#eval show TermElabM Unit from do
+  let some (.inr doc) ← findInternalDocString? (← getEnv) ``assertDisplay
+    | throwError "expected verso doc"
+  doc.text.flatMap (findInBlock ``Data.LeanTerm) |>.map docCodeStr |>.forM (IO.println ·)
+
+/--
+error: Expected Nat.zero = Nat.zero.succ, which is Nat.zero = 1, reducing to Nat.zero = 1 but they are not equal.
+-/
+#guard_msgs in
+/-! {assert'}[`Nat.zero` `Nat.succ Nat.zero`] -/
+
+/--
+error: Expected two or three code arguments: the two sides of the equality, optionally followed by their type, but got 1 arguments.
+-/
+#guard_msgs in
+/-! {assert'}[`Nat.zero`] -/
+
+end AssertRoleTests
+
+/-!
+Test that the {name}`kw` and {name}`kw?` roles store their payloads as {name}`Lean.Doc.Data.Atom`.
+-/
+
+section KwAtomPublicTests
+open Doc Elab
+
+private partial def findAtomInInline : Inline ElabInline → Array (Name × Data.Atom)
+  | .other container _ =>
+    if let some (a : Data.Atom) := container.val.get? Data.Atom then
+      #[(container.name, a)]
+    else #[]
+  | .emph xs | .bold xs | .concat xs | .link xs _ | .footnote _ xs =>
+    xs.flatMap findAtomInInline
+  | .text .. | .code .. | .math .. | .linebreak .. | .image .. => #[]
+
+private partial def findAtomInBlock : Block ElabInline ElabBlock → Array (Name × Data.Atom)
+  | .para inlines => inlines.flatMap findAtomInInline
+  | .concat blocks | .blockquote blocks => blocks.flatMap findAtomInBlock
+  | .dl items => items.flatMap fun ⟨x, y⟩ =>
+    x.flatMap findAtomInInline ++ y.flatMap findAtomInBlock
+  | .ol _ xs | .ul xs => xs.flatMap fun ⟨x⟩ => x.flatMap findAtomInBlock
+  | .other .. | .code .. => #[]
+
+/--
+{kw (cat := term)}`Type`
+-/
+def kwAtomDisplay := ()
+
+/--
+info: container: Lean.Doc.Data.Atom
+category: term
+-/
+#guard_msgs in
+#eval show TermElabM Unit from do
+  let some (.inr doc) ← findInternalDocString? (← getEnv) ``kwAtomDisplay
+    | throwError "expected verso doc"
+  for (containerName, atom) in doc.text.flatMap findAtomInBlock do
+    IO.println s!"container: {containerName}"
+    IO.println s!"category: {atom.category}"
+
+end KwAtomPublicTests