From d64a2a40710e765f6c21aff9a5d2f5b541f34009 Mon Sep 17 00:00:00 2001
From: Scott Morrison <scott.morrison@gmail.com>
Date: Tue, 10 Oct 2023 13:59:24 +1100
Subject: [PATCH 1/5] more tests

---
 test/synthesize_placeholder.expected.out  |  7 +++++++
 test/synthesize_placeholder.in            |  1 +
 test/unfinished_tactic_block.expected.out | 11 +++++++++++
 test/unfinished_tactic_block.in           |  1 +
 4 files changed, 20 insertions(+)
 create mode 100644 test/synthesize_placeholder.expected.out
 create mode 100644 test/synthesize_placeholder.in
 create mode 100644 test/unfinished_tactic_block.expected.out
 create mode 100644 test/unfinished_tactic_block.in

diff --git a/test/synthesize_placeholder.expected.out b/test/synthesize_placeholder.expected.out
new file mode 100644
index 0000000..06119ac
--- /dev/null
+++ b/test/synthesize_placeholder.expected.out
@@ -0,0 +1,7 @@
+{"sorries": [],
+ "messages":
+ [{"severity": "error",
+   "pos": {"line": 1, "column": 15},
+   "endPos": {"line": 1, "column": 16},
+   "data": "don't know how to synthesize placeholder\ncontext:\n⊢ Nat"}],
+ "env": 0}
diff --git a/test/synthesize_placeholder.in b/test/synthesize_placeholder.in
new file mode 100644
index 0000000..d054d45
--- /dev/null
+++ b/test/synthesize_placeholder.in
@@ -0,0 +1 @@
+{"cmd": "def f : Nat := _"}
diff --git a/test/unfinished_tactic_block.expected.out b/test/unfinished_tactic_block.expected.out
new file mode 100644
index 0000000..7df4f82
--- /dev/null
+++ b/test/unfinished_tactic_block.expected.out
@@ -0,0 +1,11 @@
+{"sorries": [],
+ "messages":
+ [{"severity": "error",
+   "pos": {"line": 1, "column": 17},
+   "endPos": null,
+   "data": "unexpected end of input; expected '{'"},
+  {"severity": "error",
+   "pos": {"line": 1, "column": 15},
+   "endPos": {"line": 1, "column": 17},
+   "data": "unsolved goals\n⊢ Nat"}],
+ "env": 0}
diff --git a/test/unfinished_tactic_block.in b/test/unfinished_tactic_block.in
new file mode 100644
index 0000000..0c94423
--- /dev/null
+++ b/test/unfinished_tactic_block.in
@@ -0,0 +1 @@
+{"cmd": "def f : Nat := by"}

From d681ee581e9cad11f16afb06a4fa9057d292306d Mon Sep 17 00:00:00 2001
From: Scott Morrison <scott.morrison@gmail.com>
Date: Tue, 10 Oct 2023 15:53:02 +1100
Subject: [PATCH 2/5] working prototype

---
 REPL/JSON.lean                        |  26 ++++--
 REPL/Main.lean                        | 126 +++++++++++++++++++++++---
 test.sh                               |   2 +-
 test/proof_step.expected.out          |  14 +++
 test/proof_step.in                    |   7 ++
 test/tactic_sorry.expected.out        |   3 +-
 test/term_sorry.expected.out          |   3 +-
 test/unknown_proof_state.expected.out |  12 +++
 test/unknown_proof_state.in           |   3 +
 test/unknown_tactic.expected.out      |  12 +++
 test/unknown_tactic.in                |   3 +
 11 files changed, 185 insertions(+), 26 deletions(-)
 create mode 100644 test/proof_step.expected.out
 create mode 100644 test/proof_step.in
 create mode 100644 test/unknown_proof_state.expected.out
 create mode 100644 test/unknown_proof_state.in
 create mode 100644 test/unknown_tactic.expected.out
 create mode 100644 test/unknown_tactic.in

diff --git a/REPL/JSON.lean b/REPL/JSON.lean
index 0413ef6..db2f5a7 100644
--- a/REPL/JSON.lean
+++ b/REPL/JSON.lean
@@ -6,8 +6,6 @@ Authors: Scott Morrison
 import Lean.Data.Json
 import Lean.Message
 import Lean.Elab.InfoTree.Main
-import REPL.Lean.ContextInfo
-import REPL.InfoTree
 
 open Lean Elab InfoTree
 
@@ -17,11 +15,16 @@ namespace REPL
 If `env = none`, starts a new session (in which you can use `import`).
 If `env = some n`, builds on the existing environment `n`.
 -/
-structure Run where
+structure Command where
   env : Option Nat
   cmd : String
 deriving ToJson, FromJson
 
+structure ProofStep where
+  proofState : Nat
+  tactic : String
+deriving ToJson, FromJson
+
 /-- Line and column information for error messages and sorries. -/
 structure Pos where
   line : Nat
@@ -56,28 +59,31 @@ structure Sorry where
   pos : Pos
   endPos : Pos
   goal : String
+  proofState : Option Nat
 deriving ToJson, FromJson
 
 /-- Construct the JSON representation of a Lean sorry. -/
-def Sorry.of (ctx : ContextInfo) (g : SorryType) (pos endPos : Lean.Position) :
-    IO Sorry := do pure <|
+def Sorry.of (goal : String) (pos endPos : Lean.Position) (proofState : Option Nat) : Sorry :=
   { pos := ⟨pos.line, pos.column⟩,
     endPos := ⟨endPos.line, endPos.column⟩,
-    goal := ← match g with
-    | SorryType.tactic g => do pure s!"{(← ctx.ppGoals [g])}".trim
-    | SorryType.term _ none => unreachable!
-    | SorryType.term lctx (some t) => do pure s!"⊢ {← ctx.ppExpr lctx t}" }
+    goal,
+    proofState }
 
 /--
 A response to a Lean command.
 `env` can be used in later calls, to build on the stored environment.
 -/
-structure Response where
+structure CommandResponse where
   env : Nat
   messages : List Message
   sorries : List Sorry
 deriving ToJson, FromJson
 
