From 105fb7af4bed8da4969ef1265106c239861dece9 Mon Sep 17 00:00:00 2001
From: Leni Aniva <v@leni.sh>
Date: Thu, 5 Dec 2024 14:23:55 -0800
Subject: [PATCH] feat: Goal state pickling

---
 Pantograph/Serial.lean | 89 ++++++++++++++++++++++++++++++++++++++++++
 Test/Serial.lean       | 36 +++++++++++++----
 2 files changed, 118 insertions(+), 7 deletions(-)

diff --git a/Pantograph/Serial.lean b/Pantograph/Serial.lean
index e60fc54..bd01169 100644
--- a/Pantograph/Serial.lean
+++ b/Pantograph/Serial.lean
@@ -2,6 +2,7 @@ import Lean.Environment
 import Lean.Replay
 import Init.System.IOError
 import Std.Data.HashMap
+import Pantograph.Goal
 
 /-!
 Input/Output functions
@@ -70,4 +71,92 @@ def environmentUnpickle (path : System.FilePath) : IO (Environment × CompactedR
   let env ← importModules imports {} 0
   return (← env.replay (Std.HashMap.ofList map₂.toList), region)
 
+
+open Lean.Core in
+structure CompactCoreState where
+  -- env             : Environment
+  nextMacroScope  : MacroScope     := firstFrontendMacroScope + 1
+  ngen            : NameGenerator  := {}
+  -- traceState      : TraceState     := {}
+  -- cache           : Cache     := {}
+  -- messages        : MessageLog     := {}
+  -- infoState       : Elab.InfoState := {}
+
+@[export pantograph_goal_state_pickle_m]
+def goalStatePickle (goalState : GoalState) (path : System.FilePath) : IO Unit :=
+  let {
+    savedState := {
+      term := {
+        meta := {
+          core,
+          meta,
+        }
+        «elab»,
+      },
+      tactic
+    }
+    root,
+    parentMVar?,
+    convMVar?,
+    calcPrevRhs?,
+  } := goalState
+  --let env := core.env
+  Pantograph.pickle path (
+    ({ core with } : CompactCoreState),
+    meta,
+    «elab»,
+    tactic,
+
+    root,
+    parentMVar?,
+    convMVar?,
+    calcPrevRhs?,
+  )
+
+@[export pantograph_goal_state_unpickle_m]
+def goalStateUnpickle (path : System.FilePath) (env : Environment)
+    : IO (GoalState × CompactedRegion) := unsafe do
+  let ((
+    compactCore,
+    meta,
+    «elab»,
+    tactic,
+
+    root,
+    parentMVar?,
+    convMVar?,
+    calcPrevRhs?,
+  ), region) ← Pantograph.unpickle (
+    CompactCoreState ×
+    Meta.State ×
+    Elab.Term.State ×
+    Elab.Tactic.State ×
+
+    MVarId ×
+    Option MVarId ×
+    Option (MVarId × MVarId × List MVarId) ×
+    Option (MVarId × Expr)
+  ) path
+  let goalState := {
+    savedState := {
+      term := {
+        meta := {
+          core := {
+            compactCore with
+            passedHeartbeats := 0,
+            env,
+          },
+          meta,
+        },
+        «elab»,
+      },
+      tactic,
+    },
+    root,
+    parentMVar?,
+    convMVar?,
+    calcPrevRhs?,
+  }
+  return (goalState, region)
+
 end Pantograph
diff --git a/Test/Serial.lean b/Test/Serial.lean
index d1ce661..fcdc155 100644
--- a/Test/Serial.lean
+++ b/Test/Serial.lean
@@ -31,12 +31,12 @@ def runCoreM { α } (state : Core.State) (testCoreM : TestT CoreM α) : TestM (
     return (a, state')
 
 def test_environment_pickling : TestM Unit := do
-  let stateSrc: Core.State := { env := ← getEnv }
-  let stateDst: Core.State := { env := ← getEnv }
+  let coreSrc : Core.State := { env := ← getEnv }
+  let coreDst : Core.State := { env := ← getEnv }
 
   let name := `mystery
   let envPicklePath ← tempPath
-  let ((), _) ← runCoreM stateSrc do
+  let ((), _) ← runCoreM coreSrc do
     let type: Expr := .forallE `p (.sort 0) (.forallE `h (.bvar 0) (.bvar 1) .default) .default
     let value: Expr := .lam `p (.sort 0) (.lam `h (.bvar 0) (.bvar 0) .default) .default
     let c := Lean.Declaration.defnDecl <| Lean.mkDefinitionValEx
@@ -54,7 +54,7 @@ def test_environment_pickling : TestM Unit := do
       | .ok env' => pure env'
     environmentPickle env' envPicklePath
 
-  let _ ← runCoreM stateDst do
+  let _ ← runCoreM coreDst do
     let (env', _) ← environmentUnpickle envPicklePath
     checkTrue s!"Has symbol {name}" (env'.find? name).isSome
     let anotherName := `mystery2
@@ -62,9 +62,30 @@ def test_environment_pickling : TestM Unit := do
 
   IO.FS.removeFile envPicklePath
 
+def test_goal_state_pickling_simple : TestM Unit := do
+  let coreSrc : Core.State := { env := ← getEnv }
+  let coreDst : Core.State := { env := ← getEnv }
+  let statePath ← tempPath
+
+  let type: Expr := .forallE `p (.sort 0) (.forallE `h (.bvar 0) (.bvar 1) .default) .default
+  let stateGenerate : MetaM GoalState := runTermElabMInMeta do
+    GoalState.create type
+
+  let ((), _) ← runCoreM coreSrc do
+    let state ← stateGenerate.run'
+    goalStatePickle state statePath
+
+  let ((), _) ← runCoreM coreDst do
+    let (goalState, _) ← goalStateUnpickle statePath (← getEnv)
+    let metaM : MetaM (List Expr) := do
+      goalState.goals.mapM λ goal => goalState.withContext goal goal.getType
+    let types ← metaM.run'
+    checkTrue "Goals" $ types[0]!.equal type
+
+  IO.FS.removeFile statePath
+
 structure Test where
   name : String
-  nInstances : Nat
   routine: TestM Unit
 
 protected def Test.run (test: Test) (env: Lean.Environment) : IO LSpec.TestSeq := do
@@ -76,11 +97,12 @@ protected def Test.run (test: Test) (env: Lean.Environment) : IO LSpec.TestSeq :
   match ← ((runTest $ test.routine).run' state).toBaseIO with
   | .ok e => return e
   | .error e =>
-    return LSpec.check "Emitted exception" (e.toString == "")
+    return LSpec.check s!"Emitted exception: {e.toString}" (e.toString == "")
 
 def suite (env : Lean.Environment): List (String × IO LSpec.TestSeq) :=
   let tests: List Test := [
-    { name := "environment_pickling", nInstances := 2, routine := test_environment_pickling },
+    { name := "environment_pickling", routine := test_environment_pickling, },
+    { name := "goal_state_pickling_simple", routine := test_goal_state_pickling_simple, },
   ]
   tests.map (fun test => (test.name, test.run env))