fix: correctly pass info trees in highlight generation (#18)

Fixes #17. Also adds CI to compare versions of Lean for equivalent
proof output.

Tests.lean

@@ -1,15 +1,92 @@
 import SubVerso.Highlighting
 import SubVerso.Examples
 import SubVerso.Examples.Env
-import Lean
+import Lean.Data.Json
+import Lean.Data.NameMap
 
 open SubVerso.Examples (loadExamples Example)
 open Lean.FromJson (fromJson?)
 
+open Lean in
+def exampleArray (examples : NameMap (NameMap α)) : Array α := Id.run do
+  let mut exs := #[]
+  for (_, inner) in examples do
+    for (_, x) in inner do
+      exs := exs.push x
+  exs
+
+open Lean in
+def findVersionString (examples : NameMap (NameMap Example)) : Option String := do
+  let demo ← examples.find? `Demo
+  let version ← demo.find? `version
+  let [(.information, str)] := version.messages
+    | none
+  str.trim |>.drop 1 |>.dropRight 1
+
+namespace SubVerso
+
+partial
+def Highlighting.Highlighted.countProofStates (hl : Highlighting.Highlighted) : Nat :=
+  match hl with
+  | .seq hls =>
+    hls.map countProofStates |>.foldr (· + · ) 0
+  | .span _ hl' =>
+    hl'.countProofStates
+  | .tactics _ _ _ hl' =>
+    hl'.countProofStates + 1
+  | _ => 0
+
+namespace Examples
+
+def Example.countProofStates (e : Example) : Nat :=
+  e.highlighted.map (·.countProofStates) |>.foldl (· + ·) 0
+
+end Examples
+end SubVerso
+
+def cleanupDemo : IO Unit := do
+  if ← System.FilePath.pathExists "demo/lake-manifest.json" then
+    IO.FS.removeFile "demo/lake-manifest.json"
+  if ← System.FilePath.isDir "demo/.lake" then
+    IO.FS.removeDirAll "demo/.lake"
+
+open Lean in
+def proofCount (examples : NameMap (NameMap Example)) : Nat := Id.run do
+  let mut n := 0
+  for e in exampleArray examples do
+    n := n + e.countProofStates
+  n
+
+open Lean in
+def checkVersion (expected : String) (examples : NameMap (NameMap Example)) : IO Unit := do
+  let v := findVersionString examples
+  IO.println s!"Reported version {v}"
+  if v != some expected then
+    IO.eprintln "Unexpected version!"
+    IO.Process.exit 1
+  pure ()
+
 def main : IO UInt32 := do
   IO.println "Checking that the test project generates at least some deserializable JSON"
+  cleanupDemo
   let examples ← loadExamples "demo"
   if examples.isEmpty then
     IO.eprintln "No examples found"
     return 1
+  checkVersion "4.3.0" examples
+  let proofCount1 := proofCount examples
+  IO.println s!"Found {proofCount1} proofs "
+  IO.println "Checking that the test project generates at least some deserializable JSON with a newer toolchain"
+  cleanupDemo
+  let examples' ← loadExamples "demo" (overrideToolchain := some "leanprover/lean4:nightly-2024-04-25")
+  if examples'.isEmpty then
+    IO.eprintln "No examples found with later toolchain"
+    return 1
+  checkVersion "4.8.0-nightly-2024-04-25" examples'
+  let proofCount2 := proofCount examples'
+  IO.println s!"Found {proofCount2} proofs "
+
+  if proofCount1 != proofCount2 then
+    IO.eprintln "Example proof count mismatch"
+    return 1
   pure 0

demo/Demo.lean

@@ -6,6 +6,19 @@ import SubVerso.Examples
 
 open SubVerso.Examples
 
+%example version
+#eval Lean.versionString
+%end
+
 %example xdef
 def f (x : Nat) := %ex{add}{33 + %ex{X}{x}}
 %end
+
+%example proof
+theorem test (n : Nat) : n * 1 = n := by
+  induction n with
+  | zero => rfl
+  | succ n ih =>
+    rw [← ih]
+    simp
+%end

demo/lean-toolchain

@@ -1 +1 @@
-leanprover/lean4:nightly-2024-02-26
+leanprover/lean4:v4.3.0

src/examples/SubVerso/Examples.lean

@@ -149,16 +149,21 @@ elab_rules : command
       throwErrorAt name "No highlighting found for '{name}'"
 
 open System in
-partial def loadExamples (leanProject : FilePath) : IO (NameMap (NameMap Example)) := do
+partial def loadExamples
+    (leanProject : FilePath)
+    (overrideToolchain : Option String := none) : IO (NameMap (NameMap Example)) := do
   let projectDir := ((← IO.currentDir) / leanProject).normalize
   -- Validate that the path is really a Lean project
   let lakefile := projectDir / "lakefile.lean"
   if !(← lakefile.pathExists) then
     throw <| .userError s!"File {lakefile} doesn't exist, couldn't load project"
-  let toolchainfile := projectDir / "lean-toolchain"
-  if !(← toolchainfile.pathExists) then
-    throw <| .userError s!"File {toolchainfile} doesn't exist, couldn't load project"
-  let toolchain := (← IO.FS.readFile toolchainfile).trim
+  let toolchain ← match overrideToolchain with
+    | none =>
+      let toolchainfile := projectDir / "lean-toolchain"
+      if !(← toolchainfile.pathExists) then
+        throw <| .userError s!"File {toolchainfile} doesn't exist, couldn't load project"
+      pure (← IO.FS.readFile toolchainfile).trim
+    | some override => pure override
 
   -- Kludge: remove variables introduced by Lake. Clearing out DYLD_LIBRARY_PATH and
   -- LD_LIBRARY_PATH is useful so the version selected by Elan doesn't get the wrong shared

src/highlighting/SubVerso/Highlighting/Code.lean

@@ -190,7 +190,7 @@ def infoExists [Monad m] [MonadLiftT IO m] (trees : PersistentArray InfoTree) (s
 inductive Output where
   | seq (emitted : Array Highlighted)
   | span (info : Array (Highlighted.Span.Kind × String))
-  | tactics (info : Array (Highlighted.Goal Highlighted)) (pos : Nat)
+  | tactics (info : Array (Highlighted.Goal Highlighted)) (startPos : Nat) (endPos : Nat)
 
 def Output.add (output : List Output) (hl : Highlighted) : List Output :=
   match output with
@@ -208,19 +208,16 @@ def Output.addText (output : List Output) (str : String) : List Output :=
 def Output.openSpan (output : List Output) (messages : Array (Highlighted.Span.Kind × String)) : List Output :=
   .span messages :: output
 
-def Output.openTactic (output : List Output) (info : Array (Highlighted.Goal Highlighted)) (pos : Nat) : List Output :=
-  .tactics info pos :: output
-
 def Output.inTacticState (output : List Output) (info : Array (Highlighted.Goal Highlighted)) : Bool :=
   output.any fun
-    | .tactics info' _ => info == info'
+    | .tactics info' _ _ => info == info'
     | _ => false
 
 def Output.closeSpan (output : List Output) : List Output :=
   let rec go (acc : Highlighted) : List Output → List Output
     | [] => [.seq #[acc]]
     | .span info :: more => Output.add more (.span info acc)
-    | .tactics info pos :: more => Output.add more (.tactics info pos acc)
+    | .tactics info startPos endPos :: more => Output.add more (.tactics info startPos endPos acc)
     | .seq left :: more => go (.seq (left.push acc)) more
   go .empty output
 
@@ -229,7 +226,7 @@ def Highlighted.fromOutput (output : List Output) : Highlighted :=
     | [] => acc
     | .seq left :: more => go (.seq (left.push acc)) more
     | .span info :: more => go (.span info acc) more
-    | .tactics info pos :: more => go (.tactics info pos acc) more
+    | .tactics info startPos endPos :: more => go (.tactics info startPos endPos acc) more
   go .empty output
 
 structure OpenTactic where
@@ -315,9 +312,33 @@ where
     else
       pure <| !(msg.pos.before s) && !(e.before msg.pos)
 
-
 abbrev HighlightM (α : Type) : Type := StateT HighlightState TermElabM α
 
+private def modify? (f : α → Option α) : (xs : List α) → Option (List α)
+  | [] => none
+  | x :: xs =>
+    if let some x' := f x then
+      some (x' :: xs)
+    else (x :: ·) <$> modify? f xs
+
+def HighlightState.openTactic (st : HighlightState) (info : Array (Highlighted.Goal Highlighted)) (startPos endPos : Nat) (pos : Position) : HighlightState :=
+  if let some out' := modify? update? st.output then
+    {st with output := out'}
+  else {st with
+    output := .tactics info startPos endPos :: st.output,
+    inMessages := .inr ⟨pos⟩ :: st.inMessages
+  }
+where
+  update?
+    | .tactics info' startPos' endPos' =>
+      if startPos == startPos' && endPos == endPos' then
+        some <| .tactics (info ++ info') startPos' endPos'
+      else none
+    | _ => none
+
+def HighlightM.openTactic (info : Array (Highlighted.Goal Highlighted)) (startPos endPos : Nat) (pos : Position) : HighlightM Unit :=
+  modify fun st => st.openTactic info startPos endPos pos
+
 instance : Inhabited (HighlightM α) where
   default := fun _ => default
 
@@ -421,22 +442,21 @@ partial def subsyntaxes (stx : Syntax) : Array Syntax := Id.run do
     stxs
   | _ => #[]
 
-def hasTactics (stx : Syntax) : HighlightM Bool := do
-  let trees ← getInfoTrees
+def hasTactics (stx : Syntax) (trees : PersistentArray Lean.Elab.InfoTree) : HighlightM Bool := do
   for t in trees do
     for (_, i) in infoForSyntax t stx do
       if not i.isOriginal then continue
       if let .ofTacticInfo _ := i then
         return true
   return false
 
-partial def childHasTactics (stx : Syntax) : HighlightM Bool := do
+partial def childHasTactics (stx : Syntax) (trees : PersistentArray Lean.Elab.InfoTree) : HighlightM Bool := do
   match stx with
   | .node _ _ children =>
     for s in children do
-      if ← hasTactics s then
+      if ← hasTactics s trees then
         return true
-      if ← childHasTactics s then
+      if ← childHasTactics s trees then
         return true
     pure false
   | _ => pure false
@@ -496,28 +516,44 @@ partial def _root_.Lean.Widget.TaggedText.indent (doc : TaggedText α) : TaggedT
     | .append docs => .append (docs.map indent')
   .append #[.text "  ", indent' doc]
 
-def findTactics (ids : HashMap Lsp.RefIdent Lsp.RefIdent) (stx : Syntax) : HighlightM Unit := do
+partial def findTactics
+    (ids : HashMap Lsp.RefIdent Lsp.RefIdent)
+    (trees : PersistentArray Lean.Elab.InfoTree)
+    (stx : Syntax)
+    : HighlightM (Option (Array (Highlighted.Goal Highlighted) × Nat × Nat × Position)) := do
+  let text ← getFileMap
+  let some startPos := stx.getPos?
+    | return none
+  let some endPos := stx.getTailPos?
+    | return none
+  let endPosition := text.toPosition endPos
+
+  -- Blacklisted tactics. TODO: make into an extensible table.
+  -- For now, the `by` keyword itself is blacklisted, as is its leading token
+  if stx.getKind ∈ [``Lean.Parser.Term.byTactic, ``Lean.Parser.Term.byTactic'] ||
+     stx matches .atom _ "by" then
+    return none
+
   -- 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
-    if ← childHasTactics stx then return ()
-  let trees ← getInfoTrees
-  let text ← getFileMap
+    if ← childHasTactics stx trees then return none
+
+  -- Override states - some tactics show many intermediate states, which is overwhelming in rendered
+  -- output. Get the right one to show for the whole thing, then adjust its positioning.
+  if let `(tactic|rw $[$_c:config]? [$_rs,*]%$brak $[$_l:location]?) := stx then
+    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
-        let some endPos := stx.getTailPos?
-          | continue
-        let endPosition := text.toPosition endPos
         if !tacticInfo.goalsBefore.isEmpty && tacticInfo.goalsAfter.isEmpty then
-          modify fun st => {st with
-            output := Output.openTactic st.output #[] endPos.byteIdx,
-            inMessages := .inr ⟨endPosition⟩ :: st.inMessages
-          }
-          break
+          return some (#[], startPos.byteIdx, endPos.byteIdx, endPosition)
+
         let goals := tacticInfo.goalsAfter
         if goals.isEmpty then continue
 
@@ -554,22 +590,31 @@ def findTactics (ids : HashMap Lsp.RefIdent Lsp.RefIdent) (stx : Syntax) : Highl
           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
-          modify fun st => {st with
-            output := Output.openTactic st.output goalView endPos.byteIdx,
-            inMessages := .inr ⟨endPosition⟩ :: st.inMessages
-          }
-        return
-
-partial def highlight' (ids : HashMap Lsp.RefIdent Lsp.RefIdent) (trees : PersistentArray Lean.Elab.InfoTree) (stx : Syntax) (lookingAt : Option (Name × String.Pos) := none) : HighlightM Unit := do
-  findTactics ids stx
+          return some (goalView, startPos.byteIdx, endPos.byteIdx, endPosition)
+  return none
+
+partial def highlight'
+    (ids : HashMap Lsp.RefIdent Lsp.RefIdent)
+    (trees : PersistentArray Lean.Elab.InfoTree)
+    (stx : Syntax)
+    (tactics : Bool)
+    (lookingAt : Option (Name × String.Pos) := none)
+    : HighlightM Unit := do
+  let mut tactics := tactics
+  if tactics then
+    if let some (tacticInfo, startPos, endPos, position) ← findTactics ids trees stx then
+      HighlightM.openTactic tacticInfo startPos endPos position
+      -- No nested tactics - the tactic search process should only have returned results
+      -- on "leaf" nodes anyway
+      tactics := false
   match stx with
   | `($e.%$tk$field:ident) =>
-      highlight' ids trees e
+      highlight' ids trees e tactics
       if let some ⟨pos, endPos⟩ := tk.getRange? then
         emitToken tk.getHeadInfo <| .mk .unknown <| (← getFileMap).source.extract pos endPos
       else
         emitString' "."
-      highlight' ids trees field
+      highlight' ids trees field tactics
   | _ =>
     match stx with
     | .missing => pure () -- TODO emit unhighlighted string
@@ -583,9 +628,9 @@ partial def highlight' (ids : HashMap Lsp.RefIdent Lsp.RefIdent) (trees : Persis
           match stx.identComponents (nFields? := some 1) with
           | [y, field] =>
             if (← infoExists trees field) then
-              highlight' ids trees y
+              highlight' ids trees y tactics
               emitToken' <| fakeToken .unknown "."
-              highlight' ids trees field
+              highlight' ids trees field tactics
             else
               emitToken i ⟨← identKind ids trees ⟨stx⟩, x.toString⟩
           | _ => emitToken i ⟨← identKind ids trees ⟨stx⟩, x.toString⟩
@@ -629,23 +674,23 @@ partial def highlight' (ids : HashMap Lsp.RefIdent Lsp.RefIdent) (trees : Persis
         let info := .original leading pos trailing endPos
         emitToken info ⟨← identKind ids trees ⟨stx⟩, s!".{x.toString}"⟩
       | _, _ =>
-        highlight' ids trees dot
-        highlight' ids trees name
+        highlight' ids trees dot tactics
+        highlight' ids trees name tactics
     | .node _ `choice alts =>
       -- Arbitrarily select one of the alternatives found by the parser. Otherwise, quotations of
       -- let-bindings with antiquotations as the bound variable leads to doubled bound variables,
       -- because the parser emits a choice node in the quotation. And it's never correct to show
       -- both alternatives!
       if h : alts.size > 0 then
-        highlight' ids trees alts[0]
+        highlight' ids trees alts[0]  tactics
     | stx@(.node _ k children) =>
       let pos := stx.getPos?
       for child in children do
-        highlight' ids trees child (lookingAt := pos.map (k, ·))
+        highlight' ids trees child tactics (lookingAt := pos.map (k, ·))
 
 def highlight (stx : Syntax) (messages : Array Message) (trees : PersistentArray Lean.Elab.InfoTree) : TermElabM Highlighted := do
   let modrefs := Lean.Server.findModuleRefs (← getFileMap) trees.toArray
   let ids := build modrefs
   let st ← HighlightState.ofMessages stx messages
-  let ((), {output := output, ..}) ← StateT.run (highlight' ids trees stx) st
+  let ((), {output := output, ..}) ← StateT.run (highlight' ids trees stx true) st
   pure <| .fromOutput output