feat: replay constants into an Environment (#2617)

* feat: replay constants into an Environment

src/Lean/Data/NameMap.lean

@@ -46,6 +46,13 @@ def contains (s : NameSet) (n : Name) : Bool := RBMap.contains s n
 instance : ForIn m NameSet Name :=
   inferInstanceAs (ForIn _ (RBTree ..) ..)
 
+/-- The union of two `NameSet`s. -/
+def append (s t : NameSet) : NameSet :=
+  s.mergeBy (fun _ _ _ => .unit) t
+
+instance : Append NameSet where
+  append := NameSet.append
+
 end NameSet
 
 def NameSSet := SSet Name

src/Lean/Declaration.lean

@@ -70,11 +70,11 @@ structure ConstantVal where
   name : Name
   levelParams : List Name
   type : Expr
-  deriving Inhabited
+  deriving Inhabited, BEq
 
 structure AxiomVal extends ConstantVal where
   isUnsafe : Bool
-  deriving Inhabited
+  deriving Inhabited, BEq
 
 @[export lean_mk_axiom_val]
 def mkAxiomValEx (name : Name) (levelParams : List Name) (type : Expr) (isUnsafe : Bool) : AxiomVal := {
@@ -119,7 +119,7 @@ structure TheoremVal extends ConstantVal where
     List of all (including this one) declarations in the same mutual block.
     See comment at `DefinitionVal.all`. -/
   all : List Name := [name]
-  deriving Inhabited
+  deriving Inhabited, BEq
 
 /-- Value for an opaque constant declaration `opaque x : t := e` -/
 structure OpaqueVal extends ConstantVal where
@@ -129,7 +129,7 @@ structure OpaqueVal extends ConstantVal where
     List of all (including this one) declarations in the same mutual block.
     See comment at `DefinitionVal.all`. -/
   all : List Name := [name]
-  deriving Inhabited
+  deriving Inhabited, BEq
 
 @[export lean_mk_opaque_val]
 def mkOpaqueValEx (name : Name) (levelParams : List Name) (type : Expr) (value : Expr) (isUnsafe : Bool) (all : List Name) : OpaqueVal := {
@@ -272,7 +272,7 @@ structure ConstructorVal extends ConstantVal where
   /-- Number of fields (i.e., arity - nparams) -/
   numFields : Nat
   isUnsafe : Bool
-  deriving Inhabited
+  deriving Inhabited, BEq
 
 @[export lean_mk_constructor_val]
 def mkConstructorValEx (name : Name) (levelParams : List Name) (type : Expr) (induct : Name) (cidx numParams numFields : Nat) (isUnsafe : Bool) : ConstructorVal := {
@@ -296,7 +296,7 @@ structure RecursorRule where
   nfields : Nat
   /-- Right hand side of the reduction rule -/
   rhs : Expr
-  deriving Inhabited
+  deriving Inhabited, BEq
 
 structure RecursorVal extends ConstantVal where
   /-- List of all inductive datatypes in the mutual declaration that generated this recursor -/
@@ -322,7 +322,7 @@ structure RecursorVal extends ConstantVal where
   -/
   k : Bool
   isUnsafe : Bool
-  deriving Inhabited
+  deriving Inhabited, BEq
 
 @[export lean_mk_recursor_val]
 def mkRecursorValEx (name : Name) (levelParams : List Name) (type : Expr) (all : List Name) (numParams numIndices numMotives numMinors : Nat)
@@ -441,6 +441,10 @@ def isInductive : ConstantInfo → Bool
   | inductInfo _ => true
   | _            => false
 
+def inductiveVal! : ConstantInfo → InductiveVal
+  | .inductInfo val => val
+  | _ => panic! "Expected a `ConstantInfo.inductInfo`."
+
 /--
   List of all (including this one) declarations in the same mutual block.
 -/

src/Lean/Replay.lean

@@ -0,0 +1,169 @@
+/-
+Copyright (c) 2023 Scott Morrison. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Scott Morrison
+-/
+import Lean.CoreM
+import Lean.Util.FoldConsts
+
+/-!
+# `Lean.Environment.replay`
+
+`replay env constantMap` will "replay" all the constants in `constantMap : HashMap Name ConstantInfo` into `env`,
+sending each declaration to the kernel for checking.
+
+`replay` does not send constructors or recursors in `constantMap` to the kernel,
+but rather checks that they are identical to constructors or recursors generated in the enviroment
+after replaying any inductive definitions occurring in `constantMap`.
+
+`replay` can be used either as:
+* a verifier for an `Environment`, by sending everything to the kernel, or
+* a mechanism to safely transfer constants from one `Environment` to another.
+
+-/
+
+namespace Lean.Environment
+
+namespace Replay
+
+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
+
+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}"
+
+end Replay
+
+open Replay
+
+/--
+"Replay" some constants into an `Environment`, sending them to the kernel for checking.
+
+Throws a `IO.userError` if the kernel rejects a constant,
+or if there are malformed recursors or constructors for inductive types.
+-/
+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

src/Lean/Util/FoldConsts.lean

@@ -66,8 +66,27 @@ opaque foldConsts {α : Type} (e : Expr) (init : α) (f : Name → α → α) :
 def getUsedConstants (e : Expr) : Array Name :=
   e.foldConsts #[] fun c cs => cs.push c
 
+/-- Like `Expr.getUsedConstants`, but produce a `NameSet`. -/
+def getUsedConstantsAsSet (e : Expr) : NameSet :=
+  e.foldConsts {} fun c cs => cs.insert c
+
 end Expr
 
+namespace ConstantInfo
+
+/-- Return all names appearing in the type or value of a `ConstantInfo`. -/
+def getUsedConstantsAsSet (c : ConstantInfo) : NameSet :=
+  c.type.getUsedConstantsAsSet ++ match c.value? with
+  | some v => v.getUsedConstantsAsSet
+  | none => match c with
+    | .inductInfo val => .ofList val.ctors
+    | .opaqueInfo val => val.value.getUsedConstantsAsSet
+    | .ctorInfo val => ({} : NameSet).insert val.name
+    | .recInfo val => .ofList val.all
+    | _ => {}
+
+end ConstantInfo
+
 def getMaxHeight (env : Environment) (e : Expr) : UInt32 :=
   e.foldConsts 0 fun constName max =>
     match env.find? constName with