From 495ea1ac144f524fe585979f7561470010599180 Mon Sep 17 00:00:00 2001
From: Leni Aniva <v@leni.sh>
Date: Fri, 8 Nov 2024 14:49:49 -0800
Subject: [PATCH] feat: Environment pickling

---
 Pantograph.lean             |  1 +
 Pantograph/Environment.lean |  4 +-
 Pantograph/Serial.lean      | 75 +++++++++++++++++++++++++++++++++++++
 3 files changed, 79 insertions(+), 1 deletion(-)
 create mode 100644 Pantograph/Serial.lean

diff --git a/Pantograph.lean b/Pantograph.lean
index 822824c..2c334b6 100644
--- a/Pantograph.lean
+++ b/Pantograph.lean
@@ -5,4 +5,5 @@ import Pantograph.Frontend
 import Pantograph.Goal
 import Pantograph.Library
 import Pantograph.Protocol
+import Pantograph.Serial
 import Pantograph.Version
diff --git a/Pantograph/Environment.lean b/Pantograph/Environment.lean
index 40b3386..ad21284 100644
--- a/Pantograph/Environment.lean
+++ b/Pantograph/Environment.lean
@@ -1,7 +1,9 @@
 import Pantograph.Delate
 import Pantograph.Elab
 import Pantograph.Protocol
-import Lean
+import Pantograph.Serial
+import Lean.Environment
+import Lean.Replay
 
 open Lean
 open Pantograph
diff --git a/Pantograph/Serial.lean b/Pantograph/Serial.lean
new file mode 100644
index 0000000..c6aecae
--- /dev/null
+++ b/Pantograph/Serial.lean
@@ -0,0 +1,75 @@
+import Lean.Environment
+import Lean.Replay
+import Std.Data.HashMap
+
+/-!
+Input/Output functions
+
+# Pickling and unpickling objects
+
+By abusing `saveModuleData` and `readModuleData` we can pickle and unpickle objects to disk.
+-/
+
+open Lean
+
+namespace Pantograph
+
+/--
+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 : System.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 : System.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 : System.FilePath) (f : α → m β) : m β := do
+  let (x, region) ← unpickle α path
+  let r ← f x
+  region.free
+  pure r
+end Pantograph
+
+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.
+-/
+@[export pantograph_env_pickle_m]
+def pickle (env : Environment) (path : System.FilePath) : IO Unit :=
+  Pantograph.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.
+-/
+@[export pantograph_env_unpickle_m]
+def 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