feat: show proof states after =>, and not after semicolons (#27)

InternalTests.lean

@@ -0,0 +1,80 @@
+
+import SubVerso.Examples
+import SubVerso.Highlighting.Highlighted
+
+/-! These are SubVerso tests that don't involve a subprocess, to make development easier. -/
+
+
+open SubVerso Examples
+
+partial def SubVerso.Highlighting.Highlighted.asString (hl : Highlighted) : String := Id.run do
+  let mut out := ""
+  match hl with
+  | .seq hls =>
+    for x in hls.map asString do
+      out := out ++ x
+  | .span _ hl' =>
+    out := out ++ hl'.asString
+  | .tactics _ _ _ hl' =>
+    out := out ++ hl'.asString
+  | .point .. => pure ()
+  | .text s => out := out ++ s
+  | .token t => out := out ++ t.content
+  out
+
+
+partial def SubVerso.Highlighting.Highlighted.proofStates (hl : Highlighting.Highlighted) : Array (String × Array (Goal String)) := Id.run do
+  let mut out := #[]
+  match hl with
+  | .seq hls =>
+    for x in hls.map proofStates do
+      out := out ++ x
+  | .span _ hl' =>
+    out := out ++ hl'.proofStates
+  | .tactics info _ _ hl' =>
+    out := out.push (hl'.asString, info.map (·.map (·.asString)))
+  | _ => pure ()
+  out
+
+%example proof
+theorem test (n : Nat) : n * 1 = n := by
+  induction n with
+  | zero => rfl
+  | succ n ih =>
+    rw [← ih]
+    cases n
+    next => simp
+    case' succ =>
+      skip
+    case succ =>
+      . skip; simp
+%end
+
+%dump proof into proofJson
+
+%dumpE proof into proofEx
+
+
+
+-- We don't have #guard_msgs in all supported Lean versions, so here's a dumb replacement:
+
+open Lean Elab Command in
+elab "#evalString" s:str e:term : command => do
+  let msgs ← MonadMessageState.getMessages
+  try
+    MonadMessageState.setMessages {}
+    elabCommand <| ← `(#eval $e)
+    let msgs' ← MonadMessageState.getMessages
+    let [msg] := msgs'.toList
+      | throwError "Too many messages"
+    if (← msg.toString) != s.getString then
+      throwErrorAt e "Expected {s.getString}, got {← msg.toString}"
+  finally
+    MonadMessageState.setMessages msgs
+
+#evalString "[[some `zero], [some `succ], [none], [some `succ.succ], [none]]\n"
+  (proofEx.highlighted[0].proofStates.data.filter (·.fst == "=>") |>.map (·.snd.data.map (·.name)))
+
+
+
+def main : IO Unit := pure ()

lakefile.lean

@@ -21,6 +21,11 @@ lean_exe «subverso-tests» where
   root := `Tests
   supportInterpreter := true
 
+@[default_target]
+lean_exe «subverso-internal-tests» where
+  root := `InternalTests
+  supportInterpreter := true
+
 @[default_target]
 lean_exe «subverso-extract» where
   root := `Extract

src/examples/SubVerso/Examples.lean

@@ -137,7 +137,7 @@ elab_rules : command
       logMessage msg
 
 
-scoped syntax "%dump" ident : command
+scoped syntax "%dump " ident : command
 
 elab_rules : command
   | `(%dump%$kw $name:ident) => do
@@ -148,6 +148,82 @@ elab_rules : command
     else
       throwErrorAt name "No highlighting found for '{name}'"
 
