feat: zero cost BaseIO

src/Init/Data/Format/Basic.lean

@@ -411,7 +411,6 @@ Renders a `Format` to a string.
 * `column`: begin the first line wrap `column` characters earlier than usual
   (this is useful when the output String will be printed starting at `column`)
 -/
-@[export lean_format_pretty]
 def pretty (f : Format) (width : Nat := defWidth) (indent : Nat := 0) (column := 0) : String :=
   let act : StateM State Unit := prettyM f width indent
   State.out <| act (State.mk "" column) |>.snd

src/Init/Internal/Order/Basic.lean

@@ -16,7 +16,7 @@ public import Init.System.IO  -- for `MonoBind` instance
 import all Init.Control.Except  -- for `MonoBind` instance
 import all Init.Control.StateRef  -- for `MonoBind` instance
 import all Init.Control.Option  -- for `MonoBind` instance
-import all Init.System.IO  -- for `MonoBind` instance
+import all Init.System.ST  -- for `MonoBind` instance
 
 public section
 
@@ -926,28 +926,41 @@ theorem monotone_stateTRun [PartialOrder γ]
     monotone (fun (x : γ) => StateT.run (f x) s) :=
   monotone_apply s _ hmono
 
--- TODO: axiomatize these instances (ideally without `Nonempty ε`) when EIO is opaque
-noncomputable instance [Nonempty ε] : CCPO (EIO ε α) :=
+-- TODO: axiomatize these instances (ideally without `Nonempty ε`) when EIO and friends are opaque
+
+noncomputable instance [Nonempty ε] : CCPO (EST ε σ α) :=
   inferInstanceAs (CCPO ((s : _) → FlatOrder (.error Classical.ofNonempty (Classical.choice ⟨s⟩))))
 
-noncomputable instance [Nonempty ε] : MonoBind (EIO ε) where
+noncomputable instance [Nonempty ε] : MonoBind (EST ε σ) where
   bind_mono_left {_ _ a₁ a₂ f} h₁₂ := by
     intro s
     specialize h₁₂ s
     change FlatOrder.rel (a₁.bind f s) (a₂.bind f s)
-    simp only [EStateM.bind]
+    simp only [EST.bind]
     generalize a₁ s = a₁ at h₁₂; generalize a₂ s = a₂ at h₁₂
     cases h₁₂
     · exact .bot
     · exact .refl
   bind_mono_right {_ _ a f₁ f₂} h₁₂ := by
     intro w
     change FlatOrder.rel (a.bind f₁ w) (a.bind f₂ w)
-    simp only [EStateM.bind]
+    simp only [EST.bind]
     split
     · apply h₁₂
     · exact .refl
 
+noncomputable instance [Nonempty α] : CCPO (ST σ α) :=
+  inferInstanceAs (CCPO ((s : _) → FlatOrder (.mk Classical.ofNonempty (Classical.choice ⟨s⟩))))
+
+noncomputable instance [Nonempty α] : CCPO (BaseIO α) :=
+  inferInstanceAs (CCPO (ST IO.RealWorld α))
+
+noncomputable instance [Nonempty ε] : CCPO (EIO ε α) :=
+  inferInstanceAs (CCPO (EST ε IO.RealWorld α))
+
+noncomputable instance [Nonempty ε] : MonoBind (EIO ε) :=
+  inferInstanceAs (MonoBind (EST ε IO.RealWorld))
+
 end mono_bind
 
 section implication_order

src/Init/Prelude.lean

@@ -23,6 +23,16 @@ use `PUnit` in the desugaring of `do` notation, or in the pattern match compiler
 
 universe u v w
 
+/-- Marker for information that has been erased by the code generator. -/
+unsafe axiom lcErased : Type
+
+/-- Marker for type dependency that has been erased by the code generator. -/
+unsafe axiom lcAny : Type
+
+/-- Internal representation of `Void` in the compiler. -/
+unsafe axiom lcVoid : Type
+
+
 /--
 The identity function. `id` takes an implicit argument `α : Sort u`
 (a type in any universe), and an argument `a : α`, and returns `a`.
@@ -135,15 +145,6 @@ It can be written as an empty tuple: `()`.
 -/
 @[match_pattern] abbrev Unit.unit : Unit := PUnit.unit
 
-/-- Marker for information that has been erased by the code generator. -/
-unsafe axiom lcErased : Type
-
-/-- Marker for type dependency that has been erased by the code generator. -/
-unsafe axiom lcAny : Type
-
-/-- Internal representation of `IO.RealWorld` in the compiler. -/
-unsafe axiom lcRealWorld : Type
-
 /--
 Auxiliary unsafe constant used by the Compiler when erasing proofs from code.
 

