From 31cc4a8b36f0cc675956cc782ea62f5e55974f01 Mon Sep 17 00:00:00 2001
From: David Thrane Christiansen <david@davidchristiansen.dk>
Date: Wed, 20 Nov 2024 12:36:23 +0100
Subject: [PATCH 1/4] feat: source file highlighting

Adds the ability to highlight a whole module.
---
 .github/workflows/ci.yml        |  1 +
 ExtractModule.lean              | 59 +++++++++++++++++++++++++++++++++
 lakefile.lean                   | 49 +++++++++++++++++++++++++++
 src/compat/SubVerso/Compat.lean | 12 +++++++
 4 files changed, 121 insertions(+)
 create mode 100644 ExtractModule.lean

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 7f7efa6..a3c294c 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -96,6 +96,7 @@ jobs:
           pushd demo
           lake update
           lake build :examples
+          lake build :highlighted
           popd
 
       - name: Run tests
diff --git a/ExtractModule.lean b/ExtractModule.lean
new file mode 100644
index 0000000..218104e
--- /dev/null
+++ b/ExtractModule.lean
@@ -0,0 +1,59 @@
+/-
+Copyright (c) 2023-2024 Lean FRO LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Author: David Thrane Christiansen
+-/
+import SubVerso.Compat
+import SubVerso.Examples.Env
+import Lean.Util.Paths
+
+open Lean Elab Frontend
+open Lean.Elab.Command (liftTermElabM)
+open SubVerso Examples
+open SubVerso.Highlighting (Highlighted highlight)
+
+
+def main : (args : List String) → IO UInt32
+  | [mod, outFile] => do
+    try
+      initSearchPath (← findSysroot)
+      let modName := mod.toName
+
+      let sp ← Compat.initSrcSearchPath
+      let fname ← do
+        if let some fname ← sp.findModuleWithExt "lean" modName then
+          pure fname
+        else
+          throw <| IO.userError s!"Failed to load {modName}"
+      let ictx := Parser.mkInputContext (← IO.FS.readFile fname) fname.toString
+      let (headerStx, parserState, msgs) ← Parser.parseHeader ictx
+      let imports := headerToImports headerStx
+      let env ← importModules imports {}
+      let pctx : Context := {inputCtx := ictx}
+
+      let commandState := {env, maxRecDepth := defaultMaxRecDepth, messages := msgs}
+      let cmdPos := parserState.pos
+      let cmdSt ← IO.mkRef {commandState, parserState, cmdPos}
+      processCommands pctx cmdSt
+
+      let cmdStx := (← cmdSt.get).commands
+      let infos := (← cmdSt.get).commandState.infoState.trees
+      let msgs := Compat.messageLogArray (← cmdSt.get).commandState.messages
+
+      let mut hls := Highlighted.empty
+      for cmd in #[headerStx] ++ cmdStx do
+        hls := hls ++ (← (Frontend.runCommandElabM <| liftTermElabM <| highlight cmd msgs infos) pctx cmdSt)
+      if let some p := (outFile : System.FilePath).parent then
+        IO.FS.createDirAll p
+      IO.FS.writeFile outFile (toString (ToJson.toJson hls) ++ "\n")
+      return (0 : UInt32)
+
+    catch e =>
+      IO.eprintln s!"error finding highlighted code: {toString e}"
+      return 2
+  | other => do
+    IO.eprintln s!"Didn't understand args: {other}"
+    pure 1
+where
+  relevant (mod : Name) (examples : NameMap (NameMap Json)) : List (String × Json) :=
+    examples.find? mod |>.getD {} |>.toList |>.map fun p => {p with fst := p.fst.toString (escape := false)}
diff --git a/lakefile.lean b/lakefile.lean
index 2583e19..76f6987 100644
--- a/lakefile.lean
+++ b/lakefile.lean
@@ -33,6 +33,12 @@ lean_exe «subverso-extract» where
   root := `Extract
   supportInterpreter := true
 
+@[default_target]
+lean_exe «subverso-extract-mod» where
+  root := `ExtractModule
+  supportInterpreter := true
+
+
 -- Compatibility shims for older Lake (where logging was manual) and
 -- newer Lake (where it isn't). Necessary from Lean 4.8.0 and up.
 section
@@ -45,6 +51,29 @@ open Lean Elab Command
     elabCommand <| ← `(def $(mkIdent `logInfo) [Pure $(mkIdent `m)] (message : String) : $(mkIdent `m) Unit := pure ())
 end
 
