feat: Environment save/load commands

Pantograph/Protocol.lean

@@ -183,6 +183,12 @@ structure EnvAdd where
 structure EnvAddResult where
   deriving Lean.ToJson
 
+structure EnvSaveLoad where
+  path: System.FilePath
+  deriving Lean.FromJson
+structure EnvSaveLoadResult where
+  deriving Lean.ToJson
+
 /-- Set options; See `Options` struct above for meanings -/
 structure OptionsSet where
   printJsonPretty?: Option Bool

Pantograph/Serial.lean

@@ -1,5 +1,6 @@
 import Lean.Environment
 import Lean.Replay
+import Init.System.IOError
 import Std.Data.HashMap
 
 /-!
@@ -43,9 +44,6 @@ unsafe def withUnpickle [Monad m] [MonadLiftT IO m] {α β : Type}
   let r ← f x
   region.free
   pure r
-end Pantograph
-
-namespace Lean.Environment
 
 /--
 Pickle an `Environment` to disk.
@@ -57,7 +55,7 @@ and when unpickling, we build a fresh `Environment` from the imports,
 and then add the new constants.
 -/
 @[export pantograph_env_pickle_m]
-def pickle (env : Environment) (path : System.FilePath) : IO Unit :=
+def env_pickle (env : Environment) (path : System.FilePath) : IO Unit :=
   Pantograph.pickle path (env.header.imports, env.constants.map₂)
 
 /--
@@ -67,9 +65,9 @@ We construct a fresh `Environment` with the relevant imports,
 and then replace the new constants.
 -/
 @[export pantograph_env_unpickle_m]
-def unpickle (path : System.FilePath) : IO (Environment × CompactedRegion) := unsafe do
+def env_unpickle (path : System.FilePath) : IO (Environment × CompactedRegion) := unsafe do
   let ((imports, map₂), region) ← Pantograph.unpickle (Array Import × PHashMap Name ConstantInfo) path
   let env ← importModules imports {} 0
   return (← env.replay (Std.HashMap.ofList map₂.toList), region)
 
-end Lean.Environment
+end Pantograph

Repl.lean

@@ -24,6 +24,7 @@ def runMetaInMainM { α } (metaM: Lean.MetaM α): MainM α :=
 def runTermElabInMainM { α } (termElabM: Lean.Elab.TermElabM α) : MainM α :=
   termElabM.run' (ctx := defaultElabContext) |>.run'
 
+/-- Main loop command of the REPL -/
 def execute (command: Protocol.Command): MainM Lean.Json := do
   let run { α β: Type } [Lean.FromJson α] [Lean.ToJson β] (comm: α → MainM (CR β)): MainM Lean.Json :=
     match Lean.fromJson? command.payload with
@@ -32,28 +33,35 @@ def execute (command: Protocol.Command): MainM Lean.Json := do
       | .ok result =>  return Lean.toJson result
       | .error ierror => return Lean.toJson ierror
     | .error error => return Lean.toJson $ errorCommand s!"Unable to parse json: {error}"
-  match command.cmd with
-  | "reset"         => run reset
-  | "stat"          => run stat
-  | "expr.echo"     => run expr_echo
-  | "env.catalog"   => run env_catalog
-  | "env.inspect"   => run env_inspect
-  | "env.add"       => run env_add
-  | "options.set"   => run options_set
-  | "options.print" => run options_print
-  | "goal.start"    => run goal_start
-  | "goal.tactic"   => run goal_tactic
-  | "goal.continue" => run goal_continue
-  | "goal.delete"   => run goal_delete
-  | "goal.print"    => run goal_print
-  | "frontend.process"  => run frontend_process
-  | cmd =>
-    let error: Protocol.InteractionError :=
-      errorCommand s!"Unknown command {cmd}"
-    return Lean.toJson error
+  try
+    match command.cmd with
+    | "reset"         => run reset
+    | "stat"          => run stat
+    | "expr.echo"     => run expr_echo
+    | "env.catalog"   => run env_catalog
+    | "env.inspect"   => run env_inspect
+    | "env.add"       => run env_add
+    | "env.save"      => run env_save
+    | "env.load"      => run env_load
+    | "options.set"   => run options_set
+    | "options.print" => run options_print
+    | "goal.start"    => run goal_start
+    | "goal.tactic"   => run goal_tactic
+    | "goal.continue" => run goal_continue
+    | "goal.delete"   => run goal_delete
+    | "goal.print"    => run goal_print
+    | "frontend.process"  => run frontend_process
+    | cmd =>
+      let error: Protocol.InteractionError :=
+        errorCommand s!"Unknown command {cmd}"
+      return Lean.toJson error
+  catch ex => do
+    let error ← ex.toMessageData.toString
+    return Lean.toJson $ errorIO error
   where
   errorCommand := errorI "command"
   errorIndex := errorI "index"
+  errorIO := errorI "io"
   newGoalState (goalState: GoalState) : MainM Nat := do
     let state ← get
     let stateId := state.nextId
@@ -80,6 +88,14 @@ def execute (command: Protocol.Command): MainM Lean.Json := do
     Environment.inspect args state.options
   env_add (args: Protocol.EnvAdd): MainM (CR Protocol.EnvAddResult) := do
     Environment.addDecl args
+  env_save (args: Protocol.EnvSaveLoad): MainM (CR Protocol.EnvSaveLoadResult) := do
+    let env ← Lean.MonadEnv.getEnv
+    env_pickle env args.path
+    return .ok {}
+  env_load (args: Protocol.EnvSaveLoad): MainM (CR Protocol.EnvSaveLoadResult) := do
+    let (env, _) ← env_unpickle args.path
+    Lean.setEnv env
+    return .ok {}
   expr_echo (args: Protocol.ExprEcho): MainM (CR Protocol.ExprEchoResult) := do
     let state ← get
     exprEcho args.expr (expectedType? := args.type?) (levels := args.levels.getD #[]) (options := state.options)