+structure ProofStepResponse where
+  proofState : Nat
+  goals : List String
+deriving ToJson, FromJson
+
 /-- Json wrapper for an error. -/
 structure Error where
   message : String
diff --git a/REPL/Main.lean b/REPL/Main.lean
index 22efcb1..bbad33f 100644
--- a/REPL/Main.lean
+++ b/REPL/Main.lean
@@ -8,6 +8,8 @@ import REPL.Frontend
 import REPL.InfoTree
 import REPL.Util.Path
 import REPL.Util.TermUnsafe
+import REPL.Lean.ContextInfo
+import REPL.InfoTree
 
 /-!
 # A REPL for Lean.
@@ -55,31 +57,122 @@ open Lean Elab
 
 namespace REPL
 
+-- #print Tactic.State -- goals: List MVarId
+-- #print Tactic.Context
+-- #print Term.State -- pendingMVars, syntheticMVars
+-- #print Term.Context
+-- #print Meta.State -- MetavarContext
+-- #print Meta.Context -- LocalContext, ...
+-- #print Core.State -- Environment, MessageLog, Elab.InfoState, ...
+-- #print Core.Context
+
+structure ProofState where
+  goals : List MVarId
+  mctx : MetavarContext
+  lctx : LocalContext
+  env : Environment
+
+namespace ProofState
+
+open Lean Elab Tactic
+
+def runCoreM (p : ProofState) (t : CoreM α) : IO (α × ProofState) := do
+  let ctx : Core.Context := { fileName := "", options := {}, fileMap := default }
+  let state : Core.State := { env := p.env }
+  let (a, state') ← (Lean.Core.CoreM.toIO · ctx state) do t
+  return (a, { p with env := state'.env })
+
+def runMetaM (p : ProofState) (t : MetaM α) : IO (α × ProofState) := do
+  let ctx : Meta.Context := { lctx := p.lctx }
+  let state : Meta.State := { mctx := p.mctx }
+  let ((a, state'), p') ← p.runCoreM (Lean.Meta.MetaM.run (ctx := ctx) (s := state) do t)
+  return (a, { p' with mctx := state'.mctx })
+
+def runTermElabM (p : ProofState) (t : TermElabM α) : IO (α × ProofState) := do
+  let ((a, _), p') ← p.runMetaM (Lean.Elab.Term.TermElabM.run do t)
+  return (a, p')
+
+def runTacticM (p : ProofState) (t : TacticM α) : IO (α × ProofState) := do
+  let ctx : Tactic.Context := { elaborator := .anonymous }
+  let state : Tactic.State := { goals := p.goals }
+  let ((a, state'), p') ← p.runTermElabM (t ctx |>.run state)
+  return (a, { p' with goals := state'.goals })
+
+def runTacticM' (p : ProofState) (t : TacticM α) : IO ProofState :=
+  Prod.snd <$> p.runTacticM t
+
+def runSyntax (p : ProofState) (t : Syntax) : IO ProofState :=
+  Prod.snd <$> p.runTacticM (evalTactic t)
+
+def runString (p : ProofState) (t : String) : IO ProofState :=
+  match Parser.runParserCategory p.env `tactic t with
+  | .error e => throw (IO.userError e)
+  | .ok stx => p.runSyntax stx
+
+def ppGoals (p : ProofState) : IO (List Format) :=
+  Prod.fst <$> p.runTacticM do (← getGoals).mapM (Meta.ppGoal ·)
+
+end ProofState
+
 /-- The monadic state for the Lean REPL. -/
 structure State where
-  environments : Array Environment
-  lines : Array Nat
+  environments : Array Environment := #[]
+  proofStates : Array ProofState := #[]
+  lines : Array Nat := #[]
 
 /-- The Lean REPL monad. -/
 abbrev M (m : Type → Type) := StateT State m
 
 variable [Monad m] [MonadLiftT IO m]
 
-/-- Get the next available id for a new environment. -/
-def nextId : M m Nat := do pure (← get).environments.size
+def recordEnvironment (env : Environment) : M m Nat := do
+  let id := (← get).environments.size
+  modify fun s => { s with environments := s.environments.push env }
+  return id
+
+def recordLines (lines : Nat) : M m Unit :=
+  modify fun s => { s with lines := s.lines.push lines }
+
+def recordProofState (proofState : ProofState) : M m Nat := do
+  let id := (← get).proofStates.size
+  modify fun s => { s with proofStates := s.proofStates.push proofState }
+  return id
+
+def createProofState (ctx : ContextInfo) (lctx? : Option LocalContext) (goals : List MVarId) (types : List Expr := []) :
+    IO ProofState := do
+  ctx.runMetaM (lctx?.getD {}) do pure <|
+    { goals := goals ++ (← types.mapM fun t => Expr.mvarId! <$> Meta.mkFreshExprMVar (some t))
+      mctx := ← getMCtx
+      lctx := ← getLCtx
+      env := ← getEnv }
 
 /-- Run a command, returning the id of the new environment, and any messages and sorries. -/
-def run (s : Run) : M m Response := do
+def runCommand (s : Command) : M m CommandResponse := do
   let env? := s.env.bind ((← get).environments[·]?)
   let (env, messages, trees) ← IO.processInput s.cmd env? {} ""
   let messages ← messages.mapM fun m => Message.of m
   let sorries ← trees.bind InfoTree.sorries |>.mapM
-    fun ⟨ctx, g, pos, endPos⟩ => Sorry.of ctx g pos endPos
-  let lines := s.cmd.splitOn "\n" |>.length
-  let id ← nextId
-  modify fun s => { environments := s.environments.push env, lines := s.lines.push lines }
+    fun ⟨ctx, g, pos, endPos⟩ => do
+      let (goal, proofState) ← match g with
+      | .tactic g => do
+         pure (s!"{(← ctx.ppGoals [g])}".trim, some (← createProofState ctx none [g]))
+      | .term lctx (some t) => do
+         pure (s!"⊢ {← ctx.ppExpr lctx t}", some (← createProofState ctx lctx [] [t]))
+      | .term _ none => unreachable!
+      let proofStateId ← proofState.mapM recordProofState
+      return Sorry.of goal pos endPos proofStateId
+  recordLines <| s.cmd.splitOn "\n" |>.length
+  let id ← recordEnvironment env
   pure ⟨id, messages, sorries⟩
 
+def runProofStep (s : ProofStep) : M m (ProofStepResponse ⊕ Error) := do
+  match (← get).proofStates[s.proofState]? with
+  | none => return Sum.inr ⟨"Unknown proof state."⟩
+  | some proofState =>
+    let proofState' ← proofState.runString s.tactic
+    let id ← recordProofState proofState'
+    return Sum.inl { proofState := id, goals := (← proofState'.ppGoals).map fun s => s!"{s}" }
+
 end REPL
 
 open REPL
@@ -91,9 +184,14 @@ partial def getLines : IO String := do
   | "\n" => pure "\n"
   | line => pure <| line ++ (← getLines)
 
+instance [ToJson α] [ToJson β] : ToJson (α ⊕ β) where
+  toJson x := match x with
+  | .inl a => toJson a
+  | .inr b => toJson b
+
 /-- Read-eval-print loop for Lean. -/
 partial def repl : IO Unit :=
-  StateT.run' loop ⟨#[], #[]⟩
+  StateT.run' loop {}
 where loop : M IO Unit := do
   let query ← getLines
   if query = "" then
@@ -102,11 +200,13 @@ where loop : M IO Unit := do
   match json with
   | .error e => IO.println <| toString <| toJson (⟨e⟩ : Error)
   | .ok j => match fromJson? j with
-    | .error e => IO.println <| toString <| toJson (⟨e⟩ : Error)
-    | .ok (r : Run) => IO.println <| toString <| toJson (← run r)
+    | .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)
   loop
 
-/-- Main executable function, run as `lake env lean --run Mathlib/Util/REPL.lean`. -/
+/-- Main executable function, run as `lake exe repl`. -/
 def main (_ : List String) : IO Unit := do
   searchPathRef.set compile_time_search_path%
   repl
diff --git a/test.sh b/test.sh
index 385ba5c..ae6ca2f 100755
--- a/test.sh
+++ b/test.sh
@@ -22,7 +22,7 @@ for infile in $IN_DIR/*.in; do
 
     # Run the command and store its output in a temporary file
     tmpfile=$(mktemp)
-    build/bin/repl < "$infile" > "$tmpfile"
+    build/bin/repl < "$infile" > "$tmpfile" 2>&1
 
     # Compare the output with the expected output
     if diff "$tmpfile" "$expectedfile"; then
diff --git a/test/proof_step.expected.out b/test/proof_step.expected.out
new file mode 100644
index 0000000..cb909e2
--- /dev/null
+++ b/test/proof_step.expected.out
@@ -0,0 +1,14 @@
+{"sorries":
+ [{"proofState": 0,
+   "pos": {"line": 1, "column": 18},
+   "goal": "⊢ Nat",
+   "endPos": {"line": 1, "column": 23}}],
+ "messages":
+ [{"severity": "warning",
+   "pos": {"line": 1, "column": 4},
+   "endPos": {"line": 1, "column": 5},
+   "data": "declaration uses 'sorry'"}],
+ "env": 0}
+{"proofState": 1, "goals": ["⊢ Int"]}
+{"proofState": 2, "goals": ["t : Nat\n⊢ Nat"]}
+{"proofState": 3, "goals": []}
diff --git a/test/proof_step.in b/test/proof_step.in
new file mode 100644
index 0000000..51a3f50
--- /dev/null
+++ b/test/proof_step.in
@@ -0,0 +1,7 @@
+{"cmd" : "def f : Nat := by sorry"}
+
+{"tactic": "apply Int.natAbs", "proofState": 0}
+
+{"tactic": "have t := 42", "proofState": 0}
+
+{"tactic": "exact t", "proofState": 2}
diff --git a/test/tactic_sorry.expected.out b/test/tactic_sorry.expected.out
index f12573b..43fd2cb 100644
--- a/test/tactic_sorry.expected.out
+++ b/test/tactic_sorry.expected.out
@@ -1,5 +1,6 @@
 {"sorries":
- [{"pos": {"line": 1, "column": 18},
+ [{"proofState": 0,
+   "pos": {"line": 1, "column": 18},
    "goal": "⊢ Nat",
    "endPos": {"line": 1, "column": 23}}],
  "messages":
diff --git a/test/term_sorry.expected.out b/test/term_sorry.expected.out
index b38c068..71833dc 100644
--- a/test/term_sorry.expected.out
+++ b/test/term_sorry.expected.out
@@ -1,5 +1,6 @@
 {"sorries":
- [{"pos": {"line": 1, "column": 15},
+ [{"proofState": 0,
+   "pos": {"line": 1, "column": 15},
    "goal": "⊢ Nat",
    "endPos": {"line": 1, "column": 20}}],
  "messages":
diff --git a/test/unknown_proof_state.expected.out b/test/unknown_proof_state.expected.out
new file mode 100644
index 0000000..ccc3261
--- /dev/null
+++ b/test/unknown_proof_state.expected.out
@@ -0,0 +1,12 @@
+{"sorries":
+ [{"proofState": 0,
+   "pos": {"line": 1, "column": 18},
+   "goal": "⊢ Nat",
+   "endPos": {"line": 1, "column": 23}}],
+ "messages":
+ [{"severity": "warning",
+   "pos": {"line": 1, "column": 4},
+   "endPos": {"line": 1, "column": 5},
+   "data": "declaration uses 'sorry'"}],
+ "env": 0}
+{"message": "Unknown proof state."}
diff --git a/test/unknown_proof_state.in b/test/unknown_proof_state.in
new file mode 100644
index 0000000..bdd1a68
--- /dev/null
+++ b/test/unknown_proof_state.in
@@ -0,0 +1,3 @@
+{"cmd" : "def f : Nat := by sorry"}
+
+{"tactic": "exact 42", "proofState": 1}
diff --git a/test/unknown_tactic.expected.out b/test/unknown_tactic.expected.out
new file mode 100644
index 0000000..5b60a97
--- /dev/null
+++ b/test/unknown_tactic.expected.out
@@ -0,0 +1,12 @@
+{"sorries":
+ [{"proofState": 0,
+   "pos": {"line": 1, "column": 18},
+   "goal": "⊢ Nat",
+   "endPos": {"line": 1, "column": 23}}],
+ "messages":
+ [{"severity": "warning",
+   "pos": {"line": 1, "column": 4},
+   "endPos": {"line": 1, "column": 5},
+   "data": "declaration uses 'sorry'"}],
+ "env": 0}
+uncaught exception: <input>:1:1: unknown tactic
diff --git a/test/unknown_tactic.in b/test/unknown_tactic.in
new file mode 100644
index 0000000..d1af249
--- /dev/null
+++ b/test/unknown_tactic.in
@@ -0,0 +1,3 @@
+{"cmd" : "def f : Nat := by sorry"}
+
+{"tactic": "exat 42", "proofState": 0}

From 8a429c70ce024591c2031e62969a3ab862d42adb Mon Sep 17 00:00:00 2001
From: Scott Morrison <scott.morrison@gmail.com>
Date: Wed, 11 Oct 2023 09:45:49 +1100
Subject: [PATCH 3/5] use full state

---
 REPL/Main.lean | 101 ++++++++++++++++++++++++++++++++++++-------------
 1 file changed, 75 insertions(+), 26 deletions(-)

diff --git a/REPL/Main.lean b/REPL/Main.lean
index bbad33f..54c949d 100644
--- a/REPL/Main.lean
+++ b/REPL/Main.lean
@@ -67,36 +67,76 @@ namespace REPL
 -- #print Core.Context
 
 structure ProofState where
-  goals : List MVarId
-  mctx : MetavarContext
-  lctx : LocalContext
-  env : Environment
+  tacticSavedState : Tactic.SavedState
+  tacticContext : Tactic.Context
+  termContext : Term.Context
+  metaContext : Meta.Context
+  coreContext : Core.Context
 
 namespace ProofState
 
 open Lean Elab Tactic
 
+def tacticState (p : ProofState) : Tactic.State := p.tacticSavedState.tactic
+def termState (p : ProofState) : Term.State := p.tacticSavedState.term.elab
+def metaState (p : ProofState) : Meta.State := p.tacticSavedState.term.meta.meta
+def coreState (p : ProofState) : Core.State := p.tacticSavedState.term.meta.core
+
+def withCoreState (p : ProofState) (core : Core.State) : ProofState :=
+  { p with
+    tacticSavedState :=
+    { p.tacticSavedState with
+      term :=
+      { p.tacticSavedState.term with
+        meta :=
+        { p.tacticSavedState.term.meta with
+          core }}}}
+
+def withMetaState (p : ProofState) (meta : Meta.State) : ProofState :=
+  { p with
+    tacticSavedState :=
+    { p.tacticSavedState with
+      term :=
+      { p.tacticSavedState.term with
+        meta :=
+        { p.tacticSavedState.term.meta with
+          meta }}}}
+
+def withTermState (p : ProofState) («elab» : Term.State) : ProofState :=
+  { p with
+    tacticSavedState :=
+    { p.tacticSavedState with
+      term :=
+      { p.tacticSavedState.term with
+        «elab» }}}
+
+def withTacticState (p : ProofState) (tactic : Tactic.State) : ProofState :=
+  { p with
+    tacticSavedState :=
+    { p.tacticSavedState with
+      tactic }}
+
 def runCoreM (p : ProofState) (t : CoreM α) : IO (α × ProofState) := do
-  let ctx : Core.Context := { fileName := "", options := {}, fileMap := default }
-  let state : Core.State := { env := p.env }
-  let (a, state') ← (Lean.Core.CoreM.toIO · ctx state) do t
-  return (a, { p with env := state'.env })
+  -- let ctx : Core.Context := { fileName := "", options := {}, fileMap := default }
+  -- let state : Core.State := { env := p.env }
+  let (a, state') ← (Lean.Core.CoreM.toIO · p.coreContext p.coreState) do t
+  return (a, p.withCoreState state')
 
 def runMetaM (p : ProofState) (t : MetaM α) : IO (α × ProofState) := do
-  let ctx : Meta.Context := { lctx := p.lctx }
-  let state : Meta.State := { mctx := p.mctx }
-  let ((a, state'), p') ← p.runCoreM (Lean.Meta.MetaM.run (ctx := ctx) (s := state) do t)
-  return (a, { p' with mctx := state'.mctx })
+  -- let ctx : Meta.Context := { lctx := p.lctx }
+  -- let state : Meta.State := { mctx := p.mctx }
+  let ((a, state'), p') ← p.runCoreM (Lean.Meta.MetaM.run (ctx := p.metaContext) (s := p.metaState) do t)
+  return (a, p'.withMetaState state')
 
 def runTermElabM (p : ProofState) (t : TermElabM α) : IO (α × ProofState) := do
-  let ((a, _), p') ← p.runMetaM (Lean.Elab.Term.TermElabM.run do t)
-  return (a, p')
+  let ((a, state'), p') ← p.runMetaM (Lean.Elab.Term.TermElabM.run (s := p.termState) do t)
+  return (a, p'.withTermState state')
 
 def runTacticM (p : ProofState) (t : TacticM α) : IO (α × ProofState) := do
-  let ctx : Tactic.Context := { elaborator := .anonymous }
-  let state : Tactic.State := { goals := p.goals }
-  let ((a, state'), p') ← p.runTermElabM (t ctx |>.run state)
-  return (a, { p' with goals := state'.goals })
+  -- let ctx : Tactic.Context := { elaborator := .anonymous }
+  -- let state : Tactic.State := { goals := p.goals }
+  let ((a, state'), p') ← p.runTermElabM (t p.tacticContext |>.run p.tacticState)
+  return (a, p'.withTacticState state')
 
 def runTacticM' (p : ProofState) (t : TacticM α) : IO ProofState :=
   Prod.snd <$> p.runTacticM t
@@ -105,7 +145,7 @@ def runSyntax (p : ProofState) (t : Syntax) : IO ProofState :=
   Prod.snd <$> p.runTacticM (evalTactic t)
 
 def runString (p : ProofState) (t : String) : IO ProofState :=
-  match Parser.runParserCategory p.env `tactic t with
+  match Parser.runParserCategory p.coreState.env `tactic t with
   | .error e => throw (IO.userError e)
   | .ok stx => p.runSyntax stx
 
@@ -138,13 +178,22 @@ def recordProofState (proofState : ProofState) : M m Nat := do
   modify fun s => { s with proofStates := s.proofStates.push proofState }
   return id
 
-def createProofState (ctx : ContextInfo) (lctx? : Option LocalContext) (goals : List MVarId) (types : List Expr := []) :
-    IO ProofState := do
-  ctx.runMetaM (lctx?.getD {}) do pure <|
-    { goals := goals ++ (← types.mapM fun t => Expr.mvarId! <$> Meta.mkFreshExprMVar (some t))
-      mctx := ← getMCtx
-      lctx := ← getLCtx
-      env := ← getEnv }
+def createProofState (ctx : ContextInfo) (lctx? : Option LocalContext)
+    (goals : List MVarId) (types : List Expr := []) : IO ProofState := do
+  ctx.runMetaM (lctx?.getD {}) do
+    let goals := goals ++ (← types.mapM fun t => Expr.mvarId! <$> Meta.mkFreshExprMVar (some t))
+    pure <|
+    { coreContext := ← readThe Core.Context
+      metaContext := ← readThe Meta.Context
+      termContext := {}
+      tacticContext := { elaborator := .anonymous }
+      tacticSavedState :=
+      { tactic := { goals }
+        term :=
+        { «elab» := {},
+          meta :=
+          { meta := ← getThe Meta.State
+            core := ← getThe Core.State } } } }
 
 /-- Run a command, returning the id of the new environment, and any messages and sorries. -/
 def runCommand (s : Command) : M m CommandResponse := do

From 57acea3fbf77854727ef7f75362eb76ffc219c55 Mon Sep 17 00:00:00 2001
From: Scott Morrison <scott.morrison@gmail.com>
Date: Wed, 11 Oct 2023 09:46:01 +1100
Subject: [PATCH 4/5] further tests

---
 test/proof_step2.expected.out | 14 ++++++++++++++
 test/proof_step2.in           |  7 +++++++
 test/proof_step2.lean         |  3 +++
 3 files changed, 24 insertions(+)
 create mode 100644 test/proof_step2.expected.out
 create mode 100644 test/proof_step2.in
 create mode 100644 test/proof_step2.lean

diff --git a/test/proof_step2.expected.out b/test/proof_step2.expected.out
new file mode 100644
index 0000000..cb909e2
--- /dev/null
+++ b/test/proof_step2.expected.out
@@ -0,0 +1,14 @@
+{"sorries":
+ [{"proofState": 0,
+   "pos": {"line": 1, "column": 18},
+   "goal": "⊢ Nat",
+   "endPos": {"line": 1, "column": 23}}],
+ "messages":
+ [{"severity": "warning",
+   "pos": {"line": 1, "column": 4},
+   "endPos": {"line": 1, "column": 5},
+   "data": "declaration uses 'sorry'"}],
+ "env": 0}
+{"proofState": 1, "goals": ["⊢ Int"]}
+{"proofState": 2, "goals": ["t : Nat\n⊢ Nat"]}
+{"proofState": 3, "goals": []}
diff --git a/test/proof_step2.in b/test/proof_step2.in
new file mode 100644
index 0000000..a0b5d53
--- /dev/null
+++ b/test/proof_step2.in
@@ -0,0 +1,7 @@
+{"cmd" : "def f : Nat := by sorry"}
+
+{"tactic": "apply Int.natAbs", "proofState": 0}
+
+{"tactic": "have t : Nat := 42", "proofState": 0}
+
+{"tactic": "exact t", "proofState": 2}
diff --git a/test/proof_step2.lean b/test/proof_step2.lean
new file mode 100644
index 0000000..ea06ed8
--- /dev/null
+++ b/test/proof_step2.lean
@@ -0,0 +1,3 @@
+def f : Nat := by
+  have t : Nat := 42
+  exact t

From 933547c50168e439034b109a60388dac52db83f8 Mon Sep 17 00:00:00 2001
From: Scott Morrison <scott.morrison@gmail.com>
Date: Wed, 11 Oct 2023 10:28:20 +1100
Subject: [PATCH 5/5] cleanup and documentation

---
 README.md                             |  64 +++++++++-
 REPL/InfoTree.lean                    |   4 +-
 REPL/JSON.lean                        |  11 ++
 REPL/Main.lean                        | 173 +++++++++++++-------------
 run                                   |   1 -
 test/proof_step.in                    |   2 +-
 test/proof_step2.expected.out         |  14 ---
 test/proof_step2.in                   |   7 --
 test/proof_step2.lean                 |   3 -
 test/unknown_environment.expected.out |   1 +
 test/unknown_environment.in           |   1 +
 11 files changed, 157 insertions(+), 124 deletions(-)
 delete mode 100755 run
 delete mode 100644 test/proof_step2.expected.out
 delete mode 100644 test/proof_step2.in
 delete mode 100644 test/proof_step2.lean
 create mode 100644 test/unknown_environment.expected.out
 create mode 100644 test/unknown_environment.in

diff --git a/README.md b/README.md
index f56b171..ef32b5f 100644
--- a/README.md
+++ b/README.md
@@ -1,8 +1,15 @@
 # A read-eval-print-loop for Lean 4
 
-Run using `./run`.
+Run using `lake exe repl`.
 Communicates via JSON on stdin and stdout.
 Commands should be separated by blank lines.
+
+The REPL works both in "command" mode and "tactic" mode.
+
+## Command mode
+
+In command mode, you send complete commands (e.g. declarations) to the REPL.
+
 Commands may be of the form
 
 ```json
@@ -23,13 +30,22 @@ You can only use `import` commands when you do not specify the `env` field.
 
 You can backtrack simply by using earlier values for `env`.
 
-The results are of the form
+The response includes:
+* A numeric label for the `Environment` after your command,
+  which you can use as the starting point for subsequent commands.
+* Any messages generated while processing your command.
+* A list of the `sorry`s in your command, including
+  * their expected type, and
+  * a numeric label for the proof state at the `sorry`, which you can then use in tactic mode.
+
+Example output:
 
 ```json
 {"sorries":
  [{"pos": {"line": 1, "column": 18},
    "endPos": {"line": 1, "column": 23},
-   "goal": "⊢ Nat"}],
+   "goal": "⊢ Nat",
+   "proofState": 0}],
  "messages":
  [{"severity": "error",
    "pos": {"line": 1, "column": 23},
@@ -39,7 +55,43 @@ The results are of the form
  "env": 6}
 ```
 
-showing any messages generated, or sorries with their goal states.
+showing any messages generated, and sorries with their goal states.
+
+## Tactic mode (experimental)
+
+The only way to enter tactic mode at present is to first issue a command containing a `sorry`,
+and then using the `proofState` index returned for each `sorry`.
+
+Example usage:
+```json
+{"cmd" : "def f : Nat := by sorry"}
+
+{"sorries":
+ [{"proofState": 0,
+   "pos": {"line": 1, "column": 18},
+   "goal": "⊢ Nat",
+   "endPos": {"line": 1, "column": 23}}],
+ "messages":
+ [{"severity": "warning",
+   "pos": {"line": 1, "column": 4},
+   "endPos": {"line": 1, "column": 5},
+   "data": "declaration uses 'sorry'"}],
+ "env": 0}
+
+{"tactic": "apply Int.natAbs", "proofState": 0}
+
+{"proofState": 1, "goals": ["⊢ Int"]}
+
+{"tactic": "exact -37", "proofState": 1}
+
+{"proofState": 2, "goals": []}
+```
+
+At present there is nothing you can do with a completed proof state:
+we would like to extend this so that you can replace the original `sorry` with your tactic script,
+and obtain the resulting `Environment`
+
+## Future work
 
-Information is generated for tactic mode sorries,
-but currently not for term mode sorries.
+* Replay tactic scripts from tactic mode back into the original `sorry`.
+* Serialization and deserialization of environments and proof states
diff --git a/REPL/InfoTree.lean b/REPL/InfoTree.lean
index 0f6429b..557a45a 100644
--- a/REPL/InfoTree.lean
+++ b/REPL/InfoTree.lean
@@ -6,9 +6,7 @@ Authors: Scott Morrison
 import Lean
 
 /-!
-# Code to move (eventually) to core.
-
-Mostly additional functions to deal with `InfoTree`.
+Additional functions to deal with `InfoTree`.
 -/
 
 open Lean Elab Meta
diff --git a/REPL/JSON.lean b/REPL/JSON.lean
index db2f5a7..d3cb2d8 100644
--- a/REPL/JSON.lean
+++ b/REPL/JSON.lean
@@ -20,6 +20,9 @@ structure Command where
   cmd : String
 deriving ToJson, FromJson
 
+/--
+Run a tactic in a proof state.
+-/
 structure ProofStep where
   proofState : Nat
   tactic : String
@@ -59,6 +62,10 @@ structure Sorry where
   pos : Pos
   endPos : Pos
   goal : String
+  /--
+  The index of the proof state at the sorry.
+  You can use the `ProofStep` instruction to run a tactic at this state.
+  -/
   proofState : Option Nat
 deriving ToJson, FromJson
 
@@ -79,6 +86,10 @@ structure CommandResponse where
   sorries : List Sorry
 deriving ToJson, FromJson
 
+/--
+A response to a Lean tactic.
+`proofState` can be used in later calls, to run further tactics.
+-/
 structure ProofStepResponse where
   proofState : Nat
   goals : List String
diff --git a/REPL/Main.lean b/REPL/Main.lean
index 54c949d..7cad970 100644
--- a/REPL/Main.lean
+++ b/REPL/Main.lean
@@ -57,98 +57,66 @@ open Lean Elab
 
 namespace REPL
 
--- #print Tactic.State -- goals: List MVarId
--- #print Tactic.Context
--- #print Term.State -- pendingMVars, syntheticMVars
--- #print Term.Context
--- #print Meta.State -- MetavarContext
--- #print Meta.Context -- LocalContext, ...
--- #print Core.State -- Environment, MessageLog, Elab.InfoState, ...
--- #print Core.Context
-
+/--
+Bundled structure for the `State` and `Context` objects
+for the `CoreM`, `MetaM`, `TermElabM`, and `TacticM` monads.
+-/
 structure ProofState where
-  tacticSavedState : Tactic.SavedState
+  coreState     : Core.State
+  coreContext   : Core.Context
+  metaState     : Meta.State
+  metaContext   : Meta.Context
+  termState     : Term.State
+  termContext   : Term.Context
+  tacticState   : Tactic.State
   tacticContext : Tactic.Context
-  termContext : Term.Context
-  metaContext : Meta.Context
-  coreContext : Core.Context
-
 namespace ProofState
 
 open Lean Elab Tactic
 
-def tacticState (p : ProofState) : Tactic.State := p.tacticSavedState.tactic
-def termState (p : ProofState) : Term.State := p.tacticSavedState.term.elab
-def metaState (p : ProofState) : Meta.State := p.tacticSavedState.term.meta.meta
-def coreState (p : ProofState) : Core.State := p.tacticSavedState.term.meta.core
-
-def withCoreState (p : ProofState) (core : Core.State) : ProofState :=
-  { p with
-    tacticSavedState :=
-    { p.tacticSavedState with
-      term :=
-      { p.tacticSavedState.term with
-        meta :=
-        { p.tacticSavedState.term.meta with
-          core }}}}
-
-def withMetaState (p : ProofState) (meta : Meta.State) : ProofState :=
-  { p with
-    tacticSavedState :=
-    { p.tacticSavedState with
-      term :=
-      { p.tacticSavedState.term with
-        meta :=
-        { p.tacticSavedState.term.meta with
-          meta }}}}
-
-def withTermState (p : ProofState) («elab» : Term.State) : ProofState :=
-  { p with
-    tacticSavedState :=
-    { p.tacticSavedState with
-      term :=
-      { p.tacticSavedState.term with
-        «elab» }}}
-
-def withTacticState (p : ProofState) (tactic : Tactic.State) : ProofState :=
-  { p with
-    tacticSavedState :=
-    { p.tacticSavedState with
-      tactic }}
-
+/-- Run a `CoreM` monadic function in the current `ProofState`, updating the `Core.State`. -/
 def runCoreM (p : ProofState) (t : CoreM α) : IO (α × ProofState) := do
-  -- let ctx : Core.Context := { fileName := "", options := {}, fileMap := default }
-  -- let state : Core.State := { env := p.env }
-  let (a, state') ← (Lean.Core.CoreM.toIO · p.coreContext p.coreState) do t
-  return (a, p.withCoreState state')
+  let (a, coreState) ← (Lean.Core.CoreM.toIO · p.coreContext p.coreState) do t
+  return (a, { p with coreState })
 
+/-- Run a `MetaM` monadic function in the current `ProofState`, updating the `Meta.State`. -/
 def runMetaM (p : ProofState) (t : MetaM α) : IO (α × ProofState) := do
-  -- let ctx : Meta.Context := { lctx := p.lctx }
-  -- let state : Meta.State := { mctx := p.mctx }
-  let ((a, state'), p') ← p.runCoreM (Lean.Meta.MetaM.run (ctx := p.metaContext) (s := p.metaState) do t)
-  return (a, p'.withMetaState state')
+  let ((a, metaState), p') ← p.runCoreM (Lean.Meta.MetaM.run (ctx := p.metaContext) (s := p.metaState) do t)
+  return (a, { p' with metaState })
 
+/-- Run a `TermElabM` monadic function in the current `ProofState`, updating the `Term.State`. -/
 def runTermElabM (p : ProofState) (t : TermElabM α) : IO (α × ProofState) := do
-  let ((a, state'), p') ← p.runMetaM (Lean.Elab.Term.TermElabM.run (s := p.termState) do t)
-  return (a, p'.withTermState state')
+  let ((a, termState), p') ← p.runMetaM (Lean.Elab.Term.TermElabM.run (s := p.termState) do t)
+  return (a, { p' with termState })
 
+/-- Run a `TacticM` monadic function in the current `ProofState`, updating the `Tactic.State`. -/
 def runTacticM (p : ProofState) (t : TacticM α) : IO (α × ProofState) := do
-  -- let ctx : Tactic.Context := { elaborator := .anonymous }
-  -- let state : Tactic.State := { goals := p.goals }
-  let ((a, state'), p') ← p.runTermElabM (t p.tacticContext |>.run p.tacticState)
-  return (a, p'.withTacticState state')
+  let ((a, tacticState), p') ← p.runTermElabM (t p.tacticContext |>.run p.tacticState)
+  return (a, { p' with tacticState })
 
+/--
+Run a `TacticM` monadic function in the current `ProofState`, updating the `Tactic.State`,
+and discarding the return value.
+-/
 def runTacticM' (p : ProofState) (t : TacticM α) : IO ProofState :=
   Prod.snd <$> p.runTacticM t
 
+/--
+Evaluate a `Syntax` into a `TacticM` tactic, and run it in the current `ProofState`.
+-/
 def runSyntax (p : ProofState) (t : Syntax) : IO ProofState :=
   Prod.snd <$> p.runTacticM (evalTactic t)
 
+/--
+Parse a string into a `Syntax`, evaluate it as a `TacticM` tactic,
+and run it in the current `ProofState`.
+-/
 def runString (p : ProofState) (t : String) : IO ProofState :=
   match Parser.runParserCategory p.coreState.env `tactic t with
   | .error e => throw (IO.userError e)
   | .ok stx => p.runSyntax stx
 
+/-- Pretty print the current goals in the `ProofState`. -/
 def ppGoals (p : ProofState) : IO (List Format) :=
   Prod.fst <$> p.runTacticM do (← getGoals).mapM (Meta.ppGoal ·)
 
@@ -156,48 +124,71 @@ end ProofState
 
 /-- The monadic state for the Lean REPL. -/
 structure State where
+  /--
+  Environment snapshots after complete declarations.
+  The user can run a declaration in a given environment using `{"cmd": "def f := 37", "env": 17}`.
+  -/
   environments : Array Environment := #[]
+  /--
+  Proof states after individual tactics.
+  The user can run a tactic in a given proof state using `{"tactic": "exact 42", "proofState": 5}`.
+  Declarations with containing `sorry` record a proof state at each sorry,
+  and report the numerical index for the recorded state at each sorry.
+  -/
   proofStates : Array ProofState := #[]
-  lines : Array Nat := #[]
 
-/-- The Lean REPL monad. -/
+/--
+The Lean REPL monad.
+
+We only use this with `m := IO`, but it is set up as a monad transformer for flexibility.
+-/
 abbrev M (m : Type → Type) := StateT State m
 
 variable [Monad m] [MonadLiftT IO m]
 
+/-- Record an `Environment` into the REPL state, returning its index for future use. -/
 def recordEnvironment (env : Environment) : M m Nat := do
   let id := (← get).environments.size
   modify fun s => { s with environments := s.environments.push env }
   return id
 
-def recordLines (lines : Nat) : M m Unit :=
-  modify fun s => { s with lines := s.lines.push lines }
-
+/-- Record a `ProofState` into the REPL state, returning its index for future use. -/
 def recordProofState (proofState : ProofState) : M m Nat := do
   let id := (← get).proofStates.size
   modify fun s => { s with proofStates := s.proofStates.push proofState }
   return id
 
+/--
+Construct a `ProofState` from a `ContextInfo` and optional `LocalContext`, and a list of goals.
+
+For convenience, we also allow a list of `Expr`s, and these are appended to the goals
+as fresh metavariables with the given types.
+-/
 def createProofState (ctx : ContextInfo) (lctx? : Option LocalContext)
     (goals : List MVarId) (types : List Expr := []) : IO ProofState := do
   ctx.runMetaM (lctx?.getD {}) do
     let goals := goals ++ (← types.mapM fun t => Expr.mvarId! <$> Meta.mkFreshExprMVar (some t))
     pure <|
-    { coreContext := ← readThe Core.Context
+    { coreState := ← getThe Core.State
+      coreContext := ← readThe Core.Context
+      metaState := ← getThe Meta.State
       metaContext := ← readThe Meta.Context
+      termState := {}
       termContext := {}
-      tacticContext := { elaborator := .anonymous }
-      tacticSavedState :=
-      { tactic := { goals }
-        term :=
-        { «elab» := {},
-          meta :=
-          { meta := ← getThe Meta.State
-            core := ← getThe Core.State } } } }
-
-/-- Run a command, returning the id of the new environment, and any messages and sorries. -/
-def runCommand (s : Command) : M m CommandResponse := do
-  let env? := s.env.bind ((← get).environments[·]?)
+      tacticState := { goals }
+      tacticContext := { elaborator := .anonymous } }
+
+/--
+Run a command, returning the id of the new environment, and any messages and sorries.
+-/
+def runCommand (s : Command) : M m (CommandResponse ⊕ Error) := do
+  let (env?, notFound) ← do match s.env with
+  | none => pure (none, false)
+  | some i => do match (← get).environments[i]? with
+    | some env => pure (some env, false)
+    | none => pure (none, true)
+  if notFound then
+    return .inr ⟨"Unknown environment."⟩
   let (env, messages, trees) ← IO.processInput s.cmd env? {} ""
   let messages ← messages.mapM fun m => Message.of m
   let sorries ← trees.bind InfoTree.sorries |>.mapM
@@ -210,17 +201,21 @@ def runCommand (s : Command) : M m CommandResponse := do
       | .term _ none => unreachable!
       let proofStateId ← proofState.mapM recordProofState
       return Sorry.of goal pos endPos proofStateId
-  recordLines <| s.cmd.splitOn "\n" |>.length
   let id ← recordEnvironment env
-  pure ⟨id, messages, sorries⟩
+  return .inl ⟨id, messages, sorries⟩
 
+/--
+Run a single tactic, returning the id of the new proof statement, and the new goals.
+-/
+-- TODO return messages?
+-- TODO detect sorries?
 def runProofStep (s : ProofStep) : M m (ProofStepResponse ⊕ Error) := do
   match (← get).proofStates[s.proofState]? with
-  | none => return Sum.inr ⟨"Unknown proof state."⟩
+  | none => return .inr ⟨"Unknown proof state."⟩
   | some proofState =>
     let proofState' ← proofState.runString s.tactic
     let id ← recordProofState proofState'
-    return Sum.inl { proofState := id, goals := (← proofState'.ppGoals).map fun s => s!"{s}" }
+    return .inl { proofState := id, goals := (← proofState'.ppGoals).map fun s => s!"{s}" }
 
 end REPL
 
diff --git a/run b/run
deleted file mode 100755
index f1c441f..0000000
--- a/run
+++ /dev/null
@@ -1 +0,0 @@
-lake exe repl
diff --git a/test/proof_step.in b/test/proof_step.in
index 51a3f50..a0b5d53 100644
--- a/test/proof_step.in
+++ b/test/proof_step.in
@@ -2,6 +2,6 @@
 
 {"tactic": "apply Int.natAbs", "proofState": 0}
 
-{"tactic": "have t := 42", "proofState": 0}
+{"tactic": "have t : Nat := 42", "proofState": 0}
 
 {"tactic": "exact t", "proofState": 2}
diff --git a/test/proof_step2.expected.out b/test/proof_step2.expected.out
deleted file mode 100644
index cb909e2..0000000
--- a/test/proof_step2.expected.out
+++ /dev/null
@@ -1,14 +0,0 @@
-{"sorries":
- [{"proofState": 0,
-   "pos": {"line": 1, "column": 18},
-   "goal": "⊢ Nat",
-   "endPos": {"line": 1, "column": 23}}],
- "messages":
- [{"severity": "warning",
-   "pos": {"line": 1, "column": 4},
-   "endPos": {"line": 1, "column": 5},
-   "data": "declaration uses 'sorry'"}],
- "env": 0}
-{"proofState": 1, "goals": ["⊢ Int"]}
-{"proofState": 2, "goals": ["t : Nat\n⊢ Nat"]}
-{"proofState": 3, "goals": []}
diff --git a/test/proof_step2.in b/test/proof_step2.in
deleted file mode 100644
index a0b5d53..0000000
--- a/test/proof_step2.in
+++ /dev/null
@@ -1,7 +0,0 @@
-{"cmd" : "def f : Nat := by sorry"}
-
-{"tactic": "apply Int.natAbs", "proofState": 0}
-
-{"tactic": "have t : Nat := 42", "proofState": 0}
-
-{"tactic": "exact t", "proofState": 2}
diff --git a/test/proof_step2.lean b/test/proof_step2.lean
deleted file mode 100644
index ea06ed8..0000000
--- a/test/proof_step2.lean
+++ /dev/null
@@ -1,3 +0,0 @@
-def f : Nat := by
-  have t : Nat := 42
-  exact t
diff --git a/test/unknown_environment.expected.out b/test/unknown_environment.expected.out
new file mode 100644
index 0000000..45a6e1c
--- /dev/null
+++ b/test/unknown_environment.expected.out
@@ -0,0 +1 @@
+{"message": "Unknown environment."}
diff --git a/test/unknown_environment.in b/test/unknown_environment.in
new file mode 100644
index 0000000..9cd67db
--- /dev/null
+++ b/test/unknown_environment.in
@@ -0,0 +1 @@
+{"cmd": "def f : Nat := by sorry", "env": 5}