chore: replace implementation with Environment.replay

Main.lean

@@ -4,148 +4,11 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Scott Morrison
 -/
 import Lean.CoreM
+import Lean.Replay
 import Lean4Checker.Lean
 
 open Lean
 
-structure Context where
-  newConstants : HashMap Name ConstantInfo
-
-structure State where
-  env : Environment
-  remaining : NameSet := {}
-  pending : NameSet := {}
-  postponedConstructors : NameSet := {}
-  postponedRecursors : NameSet := {}
-
-abbrev M := ReaderT Context <| StateRefT State IO
-
-/-- Check if a `Name` still needs processing. If so, move it from `remaining` to `pending`. -/
-def isTodo (name : Name) : M Bool := do
-  let r := (← get).remaining
-  if r.contains name then
-    modify fun s => { s with remaining := s.remaining.erase name, pending := s.pending.insert name }
-    return true
-  else
-    return false
-
-/-- Use the current `Environment` to throw a `KernelException`. -/
-def throwKernelException (ex : KernelException) : M Unit := do
-    let ctx := { fileName := "", options := ({} : KVMap), fileMap := default }
-    let state := { env := (← get).env }
-    Prod.fst <$> (Lean.Core.CoreM.toIO · ctx state) do Lean.throwKernelException ex
-
-/-- Add a declaration, possibly throwing a `KernelException`. -/
-def addDecl (d : Declaration) : M Unit := do
-  match (← get).env.addDecl d with
-  | .ok env => modify fun s => { s with env := env }
-  | .error ex => throwKernelException ex
-
-deriving instance BEq for ConstantVal
-deriving instance BEq for ConstructorVal
-deriving instance BEq for RecursorRule
-deriving instance BEq for RecursorVal
-
-mutual
-/--
-Check if a `Name` still needs to be processed (i.e. is in `remaining`).
-
-If so, recursively replay any constants it refers to,
-to ensure we add declarations in the right order.
-
-The construct the `Declaration` from its stored `ConstantInfo`,
-and add it to the environment.
--/
-partial def replayConstant (name : Name) : M Unit := do
-  if ← isTodo name then
-    let some ci := (← read).newConstants.find? name | unreachable!
-    replayConstants ci.getUsedConstantsAsSet
-    -- Check that this name is still pending: a mutual block may have taken care of it.
-    if (← get).pending.contains name then
-      match ci with
-      | .defnInfo   info =>
-        addDecl (Declaration.defnDecl   info)
-      | .thmInfo    info =>
-        addDecl (Declaration.thmDecl    info)
-      | .axiomInfo  info =>
-        addDecl (Declaration.axiomDecl  info)
-      | .opaqueInfo info =>
-        addDecl (Declaration.opaqueDecl info)
-      | .inductInfo info =>
-        let lparams := info.levelParams
-        let nparams := info.numParams
-        let all ← info.all.mapM fun n => do pure <| ((← read).newConstants.find! n)
-        for o in all do
-          modify fun s =>
-            { s with remaining := s.remaining.erase o.name, pending := s.pending.erase o.name }
-        let ctorInfo ← all.mapM fun ci => do
-          pure (ci, ← ci.inductiveVal!.ctors.mapM fun n => do
-            pure ((← read).newConstants.find! n))
-        -- Make sure we are really finished with the constructors.
-        for (_, ctors) in ctorInfo do
-          for ctor in ctors do
-            replayConstants ctor.getUsedConstantsAsSet
-        let types : List InductiveType := ctorInfo.map fun ⟨ci, ctors⟩ =>
-          { name := ci.name
-            type := ci.type
-            ctors := ctors.map fun ci => { name := ci.name, type := ci.type } }
-        addDecl (Declaration.inductDecl lparams nparams types false)
-      -- We postpone checking constructors,
-      -- and at the end make sure they are identical
-      -- to the constructors generated when we replay the inductives.
-      | .ctorInfo info =>
-        modify fun s => { s with postponedConstructors := s.postponedConstructors.insert info.name }
-      -- Similarly we postpone checking recursors.
-      | .recInfo info =>
-        modify fun s => { s with postponedRecursors := s.postponedRecursors.insert info.name }
-      | .quotInfo _ =>
-        addDecl (Declaration.quotDecl)
-      modify fun s => { s with pending := s.pending.erase name }
-
-/-- Replay a set of constants one at a time. -/
-partial def replayConstants (names : NameSet) : M Unit := do
-  for n in names do replayConstant n
-
-end
-
-/--
-Check that all postponed constructors are identical to those generated
-when we replayed the inductives.
--/
-def checkPostponedConstructors : M Unit := do
-  for ctor in (← get).postponedConstructors do
-    match (← get).env.constants.find? ctor, (← read).newConstants.find? ctor with
-    | some (.ctorInfo info), some (.ctorInfo info') =>
-      if ! (info == info') then throw <| IO.userError s!"Invalid constructor {ctor}"
-    | _, _ => throw <| IO.userError s!"No such constructor {ctor}"
-
-/--
-Check that all postponed recursors are identical to those generated
-when we replayed the inductives.
--/
-def checkPostponedRecursors : M Unit := do
-  for ctor in (← get).postponedRecursors do
-    match (← get).env.constants.find? ctor, (← read).newConstants.find? ctor with
-    | some (.recInfo info), some (.recInfo info') =>
-      if ! (info == info') then throw <| IO.userError s!"Invalid recursor {ctor}"
-    | _, _ => throw <| IO.userError s!"No such recursor {ctor}"
-
-/-- "Replay" some constants into an `Environment`, sending them to the kernel for checking. -/
-def replay (newConstants : HashMap Name ConstantInfo) (env : Environment) : IO Environment := do
-  let mut remaining : NameSet := ∅
-  for (n, ci) in newConstants.toList do
-    -- We skip unsafe constants, and also partial constants.
-    -- Later we may want to handle partial constants.
-    if !ci.isUnsafe && !ci.isPartial then
-      remaining := remaining.insert n
-  let (_, s) ← StateRefT'.run (s := { env, remaining }) do
-    ReaderT.run (r := { newConstants }) do
-      for n in remaining do
-        replayConstant n
-      checkPostponedConstructors
-      checkPostponedRecursors
-  return s.env
-
 unsafe def replayFromImports (module : Name) : IO Unit := do
   let mFile ← findOLean module
   unless (← mFile.pathExists) do
@@ -157,13 +20,13 @@ unsafe def replayFromImports (module : Name) : IO Unit := do
   let mut newConstants := {}
   for name in mod.constNames, ci in mod.constants do
     newConstants := newConstants.insert name ci
-  let env' ← replay newConstants env
+  let env' ← env.replay newConstants
   env'.freeRegions
   region.free
 
 unsafe def replayFromFresh (module : Name) : IO Unit := do
   Lean.withImportModules #[{module}] {} 0 fun env => do
-    discard <| replay env.constants.map₁ (← mkEmptyEnvironment)
+    discard <| (← mkEmptyEnvironment).replay env.constants.map₁
 
 /--
 Run as e.g. `lake exe lean4checker` to check everything on the Lean search path,