bookshelf-doc/DocGen4/Load.lean

import Lean
import DocGen4.Process
import Std.Data.HashMap

namespace DocGen4

open Lean System Std IO

def getLakePath : IO FilePath := do
  match (← IO.getEnv "LAKE") with
  | some path => System.FilePath.mk path
  | none =>
    let lakePath := (←findSysroot?) / "bin" / "lake"
    lakePath.withExtension System.FilePath.exeExtension

-- Modified from the LSP Server
def lakeSetupSearchPath (lakePath : System.FilePath) (imports : Array String) : IO Lean.SearchPath := do
  let args := #["print-paths"] ++ imports
  let cmdStr := " ".intercalate (toString lakePath :: args.toList)
  let lakeProc ← Process.spawn {
    stdin  := Process.Stdio.null
    stdout := Process.Stdio.piped
    stderr := Process.Stdio.piped
    cmd    := lakePath.toString
    args
  }
  let stdout := String.trim (← lakeProc.stdout.readToEnd)
  let stderr := String.trim (← lakeProc.stderr.readToEnd)
  match (← lakeProc.wait) with
  | 0 =>
    let stdout := stdout.split (· == '\n') |>.getLast!
    let Except.ok (paths : LeanPaths) ← pure (Json.parse stdout >>= fromJson?)
      | throw $ userError s!"invalid output from `{cmdStr}`:\n{stdout}\nstderr:\n{stderr}"
    initSearchPath (← findSysroot?) paths.oleanPath
    paths.oleanPath.mapM realPathNormalized
  | 2 => pure []  -- no lakefile.lean
  | _ => throw $ userError s!"`{cmdStr}` failed:\n{stdout}\nstderr:\n{stderr}"

def load (imports : List Name) : IO (HashMap Name Module) := do
  let env ← importModules (List.map (Import.mk · false) imports) Options.empty
  -- TODO parameterize maxHeartbeats
  let doc ← Prod.fst <$> (Meta.MetaM.toIO process { maxHeartbeats := 100000000} { env := env} {} {})
  for (_, module) in doc.toList do
    let s ← Core.CoreM.toIO module.prettyPrint {} { env := env }
    IO.println s.fst
  IO.println s!"Processed {List.foldl (λ a (_, b) => a + b.members.size) 0 doc.toList} declarations"
  IO.println s!"Processed {doc.size} modules"
  return doc

end DocGen4