feat: Goal State IO in REPL

Pantograph/Protocol.lean

@@ -289,6 +289,19 @@ structure GoalDiag where
   instantiate: Bool := true
   printSexp: Bool := false
 
+structure GoalSave where
+  id: Nat
+  path: System.FilePath
+  deriving Lean.FromJson
+structure GoalSaveResult where
+  deriving Lean.ToJson
+structure GoalLoad where
+  path: System.FilePath
+  deriving Lean.FromJson
+structure GoalLoadResult where
+  id: Nat
+  deriving Lean.ToJson
+
 
 /-- Executes the Lean compiler on a single file -/
 structure FrontendProcess where

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

Repl.lean

@@ -15,6 +15,16 @@ structure State where
 /-- Main state monad for executing commands -/
 abbrev MainM := ReaderT Context (StateT State Lean.CoreM)
 
+def newGoalState (goalState: GoalState) : MainM Nat := do
+  let state ← get
+  let stateId := state.nextId
+  set { state with
+    goalStates := state.goalStates.insert stateId goalState,
+    nextId := state.nextId + 1
+  }
+  return stateId
+
+
 -- HACK: For some reason writing `CommandM α := MainM (Except ... α)` disables
 -- certain monadic features in `MainM`
 abbrev CR α := Except Protocol.InteractionError α
@@ -50,6 +60,8 @@ def execute (command: Protocol.Command): MainM Lean.Json := do
     | "goal.continue" => run goal_continue
     | "goal.delete"   => run goal_delete
     | "goal.print"    => run goal_print
+    | "goal.save"     => run goal_save
+    | "goal.load"     => run goal_load
     | "frontend.process"  => run frontend_process
     | cmd =>
       let error: Protocol.InteractionError :=
@@ -62,14 +74,6 @@ def execute (command: Protocol.Command): MainM Lean.Json := do
   errorCommand := errorI "command"
   errorIndex := errorI "index"
   errorIO := errorI "io"
-  newGoalState (goalState: GoalState) : MainM Nat := do
-    let state ← get
-    let stateId := state.nextId
-    set { state with
-      goalStates := state.goalStates.insert stateId goalState,
-      nextId := state.nextId + 1
-    }
-    return stateId
   -- Command Functions
   reset (_: Protocol.Reset): MainM (CR Protocol.StatResult) := do
     let state ← get
@@ -203,11 +207,7 @@ def execute (command: Protocol.Command): MainM Lean.Json := do
     match nextState? with
     | .error error => return .error <| errorI "structure" error
     | .ok nextGoalState =>
-      let nextStateId := state.nextId
-      set { state with
-        goalStates := state.goalStates.insert nextStateId nextGoalState,
-        nextId := state.nextId + 1
-      }
+      let nextStateId ← newGoalState nextGoalState
       let goals ← goalSerialize nextGoalState (options := state.options)
       return .ok {
         nextStateId,
@@ -224,6 +224,16 @@ def execute (command: Protocol.Command): MainM Lean.Json := do
       return .error $ errorIndex s!"Invalid state index {args.stateId}"
     let result ← runMetaInMainM <| goalPrint goalState state.options
     return .ok result
+  goal_save (args: Protocol.GoalSave): MainM (CR Protocol.GoalSaveResult) := do
+    let state ← get
+    let .some goalState := state.goalStates[args.id]? |
+      return .error $ errorIndex s!"Invalid state index {args.id}"
+    goalStatePickle goalState args.path
+    return .ok {}
+  goal_load (args: Protocol.GoalLoad): MainM (CR Protocol.GoalLoadResult) := do
+    let (goalState, _) ← goalStateUnpickle args.path (← Lean.MonadEnv.getEnv)
+    let id ← newGoalState goalState
+    return .ok { id }
   frontend_process (args: Protocol.FrontendProcess): MainM (CR Protocol.FrontendProcessResult) := do
     let options := (← get).options
     try

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))