+scoped syntax "%dump " ident &" into " ident: command
+scoped syntax "%dumpE " ident &" into " ident: command
+
+open Syntax in
+private scoped instance : Quote Int where
+  quote
+    | .ofNat n => mkCApp ``Int.ofNat #[quote n]
+    | .negSucc n => mkCApp ``Int.negSucc #[quote n]
+
+open Syntax in
+instance : Quote JsonNumber where
+  quote
+    | ⟨mantissa, exponent⟩ => mkCApp ``JsonNumber.mk #[quote mantissa, quote exponent]
+
+open Syntax in
+partial instance : Quote Json where
+  quote := q
+where
+  quoteArray {α : _} (_inst : Quote α) (xs : Array α) : TSyntax `term :=
+    mkCApp ``List.toArray #[quote xs.toList]
+  quoteField {α : _} (_inst : Quote α) (f : (_ : String) × α) : TSyntax `term :=
+    quote (f.fst, f.snd)
+  q
+    | .null => mkCApp ``Json.null #[]
+    | .str s => mkCApp ``Json.str #[quote s]
+    | .bool b => mkCApp ``Json.bool #[quote b]
+    | .num n => mkCApp ``Json.num #[quote n]
+    | .arr xs => mkCApp ``Json.arr #[(quoteArray ⟨q⟩ xs)]
+    | .obj fields =>
+      let _fieldInst : Quote ((_ : String) × Json) := ⟨quoteField ⟨q⟩⟩
+      let fieldList := quote fields.toArray.toList
+      mkCApp ``Json.mkObj #[fieldList]
+
+elab_rules : command
+  | `(%dump $name:ident into $x) => do
+    let mod ← getMainModule
+    let st := highlighted.getState (← getEnv) |>.find? mod |>.getD {}
+    if let some json := st.find? name.getId then
+      elabCommand <| ← `(def $x : Json := $(quote json))
+    else
+      throwErrorAt name "No highlighting found for '{name}'"
+
+open Syntax in
+instance : Quote MessageSeverity where
+  quote s :=
+    let n :=
+      match s with
+      | .error => ``MessageSeverity.error
+      | .information => ``MessageSeverity.information
+      | .warning => ``MessageSeverity.warning
+    mkCApp n #[]
+
+open Syntax in
+instance : Quote Lean.Position where
+  quote s :=
+    mkCApp ``Lean.Position.mk #[quote s.line, quote s.column]
+
+open Syntax in
+instance : Quote Example where
+  quote ex :=
+    mkCApp ``Example.mk #[quote ex.highlighted, quote ex.messages, quote ex.original, quote ex.start, quote ex.stop]
+
+elab_rules : command
+  | `(%dumpE $name:ident into $x) => do
+    let mod ← getMainModule
+    let st := highlighted.getState (← getEnv) |>.find? mod |>.getD {}
+    if let some json := st.find? name.getId then
+      match FromJson.fromJson? json with
+      | .ok (e : Example) =>
+        elabCommand <| ← `(def $x : Example := $(quote e))
+      | .error err =>
+        throwErrorAt name "Couldn't deserialize JSON: {err}"
+    else
+      throwErrorAt name "No highlighting found for '{name}'"
+
+
 namespace Internals
 scoped syntax "%show_name" ident : term
 elab_rules : term

src/examples/SubVerso/Examples/Env.lean

@@ -32,13 +32,15 @@ instance : FromJson MessageSeverity where
 theorem MessageSeverity.fromJson_toJson_ok (s : MessageSeverity) : fromJson? (toJson s) = .ok s := by
   cases s <;> simp [toJson, fromJson?]
 
+deriving instance Repr for MessageSeverity
+
 structure Example where
   highlighted : Array Highlighted
   messages : List (MessageSeverity × String)
   original : String
   start : Lean.Position
   stop : Lean.Position