src/Init/System/IO.lean

@@ -1,7 +1,7 @@
 /-
 Copyright (c) 2017 Microsoft Corporation. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
-Authors: Luke Nelson, Jared Roesch, Leonardo de Moura, Sebastian Ullrich, Mac Malone
+Authors: Luke Nelson, Jared Roesch, Leonardo de Moura, Sebastian Ullrich, Mac Malone, Henrik Böving
 -/
 module
 
@@ -18,7 +18,6 @@ public section
 open System
 
 opaque IO.RealWorld.nonemptyType : NonemptyType.{0}
-
 /--
 A representation of “the real world” that's used in `IO` monads to ensure that `IO` actions are not
 reordered.
@@ -28,6 +27,18 @@ reordered.
 instance IO.RealWorld.instNonempty : Nonempty IO.RealWorld :=
   by exact IO.RealWorld.nonemptyType.property
 
+/--
+An `IO` monad that cannot throw exceptions.
+-/
+@[expose] def BaseIO (α : Type) := ST IO.RealWorld α
+
+instance : Monad BaseIO := inferInstanceAs (Monad (ST IO.RealWorld))
+instance : MonadFinally BaseIO := inferInstanceAs (MonadFinally (ST IO.RealWorld))
+
+@[always_inline, inline]
+def BaseIO.map (f : α → β) (x : BaseIO α) : BaseIO β :=
+  f <$> x
+
 /--
 A monad that can have side effects on the external world or throw exceptions of type `ε`.
 
