feat: pickle proof states (#12)

* feat: pickle proof states as well

* cleanup and doc-strings

* update README

* set search path correctly

.gitignore

@@ -2,3 +2,4 @@
 /lake-packages/*
 /lakefile.olean
 /test/*.olean
+/test/*.olean.tmp

README.md

@@ -91,7 +91,33 @@ At present there is nothing you can do with a completed proof state:
 we would like to extend this so that you can replace the original `sorry` with your tactic script,
 and obtain the resulting `Environment`
 
+## Pickling
+
+The REPL supports pickling environments and proof states to disk as `.olean` files.
+As long as the same imports are available, it should be possible to move such an `.olean` file
+to another machine and unpickle into a new REPL session.
+
+The commands are
+
+```json
+{"pickleTo": "path/to/file.olean", "env": 7}
+
+{"pickleTo": "path/to/file.olean", "proofState": 17}
+
+{"unpickleEnvFrom": "path/to/file.olean"}
+
+{"unpickleProofStateFrom": "path/to/file.olean"}
+```
+
+The unpickling commands will report the new "env" or "proofState" identifier that
+you can use in subsequent commands.
+
+Pickling is quite efficient:
+* we don't record full `Environment`s, only the changes relative to imports
+* unpickling uses memory mapping
+* file sizes are generally small, but see https://github.com/digama0/leangz if compression is
+  desirable
+
 ## Future work
 
 * Replay tactic scripts from tactic mode back into the original `sorry`.
-* Serialization and deserialization of environments and proof states

REPL/Frontend.lean

@@ -34,13 +34,12 @@ Returns the resulting environment, along with a list of messages and info trees.
 def processInput (input : String) (env? : Option Environment)
     (opts : Options) (fileName : Option String := none) :
     IO (Environment × List Message × List InfoTree) := unsafe do
+  Lean.initSearchPath (← Lean.findSysroot)
+  enableInitializersExecution
   let fileName   := fileName.getD "<input>"
   let inputCtx   := Parser.mkInputContext input fileName
   let (parserState, commandState) ← match env? with
   | none => do
-    let leanPath ← Lean.findSysroot
-    Lean.initSearchPath leanPath
-    enableInitializersExecution
     let (header, parserState, messages) ← Parser.parseHeader inputCtx
     let (env, messages) ← processHeader header opts messages inputCtx
     pure (parserState, (Command.mkState env messages opts))

REPL/JSON.lean

@@ -106,7 +106,16 @@ structure PickleEnvironment where
 deriving ToJson, FromJson
 
 structure UnpickleEnvironment where
-  unpickleFrom : System.FilePath
+  unpickleEnvFrom : System.FilePath
+deriving ToJson, FromJson
+
+structure PickleProofState where
+  proofState : Nat
+  pickleTo : System.FilePath
+deriving ToJson, FromJson
+
+structure UnpickleProofState where
+  unpickleProofStateFrom : System.FilePath
 deriving ToJson, FromJson
 
 end REPL

REPL/Main.lean

@@ -11,6 +11,7 @@ import REPL.Util.TermUnsafe
 import REPL.Lean.ContextInfo
 import REPL.Lean.Environment
 import REPL.InfoTree
+import REPL.ProofState
 
 /-!
 # A REPL for Lean.
@@ -58,71 +59,6 @@ open Lean Elab
 
 namespace REPL
 
-/--
-Bundled structure for the `State` and `Context` objects
-for the `CoreM`, `MetaM`, `TermElabM`, and `TacticM` monads.
--/
-structure ProofState where
-  coreState     : Core.State
-  coreContext   : Core.Context
-  metaState     : Meta.State
-  metaContext   : Meta.Context
-  termState     : Term.State
-  termContext   : Term.Context
-  tacticState   : Tactic.State
-  tacticContext : Tactic.Context
-namespace ProofState
-
-open Lean Elab Tactic
-
-/-- Run a `CoreM` monadic function in the current `ProofState`, updating the `Core.State`. -/
-def runCoreM (p : ProofState) (t : CoreM α) : IO (α × ProofState) := do
-  let (a, coreState) ← (Lean.Core.CoreM.toIO · p.coreContext p.coreState) do t
-  return (a, { p with coreState })
-
-/-- Run a `MetaM` monadic function in the current `ProofState`, updating the `Meta.State`. -/
-def runMetaM (p : ProofState) (t : MetaM α) : IO (α × ProofState) := do
-  let ((a, metaState), p') ← p.runCoreM (Lean.Meta.MetaM.run (ctx := p.metaContext) (s := p.metaState) do t)
-  return (a, { p' with metaState })
-
-/-- Run a `TermElabM` monadic function in the current `ProofState`, updating the `Term.State`. -/
-def runTermElabM (p : ProofState) (t : TermElabM α) : IO (α × ProofState) := do
-  let ((a, termState), p') ← p.runMetaM (Lean.Elab.Term.TermElabM.run (s := p.termState) do t)
-  return (a, { p' with termState })
-
-/-- Run a `TacticM` monadic function in the current `ProofState`, updating the `Tactic.State`. -/
-def runTacticM (p : ProofState) (t : TacticM α) : IO (α × ProofState) := do
-  let ((a, tacticState), p') ← p.runTermElabM (t p.tacticContext |>.run p.tacticState)
-  return (a, { p' with tacticState })
-
-/--
-Run a `TacticM` monadic function in the current `ProofState`, updating the `Tactic.State`,
-and discarding the return value.
--/
-def runTacticM' (p : ProofState) (t : TacticM α) : IO ProofState :=
-  Prod.snd <$> p.runTacticM t
-
-/--
-Evaluate a `Syntax` into a `TacticM` tactic, and run it in the current `ProofState`.
--/
-def runSyntax (p : ProofState) (t : Syntax) : IO ProofState :=
-  Prod.snd <$> p.runTacticM (evalTactic t)
-
-/--
-Parse a string into a `Syntax`, evaluate it as a `TacticM` tactic,
-and run it in the current `ProofState`.
--/
-def runString (p : ProofState) (t : String) : IO ProofState :=
-  match Parser.runParserCategory p.coreState.env `tactic t with
-  | .error e => throw (IO.userError e)
-  | .ok stx => p.runSyntax stx
-
-/-- Pretty print the current goals in the `ProofState`. -/
-def ppGoals (p : ProofState) : IO (List Format) :=
-  Prod.fst <$> p.runTacticM do (← getGoals).mapM (Meta.ppGoal ·)
-
-end ProofState
-
 /-- The monadic state for the Lean REPL. -/
 structure State where
   /--
@@ -159,38 +95,39 @@ def recordProofState (proofState : ProofState) : M m Nat := do
   modify fun s => { s with proofStates := s.proofStates.push proofState }
   return id
 
-/--
-Construct a `ProofState` from a `ContextInfo` and optional `LocalContext`, and a list of goals.
-
-For convenience, we also allow a list of `Expr`s, and these are appended to the goals
-as fresh metavariables with the given types.
--/
-def createProofState (ctx : ContextInfo) (lctx? : Option LocalContext)
-    (goals : List MVarId) (types : List Expr := []) : IO ProofState := do
-  ctx.runMetaM (lctx?.getD {}) do
-    let goals := goals ++ (← types.mapM fun t => Expr.mvarId! <$> Meta.mkFreshExprMVar (some t))
-    pure <|
-    { coreState := ← getThe Core.State
-      coreContext := ← readThe Core.Context
-      metaState := ← getThe Meta.State
-      metaContext := ← readThe Meta.Context
-      termState := {}
-      termContext := {}
-      tacticState := { goals }
-      tacticContext := { elaborator := .anonymous } }
+/-- Record a `ProofState` and generate a JSON response for it. -/
+def createProofStepReponse (p : ProofState): M m ProofStepResponse := do
+  let id ← recordProofState p
+  return { proofState := id, goals := (← p.ppGoals).map fun s => s!"{s}" }
 
+/-- Pickle an `Environment`, generating a JSON response. -/
 def pickleEnvironment (n : PickleEnvironment) : M m (CommandResponse ⊕ Error) := do
   match (← get).environments[n.env]? with
   | none => return .inr ⟨"Unknown environment."⟩
   | some env =>
     discard <| env.pickle n.pickleTo
     return .inl { env := n.env }
 
+/-- Unpickle an `Environment`, generating a JSON response. -/
 def unpickleEnvironment (n : UnpickleEnvironment) : M IO CommandResponse := do
-  let (env, _) ← Environment.unpickle n.unpickleFrom
+  let (env, _) ← Environment.unpickle n.unpickleEnvFrom
   let env ← recordEnvironment env
   return { env }
 
+/-- Pickle a `ProofState`, generating a JSON response. -/
+-- This generates a new identifier, which perhaps is not what we want?
+def pickleProofState (n : PickleProofState) : M m (ProofStepResponse ⊕ Error) := do
+  match (← get).proofStates[n.proofState]? with
+  | none => return .inr ⟨"Unknown proof State."⟩
+  | some proofState =>
+    discard <| proofState.pickle n.pickleTo
+    return .inl (← createProofStepReponse proofState)
+
+/-- Unpickle a `ProofState`, generating a JSON response. -/
+def unpickleProofState (n : UnpickleProofState) : M IO ProofStepResponse := do
+  let (proofState, _) ← ProofState.unpickle n.unpickleProofStateFrom
+  createProofStepReponse proofState
+
 /--
 Run a command, returning the id of the new environment, and any messages and sorries.
 -/
@@ -208,9 +145,9 @@ def runCommand (s : Command) : M m (CommandResponse ⊕ Error) := do
     fun ⟨ctx, g, pos, endPos⟩ => do
       let (goal, proofState) ← match g with
       | .tactic g => do
-         pure (s!"{(← ctx.ppGoals [g])}".trim, some (← createProofState ctx none [g]))
+         pure (s!"{(← ctx.ppGoals [g])}".trim, some (← ProofState.create ctx none [g]))
       | .term lctx (some t) => do
-         pure (s!"⊢ {← ctx.ppExpr lctx t}", some (← createProofState ctx lctx [] [t]))
+         pure (s!"⊢ {← ctx.ppExpr lctx t}", some (← ProofState.create ctx lctx [] [t]))
       | .term _ none => unreachable!
       let proofStateId ← proofState.mapM recordProofState
       return Sorry.of goal pos endPos proofStateId
@@ -227,8 +164,7 @@ def runProofStep (s : ProofStep) : M m (ProofStepResponse ⊕ Error) := do
   | none => return .inr ⟨"Unknown proof state."⟩
   | some proofState =>
     let proofState' ← proofState.runString s.tactic
-    let id ← recordProofState proofState'
-    return .inl { proofState := id, goals := (← proofState'.ppGoals).map fun s => s!"{s}" }
+    return .inl (← createProofStepReponse proofState')
 
 end REPL
 
@@ -246,12 +182,16 @@ instance [ToJson α] [ToJson β] : ToJson (α ⊕ β) where
   | .inl a => toJson a
   | .inr b => toJson b
 
+/-- Commands accepted by the REPL. -/
 inductive Input
 | command : REPL.Command → Input
 | proofStep : REPL.ProofStep → Input
 | pickleEnvironment : REPL.PickleEnvironment → Input
 | unpickleEnvironment : REPL.UnpickleEnvironment → Input
+| pickleProofState : REPL.PickleProofState → Input
+| unpickleProofState : REPL.UnpickleProofState → Input
 
+/-- Parse a user input string to an input command. -/
 def parse (query : String) : IO Input := do
   let json := Json.parse query
   match json with
@@ -263,6 +203,10 @@ def parse (query : String) : IO Input := do
     | .error _ => match fromJson? j with
     | .ok (r : REPL.UnpickleEnvironment) => return .unpickleEnvironment r
     | .error _ => match fromJson? j with
+    | .ok (r : REPL.PickleProofState) => return .pickleProofState r
+    | .error _ => match fromJson? j with
+    | .ok (r : REPL.UnpickleProofState) => return .unpickleProofState r
+    | .error _ => match fromJson? j with
     | .ok (r : REPL.Command) => return .command r
     | .error e => throw <| IO.userError <| toString <| toJson (⟨e⟩ : Error)
 
@@ -278,9 +222,11 @@ where loop : M IO Unit := do
   | .proofStep r => return toJson (← runProofStep r)
   | .pickleEnvironment r => return toJson (← pickleEnvironment r)
   | .unpickleEnvironment r => return toJson (← unpickleEnvironment r)
+  | .pickleProofState r => return toJson (← pickleProofState r)
+  | .unpickleProofState r => return toJson (← unpickleProofState r)
   loop
 
 /-- Main executable function, run as `lake exe repl`. -/
 unsafe def main (_ : List String) : IO Unit := do
-  searchPathRef.set compile_time_search_path%
+  initSearchPath (← Lean.findSysroot)
   repl