feat: pickle and unpickle Environments

.gitignore

@@ -1,3 +1,4 @@
 /build
 /lake-packages/*
-/lakefile.olean
+/lakefile.olean
+/test/*.olean

REPL/JSON.lean

@@ -82,8 +82,8 @@ A response to a Lean command.
 -/
 structure CommandResponse where
   env : Nat
-  messages : List Message
-  sorries : List Sorry
+  messages : List Message := []
+  sorries : List Sorry := []
 deriving ToJson, FromJson
 
 /--
@@ -100,4 +100,13 @@ structure Error where
   message : String
 deriving ToJson, FromJson
 
+structure PickleEnvironment where
+  env : Nat
+  pickleTo : System.FilePath
+deriving ToJson, FromJson
+
+structure UnpickleEnvironment where
+  unpickleFrom : System.FilePath
+deriving ToJson, FromJson
+
 end REPL

REPL/Lean/Environment.lean

@@ -0,0 +1,30 @@
+import REPL.Util.Pickle
+
+open System (FilePath)
+
+namespace Lean.Environment
+
+/--
+Pickle an `Environment` to disk.
+
+We only store:
+* the list of imports
+* the new constants from `Environment.constants`
+and when unpickling, we build a fresh `Environment` from the imports,
+and then add the new constants.
+-/
+def pickle (env : Environment) (path : FilePath) : IO Unit :=
+  _root_.pickle path (env.header.imports, env.constants.map₂)
+
+/--
+Unpickle an `Environment` from disk.
+
+We construct a fresh `Environment` with the relevant imports,
+and then replace the new constants.
+-/
+def unpickle (path : FilePath) : IO (Environment × CompactedRegion) := unsafe do
+  let ((imports, map₂), region) ← _root_.unpickle (Array Import × PHashMap Name ConstantInfo) path
+  let env ← importModules imports {} 0
+  return ({ env with constants := { env.constants with map₂ }}, region)
+
+end Lean.Environment

REPL/Main.lean

@@ -9,6 +9,7 @@ import REPL.InfoTree
 import REPL.Util.Path
 import REPL.Util.TermUnsafe
 import REPL.Lean.ContextInfo
+import REPL.Lean.Environment
 import REPL.InfoTree
 
 /-!
@@ -178,6 +179,18 @@ def createProofState (ctx : ContextInfo) (lctx? : Option LocalContext)
       tacticState := { goals }
       tacticContext := { elaborator := .anonymous } }
 
+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 }
+
+def unpickleEnvironment (n : UnpickleEnvironment) : M IO CommandResponse := do
+  let (env, _) ← Environment.unpickle n.unpickleFrom
+  let env ← recordEnvironment env
+  return { env }
+
 /--
 Run a command, returning the id of the new environment, and any messages and sorries.
 -/
@@ -233,24 +246,41 @@ instance [ToJson α] [ToJson β] : ToJson (α ⊕ β) where
   | .inl a => toJson a
   | .inr b => toJson b
 
+inductive Input
+| command : REPL.Command → Input
+| proofStep : REPL.ProofStep → Input
+| pickleEnvironment : REPL.PickleEnvironment → Input
+| unpickleEnvironment : REPL.UnpickleEnvironment → Input
+
+def parse (query : String) : IO Input := do
+  let json := Json.parse query
+  match json with
+  | .error e => throw <| IO.userError <| toString <| toJson (⟨e⟩ : Error)
+  | .ok j => match fromJson? j with
+    | .ok (r : REPL.ProofStep) => return .proofStep r
+    | .error _ => match fromJson? j with
+    | .ok (r : REPL.PickleEnvironment) => return .pickleEnvironment r
+    | .error _ => match fromJson? j with
+    | .ok (r : REPL.UnpickleEnvironment) => return .unpickleEnvironment r
+    | .error _ => match fromJson? j with
+    | .ok (r : REPL.Command) => return .command r
+    | .error e => throw <| IO.userError <| toString <| toJson (⟨e⟩ : Error)
+
 /-- Read-eval-print loop for Lean. -/
-partial def repl : IO Unit :=
+unsafe def repl : IO Unit :=
   StateT.run' loop {}
 where loop : M IO Unit := do
   let query ← getLines
   if query = "" then
     return ()
-  let json := Json.parse query
-  match json with
-  | .error e => IO.println <| toString <| toJson (⟨e⟩ : Error)
-  | .ok j => match fromJson? j with
-    | .ok (r : REPL.ProofStep) => IO.println <| toString <| toJson (← runProofStep r)
-    | .error _ => match fromJson? j with
-      | .ok (r : REPL.Command) => IO.println <| toString <| toJson (← runCommand r)
-      | .error e => IO.println <| toString <| toJson (⟨e⟩ : Error)
+  IO.println <| toString <| ← match ← parse query with
+  | .command r => return toJson (← runCommand r)
+  | .proofStep r => return toJson (← runProofStep r)
+  | .pickleEnvironment r => return toJson (← pickleEnvironment r)
+  | .unpickleEnvironment r => return toJson (← unpickleEnvironment r)
   loop
 
 /-- Main executable function, run as `lake exe repl`. -/
-def main (_ : List String) : IO Unit := do
+unsafe def main (_ : List String) : IO Unit := do
   searchPathRef.set compile_time_search_path%
   repl

REPL/Util/Pickle.lean

@@ -0,0 +1,44 @@
+/-
+Copyright (c) 2023 Mario Carneiro. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Mario Carneiro
+-/
+import REPL.Util.TermUnsafe
+
+/-!
+# Pickling and unpickling objects
+
+By abusing `saveModuleData` and `readModuleData` we can pickle and unpickle objects to disk.
+-/
+
+open Lean System
+
+/--
+Save an object to disk.
+If you need to write multiple objects from within a single declaration,
+you will need to provide a unique `key` for each.
+-/
+def pickle {α : Type} (path : FilePath) (x : α) (key : Name := by exact decl_name%) : IO Unit :=
+  saveModuleData path key (unsafe unsafeCast x)
+
+/--
+Load an object from disk.
+Note: The returned `CompactedRegion` can be used to free the memory behind the value
+of type `α`, using `CompactedRegion.free` (which is only safe once all references to the `α` are
+released). Ignoring the `CompactedRegion` results in the data being leaked.
+Use `withUnpickle` to call `CompactedRegion.free` automatically.
+
+This function is unsafe because the data being loaded may not actually have type `α`, and this
+may cause crashes or other bad behavior.
+-/
+unsafe def unpickle (α : Type) (path : FilePath) : IO (α × CompactedRegion) := do
+  let (x, region) ← readModuleData path
+  pure (unsafeCast x, region)
+
+/-- Load an object from disk and run some continuation on it, freeing memory afterwards. -/
+unsafe def withUnpickle [Monad m] [MonadLiftT IO m] {α β : Type}
+    (path : FilePath) (f : α → m β) : m β := do
+  let (x, region) ← unpickle α path
+  let r ← f x
+  region.free
+  pure r

test/pickle_environment.expected.out

@@ -0,0 +1,4 @@
+{"sorries": [], "messages": [], "env": 0}
+{"sorries": [], "messages": [], "env": 0}
+{"sorries": [], "messages": [], "env": 1}
+{"sorries": [], "messages": [], "env": 2}

test/pickle_environment.in

@@ -0,0 +1,7 @@
+{"cmd": "def f := 42"}
+
+{"pickleTo": "test/a.olean", "env": 0}
+
+{"unpickleFrom": "test/a.olean"}
+
+{"cmd": "example : f = 42 := by rfl", "env": 1}

test/pickle_environment_with_imports.expected.out

@@ -0,0 +1,5 @@
+{"sorries": [], "messages": [], "env": 0}
+{"sorries": [], "messages": [], "env": 1}
+{"sorries": [], "messages": [], "env": 1}
+{"sorries": [], "messages": [], "env": 2}
+{"sorries": [], "messages": [], "env": 3}

test/pickle_environment_with_imports.in

@@ -0,0 +1,9 @@
+{"cmd": "import Lean"}
+
+{"cmd": "def f := 42", "env": 0}
+
+{"pickleTo": "test/b.olean", "env": 1}
+
+{"unpickleFrom": "test/b.olean"}
+
+{"cmd": "example : f = 42 := by rfl", "env": 2}