feat: source file highlighting (#61)

Adds the ability to highlight a whole module.

.github/workflows/ci.yml

@@ -96,6 +96,7 @@ jobs:
           pushd demo
           lake update
           lake build :examples
+          lake build :highlighted
           popd
 
       - name: Run tests

ExtractModule.lean

@@ -0,0 +1,60 @@
+/-
+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 sp : SearchPath := (sp : List System.FilePath) ++ [("." : System.FilePath)]
+      let fname ← do
+        if let some fname ← sp.findModuleWithExt "lean" modName then
+          pure fname
+        else
+          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
+      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)}

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,18 @@ 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 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
+    logInfo s!"Highlighted code written to '{hlDir}'"
+    pure (hlDir, trace)
+
 package_facet examples pkg : FilePath := do
   let ws ← getWorkspace
   let libs := pkg.leanLibs

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