From ca41db6da1a7ca1e42e0c308eb081aa9b833075a Mon Sep 17 00:00:00 2001
From: David Thrane Christiansen <david@davidchristiansen.dk>
Date: Mon, 29 Apr 2024 07:39:03 +0200
Subject: [PATCH] fix: correctly pass info trees in highlight generation

Fixes #17. Also adds CI to compare versions of Lean for equivalent
proof output.
---
 Tests.lean                                    | 79 ++++++++++++++++++-
 demo/Demo.lean                                | 13 +++
 demo/lean-toolchain                           |  2 +-
 src/examples/SubVerso/Examples.lean           | 15 ++--
 .../SubVerso/Highlighting/Code.lean           |  8 +-
 5 files changed, 107 insertions(+), 10 deletions(-)

diff --git a/Tests.lean b/Tests.lean
index 99bfd23..24bcaa9 100644
--- a/Tests.lean
+++ b/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
diff --git a/demo/Demo.lean b/demo/Demo.lean
index 19206f7..8a9e226 100644
--- a/demo/Demo.lean
+++ b/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
diff --git a/demo/lean-toolchain b/demo/lean-toolchain
index 6ff4890..5cadc9d 100644
--- a/demo/lean-toolchain
+++ b/demo/lean-toolchain
@@ -1 +1 @@
-leanprover/lean4:nightly-2024-02-26
+leanprover/lean4:v4.3.0
diff --git a/src/examples/SubVerso/Examples.lean b/src/examples/SubVerso/Examples.lean
index 5cd2acb..20a6c23 100644
--- a/src/examples/SubVerso/Examples.lean
+++ b/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
diff --git a/src/highlighting/SubVerso/Highlighting/Code.lean b/src/highlighting/SubVerso/Highlighting/Code.lean
index e14901f..1181f35 100644
--- a/src/highlighting/SubVerso/Highlighting/Code.lean
+++ b/src/highlighting/SubVerso/Highlighting/Code.lean
@@ -496,11 +496,13 @@ 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
+def findTactics
+    (ids : HashMap Lsp.RefIdent Lsp.RefIdent)
+    (trees : PersistentArray Lean.Elab.InfoTree)
+    (stx : Syntax) : HighlightM Unit := do
   -- 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
   for t in trees do
     -- The info is reversed here so that the _last_ state computed is shown.
@@ -561,7 +563,7 @@ def findTactics (ids : HashMap Lsp.RefIdent Lsp.RefIdent) (stx : Syntax) : Highl
         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
+  findTactics ids trees stx
   match stx with
   | `($e.%$tk$field:ident) =>
       highlight' ids trees e