+module_facet highlighted mod : FilePath := do
+  let ws ← getWorkspace
+  let some extract ← findLeanExe? `«subverso-extract-mod»
+    | error "The subverso-extract-mod executable was not found"
+
+  let exeJob ← extract.exe.fetch
+  let modJob ← mod.olean.fetch
+
+  let buildDir := ws.root.buildDir
+  let hlFile := mod.filePath (buildDir / "highlighted") "json"
+
+  exeJob.bindAsync fun exeFile exeTrace =>
+    modJob.bindSync fun _oleanPath modTrace => do
+      let depTrace := mixTrace exeTrace modTrace
+      let trace ← buildFileUnlessUpToDate hlFile depTrace do
+        logStep s!"Exporting highlighted source file JSON for '{mod.name}'"
+        proc {
+          cmd := exeFile.toString
+          args := #[mod.name.toString, hlFile.toString]
+          env := ← getAugmentedEnv
+        }
+      pure (hlFile, trace)
+
 module_facet examples mod : FilePath := do
   let ws ← getWorkspace
   let some extract ← findLeanExe? `«subverso-extract»
@@ -68,6 +97,16 @@ module_facet examples mod : FilePath := do
         }
       pure (hlFile, trace)
 
+library_facet highlighted lib : FilePath := do
+  let ws ← getWorkspace
+  let mods ← lib.modules.fetch
+  let moduleJobs ← BuildJob.mixArray <| ← mods.mapM (fetch <| ·.facet `highlighted)
+  let buildDir := ws.root.buildDir
+  let hlDir := buildDir / "highlighted"
+  moduleJobs.bindSync fun () trace => do
+    pure (hlDir, trace)
+
+
 library_facet examples lib : FilePath := do
   let ws ← getWorkspace
   let mods ← lib.modules.fetch
@@ -77,6 +116,16 @@ library_facet examples lib : FilePath := do
   moduleJobs.bindSync fun () trace => do
     pure (hlDir, trace)
 
+package_facet highlighted pkg : FilePath := do
+  let ws ← getWorkspace
+  let libs := pkg.leanLibs
+  let libJobs ← BuildJob.mixArray <| ← libs.mapM (fetch <| ·.facet `highlighted)
+  let buildDir := ws.root.buildDir
+  let hlDir := buildDir / "highlighted"
+  libJobs.bindSync fun () trace => do
+    logInfo s!"Highlighted code written to '{hlDir}'"
+    pure (hlDir, trace)
+
 package_facet examples pkg : FilePath := do
   let ws ← getWorkspace
   let libs := pkg.leanLibs
diff --git a/src/compat/SubVerso/Compat.lean b/src/compat/SubVerso/Compat.lean
index 8064bcc..b745ec1 100644
--- a/src/compat/SubVerso/Compat.lean
+++ b/src/compat/SubVerso/Compat.lean
@@ -4,6 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Author: David Thrane Christiansen
 -/
 import Lean.Elab
+import Lean.Util.Paths
 
 open Lean Elab Term
 
@@ -116,6 +117,17 @@ def rwTacticRightBracket? (stx : Syntax) : Option Syntax := Id.run do
 def getDeclarationRange? [Monad m] [MonadFileMap m] (stx : Syntax) : m (Option DeclarationRange) :=
   %first_succeeding [Lean.Elab.getDeclarationRange? stx, some <$> Lean.Elab.getDeclarationRange stx]
 
