From 29a249e8fdd640c7367415c063d817ad1aebe66d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Henrik=20B=C3=B6ving?= <hargonix@gmail.com>
Date: Sun, 5 Dec 2021 14:56:25 +0100
Subject: [PATCH] feat: Automatic search path + modules as CLI arguments

---
 DocGen4/Load.lean | 36 ++++++++++++++++++++++++++++++------
 Main.lean         | 12 ++++++------
 2 files changed, 36 insertions(+), 12 deletions(-)

diff --git a/DocGen4/Load.lean b/DocGen4/Load.lean
index 2df73fd..e44cd95 100644
--- a/DocGen4/Load.lean
+++ b/DocGen4/Load.lean
@@ -4,18 +4,42 @@ import Std.Data.HashMap
 
 namespace DocGen4
 
-open Lean System Std
+open Lean System Std IO
 
-def printSearchPath : IO PUnit := do
-  IO.println s!"{←searchPathRef.get}"
+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
 
-def setSearchPath (path : List FilePath) : IO PUnit := do
-    searchPathRef.set path
+-- 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} {} {})
+  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
diff --git a/Main.lean b/Main.lean
index 639455f..a4bab64 100644
--- a/Main.lean
+++ b/Main.lean
@@ -1,11 +1,11 @@
 import DocGen4
 import Lean
 
-open DocGen4 Lean
+open DocGen4 Lean IO
 
-
-def main : IO Unit := do
-  -- This should be set by lake at some point
-  setSearchPath ["/home/nix/Desktop/formal_verification/lean/mathlib4/build/lib", "/home/nix/.elan/toolchains/leanprover--lean4---nightly-2021-11-24/lib/lean"]
-  let doc ← load [`Mathlib]
+def main (args : List String) : IO Unit := do
+  let modules := args
+  let path ← lakeSetupSearchPath (←getLakePath) modules.toArray
+  IO.println s!"Loading modules from: {path}"
+  let doc ← load $ modules.map Name.mkSimple
   return ()