perf: lake: fix LogIO inling/lifting (#4351)

The current manner of lifting `LogIO` into `CliM` produces excessive
specializations (due to a nested inlined `forM`). There was also a bug
where `IO` was lifted into `CliM` via `LogIO` rather than directly
through `MainM`.

src/lake/Lake/CLI/Main.lean

@@ -99,8 +99,11 @@ def CliM.run (self : CliM α) (args : List String) : BaseIO ExitCode := do
   let main := main.run >>= fun | .ok a => pure a | .error e => error e.toString
   main.run
 
-instance : MonadLift LogIO CliStateM :=
-  ⟨fun x => do MainM.runLogIO x (← get).verbosity.minLogLv (← get).ansiMode⟩
+@[inline] def CliStateM.runLogIO (x : LogIO α) : CliStateM α := do
+  let opts ← get
+  MainM.runLogIO x opts.verbosity.minLogLv opts.ansiMode
+
+instance (priority := low) : MonadLift LogIO CliStateM := ⟨CliStateM.runLogIO⟩
 
 /-! ## Argument Parsing -/
 
@@ -273,7 +276,7 @@ protected def doc : CliM PUnit := do
     | none => throw <| CliError.missingScriptDoc script.name
 
 protected def help : CliM PUnit := do
-  IO.println <| helpScript <| (← takeArg?).getD ""
+  IO.println <| helpScript <| ← takeArgD ""
 
 end script
 
@@ -290,14 +293,14 @@ protected def new : CliM PUnit := do
   processOptions lakeOption
   let opts ← getThe LakeOptions
   let name ← takeArg "package name"
-  let (tmp, lang) ← parseTemplateLangSpec <| (← takeArg?).getD ""
+  let (tmp, lang) ← parseTemplateLangSpec <| ← takeArgD ""
   noArgsRem do new name tmp lang (← opts.computeEnv) opts.rootDir
 
 protected def init : CliM PUnit := do
   processOptions lakeOption
   let opts ← getThe LakeOptions
-  let name := (← takeArg?).getD "."
-  let (tmp, lang) ← parseTemplateLangSpec <| (← takeArg?).getD ""
+  let name := ← takeArgD "."
+  let (tmp, lang) ← parseTemplateLangSpec <| ← takeArgD ""
   noArgsRem do init name tmp lang (← opts.computeEnv) opts.rootDir
 
 protected def build : CliM PUnit := do
@@ -480,7 +483,7 @@ protected def selfCheck : CliM PUnit := do
   noArgsRem do verifyInstall (← getThe LakeOptions)
 
 protected def help : CliM PUnit := do
-  IO.println <| help <| (← takeArg?).getD ""
+  IO.println <| help <| ← takeArgD ""
 
 end lake
 

src/lake/Lake/Util/Cli.lean

@@ -49,6 +49,10 @@ variable [Monad m] [MonadStateOf ArgList m]
 @[inline] def takeArg? : m (Option String) :=
   modifyGet fun | [] => (none, []) | arg :: args => (some arg, args)
 
+/-- Take the head of the remaining argument list (or `default` if none). -/
+@[inline] def takeArgD (default : String) : m String :=
+  modifyGet fun | [] => (default, []) | arg :: args => (arg, args)
+
 /-- Take the remaining argument list, leaving only an empty list. -/
 @[inline] def takeArgs : m (List String) :=
   modifyGet fun args => (args, [])

src/lake/Lake/Util/Lift.lean

@@ -5,6 +5,9 @@ Authors: Mac Malone
 -/
 namespace Lake
 
+/-- Ensure direct lifts are preferred over indirect ones. -/
+instance (priority := high) [MonadLift α β] : MonadLiftT α β := ⟨MonadLift.monadLift⟩
+
 instance [Pure m] : MonadLiftT Id m where
   monadLift act := pure act.run
 

src/lake/Lake/Util/Log.lean

@@ -541,6 +541,13 @@ Prints log entries of at least `minLv` to `out`.
 @[inline] def toBaseIO (self : LogIO α)
   (minLv := LogLevel.info) (ansiMode := AnsiMode.auto) (out := OutStream.stderr)
 : BaseIO (Option α) := do
-  self.replayLog? (logger := ← out.getLogger minLv ansiMode)
+  let logger ← out.getLogger minLv ansiMode
+  let (a?, log) ← self.run? {}
+  replay log logger
+  return a?
+where
+  -- avoid specialization of this call at each call site
+  replay (log : Log) (logger : MonadLog BaseIO) : BaseIO Unit :=
+    log.replay (logger := logger)
 
 end LogIO

src/lake/Lake/Util/MainM.lean

@@ -79,6 +79,13 @@ instance : MonadLift IO MainM := ⟨MonadError.runIO⟩
 @[inline] def runLogIO (x : LogIO α)
   (minLv := LogLevel.info) (ansiMode := AnsiMode.auto) (out := OutStream.stderr)
 : MainM α := do
-  x.replayLog (logger := ← out.getLogger minLv ansiMode)
+  let logger ← out.getLogger minLv ansiMode
+  match (← x {}) with
+  | .ok a  log => replay log logger; return a
+  | .error _ log => replay log logger; exit 1
+where
+  -- avoid specialization of this call at each call site
+  replay (log : Log) (logger : MonadLog BaseIO) : BaseIO Unit :=
+    log.replay (logger := logger)
 
-instance : MonadLift LogIO MainM := ⟨runLogIO⟩
+instance (priority := low) : MonadLift LogIO MainM := ⟨runLogIO⟩