From c54ce93ef5eb745a3db7668886f35ed034f42afb Mon Sep 17 00:00:00 2001
From: Leni Aniva <v@leni.sh>
Date: Thu, 5 Dec 2024 14:31:43 -0800
Subject: [PATCH] feat: Goal State IO in REPL

---
 Pantograph/Protocol.lean | 13 +++++++++++++
 Repl.lean                | 36 +++++++++++++++++++++++-------------
 2 files changed, 36 insertions(+), 13 deletions(-)

diff --git a/Pantograph/Protocol.lean b/Pantograph/Protocol.lean
index fcd5ebe..0cb6cac 100644
--- a/Pantograph/Protocol.lean
+++ b/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
diff --git a/Repl.lean b/Repl.lean
index f0572a1..3f8a3c6 100644
--- a/Repl.lean
+++ b/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