+def messageLogArray (msgs : Lean.MessageLog) : Array Lean.Message := %first_succeeding [msgs.toArray, msgs.msgs.toArray]
+
+def initSrcSearchPath (pkgSearchPath : SearchPath := ∅) : IO SearchPath := do
+  %first_succeeding [
+    Lean.initSrcSearchPath (pkgSearchPath := pkgSearchPath),
+    Lean.initSrcSearchPath (sp := pkgSearchPath),
+    -- leanSysRoot seems to never have been used by this function
+    Lean.initSrcSearchPath (leanSysroot := "") (sp := pkgSearchPath),
+    Lean.initSrcSearchPath (_leanSysroot := "") (sp := pkgSearchPath)
+  ]
+
 namespace NameSet
 def union (xs ys : NameSet) : NameSet :=
   xs.mergeBy (fun _ _ _ => .unit) ys

From 3537a46c54572d3a3a99a8bf7154d2f08990ee39 Mon Sep 17 00:00:00 2001
From: David Thrane Christiansen <david@davidchristiansen.dk>
Date: Thu, 21 Nov 2024 15:57:53 +0100
Subject: [PATCH 2/4] feat: also highlight executables

---
 lakefile.lean | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/lakefile.lean b/lakefile.lean
index 76f6987..c724fc1 100644
--- a/lakefile.lean
+++ b/lakefile.lean
@@ -116,10 +116,12 @@ library_facet examples lib : FilePath := do
   moduleJobs.bindSync fun () trace => do
     pure (hlDir, trace)
 
+
 package_facet highlighted pkg : FilePath := do
   let ws ← getWorkspace
   let libs := pkg.leanLibs
-  let libJobs ← BuildJob.mixArray <| ← libs.mapM (fetch <| ·.facet `highlighted)
+  let exes := pkg.leanExes.map (·.toLeanLib)
+  let libJobs ← BuildJob.mixArray <| ← (libs ++ exes).mapM (fetch <| ·.facet `highlighted)
   let buildDir := ws.root.buildDir
   let hlDir := buildDir / "highlighted"
   libJobs.bindSync fun () trace => do

From 0f51d9fb51a866e28f0bd4bfce26ab2c9f54f988 Mon Sep 17 00:00:00 2001
From: David Thrane Christiansen <david@davidchristiansen.dk>
Date: Thu, 21 Nov 2024 16:29:13 +0100
Subject: [PATCH 3/4] fix: more informative error message

---
 ExtractModule.lean | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ExtractModule.lean b/ExtractModule.lean
index 218104e..034e91c 100644
--- a/ExtractModule.lean
+++ b/ExtractModule.lean
@@ -24,7 +24,7 @@ def main : (args : List String) → IO UInt32
         if let some fname ← sp.findModuleWithExt "lean" modName then
           pure fname
         else
-          throw <| IO.userError s!"Failed to load {modName}"
+          throw <| IO.userError s!"Failed to load {modName} from {sp}"
       let ictx := Parser.mkInputContext (← IO.FS.readFile fname) fname.toString
       let (headerStx, parserState, msgs) ← Parser.parseHeader ictx
       let imports := headerToImports headerStx

From 179e931b437d7b564ddf44d68b60e143c18940c4 Mon Sep 17 00:00:00 2001
From: David Thrane Christiansen <david@davidchristiansen.dk>
Date: Thu, 21 Nov 2024 16:36:14 +0100
Subject: [PATCH 4/4] fix: look here for module as well

---
 ExtractModule.lean | 1 +
 1 file changed, 1 insertion(+)

diff --git a/ExtractModule.lean b/ExtractModule.lean
index 034e91c..2d4d6de 100644
--- a/ExtractModule.lean
+++ b/ExtractModule.lean
@@ -20,6 +20,7 @@ def main : (args : List String) → IO UInt32
       let modName := mod.toName
 
       let sp ← Compat.initSrcSearchPath
+      let sp : SearchPath := (sp : List System.FilePath) ++ [("." : System.FilePath)]
       let fname ← do
         if let some fname ← sp.findModuleWithExt "lean" modName then
           pure fname