@@ -41,21 +52,7 @@ A monad that can have side effects on the external world or throw exceptions of
    def getWorld : IO (IO.RealWorld) := get
    ```
 -/
-@[expose] def EIO (ε : Type) : Type → Type := EStateM ε IO.RealWorld
-
-instance : Monad (EIO ε) := inferInstanceAs (Monad (EStateM ε IO.RealWorld))
-instance : MonadFinally (EIO ε) := inferInstanceAs (MonadFinally (EStateM ε IO.RealWorld))
-instance : MonadExceptOf ε (EIO ε) := inferInstanceAs (MonadExceptOf ε (EStateM ε IO.RealWorld))
-instance : OrElse (EIO ε α) := ⟨MonadExcept.orElse⟩
-instance [Inhabited ε] : Inhabited (EIO ε α) := inferInstanceAs (Inhabited (EStateM ε IO.RealWorld α))
-
-/--
-An `IO` monad that cannot throw exceptions.
--/
-@[expose] def BaseIO := EIO Empty
-
-instance : Monad BaseIO := inferInstanceAs (Monad (EIO Empty))
-instance : MonadFinally BaseIO := inferInstanceAs (MonadFinally (EIO Empty))
+@[expose] def EIO (ε : Type) (α : Type) : Type := EST ε IO.RealWorld α
 
 /--
 Runs a `BaseIO` action, which cannot throw an exception, in any other `EIO` monad.
@@ -66,7 +63,7 @@ lifting](lean-manual://section/lifting-monads) rather being than called explicit
 @[always_inline, inline]
 def BaseIO.toEIO (act : BaseIO α) : EIO ε α :=
   fun s => match act s with
-  | EStateM.Result.ok a s => EStateM.Result.ok a s
+  | .mk a s => .ok a s
 
 instance : MonadLift BaseIO (EIO ε) := ⟨BaseIO.toEIO⟩
 
@@ -77,8 +74,8 @@ action that returns an `Except` value.
 @[always_inline, inline]
 def EIO.toBaseIO (act : EIO ε α) : BaseIO (Except ε α) :=
   fun s => match act s with
-  | EStateM.Result.ok a s     => EStateM.Result.ok (Except.ok a) s
-  | EStateM.Result.error ex s => EStateM.Result.ok (Except.error ex) s
+  | .ok a s     => .mk (.ok a) s
+  | .error ex s => .mk (.error ex) s
 
 /--
 Handles any exception that might be thrown by an `EIO ε` action, transforming it into an
@@ -87,8 +84,25 @@ exception-free `BaseIO` action.
 @[always_inline, inline]
 def EIO.catchExceptions (act : EIO ε α) (h : ε → BaseIO α) : BaseIO α :=
   fun s => match act s with
-  | EStateM.Result.ok a s     => EStateM.Result.ok a s
-  | EStateM.Result.error ex s => h ex s
+  | .ok a s     => .mk a s
+  | .error ex s => h ex s
+
+instance : Monad (EIO ε) := inferInstanceAs (Monad (EST ε IO.RealWorld))
+instance : MonadFinally (EIO ε) := inferInstanceAs (MonadFinally (EST ε IO.RealWorld))
+instance : MonadExceptOf ε (EIO ε) := inferInstanceAs (MonadExceptOf ε (EST ε IO.RealWorld))
+instance : OrElse (EIO ε α) := ⟨MonadExcept.orElse⟩
+instance [Inhabited ε] : Inhabited (EIO ε α) := inferInstanceAs (Inhabited (EST ε IO.RealWorld α))
+
+@[always_inline, inline]
+def EIO.map (f : α → β) (x : EIO ε α) : EIO ε β :=
+  f <$> x
+
+@[always_inline, inline]
+protected def EIO.throw (e : ε) : EIO ε α := throw e
+
+@[always_inline, inline]
+protected def EIO.tryCatch (x : EIO ε α) (handle : ε → EIO ε α) : EIO ε α :=
+  MonadExceptOf.tryCatch x handle
 
 /--
 Converts an `Except ε` action into an `EIO ε` action.
@@ -101,6 +115,15 @@ def EIO.ofExcept (e : Except ε α) : EIO ε α :=
   | Except.ok a    => pure a
   | Except.error e => throw e
 
+@[always_inline, inline]
+def EIO.adapt (f : ε → ε') (m : EIO ε α) : EIO ε' α :=
+  fun s => match m s with
+  | .ok a s => .ok a s
+  | .error e s => .error (f e) s
+
+@[deprecated EIO.adapt (since := "2025-09-29"), always_inline, inline]
+def EIO.adaptExcept (f : ε → ε') (m : EIO ε α) : EIO ε' α := EIO.adapt f m
+
 open IO (Error) in
 /--
 A monad that supports arbitrary side effects and throwing exceptions of type `IO.Error`.
@@ -121,7 +144,7 @@ Converts an `EIO ε` action into an `IO` action by translating any exceptions th
 `IO.Error`s using `f`.
 -/
 @[inline] def EIO.toIO (f : ε → IO.Error) (act : EIO ε α) : IO α :=
-  act.adaptExcept f
+  act.adapt f
 
 /--
 Converts an `EIO ε` action that might throw an exception of type `ε` into an exception-free `IO`
@@ -134,7 +157,7 @@ action that returns an `Except` value.
 Runs an `IO` action in some other `EIO` monad, using `f` to translate `IO` exceptions.
 -/
 @[inline] def IO.toEIO (f : IO.Error → ε) (act : IO α) : EIO ε α :=
-  act.adaptExcept f
+  act.adapt f
 
 /- After we inline `EState.run'`, the closed term `((), ())` is generated, where the second `()`
    represents the "initial world". We don't want to cache this closed term. So, we disable
@@ -154,8 +177,8 @@ duplicate, or delete calls to this function. The side effect may even be hoisted
 causing the side effect to occur at initialization time, even if it would otherwise never be called.
 -/
 @[noinline] unsafe def unsafeBaseIO (fn : BaseIO α) : α :=
-  match fn.run (unsafeCast Unit.unit) with
-  | EStateM.Result.ok a _ => a
+  match fn (unsafeCast Unit.unit) with
+  | .mk a _ => a
 
 /--
 Executes arbitrary side effects in a pure context, with exceptions indicated via `Except`. This a
@@ -400,7 +423,7 @@ Pauses execution for the specified number of milliseconds.
 -/
 opaque sleep (ms : UInt32) : BaseIO Unit :=
   -- TODO: add a proper primitive for IO.sleep
-  fun s => dbgSleep ms fun _ => EStateM.Result.ok () s
+  fun s => dbgSleep ms fun _ => .mk () s
 
 /--
 Runs `act` in a separate `Task`, with priority `prio`. Because `IO` actions may throw an exception
@@ -1609,7 +1632,10 @@ the `IO` monad.
 -/
 abbrev Ref (α : Type) := ST.Ref IO.RealWorld α
 
-instance : MonadLift (ST IO.RealWorld) BaseIO := ⟨id⟩
+instance : MonadLift (ST IO.RealWorld) BaseIO where
+  monadLift mx := fun s =>
+    match mx s with
+    | .mk s a => .mk s a
 
 /--
 Creates a new mutable reference cell that contains `a`.