feat: unpickle proof states using an existing environment

Author: kim-em

REPL/JSON.lean

@@ -164,6 +164,7 @@ deriving ToJson, FromJson
 
 structure UnpickleProofState where
   unpickleProofStateFrom : System.FilePath
+  env : Option Nat
 deriving ToJson, FromJson
 
 end REPL

REPL/Main.lean

@@ -177,9 +177,16 @@ def pickleProofSnapshot (n : PickleProofState) : M m (ProofStepResponse ⊕ Erro
     return .inl (← createProofStepReponse proofState)
 
 /-- Unpickle a `ProofSnapshot`, generating a JSON response. -/
-def unpickleProofSnapshot (n : UnpickleProofState) : M IO ProofStepResponse := do
-  let (proofState, _) ← ProofSnapshot.unpickle n.unpickleProofStateFrom
-  createProofStepReponse proofState
+def unpickleProofSnapshot (n : UnpickleProofState) : M IO (ProofStepResponse ⊕ Error) := do
+  let (cmdSnapshot?, notFound) ← do match n.env with
+  | none => pure (none, false)
+  | some i => do match (← get).cmdStates[i]? with
+    | some env => pure (some env, false)
+    | none => pure (none, true)
+  if notFound then
+    return .inr ⟨"Unknown environment."⟩
+  let (proofState, _) ← ProofSnapshot.unpickle n.unpickleProofStateFrom cmdSnapshot?
+  Sum.inl <$> createProofStepReponse proofState
 
 /--
 Run a command, returning the id of the new environment, and any messages and sorries.

REPL/Snapshots.lean

@@ -259,14 +259,19 @@ def pickle (p : ProofSnapshot) (path : FilePath) : IO Unit := do
 /--
 Unpickle a `ProofSnapshot`.
 -/
-def unpickle (path : FilePath) : IO (ProofSnapshot × CompactedRegion) := unsafe do
+def unpickle (path : FilePath) (cmd? : Option CommandSnapshot) :
+    IO (ProofSnapshot × CompactedRegion) := unsafe do
   let ((imports, map₂, coreState, coreContext, metaState, metaContext, termState, termContext,
     tacticState, tacticContext), region) ←
     _root_.unpickle (Array Import × PHashMap Name ConstantInfo × CompactableCoreState ×
       Core.Context × Meta.State × CompactableMetaContext × Term.State × CompactableTermContext ×
       Tactic.State × Tactic.Context) path
-  enableInitializersExecution
-  let env ← (← importModules imports {} 0).replay (HashMap.ofList map₂.toList)
+  let env ← match cmd? with
+  | none =>
+    enableInitializersExecution
+    (← importModules imports {} 0).replay (HashMap.ofList map₂.toList)
+  | some cmd =>
+    cmd.cmdState.env.replay (HashMap.ofList map₂.toList)
   let p' : ProofSnapshot :=
   { coreState := { coreState with env }
     coreContext

test/import_lean.in

@@ -0,0 +1 @@
+{"cmd" : "import Lean"}

test/pickle_proof_state_env.expected.out

@@ -0,0 +1,26 @@
+{"env": 0}
+
+{"sorries":
+ [{"proofState": 0,
+   "pos": {"line": 1, "column": 18},
+   "goal": "⊢ Nat",
+   "endPos": {"line": 1, "column": 23}}],
+ "messages":
+ [{"severity": "warning",
+   "pos": {"line": 1, "column": 4},
+   "endPos": {"line": 1, "column": 5},
+   "data": "declaration uses 'sorry'"}],
+ "env": 1}
+
+{"proofState": 1, "goals": ["⊢ Nat"]}
+
+{"proofState": 2, "goals": ["⊢ Nat"]}
+
+{"proofState": 3, "goals": ["t : Nat\n⊢ Nat"]}
+
+{"proofState": 4, "goals": ["t : Nat\n⊢ Nat"]}
+
+{"proofState": 5, "goals": ["t : Nat\n⊢ Nat"]}
+
+{"proofState": 6, "goals": []}
+

test/pickle_proof_state_env.in

@@ -0,0 +1,15 @@
+{"cmd" : "import Lean"}
+
+{"cmd" : "def f : Nat := by sorry", "env": 0}
+
+{"pickleTo": "test/c.olean", "proofState": 0}
+
+{"unpickleProofStateFrom": "test/c.olean", "env": 0}
+
+{"tactic": "have t : Nat := 42", "proofState": 2}
+
+{"pickleTo": "test/d.olean", "proofState": 3}
+
+{"unpickleProofStateFrom": "test/d.olean", "env": 0}
+
+{"tactic": "exact t", "proofState": 5}