-deriving ToJson, FromJson
+deriving ToJson, FromJson, Repr
 
 initialize highlighted : PersistentEnvExtension (NameMap (NameMap Json)) (Name × Name × Example) (NameMap (NameMap Json)) ←
   registerPersistentEnvExtension {

src/highlighting/SubVerso/Highlighting/Code.lean

@@ -29,6 +29,20 @@ def infoForSyntax (t : InfoTree) (stx : Syntax) : List (ContextInfo × Info) :=
       (ci, info) :: soFar
     else soFar
 
+-- Find the nodes whose canonical span includes the given syntax
+def infoIncludingSyntax (t : InfoTree) (stx : Syntax) : List (ContextInfo × Info) :=
+  t.collectNodesBottomUp fun ci info _ soFar =>
+    if let some true := do
+      let start ← info.stx.getPos? true
+      let start' ← stx.getPos? true
+      let stop ← info.stx.getTailPos? true
+      let stop' ← stx.getTailPos? true
+      pure <| start ≤ start' && stop ≥ stop'
+    then
+      (ci, info) :: soFar
+    else soFar
+
+
 
 partial def bestD (sems : Array Token.Kind) (default : Token.Kind) : Token.Kind :=
   if let some m := sems.back? then
@@ -516,6 +530,71 @@ partial def _root_.Lean.Widget.TaggedText.indent (doc : TaggedText α) : TaggedT
     | .append docs => .append (docs.map indent')
   .append #[.text "  ", indent' doc]
 
+def highlightGoals
+    (ids : HashMap Lsp.RefIdent Lsp.RefIdent)
+    (ci : ContextInfo)
+    (goals : List MVarId)
+    : HighlightM (Array (Highlighted.Goal Highlighted)) := do
+  let mut goalView := #[]
+  for g in goals do
+    let mut hyps := #[]
+    let some mvDecl := ci.mctx.findDecl? g
+      | continue
+    let name := if mvDecl.userName.isAnonymous then none else some mvDecl.userName
+    let lctx := mvDecl.lctx |>.sanitizeNames.run' {options := (← getOptions)}
+
+    -- Tell the delaborator to tag functions that are being applied. Otherwise,
+    -- functions have no tooltips or binding info in tactics.
+    -- cf https://leanprover.zulipchat.com/#narrow/stream/270676-lean4/topic/Function.20application.20delaboration/near/265800665
+    let ci' := {ci with options := ci.options.set `pp.tagAppFns true}
+    let runMeta {α} (act : MetaM α) : HighlightM α := ci'.runMetaM lctx act
+    for c in lctx.decls do
+      let some decl := c
+        | continue
+      if decl.isAuxDecl || decl.isImplementationDetail then continue
+      match decl with
+      | .cdecl _index fvar name type _ _ =>
+        let nk ← exprKind ids ci lctx (.fvar fvar)
+        let tyStr ← renderTagged ids (← runMeta (ppExprTagged =<< instantiateMVars type))
+        hyps := hyps.push (name, nk.getD .unknown, tyStr)
+      | .ldecl _index fvar name type val _ _ =>
+        let nk ← exprKind ids ci lctx (.fvar fvar)
+        let tyDoc ← runMeta (ppExprTagged =<< instantiateMVars type)
+        let valDoc ← runMeta (ppExprTagged =<< instantiateMVars val)
+        let tyValStr ← renderTagged ids <| .append <| #[tyDoc].append <|
+          if tyDoc.oneLine && valDoc.oneLine then #[.text " := ", valDoc]
+          else #[.text " := \n", valDoc.indent]
+        hyps := hyps.push (name, nk.getD .unknown, tyValStr)
+    let concl ← renderTagged ids (← runMeta <| ppExprTagged =<< instantiateMVars mvDecl.type)
+    goalView := goalView.push ⟨name, Meta.getGoalPrefix mvDecl, hyps, concl⟩
+  pure goalView
+
+partial def findTactics'
+    (ids : HashMap Lsp.RefIdent Lsp.RefIdent)
+    (trees : PersistentArray Lean.Elab.InfoTree)
+    (stx : Syntax)
+    (startPos endPos : String.Pos)
+    (endPosition : Position)
+    (before : Bool := false)
+    : HighlightM (Option (Array (Highlighted.Goal Highlighted) × Nat × Nat × Position)) := do
+  for t in trees do
+    -- The info is reversed here so that the _last_ state computed is shown.
+    -- This makes `repeat` do the right thing, rather than either spamming
+    -- states or showing the first one.
+    for i in infoForSyntax t stx |>.reverse do
+      if not i.2.isOriginal then continue
+      if let (ci, .ofTacticInfo tacticInfo) := i then
+        if !before && !tacticInfo.goalsBefore.isEmpty && tacticInfo.goalsAfter.isEmpty then
+          return some (#[], startPos.byteIdx, endPos.byteIdx, endPosition)
+
+        let goals := if before then tacticInfo.goalsBefore else tacticInfo.goalsAfter
+        if goals.isEmpty then continue
+        let goalView ← highlightGoals ids ci goals
+
+        if !Output.inTacticState (← get).output goalView then
+          return some (goalView, startPos.byteIdx, endPos.byteIdx, endPosition)
+  return none
+
 partial def findTactics
     (ids : HashMap Lsp.RefIdent Lsp.RefIdent)
     (trees : PersistentArray Lean.Elab.InfoTree)
@@ -533,6 +612,23 @@ partial def findTactics
   if stx.getKind ∈ [``Lean.Parser.Term.byTactic, ``Lean.Parser.Term.byTactic'] ||
      stx matches .atom _ "by" then
     return none
+  -- `;` is blacklisted as well - no need to highlight states identically
+  if stx matches .atom _ ";" then return none
+
+  -- Special handling for =>: show the _before state_
+  if stx matches .atom _ "=>" then
+    for t in trees do
+      for i in infoIncludingSyntax t stx |>.reverse do
+        if not i.2.isOriginal then continue
+        if let (_, .ofTacticInfo tacticInfo) := i then
+          match tacticInfo.stx with
+          | `(Lean.Parser.Tactic.inductionAlt| $_lhs =>%$tk $rhs )
+          | `(tactic| next $_* =>%$tk $rhs )
+          | `(tactic| case $_ $_* =>%$tk $rhs )
+          | `(tactic| case' $_ $_* =>%$tk $rhs ) =>
+            if tk == stx then
+              return (← findTactics' ids trees rhs startPos endPos endPosition (before := true))
+          | _ => continue
 
   -- Only show tactic output for the most specific source spans possible, with a few exceptions
   if stx.getKind ∉ [``Lean.Parser.Tactic.rwSeq,``Lean.Parser.Tactic.simp] then
@@ -544,54 +640,8 @@ partial def findTactics
     if let some (goals, _startPos, _endPos, _endPosition) ← findTactics ids trees brak then
       return some (goals, startPos.byteIdx, endPos.byteIdx, endPosition)
 
-  for t in trees do
-    -- The info is reversed here so that the _last_ state computed is shown.
-    -- This makes `repeat` do the right thing, rather than either spamming
-    -- states or showing the first one.
-    for i in infoForSyntax t stx |>.reverse do
-      if not i.2.isOriginal then continue
-      if let (ci, .ofTacticInfo tacticInfo) := i then
-        if !tacticInfo.goalsBefore.isEmpty && tacticInfo.goalsAfter.isEmpty then
-          return some (#[], startPos.byteIdx, endPos.byteIdx, endPosition)
+  findTactics' ids trees stx startPos endPos endPosition
 
-        let goals := tacticInfo.goalsAfter
-        if goals.isEmpty then continue
-
-        let mut goalView := #[]
-        for g in goals do
-          let mut hyps := #[]
-          let some mvDecl := ci.mctx.findDecl? g
-            | continue
-          let name := if mvDecl.userName.isAnonymous then none else some mvDecl.userName
-          let lctx := mvDecl.lctx |>.sanitizeNames.run' {options := (← getOptions)}
-
-          -- Tell the delaborator to tag functions that are being applied. Otherwise,
-          -- functions have no tooltips or binding info in tactics.
-          -- cf https://leanprover.zulipchat.com/#narrow/stream/270676-lean4/topic/Function.20application.20delaboration/near/265800665
-          let ci' := {ci with options := ci.options.set `pp.tagAppFns true}
-          let runMeta {α} (act : MetaM α) : HighlightM α := ci'.runMetaM lctx act
-          for c in lctx.decls do
-            let some decl := c
-              | continue
-            if decl.isAuxDecl || decl.isImplementationDetail then continue
-            match decl with
-            | .cdecl _index fvar name type _ _ =>
-              let nk ← exprKind ids ci lctx (.fvar fvar)
-              let tyStr ← renderTagged ids (← runMeta (ppExprTagged =<< instantiateMVars type))
-              hyps := hyps.push (name, nk.getD .unknown, tyStr)
-            | .ldecl _index fvar name type val _ _ =>
-              let nk ← exprKind ids ci lctx (.fvar fvar)
-              let tyDoc ← runMeta (ppExprTagged =<< instantiateMVars type)
-              let valDoc ← runMeta (ppExprTagged =<< instantiateMVars val)
-              let tyValStr ← renderTagged ids <| .append <| #[tyDoc].append <|
-                if tyDoc.oneLine && valDoc.oneLine then #[.text " := ", valDoc]
-                else #[.text " := \n", valDoc.indent]
-              hyps := hyps.push (name, nk.getD .unknown, tyValStr)
-          let concl ← renderTagged ids (← runMeta <| ppExprTagged =<< instantiateMVars mvDecl.type)
-          goalView := goalView.push ⟨name, Meta.getGoalPrefix mvDecl, hyps, concl⟩
-        if !Output.inTacticState (← get).output goalView then
-          return some (goalView, startPos.byteIdx, endPos.byteIdx, endPosition)
-  return none
 
 partial def highlight'
     (ids : HashMap Lsp.RefIdent Lsp.RefIdent)

src/highlighting/SubVerso/Highlighting/Highlighted.lean

@@ -74,6 +74,11 @@ structure Highlighted.Goal (expr) where
   conclusion : expr
 deriving Repr, BEq, Hashable, ToJson, FromJson
 
+def Highlighted.Goal.map (f : α → β) (g : Goal α) : Goal β :=
+  {g with
+    hypotheses := g.hypotheses.map (fun (x, k, e) => (x, k, f e))
+    conclusion := f g.conclusion}
+
 instance [Quote expr] : Quote (Highlighted.Goal expr) where
   quote
     | {name, goalPrefix, hypotheses, conclusion} =>