From aebf027711e550ee5ac7ec97f5f9f9d80b080f24 Mon Sep 17 00:00:00 2001
From: Qiu233 <lopi258@outlook.com>
Date: Tue, 15 Jul 2025 10:52:03 +0800
Subject: [PATCH 1/4] chore: collect actions

---
 REPL/Actions.lean         |   4 +
 REPL/Actions/Basic.lean   | 109 ++++++++++++
 REPL/Actions/Command.lean | 129 ++++++++++++++
 REPL/Actions/File.lean    |  20 +++
 REPL/Actions/Tactic.lean  | 238 ++++++++++++++++++++++++++
 REPL/Basic.lean           |  30 ++++
 REPL/JSON.lean            | 192 ---------------------
 REPL/Main.lean            | 346 +++-----------------------------------
 8 files changed, 550 insertions(+), 518 deletions(-)
 create mode 100644 REPL/Actions.lean
 create mode 100644 REPL/Actions/Basic.lean
 create mode 100644 REPL/Actions/Command.lean
 create mode 100644 REPL/Actions/File.lean
 create mode 100644 REPL/Actions/Tactic.lean
 create mode 100644 REPL/Basic.lean
 delete mode 100644 REPL/JSON.lean

diff --git a/REPL/Actions.lean b/REPL/Actions.lean
new file mode 100644
index 00000000..2ee80f3e
--- /dev/null
+++ b/REPL/Actions.lean
@@ -0,0 +1,4 @@
+import REPL.Actions.Basic
+import REPL.Actions.Tactic
+import REPL.Actions.Command
+import REPL.Actions.File
diff --git a/REPL/Actions/Basic.lean b/REPL/Actions/Basic.lean
new file mode 100644
index 00000000..231e41b7
--- /dev/null
+++ b/REPL/Actions/Basic.lean
@@ -0,0 +1,109 @@
+import REPL.Basic
+import REPL.Frontend
+import REPL.Util.Path
+import REPL.Lean.ContextInfo
+import REPL.Lean.Environment
+import REPL.Lean.InfoTree
+import REPL.Lean.InfoTree.ToJson
+import REPL.Snapshots
+
+open Lean Elab InfoTree
+
+namespace REPL.Actions
+
+/-- Line and column information for error messages and sorries. -/
+structure Pos where
+  line : Nat
+  column : Nat
+deriving ToJson, FromJson
+
+/-- Severity of a message. -/
+inductive Severity
+  | trace | info | warning | error
+deriving ToJson, FromJson
+
+/-- A Lean message. -/
+structure Message where
+  pos : Pos
+  endPos : Option Pos
+  severity : Severity
+  data : String
+deriving ToJson, FromJson
+
+/-- Construct the JSON representation of a Lean message. -/
+def Message.of (m : Lean.Message) : IO Message := do pure <|
+  { pos := ⟨m.pos.line, m.pos.column⟩,
+    endPos := m.endPos.map fun p => ⟨p.line, p.column⟩,
+    severity := match m.severity with
+    | .information => .info
+    | .warning => .warning
+    | .error => .error,
+    data := (← m.data.toString).trim }
+
+/-- A Lean `sorry`. -/
+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 FromJson
+
+instance : ToJson Sorry where
+  toJson r := Json.mkObj <| .flatten [
+    [("goal", r.goal)],
+    [("proofState", toJson r.proofState)],
+    if r.pos.line ≠ 0 then [("pos", toJson r.pos)] else [],
+    if r.endPos.line ≠ 0 then [("endPos", toJson r.endPos)] else [],
+  ]
+
+/-- Construct the JSON representation of a Lean sorry. -/
+def Sorry.of (goal : String) (pos endPos : Lean.Position) (proofState : Option Nat) : Sorry :=
+  { pos := ⟨pos.line, pos.column⟩,
+    endPos := ⟨endPos.line, endPos.column⟩,
+    goal,
+    proofState }
+
+def Json.nonemptyList [ToJson α] (k : String) : List α → List (String × Json)
+  | [] => []
+  | l  => [⟨k, toJson l⟩]
+
+/-- Json wrapper for an error. -/
+structure Error where
+  message : String
+deriving ToJson, FromJson
+
+variable [Monad m] [MonadLiftT IO m]
+
+def collectFVarsAux : Expr → NameSet
+  | .fvar fvarId => NameSet.empty.insert fvarId.name
+  | .app fm arg => (collectFVarsAux fm).union $ collectFVarsAux arg
+  | .lam _ binderType body _ => (collectFVarsAux binderType).union $ collectFVarsAux body
+  | .forallE _ binderType body _ => (collectFVarsAux binderType).union $ collectFVarsAux body
+  | .letE _ type value body _ => ((collectFVarsAux type).union $ collectFVarsAux value).union $ collectFVarsAux body
+  | .mdata _ expr => collectFVarsAux expr
+  | .proj _ _ struct => collectFVarsAux struct
+  | _ => NameSet.empty
+
+/-- Collect all fvars in the expression, and return their names. -/
+def collectFVars (e : Expr) : MetaM (Array Expr) := do
+  let names := collectFVarsAux e
+  let mut fvars := #[]
+  for ldecl in ← getLCtx do
+    if ldecl.isImplementationDetail then
+      continue
+    if names.contains ldecl.fvarId.name then
+      fvars := fvars.push $ .fvar ldecl.fvarId
+  return fvars
+
+def abstractAllLambdaFVars (e : Expr) : MetaM Expr := do
+  let mut e' := e
+  while e'.hasFVar do
+    let fvars ← collectFVars e'
+    if fvars.isEmpty then
+      break
+    e' ← Meta.mkLambdaFVars fvars e'
+  return e'
diff --git a/REPL/Actions/Command.lean b/REPL/Actions/Command.lean
new file mode 100644
index 00000000..1fb4f9cb
--- /dev/null
+++ b/REPL/Actions/Command.lean
@@ -0,0 +1,129 @@
+import REPL.Basic
+import REPL.Actions.Basic
+import REPL.Actions.Tactic
+import REPL.Frontend
+
+open Lean Elab InfoTree
+
+namespace REPL.Actions
+
+structure CommandOptions where
+  allTactics : Option Bool := none
+  rootGoals : Option Bool := none
+  /--
+  Should be "full", "tactics", "original", or "substantive".
+  Anything else is ignored.
+  -/
+  infotree : Option String
+
+/-- Run Lean commands.
+If `env = none`, starts a new session (in which you can use `import`).
+If `env = some n`, builds on the existing environment `n`.
+-/
+structure Command extends CommandOptions where
+  env : Option Nat
+  cmd : String
+deriving ToJson, FromJson
+
+/--
+A response to a Lean command.
+`env` can be used in later calls, to build on the stored environment.
+-/
+structure CommandResponse where
+  env : Nat
+  messages : List Message := []
+  sorries : List Sorry := []
+  tactics : List Tactic := []
+  infotree : Option Json := none
+deriving FromJson
+
+instance : ToJson CommandResponse where
+  toJson r := Json.mkObj <| .flatten [
+    [("env", r.env)],
+    Json.nonemptyList "messages" r.messages,
+    Json.nonemptyList "sorries" r.sorries,
+    Json.nonemptyList "tactics" r.tactics,
+    match r.infotree with | some j => [("infotree", j)] | none => []
+  ]
+
+structure PickleEnvironment where
+  env : Nat
+  pickleTo : System.FilePath
+deriving ToJson, FromJson
+
+structure UnpickleEnvironment where
+  unpickleEnvFrom : System.FilePath
+deriving ToJson, FromJson
+
+variable [Monad m] [MonadLiftT IO m]
+
+/-- Record an `CommandSnapshot` into the REPL state, returning its index for future use. -/
+def recordCommandSnapshot (state : CommandSnapshot) : M m Nat := do
+  let id := (← get).cmdStates.size
+  modify fun s => { s with cmdStates := s.cmdStates.push state }
+  return id
+
+/--
+Run a command, returning the id of the new environment, and any messages and sorries.
+-/
+def runCommand (s : Command) : M IO (CommandResponse ⊕ Error) := do
+  let (cmdSnapshot?, notFound) ← do match s.env with
+  | none => pure (none, false)
+  | some i => do match (← get).cmdStates[i]? with
+    | some env => pure (some env, false)
+    | none => pure (none, true)
+  if notFound then
+    return .inr ⟨"Unknown environment."⟩
+  let initialCmdState? := cmdSnapshot?.map fun c => c.cmdState
+  let (initialCmdState, cmdState, messages, trees) ← try
+    IO.processInput s.cmd initialCmdState?
+  catch ex =>
+    return .inr ⟨ex.toString⟩
+  let messages ← messages.mapM fun m => Message.of m
+  -- For debugging purposes, sometimes we print out the trees here:
+  -- trees.forM fun t => do IO.println (← t.format)
+  let sorries ← sorries trees initialCmdState.env none
+  let sorries ← match s.rootGoals with
+  | some true => pure (sorries ++ (← collectRootGoalsAsSorries trees initialCmdState.env))
+  | _ => pure sorries
+  let tactics ← match s.allTactics with
+  | some true => tactics trees initialCmdState.env
+  | _ => pure []
+  let cmdSnapshot :=
+  { cmdState
+    cmdContext := (cmdSnapshot?.map fun c => c.cmdContext).getD
+      { fileName := "",
+        fileMap := default,
+        snap? := none,
+        cancelTk? := none } }
+  let env ← recordCommandSnapshot cmdSnapshot
+  let jsonTrees := match s.infotree with
+  | some "full" => trees
+  | some "tactics" => trees.flatMap InfoTree.retainTacticInfo
+  | some "original" => trees.flatMap InfoTree.retainTacticInfo |>.flatMap InfoTree.retainOriginal
+  | some "substantive" => trees.flatMap InfoTree.retainTacticInfo |>.flatMap InfoTree.retainSubstantive
+  | _ => []
+  let infotree ← if jsonTrees.isEmpty then
+    pure none
+  else
+    pure <| some <| Json.arr (← jsonTrees.toArray.mapM fun t => t.toJson none)
+  return .inl
+    { env,
+      messages,
+      sorries,
+      tactics
+      infotree }
+
+/-- Pickle a `CommandSnapshot`, generating a JSON response. -/
+def pickleCommandSnapshot (n : PickleEnvironment) : M m (CommandResponse ⊕ Error) := do
+  match (← get).cmdStates[n.env]? with
+  | none => return .inr ⟨"Unknown environment."⟩
+  | some env =>
+    discard <| env.pickle n.pickleTo
+    return .inl { env := n.env }
+
+/-- Unpickle a `CommandSnapshot`, generating a JSON response. -/
+def unpickleCommandSnapshot (n : UnpickleEnvironment) : M IO CommandResponse := do
+  let (env, _) ← CommandSnapshot.unpickle n.unpickleEnvFrom
+  let env ← recordCommandSnapshot env
+  return { env }
diff --git a/REPL/Actions/File.lean b/REPL/Actions/File.lean
new file mode 100644
index 00000000..572f01ed
--- /dev/null
+++ b/REPL/Actions/File.lean
@@ -0,0 +1,20 @@
+import REPL.Basic
+import REPL.Actions.Basic
+import REPL.Actions.Command
+
+open Lean Elab InfoTree
+
+namespace REPL.Actions
+
+/-- Process a Lean file in a fresh environment if `env` is not provided. -/
+structure File extends CommandOptions where
+  env : Option Nat
+  path : System.FilePath
+deriving ToJson, FromJson
+
+def processFile (s : File) : M IO (CommandResponse ⊕ Error) := do
+  try
+    let cmd ← IO.FS.readFile s.path
+    runCommand { s with env := s.env, cmd }
+  catch e =>
+    pure <| .inr ⟨e.toString⟩
diff --git a/REPL/Actions/Tactic.lean b/REPL/Actions/Tactic.lean
new file mode 100644
index 00000000..dc54b32e
--- /dev/null
+++ b/REPL/Actions/Tactic.lean
@@ -0,0 +1,238 @@
+import REPL.Basic
+import REPL.Actions.Basic
+
+open Lean Elab InfoTree
+
+namespace REPL.Actions
+
+structure Tactic where
+  pos : Pos
+  endPos : Pos
+  goals : String
+  tactic : String
+  proofState : Option Nat
+  usedConstants : Array Name
+deriving ToJson, FromJson
+
+/-- Construct the JSON representation of a Lean tactic. -/
+def Tactic.of (goals tactic : String) (pos endPos : Lean.Position) (proofState : Option Nat) (usedConstants : Array Name) : Tactic :=
+  { pos := ⟨pos.line, pos.column⟩,
+    endPos := ⟨endPos.line, endPos.column⟩,
+    goals,
+    tactic,
+    proofState,
+    usedConstants }
+
+/--
+Run a tactic in a proof state.
+-/
+structure ProofStep where
+  proofState : Nat
+  tactic : String
+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
+  messages : List Message := []
+  sorries : List Sorry := []
+  traces : List String
+  proofStatus : String
+deriving ToJson, FromJson
+
+instance : ToJson ProofStepResponse where
+  toJson r := Json.mkObj <| .flatten [
+    [("proofState", r.proofState)],
+    [("goals", toJson r.goals)],
+    Json.nonemptyList "messages" r.messages,
+    Json.nonemptyList "sorries" r.sorries,
+    Json.nonemptyList "traces" r.traces,
+    [("proofStatus", r.proofStatus)]
+  ]
+
+structure PickleProofState where
+  proofState : Nat
+  pickleTo : System.FilePath
+deriving ToJson, FromJson
+
+structure UnpickleProofState where
+  unpickleProofStateFrom : System.FilePath
+  env : Option Nat
+deriving ToJson, FromJson
+
+variable [Monad m] [MonadLiftT IO m]
+
+/-- Record a `ProofSnapshot` into the REPL state, returning its index for future use. -/
+def recordProofSnapshot (proofState : ProofSnapshot) : M m Nat := do
+  let id := (← get).proofStates.size
+  modify fun s => { s with proofStates := s.proofStates.push proofState }
+  return id
+
+def sorries (trees : List InfoTree) (env? : Option Environment) (rootGoals? : Option (List MVarId))
+: M m (List Sorry) :=
+  trees.flatMap InfoTree.sorries |>.filter (fun t => match t.2.1 with
+    | .term _ none => false
+    | _ => true ) |>.mapM
+      fun ⟨ctx, g, pos, endPos⟩ => do
+        let (goal, proofState) ← match g with
+        | .tactic g => do
+          let lctx ← ctx.runMetaM {} do
+              match ctx.mctx.findDecl? g with
+              | some decl => return decl.lctx
+              | none => throwError "unknown metavariable '{g}'"
+          let s ← ProofSnapshot.create ctx lctx env? [g] rootGoals?
+          pure ("\n".intercalate <| (← s.ppGoals).map fun s => s!"{s}", some s)
+        | .term lctx (some t) => do
+          let s ← ProofSnapshot.create ctx lctx env? [] rootGoals? [t]
+          pure ("\n".intercalate <| (← s.ppGoals).map fun s => s!"{s}", some s)
+        | .term _ none => unreachable!
+        let proofStateId ← proofState.mapM recordProofSnapshot
+        return Sorry.of goal pos endPos proofStateId
+
+def ppTactic (ctx : ContextInfo) (stx : Syntax) : IO Format :=
+  ctx.runMetaM {} try
+    Lean.PrettyPrinter.ppTactic ⟨stx⟩
+  catch _ =>
+    pure "<failed to pretty print>"
+
+def tactics (trees : List InfoTree) (env? : Option Environment) : M m (List Tactic) :=
+  trees.flatMap InfoTree.tactics |>.mapM
+    fun ⟨ctx, stx, rootGoals, goals, pos, endPos, ns⟩ => do
+      let proofState := some (← ProofSnapshot.create ctx none env? goals rootGoals)
+      let goals := s!"{(← ctx.ppGoals goals)}".trim
+      let tactic := Format.pretty (← ppTactic ctx stx)
+      let proofStateId ← proofState.mapM recordProofSnapshot
+      return Tactic.of goals tactic pos endPos proofStateId ns
+
+def collectRootGoalsAsSorries (trees : List InfoTree) (env? : Option Environment) : M m (List Sorry) := do
+  trees.flatMap InfoTree.rootGoals |>.mapM
+    fun ⟨ctx, goals, pos⟩ => do
+      let proofState := some (← ProofSnapshot.create ctx none env? goals goals)
+      let goals := s!"{(← ctx.ppGoals goals)}".trim
+      let proofStateId ← proofState.mapM recordProofSnapshot
+      return Sorry.of goals pos pos proofStateId
+
+/--
+Evaluates the current status of a proof, returning a string description.
+Main states include:
+- "Completed": Proof is complete and type checks successfully
+- "Incomplete": When goals remain, or proof contains sorry/metavariables
+- "Error": When kernel type checking errors occur
+
+Inspired by LeanDojo REPL's status tracking.
+-/
+def getProofStatus (proofState : ProofSnapshot) : M m String := do
+  match proofState.tacticState.goals with
+    | [] =>
+      let res := proofState.runMetaM do
+        match proofState.rootGoals with
+        | [goalId] =>
+          goalId.withContext do
+          match proofState.metaState.mctx.getExprAssignmentCore? goalId with
+          | none => return "Error: Goal not assigned"
+          | some pf => do
+            let pf ← instantiateMVars pf
+
+            -- First check that the proof has the expected type
+            let pft ← Meta.inferType pf >>= instantiateMVars
+            let expectedType ← Meta.inferType (mkMVar goalId) >>= instantiateMVars
+            unless (← Meta.isDefEq pft expectedType) do
+              return s!"Error: proof has type {pft} but root goal has type {expectedType}"
+
+            let pf ← abstractAllLambdaFVars pf
+            let pft ← Meta.inferType pf >>= instantiateMVars
+
+            if pf.hasExprMVar then
+              return "Incomplete: contains metavariable(s)"
+
+            -- Find all level parameters
+            let usedLevels := collectLevelParams {} pft
+            let usedLevels := collectLevelParams usedLevels pf
+
+            let decl := Declaration.defnDecl {
+              name := Name.anonymous,
+              type := pft,
+              value := pf,
+              levelParams := usedLevels.params.toList,
+              hints := ReducibilityHints.opaque,
+              safety := DefinitionSafety.safe
+            }
+
+            try
+              let _ ← addDecl decl
+            catch ex =>
+              return s!"Error: kernel type check failed: {← ex.toMessageData.toString}"
+
+            if pf.hasSorry then
+              return "Incomplete: contains sorry"
+            return "Completed"
+
+        | _ => return "Not verified: more than one initial goal"
+      return (← res).fst
+
+    | _ => return "Incomplete: open goals remain"
+
+
+/-- Record a `ProofSnapshot` and generate a JSON response for it. -/
+def createProofStepReponse (proofState : ProofSnapshot) (old? : Option ProofSnapshot := none) :
+    M m ProofStepResponse := do
+  let messages := proofState.newMessages old?
+  let messages ← messages.mapM fun m => Message.of m
+  let traces ← proofState.newTraces old?
+  let trees := proofState.newInfoTrees old?
+  let trees ← match old? with
+  | some old => do
+    let (ctx, _) ← old.runMetaM do return { ← CommandContextInfo.save with }
+    let ctx := PartialContextInfo.commandCtx ctx
+    pure <| trees.map fun t => InfoTree.context ctx t
+  | none => pure trees
+  -- For debugging purposes, sometimes we print out the trees here:
+  -- trees.forM fun t => do IO.println (← t.format)
+  let sorries ← sorries trees none (some proofState.rootGoals)
+  let id ← recordProofSnapshot proofState
+  return {
+    proofState := id
+    goals := (← proofState.ppGoals).map fun s => s!"{s}"
+    messages
+    sorries
+    traces
+    proofStatus := (← getProofStatus proofState) }
+
+/-- Pickle a `ProofSnapshot`, generating a JSON response. -/
+-- This generates a new identifier, which perhaps is not what we want?
+def pickleProofSnapshot (n : PickleProofState) : M m (ProofStepResponse ⊕ Error) := do
+  match (← get).proofStates[n.proofState]? with
+  | none => return .inr ⟨"Unknown proof State."⟩
+  | some proofState =>
+    discard <| proofState.pickle n.pickleTo
+    return .inl (← createProofStepReponse proofState)
+
+/-- Unpickle a `ProofSnapshot`, generating a JSON response. -/
+def unpickleProofSnapshot (n : UnpickleProofState) : M IO (ProofStepResponse ⊕ Error) := do
+  let (cmdSnapshot?, notFound) ← do match n.env with
+  | none => pure (none, false)
+  | some i => do match (← get).cmdStates[i]? with
+    | some env => pure (some env, false)
+    | none => pure (none, true)
+  if notFound then
+    return .inr ⟨"Unknown environment."⟩
+  let (proofState, _) ← ProofSnapshot.unpickle n.unpickleProofStateFrom cmdSnapshot?
+  Sum.inl <$> createProofStepReponse proofState
+
+/--
+Run a single tactic, returning the id of the new proof statement, and the new goals.
+-/
+-- TODO detect sorries?
+def runProofStep (s : ProofStep) : M IO (ProofStepResponse ⊕ Error) := do
+  match (← get).proofStates[s.proofState]? with
+  | none => return .inr ⟨"Unknown proof state."⟩
+  | some proofState =>
+    try
+      let proofState' ← proofState.runString s.tactic
+      return .inl (← createProofStepReponse proofState' proofState)
+    catch ex =>
+      return .inr ⟨"Lean error:\n" ++ ex.toString⟩
diff --git a/REPL/Basic.lean b/REPL/Basic.lean
new file mode 100644
index 00000000..328853bb
--- /dev/null
+++ b/REPL/Basic.lean
@@ -0,0 +1,30 @@
+import Lean
+import REPL.Snapshots
+
+open Lean Elab
+
+namespace REPL
+
+/-- 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}`.
+  -/
+  cmdStates : Array CommandSnapshot := #[]
+  /--
+  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 ProofSnapshot := #[]
+
+/--
+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]
diff --git a/REPL/JSON.lean b/REPL/JSON.lean
deleted file mode 100644
index ebec4c56..00000000
--- a/REPL/JSON.lean
+++ /dev/null
@@ -1,192 +0,0 @@
-/-
-Copyright (c) 2023 Scott Morrison. All rights reserved.
-Released under Apache 2.0 license as described in the file LICENSE.
-Authors: Scott Morrison
--/
-import Lean.Data.Json
-import Lean.Message
-import Lean.Elab.InfoTree.Main
-
-open Lean Elab InfoTree
-
-namespace REPL
-
-structure CommandOptions where
-  allTactics : Option Bool := none
-  rootGoals : Option Bool := none
-  /--
-  Should be "full", "tactics", "original", or "substantive".
-  Anything else is ignored.
-  -/
-  infotree : Option String
-
-/-- Run Lean commands.
-If `env = none`, starts a new session (in which you can use `import`).
-If `env = some n`, builds on the existing environment `n`.
--/
-structure Command extends CommandOptions where
-  env : Option Nat
-  cmd : String
-deriving ToJson, FromJson
-
-/-- Process a Lean file in a fresh environment if `env` is not provided. -/
-structure File extends CommandOptions where
-  env : Option Nat
-  path : System.FilePath
-deriving ToJson, FromJson
-
-/--
-Run a tactic in a proof state.
--/
-structure ProofStep where
-  proofState : Nat
-  tactic : String
-deriving ToJson, FromJson
-
-/-- Line and column information for error messages and sorries. -/
-structure Pos where
-  line : Nat
-  column : Nat
-deriving ToJson, FromJson
-
-/-- Severity of a message. -/
-inductive Severity
-  | trace | info | warning | error
-deriving ToJson, FromJson
-
-/-- A Lean message. -/
-structure Message where
-  pos : Pos
-  endPos : Option Pos
-  severity : Severity
-  data : String
-deriving ToJson, FromJson
-
-/-- Construct the JSON representation of a Lean message. -/
-def Message.of (m : Lean.Message) : IO Message := do pure <|
-  { pos := ⟨m.pos.line, m.pos.column⟩,
-    endPos := m.endPos.map fun p => ⟨p.line, p.column⟩,
-    severity := match m.severity with
-    | .information => .info
-    | .warning => .warning
-    | .error => .error,
-    data := (← m.data.toString).trim }
-
-/-- A Lean `sorry`. -/
-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 FromJson
-
-instance : ToJson Sorry where
-  toJson r := Json.mkObj <| .flatten [
-    [("goal", r.goal)],
-    [("proofState", toJson r.proofState)],
-    if r.pos.line ≠ 0 then [("pos", toJson r.pos)] else [],
-    if r.endPos.line ≠ 0 then [("endPos", toJson r.endPos)] else [],
-  ]
-
-/-- Construct the JSON representation of a Lean sorry. -/
-def Sorry.of (goal : String) (pos endPos : Lean.Position) (proofState : Option Nat) : Sorry :=
-  { pos := ⟨pos.line, pos.column⟩,
-    endPos := ⟨endPos.line, endPos.column⟩,
-    goal,
-    proofState }
-
-structure Tactic where
-  pos : Pos
-  endPos : Pos
-  goals : String
-  tactic : String
-  proofState : Option Nat
-  usedConstants : Array Name
-deriving ToJson, FromJson
-
-/-- Construct the JSON representation of a Lean tactic. -/
-def Tactic.of (goals tactic : String) (pos endPos : Lean.Position) (proofState : Option Nat) (usedConstants : Array Name) : Tactic :=
-  { pos := ⟨pos.line, pos.column⟩,
-    endPos := ⟨endPos.line, endPos.column⟩,
-    goals,
-    tactic,
-    proofState,
-    usedConstants }
-
-/--
-A response to a Lean command.
-`env` can be used in later calls, to build on the stored environment.
--/
-structure CommandResponse where
-  env : Nat
-  messages : List Message := []
-  sorries : List Sorry := []
-  tactics : List Tactic := []
-  infotree : Option Json := none
-deriving FromJson
-
-def Json.nonemptyList [ToJson α] (k : String) : List α → List (String × Json)
-  | [] => []
-  | l  => [⟨k, toJson l⟩]
-
-instance : ToJson CommandResponse where
-  toJson r := Json.mkObj <| .flatten [
-    [("env", r.env)],
-    Json.nonemptyList "messages" r.messages,
-    Json.nonemptyList "sorries" r.sorries,
-    Json.nonemptyList "tactics" r.tactics,
-    match r.infotree with | some j => [("infotree", j)] | none => []
-  ]
-
-/--
-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
-  messages : List Message := []
-  sorries : List Sorry := []
-  traces : List String
-  proofStatus : String
-deriving ToJson, FromJson
-
-instance : ToJson ProofStepResponse where
-  toJson r := Json.mkObj <| .flatten [
-    [("proofState", r.proofState)],
-    [("goals", toJson r.goals)],
-    Json.nonemptyList "messages" r.messages,
-    Json.nonemptyList "sorries" r.sorries,
-    Json.nonemptyList "traces" r.traces,
-    [("proofStatus", r.proofStatus)]
-  ]
-
-/-- Json wrapper for an error. -/
-structure Error where
-  message : String
-deriving ToJson, FromJson
-
-structure PickleEnvironment where
-  env : Nat
-  pickleTo : System.FilePath
-deriving ToJson, FromJson
-
-structure UnpickleEnvironment where
-  unpickleEnvFrom : System.FilePath
-deriving ToJson, FromJson
-
-structure PickleProofState where
-  proofState : Nat
-  pickleTo : System.FilePath
-deriving ToJson, FromJson
-
-structure UnpickleProofState where
-  unpickleProofStateFrom : System.FilePath
-  env : Option Nat
-deriving ToJson, FromJson
-
-end REPL
diff --git a/REPL/Main.lean b/REPL/Main.lean
index d181080f..db4f87b9 100644
--- a/REPL/Main.lean
+++ b/REPL/Main.lean
@@ -3,7 +3,8 @@ Copyright (c) 2023 Scott Morrison. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Scott Morrison
 -/
-import REPL.JSON
+import REPL.Basic
+-- import REPL.JSON
 import REPL.Frontend
 import REPL.Util.Path
 import REPL.Lean.ContextInfo
@@ -11,6 +12,7 @@ import REPL.Lean.Environment
 import REPL.Lean.InfoTree
 import REPL.Lean.InfoTree.ToJson
 import REPL.Snapshots
+import REPL.Actions
 
 /-!
 # A REPL for Lean.
@@ -56,315 +58,7 @@ The results are of the form
 
 open Lean Elab
 
-namespace REPL
-
-/-- 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}`.
-  -/
-  cmdStates : Array CommandSnapshot := #[]
-  /--
-  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 ProofSnapshot := #[]
-
-/--
-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 `CommandSnapshot` into the REPL state, returning its index for future use. -/
-def recordCommandSnapshot (state : CommandSnapshot) : M m Nat := do
-  let id := (← get).cmdStates.size
-  modify fun s => { s with cmdStates := s.cmdStates.push state }
-  return id
-
-/-- Record a `ProofSnapshot` into the REPL state, returning its index for future use. -/
-def recordProofSnapshot (proofState : ProofSnapshot) : M m Nat := do
-  let id := (← get).proofStates.size
-  modify fun s => { s with proofStates := s.proofStates.push proofState }
-  return id
-
-def sorries (trees : List InfoTree) (env? : Option Environment) (rootGoals? : Option (List MVarId))
-: M m (List Sorry) :=
-  trees.flatMap InfoTree.sorries |>.filter (fun t => match t.2.1 with
-    | .term _ none => false
-    | _ => true ) |>.mapM
-      fun ⟨ctx, g, pos, endPos⟩ => do
-        let (goal, proofState) ← match g with
-        | .tactic g => do
-          let lctx ← ctx.runMetaM {} do
-              match ctx.mctx.findDecl? g with
-              | some decl => return decl.lctx
-              | none => throwError "unknown metavariable '{g}'"
-          let s ← ProofSnapshot.create ctx lctx env? [g] rootGoals?
-          pure ("\n".intercalate <| (← s.ppGoals).map fun s => s!"{s}", some s)
-        | .term lctx (some t) => do
-          let s ← ProofSnapshot.create ctx lctx env? [] rootGoals? [t]
-          pure ("\n".intercalate <| (← s.ppGoals).map fun s => s!"{s}", some s)
-        | .term _ none => unreachable!
-        let proofStateId ← proofState.mapM recordProofSnapshot
-        return Sorry.of goal pos endPos proofStateId
-
-def ppTactic (ctx : ContextInfo) (stx : Syntax) : IO Format :=
-  ctx.runMetaM {} try
-    Lean.PrettyPrinter.ppTactic ⟨stx⟩
-  catch _ =>
-    pure "<failed to pretty print>"
-
-def tactics (trees : List InfoTree) (env? : Option Environment) : M m (List Tactic) :=
-  trees.flatMap InfoTree.tactics |>.mapM
-    fun ⟨ctx, stx, rootGoals, goals, pos, endPos, ns⟩ => do
-      let proofState := some (← ProofSnapshot.create ctx none env? goals rootGoals)
-      let goals := s!"{(← ctx.ppGoals goals)}".trim
-      let tactic := Format.pretty (← ppTactic ctx stx)
-      let proofStateId ← proofState.mapM recordProofSnapshot
-      return Tactic.of goals tactic pos endPos proofStateId ns
-
-def collectRootGoalsAsSorries (trees : List InfoTree) (env? : Option Environment) : M m (List Sorry) := do
-  trees.flatMap InfoTree.rootGoals |>.mapM
-    fun ⟨ctx, goals, pos⟩ => do
-      let proofState := some (← ProofSnapshot.create ctx none env? goals goals)
-      let goals := s!"{(← ctx.ppGoals goals)}".trim
-      let proofStateId ← proofState.mapM recordProofSnapshot
-      return Sorry.of goals pos pos proofStateId
-
-
-private def collectFVarsAux : Expr → NameSet
-  | .fvar fvarId => NameSet.empty.insert fvarId.name
-  | .app fm arg => (collectFVarsAux fm).union $ collectFVarsAux arg
-  | .lam _ binderType body _ => (collectFVarsAux binderType).union $ collectFVarsAux body
-  | .forallE _ binderType body _ => (collectFVarsAux binderType).union $ collectFVarsAux body
-  | .letE _ type value body _ => ((collectFVarsAux type).union $ collectFVarsAux value).union $ collectFVarsAux body
-  | .mdata _ expr => collectFVarsAux expr
-  | .proj _ _ struct => collectFVarsAux struct
-  | _ => NameSet.empty
-
-/-- Collect all fvars in the expression, and return their names. -/
-private def collectFVars (e : Expr) : MetaM (Array Expr) := do
-  let names := collectFVarsAux e
-  let mut fvars := #[]
-  for ldecl in ← getLCtx do
-    if ldecl.isImplementationDetail then
-      continue
-    if names.contains ldecl.fvarId.name then
-      fvars := fvars.push $ .fvar ldecl.fvarId
-  return fvars
-
-
-private def abstractAllLambdaFVars (e : Expr) : MetaM Expr := do
-  let mut e' := e
-  while e'.hasFVar do
-    let fvars ← collectFVars e'
-    if fvars.isEmpty then
-      break
-    e' ← Meta.mkLambdaFVars fvars e'
-  return e'
-
-/--
-Evaluates the current status of a proof, returning a string description.
-Main states include:
-- "Completed": Proof is complete and type checks successfully
-- "Incomplete": When goals remain, or proof contains sorry/metavariables
-- "Error": When kernel type checking errors occur
-
-Inspired by LeanDojo REPL's status tracking.
--/
-def getProofStatus (proofState : ProofSnapshot) : M m String := do
-  match proofState.tacticState.goals with
-    | [] =>
-      let res := proofState.runMetaM do
-        match proofState.rootGoals with
-        | [goalId] =>
-          goalId.withContext do
-          match proofState.metaState.mctx.getExprAssignmentCore? goalId with
-          | none => return "Error: Goal not assigned"
-          | some pf => do
-            let pf ← instantiateMVars pf
-
-            -- First check that the proof has the expected type
-            let pft ← Meta.inferType pf >>= instantiateMVars
-            let expectedType ← Meta.inferType (mkMVar goalId) >>= instantiateMVars
-            unless (← Meta.isDefEq pft expectedType) do
-              return s!"Error: proof has type {pft} but root goal has type {expectedType}"
-
-            let pf ← abstractAllLambdaFVars pf
-            let pft ← Meta.inferType pf >>= instantiateMVars
-
-            if pf.hasExprMVar then
-              return "Incomplete: contains metavariable(s)"
-
-            -- Find all level parameters
-            let usedLevels := collectLevelParams {} pft
-            let usedLevels := collectLevelParams usedLevels pf
-
-            let decl := Declaration.defnDecl {
-              name := Name.anonymous,
-              type := pft,
-              value := pf,
-              levelParams := usedLevels.params.toList,
-              hints := ReducibilityHints.opaque,
-              safety := DefinitionSafety.safe
-            }
-
-            try
-              let _ ← addDecl decl
-            catch ex =>
-              return s!"Error: kernel type check failed: {← ex.toMessageData.toString}"
-
-            if pf.hasSorry then
-              return "Incomplete: contains sorry"
-            return "Completed"
-
-        | _ => return "Not verified: more than one initial goal"
-      return (← res).fst
-
-    | _ => return "Incomplete: open goals remain"
-
-/-- Record a `ProofSnapshot` and generate a JSON response for it. -/
-def createProofStepReponse (proofState : ProofSnapshot) (old? : Option ProofSnapshot := none) :
-    M m ProofStepResponse := do
-  let messages := proofState.newMessages old?
-  let messages ← messages.mapM fun m => Message.of m
-  let traces ← proofState.newTraces old?
-  let trees := proofState.newInfoTrees old?
-  let trees ← match old? with
-  | some old => do
-    let (ctx, _) ← old.runMetaM do return { ← CommandContextInfo.save with }
-    let ctx := PartialContextInfo.commandCtx ctx
-    pure <| trees.map fun t => InfoTree.context ctx t
-  | none => pure trees
-  -- For debugging purposes, sometimes we print out the trees here:
-  -- trees.forM fun t => do IO.println (← t.format)
-  let sorries ← sorries trees none (some proofState.rootGoals)
-  let id ← recordProofSnapshot proofState
-  return {
-    proofState := id
-    goals := (← proofState.ppGoals).map fun s => s!"{s}"
-    messages
-    sorries
-    traces
-    proofStatus := (← getProofStatus proofState) }
-
-/-- Pickle a `CommandSnapshot`, generating a JSON response. -/
-def pickleCommandSnapshot (n : PickleEnvironment) : M m (CommandResponse ⊕ Error) := do
-  match (← get).cmdStates[n.env]? with
-  | none => return .inr ⟨"Unknown environment."⟩
-  | some env =>
-    discard <| env.pickle n.pickleTo
-    return .inl { env := n.env }
-
-/-- Unpickle a `CommandSnapshot`, generating a JSON response. -/
-def unpickleCommandSnapshot (n : UnpickleEnvironment) : M IO CommandResponse := do
-  let (env, _) ← CommandSnapshot.unpickle n.unpickleEnvFrom
-  let env ← recordCommandSnapshot env
-  return { env }
-
-/-- Pickle a `ProofSnapshot`, generating a JSON response. -/
--- This generates a new identifier, which perhaps is not what we want?
-def pickleProofSnapshot (n : PickleProofState) : M m (ProofStepResponse ⊕ Error) := do
-  match (← get).proofStates[n.proofState]? with
-  | none => return .inr ⟨"Unknown proof State."⟩
-  | some proofState =>
-    discard <| proofState.pickle n.pickleTo
-    return .inl (← createProofStepReponse proofState)
-
-/-- Unpickle a `ProofSnapshot`, generating a JSON response. -/
-def unpickleProofSnapshot (n : UnpickleProofState) : M IO (ProofStepResponse ⊕ Error) := do
-  let (cmdSnapshot?, notFound) ← do match n.env with
-  | none => pure (none, false)
-  | some i => do match (← get).cmdStates[i]? with
-    | some env => pure (some env, false)
-    | none => pure (none, true)
-  if notFound then
-    return .inr ⟨"Unknown environment."⟩
-  let (proofState, _) ← ProofSnapshot.unpickle n.unpickleProofStateFrom cmdSnapshot?
-  Sum.inl <$> createProofStepReponse proofState
-
-/--
-Run a command, returning the id of the new environment, and any messages and sorries.
--/
-def runCommand (s : Command) : M IO (CommandResponse ⊕ Error) := do
-  let (cmdSnapshot?, notFound) ← do match s.env with
-  | none => pure (none, false)
-  | some i => do match (← get).cmdStates[i]? with
-    | some env => pure (some env, false)
-    | none => pure (none, true)
-  if notFound then
-    return .inr ⟨"Unknown environment."⟩
-  let initialCmdState? := cmdSnapshot?.map fun c => c.cmdState
-  let (initialCmdState, cmdState, messages, trees) ← try
-    IO.processInput s.cmd initialCmdState?
-  catch ex =>
-    return .inr ⟨ex.toString⟩
-  let messages ← messages.mapM fun m => Message.of m
-  -- For debugging purposes, sometimes we print out the trees here:
-  -- trees.forM fun t => do IO.println (← t.format)
-  let sorries ← sorries trees initialCmdState.env none
-  let sorries ← match s.rootGoals with
-  | some true => pure (sorries ++ (← collectRootGoalsAsSorries trees initialCmdState.env))
-  | _ => pure sorries
-  let tactics ← match s.allTactics with
-  | some true => tactics trees initialCmdState.env
-  | _ => pure []
-  let cmdSnapshot :=
-  { cmdState
-    cmdContext := (cmdSnapshot?.map fun c => c.cmdContext).getD
-      { fileName := "",
-        fileMap := default,
-        snap? := none,
-        cancelTk? := none } }
-  let env ← recordCommandSnapshot cmdSnapshot
-  let jsonTrees := match s.infotree with
-  | some "full" => trees
-  | some "tactics" => trees.flatMap InfoTree.retainTacticInfo
-  | some "original" => trees.flatMap InfoTree.retainTacticInfo |>.flatMap InfoTree.retainOriginal
-  | some "substantive" => trees.flatMap InfoTree.retainTacticInfo |>.flatMap InfoTree.retainSubstantive
-  | _ => []
-  let infotree ← if jsonTrees.isEmpty then
-    pure none
-  else
-    pure <| some <| Json.arr (← jsonTrees.toArray.mapM fun t => t.toJson none)
-  return .inl
-    { env,
-      messages,
-      sorries,
-      tactics
-      infotree }
-
-def processFile (s : File) : M IO (CommandResponse ⊕ Error) := do
-  try
-    let cmd ← IO.FS.readFile s.path
-    runCommand { s with env := s.env, cmd }
-  catch e =>
-    pure <| .inr ⟨e.toString⟩
-
-/--
-Run a single tactic, returning the id of the new proof statement, and the new goals.
--/
--- TODO detect sorries?
-def runProofStep (s : ProofStep) : M IO (ProofStepResponse ⊕ Error) := do
-  match (← get).proofStates[s.proofState]? with
-  | none => return .inr ⟨"Unknown proof state."⟩
-  | some proofState =>
-    try
-      let proofState' ← proofState.runString s.tactic
-      return .inl (← createProofStepReponse proofState' proofState)
-    catch ex =>
-      return .inr ⟨"Lean error:\n" ++ ex.toString⟩
-
-end REPL
-
-open REPL
+open REPL Actions
 
 /-- Get lines from stdin until a blank line is entered. -/
 partial def getLines : IO String := do
@@ -381,36 +75,36 @@ instance [ToJson α] [ToJson β] : ToJson (α ⊕ β) where
 
 /-- Commands accepted by the REPL. -/
 inductive Input
-| command : REPL.Command → Input
-| file : REPL.File → Input
-| proofStep : REPL.ProofStep → Input
-| pickleEnvironment : REPL.PickleEnvironment → Input
-| unpickleEnvironment : REPL.UnpickleEnvironment → Input
-| pickleProofSnapshot : REPL.PickleProofState → Input
-| unpickleProofSnapshot : REPL.UnpickleProofState → Input
+| command : Actions.Command → Input
+| file : Actions.File → Input
+| proofStep : Actions.ProofStep → Input
+| pickleEnvironment : Actions.PickleEnvironment → Input
+| unpickleEnvironment : Actions.UnpickleEnvironment → Input
+| pickleProofSnapshot : Actions.PickleProofState → Input
+| unpickleProofSnapshot : Actions.UnpickleProofState → Input
 
 /-- Parse a user input string to an input command. -/
 def parse (query : String) : IO Input := do
   let json := Json.parse query
   match json with
   | .error e => throw <| IO.userError <| toString <| toJson <|
-      (⟨"Could not parse JSON:\n" ++ e⟩ : Error)
+      (⟨"Could not parse JSON:\n" ++ e⟩ : Actions.Error)
   | .ok j => match fromJson? j with
-    | .ok (r : REPL.ProofStep) => return .proofStep r
+    | .ok (r : Actions.ProofStep) => return .proofStep r
     | .error _ => match fromJson? j with
-    | .ok (r : REPL.PickleEnvironment) => return .pickleEnvironment r
+    | .ok (r : Actions.PickleEnvironment) => return .pickleEnvironment r
     | .error _ => match fromJson? j with
-    | .ok (r : REPL.UnpickleEnvironment) => return .unpickleEnvironment r
+    | .ok (r : Actions.UnpickleEnvironment) => return .unpickleEnvironment r
     | .error _ => match fromJson? j with
-    | .ok (r : REPL.PickleProofState) => return .pickleProofSnapshot r
+    | .ok (r : Actions.PickleProofState) => return .pickleProofSnapshot r
     | .error _ => match fromJson? j with
-    | .ok (r : REPL.UnpickleProofState) => return .unpickleProofSnapshot r
+    | .ok (r : Actions.UnpickleProofState) => return .unpickleProofSnapshot r
     | .error _ => match fromJson? j with
-    | .ok (r : REPL.Command) => return .command r
+    | .ok (r : Actions.Command) => return .command r
     | .error _ => match fromJson? j with
-    | .ok (r : REPL.File) => return .file r
+    | .ok (r : Actions.File) => return .file r
     | .error e => throw <| IO.userError <| toString <| toJson <|
-        (⟨"Could not parse as a valid JSON command:\n" ++ e⟩ : Error)
+        (⟨"Could not parse as a valid JSON command:\n" ++ e⟩ : Actions.Error)
 
 /-- Avoid buffering the output. -/
 def printFlush [ToString α] (s : α) : IO Unit := do

From 2bc343354f8ee3546acf721308f534295150fd79 Mon Sep 17 00:00:00 2001
From: Qiu233 <lopi258@outlook.com>
Date: Tue, 15 Jul 2025 11:48:07 +0800
Subject: [PATCH 2/4] chore: replace StateT to improve performance

---
 REPL/Actions/Basic.lean   |  2 --
 REPL/Actions/Command.lean | 18 +++++++++---------
 REPL/Actions/File.lean    |  2 +-
 REPL/Actions/Tactic.lean  | 36 +++++++++++++++++++++--------------
 REPL/Basic.lean           | 40 +++++++++++++++++++++++++++++++++------
 REPL/Main.lean            |  4 ++--
 6 files changed, 68 insertions(+), 34 deletions(-)

diff --git a/REPL/Actions/Basic.lean b/REPL/Actions/Basic.lean
index 231e41b7..53d1869c 100644
--- a/REPL/Actions/Basic.lean
+++ b/REPL/Actions/Basic.lean
@@ -76,8 +76,6 @@ structure Error where
   message : String
 deriving ToJson, FromJson
 
-variable [Monad m] [MonadLiftT IO m]
-
 def collectFVarsAux : Expr → NameSet
   | .fvar fvarId => NameSet.empty.insert fvarId.name
   | .app fm arg => (collectFVarsAux fm).union $ collectFVarsAux arg
diff --git a/REPL/Actions/Command.lean b/REPL/Actions/Command.lean
index 1fb4f9cb..b771e2b4 100644
--- a/REPL/Actions/Command.lean
+++ b/REPL/Actions/Command.lean
@@ -55,18 +55,17 @@ structure UnpickleEnvironment where
   unpickleEnvFrom : System.FilePath
 deriving ToJson, FromJson
 
-variable [Monad m] [MonadLiftT IO m]
+variable [Monad m] [MonadREPL m] [MonadLiftT IO m]
 
 /-- Record an `CommandSnapshot` into the REPL state, returning its index for future use. -/
-def recordCommandSnapshot (state : CommandSnapshot) : M m Nat := do
-  let id := (← get).cmdStates.size
-  modify fun s => { s with cmdStates := s.cmdStates.push state }
-  return id
+@[specialize]
+def recordCommandSnapshot (state : CommandSnapshot) : m Nat := do
+  modifyGetCmdSnaps fun s => (s.size, s.push state)
 
 /--
 Run a command, returning the id of the new environment, and any messages and sorries.
 -/
-def runCommand (s : Command) : M IO (CommandResponse ⊕ Error) := do
+def runCommand (s : Command) : M (CommandResponse ⊕ Error) := do
   let (cmdSnapshot?, notFound) ← do match s.env with
   | none => pure (none, false)
   | some i => do match (← get).cmdStates[i]? with
@@ -115,15 +114,16 @@ def runCommand (s : Command) : M IO (CommandResponse ⊕ Error) := do
       infotree }
 
 /-- Pickle a `CommandSnapshot`, generating a JSON response. -/
-def pickleCommandSnapshot (n : PickleEnvironment) : M m (CommandResponse ⊕ Error) := do
-  match (← get).cmdStates[n.env]? with
+@[specialize]
+def pickleCommandSnapshot (n : PickleEnvironment) : m (CommandResponse ⊕ Error) := do
+  match (← getCmdSnaps)[n.env]? with
   | none => return .inr ⟨"Unknown environment."⟩
   | some env =>
     discard <| env.pickle n.pickleTo
     return .inl { env := n.env }
 
 /-- Unpickle a `CommandSnapshot`, generating a JSON response. -/
-def unpickleCommandSnapshot (n : UnpickleEnvironment) : M IO CommandResponse := do
+def unpickleCommandSnapshot (n : UnpickleEnvironment) : M CommandResponse := do
   let (env, _) ← CommandSnapshot.unpickle n.unpickleEnvFrom
   let env ← recordCommandSnapshot env
   return { env }
diff --git a/REPL/Actions/File.lean b/REPL/Actions/File.lean
index 572f01ed..09abf74e 100644
--- a/REPL/Actions/File.lean
+++ b/REPL/Actions/File.lean
@@ -12,7 +12,7 @@ structure File extends CommandOptions where
   path : System.FilePath
 deriving ToJson, FromJson
 
-def processFile (s : File) : M IO (CommandResponse ⊕ Error) := do
+def processFile (s : File) : M (CommandResponse ⊕ Error) := do
   try
     let cmd ← IO.FS.readFile s.path
     runCommand { s with env := s.env, cmd }
diff --git a/REPL/Actions/Tactic.lean b/REPL/Actions/Tactic.lean
index dc54b32e..008e54d5 100644
--- a/REPL/Actions/Tactic.lean
+++ b/REPL/Actions/Tactic.lean
@@ -64,16 +64,16 @@ structure UnpickleProofState where
   env : Option Nat
 deriving ToJson, FromJson
 
-variable [Monad m] [MonadLiftT IO m]
+variable [Monad m] [MonadREPL m] [MonadLiftT IO m]
 
 /-- Record a `ProofSnapshot` into the REPL state, returning its index for future use. -/
-def recordProofSnapshot (proofState : ProofSnapshot) : M m Nat := do
-  let id := (← get).proofStates.size
-  modify fun s => { s with proofStates := s.proofStates.push proofState }
-  return id
+@[specialize]
+def recordProofSnapshot (proofState : ProofSnapshot) : m Nat := do
+  modifyGetProofSnaps fun s => (s.size, s.push proofState)
 
+@[specialize]
 def sorries (trees : List InfoTree) (env? : Option Environment) (rootGoals? : Option (List MVarId))
-: M m (List Sorry) :=
+: m (List Sorry) :=
   trees.flatMap InfoTree.sorries |>.filter (fun t => match t.2.1 with
     | .term _ none => false
     | _ => true ) |>.mapM
@@ -93,13 +93,15 @@ def sorries (trees : List InfoTree) (env? : Option Environment) (rootGoals? : Op
         let proofStateId ← proofState.mapM recordProofSnapshot
         return Sorry.of goal pos endPos proofStateId
 
+@[specialize]
 def ppTactic (ctx : ContextInfo) (stx : Syntax) : IO Format :=
   ctx.runMetaM {} try
     Lean.PrettyPrinter.ppTactic ⟨stx⟩
   catch _ =>
     pure "<failed to pretty print>"
 
-def tactics (trees : List InfoTree) (env? : Option Environment) : M m (List Tactic) :=
+@[specialize]
+def tactics (trees : List InfoTree) (env? : Option Environment) : m (List Tactic) :=
   trees.flatMap InfoTree.tactics |>.mapM
     fun ⟨ctx, stx, rootGoals, goals, pos, endPos, ns⟩ => do
       let proofState := some (← ProofSnapshot.create ctx none env? goals rootGoals)
@@ -108,7 +110,8 @@ def tactics (trees : List InfoTree) (env? : Option Environment) : M m (List Tact
       let proofStateId ← proofState.mapM recordProofSnapshot
       return Tactic.of goals tactic pos endPos proofStateId ns
 
-def collectRootGoalsAsSorries (trees : List InfoTree) (env? : Option Environment) : M m (List Sorry) := do
+@[specialize]
+def collectRootGoalsAsSorries (trees : List InfoTree) (env? : Option Environment) : m (List Sorry) := do
   trees.flatMap InfoTree.rootGoals |>.mapM
     fun ⟨ctx, goals, pos⟩ => do
       let proofState := some (← ProofSnapshot.create ctx none env? goals goals)
@@ -125,7 +128,8 @@ Main states include:
 
 Inspired by LeanDojo REPL's status tracking.
 -/
-def getProofStatus (proofState : ProofSnapshot) : M m String := do
+@[specialize]
+def getProofStatus (proofState : ProofSnapshot) : m String := do
   match proofState.tacticState.goals with
     | [] =>
       let res := proofState.runMetaM do
@@ -178,8 +182,9 @@ def getProofStatus (proofState : ProofSnapshot) : M m String := do
 
 
 /-- Record a `ProofSnapshot` and generate a JSON response for it. -/
+@[specialize]
 def createProofStepReponse (proofState : ProofSnapshot) (old? : Option ProofSnapshot := none) :
-    M m ProofStepResponse := do
+    m ProofStepResponse := do
   let messages := proofState.newMessages old?
   let messages ← messages.mapM fun m => Message.of m
   let traces ← proofState.newTraces old?
@@ -204,15 +209,17 @@ def createProofStepReponse (proofState : ProofSnapshot) (old? : Option ProofSnap
 
 /-- Pickle a `ProofSnapshot`, generating a JSON response. -/
 -- This generates a new identifier, which perhaps is not what we want?
-def pickleProofSnapshot (n : PickleProofState) : M m (ProofStepResponse ⊕ Error) := do
-  match (← get).proofStates[n.proofState]? with
+@[specialize]
+def pickleProofSnapshot (n : PickleProofState) : m (ProofStepResponse ⊕ Error) := do
+  match (← getProofSnaps)[n.proofState]? with
   | none => return .inr ⟨"Unknown proof State."⟩
   | some proofState =>
     discard <| proofState.pickle n.pickleTo
     return .inl (← createProofStepReponse proofState)
 
 /-- Unpickle a `ProofSnapshot`, generating a JSON response. -/
-def unpickleProofSnapshot (n : UnpickleProofState) : M IO (ProofStepResponse ⊕ Error) := do
+@[specialize]
+def unpickleProofSnapshot (n : UnpickleProofState) : M (ProofStepResponse ⊕ Error) := do
   let (cmdSnapshot?, notFound) ← do match n.env with
   | none => pure (none, false)
   | some i => do match (← get).cmdStates[i]? with
@@ -227,7 +234,8 @@ def unpickleProofSnapshot (n : UnpickleProofState) : M IO (ProofStepResponse ⊕
 Run a single tactic, returning the id of the new proof statement, and the new goals.
 -/
 -- TODO detect sorries?
-def runProofStep (s : ProofStep) : M IO (ProofStepResponse ⊕ Error) := do
+@[specialize]
+def runProofStep (s : ProofStep) : M (ProofStepResponse ⊕ Error) := do
   match (← get).proofStates[s.proofState]? with
   | none => return .inr ⟨"Unknown proof state."⟩
   | some proofState =>
diff --git a/REPL/Basic.lean b/REPL/Basic.lean
index 328853bb..dcaf42bb 100644
--- a/REPL/Basic.lean
+++ b/REPL/Basic.lean
@@ -5,6 +5,23 @@ open Lean Elab
 
 namespace REPL
 
+class MonadCommandSnapshots (m : Type → Type) where
+  getCmdSnaps : m (Array CommandSnapshot)
+  setCmdSnaps : Array CommandSnapshot → m Unit
+  modifyGetCmdSnaps {α} : (Array CommandSnapshot → (α × Array CommandSnapshot)) → m α
+
+class MonadProofSnapshots (m : Type → Type) where
+  getProofSnaps : m (Array ProofSnapshot)
+  setProofSnaps : Array ProofSnapshot → m Unit
+  modifyGetProofSnaps {α} : (Array ProofSnapshot → (α × Array ProofSnapshot)) → m α
+
+attribute [specialize] MonadCommandSnapshots.modifyGetCmdSnaps MonadProofSnapshots.modifyGetProofSnaps
+
+class MonadREPL (m : Type → Type) extends MonadCommandSnapshots m, MonadProofSnapshots m where
+
+export MonadCommandSnapshots (getCmdSnaps setCmdSnaps modifyGetCmdSnaps)
+export MonadProofSnapshots (getProofSnaps setProofSnaps modifyGetProofSnaps)
+
 /-- The monadic state for the Lean REPL. -/
 structure State where
   /--
@@ -20,11 +37,22 @@ structure State where
   -/
   proofStates : Array ProofSnapshot := #[]
 
-/--
-The Lean REPL monad.
+-- /--
+-- The Lean REPL monad.
+-- -/
+abbrev M : Type → Type := StateRefT State IO
+
+@[inline]
+instance : MonadREPL M where
+  getCmdSnaps := State.cmdStates <$> get
+  setCmdSnaps snaps := modify ({· with cmdStates := snaps})
+  modifyGetCmdSnaps f := modifyGet fun x => let (a, v) := f x.cmdStates; (a, {x with cmdStates := v})
+  getProofSnaps := State.proofStates <$> get
+  setProofSnaps snaps := modify ({· with proofStates := snaps})
+  modifyGetProofSnaps f := modifyGet fun x => let (a, v) := f x.proofStates; (a, {x with proofStates := v})
 
-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
+@[always_inline, specialize]
+def modifyCmdSnaps [MonadCommandSnapshots m] : (Array CommandSnapshot → Array CommandSnapshot) → m Unit := fun f => modifyGetCmdSnaps fun x => ((), f x)
 
-variable [Monad m] [MonadLiftT IO m]
+@[always_inline, specialize]
+def modifyProofSnaps [MonadProofSnapshots m] : (Array ProofSnapshot → Array ProofSnapshot) → m Unit := fun f => modifyGetProofSnaps fun x => ((), f x)
diff --git a/REPL/Main.lean b/REPL/Main.lean
index db4f87b9..342dbd84 100644
--- a/REPL/Main.lean
+++ b/REPL/Main.lean
@@ -114,8 +114,8 @@ def printFlush [ToString α] (s : α) : IO Unit := do
 
 /-- Read-eval-print loop for Lean. -/
 unsafe def repl : IO Unit :=
-  StateT.run' loop {}
-where loop : M IO Unit := do
+  loop.run' {}
+where loop : M Unit := do
   let query ← getLines
   if query = "" then
     return ()

From f0b4a59a117e73c3fb7c9e292c37c218f35c818e Mon Sep 17 00:00:00 2001
From: Qiu233 <lopi258@outlook.com>
Date: Tue, 15 Jul 2025 12:38:42 +0800
Subject: [PATCH 3/4] chore: greatly improve readability of effectful actions

---
 REPL/Actions/Basic.lean   |  2 ++
 REPL/Actions/Command.lean | 72 ++++++++++++++++++---------------------
 REPL/Actions/File.lean    |  6 ++--
 REPL/Actions/Tactic.lean  | 41 +++++++++-------------
 REPL/Basic.lean           |  9 +++++
 REPL/Main.lean            | 36 ++++++++++----------
 6 files changed, 82 insertions(+), 84 deletions(-)

diff --git a/REPL/Actions/Basic.lean b/REPL/Actions/Basic.lean
index 53d1869c..34806c43 100644
--- a/REPL/Actions/Basic.lean
+++ b/REPL/Actions/Basic.lean
@@ -76,6 +76,8 @@ structure Error where
   message : String
 deriving ToJson, FromJson
 
+abbrev ResultT (m : Type → Type) (α : Type) := ExceptT Error m α
+
 def collectFVarsAux : Expr → NameSet
   | .fvar fvarId => NameSet.empty.insert fvarId.name
   | .app fm arg => (collectFVarsAux fm).union $ collectFVarsAux arg
diff --git a/REPL/Actions/Command.lean b/REPL/Actions/Command.lean
index b771e2b4..59dbca04 100644
--- a/REPL/Actions/Command.lean
+++ b/REPL/Actions/Command.lean
@@ -65,48 +65,46 @@ def recordCommandSnapshot (state : CommandSnapshot) : m Nat := do
 /--
 Run a command, returning the id of the new environment, and any messages and sorries.
 -/
-def runCommand (s : Command) : M (CommandResponse ⊕ Error) := do
-  let (cmdSnapshot?, notFound) ← do match s.env with
-  | none => pure (none, false)
-  | some i => do match (← get).cmdStates[i]? with
-    | some env => pure (some env, false)
-    | none => pure (none, true)
-  if notFound then
-    return .inr ⟨"Unknown environment."⟩
+def runCommand (s : Command) : ResultT M CommandResponse := do
+  let cmdSnapshot? ← s.env.mapM fun i => do
+    let some env := (← get).cmdStates[i]? | throw ⟨"Unknown environment."⟩
+    return env
   let initialCmdState? := cmdSnapshot?.map fun c => c.cmdState
-  let (initialCmdState, cmdState, messages, trees) ← try
-    IO.processInput s.cmd initialCmdState?
-  catch ex =>
-    return .inr ⟨ex.toString⟩
+  let (initialCmdState, cmdState, messages, trees) ← do
+    try
+      IO.processInput s.cmd initialCmdState?
+    catch ex : IO.Error =>
+      throw ⟨ex.toString⟩
   let messages ← messages.mapM fun m => Message.of m
   -- For debugging purposes, sometimes we print out the trees here:
   -- trees.forM fun t => do IO.println (← t.format)
   let sorries ← sorries trees initialCmdState.env none
   let sorries ← match s.rootGoals with
-  | some true => pure (sorries ++ (← collectRootGoalsAsSorries trees initialCmdState.env))
-  | _ => pure sorries
+    | some true => pure (sorries ++ (← collectRootGoalsAsSorries trees initialCmdState.env))
+    | _ => pure sorries
   let tactics ← match s.allTactics with
-  | some true => tactics trees initialCmdState.env
-  | _ => pure []
+    | some true => tactics trees initialCmdState.env
+    | _ => pure []
   let cmdSnapshot :=
-  { cmdState
-    cmdContext := (cmdSnapshot?.map fun c => c.cmdContext).getD
-      { fileName := "",
-        fileMap := default,
-        snap? := none,
-        cancelTk? := none } }
+    { cmdState
+      cmdContext := (cmdSnapshot?.map fun c => c.cmdContext).getD
+        { fileName := "",
+          fileMap := default,
+          snap? := none,
+          cancelTk? := none } }
   let env ← recordCommandSnapshot cmdSnapshot
   let jsonTrees := match s.infotree with
-  | some "full" => trees
-  | some "tactics" => trees.flatMap InfoTree.retainTacticInfo
-  | some "original" => trees.flatMap InfoTree.retainTacticInfo |>.flatMap InfoTree.retainOriginal
-  | some "substantive" => trees.flatMap InfoTree.retainTacticInfo |>.flatMap InfoTree.retainSubstantive
-  | _ => []
-  let infotree ← if jsonTrees.isEmpty then
-    pure none
-  else
-    pure <| some <| Json.arr (← jsonTrees.toArray.mapM fun t => t.toJson none)
-  return .inl
+    | some "full" => trees
+    | some "tactics" => trees.flatMap InfoTree.retainTacticInfo
+    | some "original" => trees.flatMap InfoTree.retainTacticInfo |>.flatMap InfoTree.retainOriginal
+    | some "substantive" => trees.flatMap InfoTree.retainTacticInfo |>.flatMap InfoTree.retainSubstantive
+    | _ => []
+  let infotree ←
+    if jsonTrees.isEmpty then
+      pure none
+    else
+      pure <| some <| Json.arr (← jsonTrees.toArray.mapM fun t => t.toJson none)
+  return show CommandResponse from
     { env,
       messages,
       sorries,
@@ -115,12 +113,10 @@ def runCommand (s : Command) : M (CommandResponse ⊕ Error) := do
 
 /-- Pickle a `CommandSnapshot`, generating a JSON response. -/
 @[specialize]
-def pickleCommandSnapshot (n : PickleEnvironment) : m (CommandResponse ⊕ Error) := do
-  match (← getCmdSnaps)[n.env]? with
-  | none => return .inr ⟨"Unknown environment."⟩
-  | some env =>
-    discard <| env.pickle n.pickleTo
-    return .inl { env := n.env }
+def pickleCommandSnapshot (n : PickleEnvironment) : ResultT m CommandResponse := do
+  let some env := (← getCmdSnaps)[n.env]? | throw ⟨"Unknown environment."⟩
+  discard <| env.pickle n.pickleTo
+  return { env := n.env }
 
 /-- Unpickle a `CommandSnapshot`, generating a JSON response. -/
 def unpickleCommandSnapshot (n : UnpickleEnvironment) : M CommandResponse := do
diff --git a/REPL/Actions/File.lean b/REPL/Actions/File.lean
index 09abf74e..db97df38 100644
--- a/REPL/Actions/File.lean
+++ b/REPL/Actions/File.lean
@@ -12,9 +12,9 @@ structure File extends CommandOptions where
   path : System.FilePath
 deriving ToJson, FromJson
 
-def processFile (s : File) : M (CommandResponse ⊕ Error) := do
+def processFile (s : File) : ResultT M CommandResponse := do
   try
     let cmd ← IO.FS.readFile s.path
     runCommand { s with env := s.env, cmd }
-  catch e =>
-    pure <| .inr ⟨e.toString⟩
+  catch e : IO.Error =>
+    throw ⟨e.toString⟩
diff --git a/REPL/Actions/Tactic.lean b/REPL/Actions/Tactic.lean
index 008e54d5..6f7e2de4 100644
--- a/REPL/Actions/Tactic.lean
+++ b/REPL/Actions/Tactic.lean
@@ -210,37 +210,28 @@ def createProofStepReponse (proofState : ProofSnapshot) (old? : Option ProofSnap
 /-- Pickle a `ProofSnapshot`, generating a JSON response. -/
 -- This generates a new identifier, which perhaps is not what we want?
 @[specialize]
-def pickleProofSnapshot (n : PickleProofState) : m (ProofStepResponse ⊕ Error) := do
-  match (← getProofSnaps)[n.proofState]? with
-  | none => return .inr ⟨"Unknown proof State."⟩
-  | some proofState =>
-    discard <| proofState.pickle n.pickleTo
-    return .inl (← createProofStepReponse proofState)
+def pickleProofSnapshot (n : PickleProofState) : ResultT m ProofStepResponse := do
+  let some proofState := (← getProofSnaps)[n.proofState]? | throw ⟨"Unknown proof State."⟩
+  discard <| proofState.pickle n.pickleTo
+  createProofStepReponse proofState
 
 /-- Unpickle a `ProofSnapshot`, generating a JSON response. -/
-@[specialize]
-def unpickleProofSnapshot (n : UnpickleProofState) : M (ProofStepResponse ⊕ Error) := do
-  let (cmdSnapshot?, notFound) ← do match n.env with
-  | none => pure (none, false)
-  | some i => do match (← get).cmdStates[i]? with
-    | some env => pure (some env, false)
-    | none => pure (none, true)
-  if notFound then
-    return .inr ⟨"Unknown environment."⟩
+def unpickleProofSnapshot (n : UnpickleProofState) : ResultT M ProofStepResponse := do
+  let cmdSnapshot? ← n.env.mapM fun i => do
+    let some env := (← get).cmdStates[i]? | throw ⟨"Unknown environment."⟩
+    return env
   let (proofState, _) ← ProofSnapshot.unpickle n.unpickleProofStateFrom cmdSnapshot?
-  Sum.inl <$> createProofStepReponse proofState
+  createProofStepReponse proofState
 
 /--
 Run a single tactic, returning the id of the new proof statement, and the new goals.
 -/
 -- TODO detect sorries?
 @[specialize]
-def runProofStep (s : ProofStep) : M (ProofStepResponse ⊕ Error) := do
-  match (← get).proofStates[s.proofState]? with
-  | none => return .inr ⟨"Unknown proof state."⟩
-  | some proofState =>
-    try
-      let proofState' ← proofState.runString s.tactic
-      return .inl (← createProofStepReponse proofState' proofState)
-    catch ex =>
-      return .inr ⟨"Lean error:\n" ++ ex.toString⟩
+def runProofStep (s : ProofStep) : ResultT M ProofStepResponse := do
+  let some proofState := (← get).proofStates[s.proofState]? | throw ⟨"Unknown proof state."⟩
+  try
+    let proofState' ← proofState.runString s.tactic
+    createProofStepReponse proofState' proofState
+  catch ex : IO.Error => -- catch the outer IO error
+    throw ⟨"Lean error:\n" ++ ex.toString⟩
diff --git a/REPL/Basic.lean b/REPL/Basic.lean
index dcaf42bb..433d8db8 100644
--- a/REPL/Basic.lean
+++ b/REPL/Basic.lean
@@ -56,3 +56,12 @@ def modifyCmdSnaps [MonadCommandSnapshots m] : (Array CommandSnapshot → Array
 
 @[always_inline, specialize]
 def modifyProofSnaps [MonadProofSnapshots m] : (Array ProofSnapshot → Array ProofSnapshot) → m Unit := fun f => modifyGetProofSnaps fun x => ((), f x)
+
+@[inline]
+instance {m n} [MonadLift m n] [MonadREPL m] : MonadREPL n where
+  getCmdSnaps := MonadLift.monadLift (m := m) (n := n) getCmdSnaps
+  setCmdSnaps snaps := MonadLift.monadLift (m := m) (n := n) <| setCmdSnaps snaps
+  modifyGetCmdSnaps f := MonadLift.monadLift (m := m) (n := n) <| modifyGetCmdSnaps f
+  getProofSnaps := MonadLift.monadLift (m := m) (n := n) getProofSnaps
+  setProofSnaps snaps := MonadLift.monadLift (m := m) (n := n) <| setProofSnaps snaps
+  modifyGetProofSnaps f := MonadLift.monadLift (m := m) (n := n) <| modifyGetProofSnaps f
diff --git a/REPL/Main.lean b/REPL/Main.lean
index 342dbd84..1b871e77 100644
--- a/REPL/Main.lean
+++ b/REPL/Main.lean
@@ -68,20 +68,20 @@ partial def getLines : IO String := do
   else
     return line.trimRight ++ (← getLines)
 
-instance [ToJson α] [ToJson β] : ToJson (α ⊕ β) where
-  toJson x := match x with
-  | .inl a => toJson a
-  | .inr b => toJson b
+local instance [ToJson ε] [ToJson α] : ToJson (Except ε α) where
+  toJson
+  | .error e => toJson e
+  | .ok a => toJson a
 
 /-- Commands accepted by the REPL. -/
 inductive Input
-| command : Actions.Command → Input
-| file : Actions.File → Input
-| proofStep : Actions.ProofStep → Input
-| pickleEnvironment : Actions.PickleEnvironment → Input
-| unpickleEnvironment : Actions.UnpickleEnvironment → Input
-| pickleProofSnapshot : Actions.PickleProofState → Input
-| unpickleProofSnapshot : Actions.UnpickleProofState → Input
+  | command : Actions.Command → Input
+  | file : Actions.File → Input
+  | proofStep : Actions.ProofStep → Input
+  | pickleEnvironment : Actions.PickleEnvironment → Input
+  | unpickleEnvironment : Actions.UnpickleEnvironment → Input
+  | pickleProofSnapshot : Actions.PickleProofState → Input
+  | unpickleProofSnapshot : Actions.UnpickleProofState → Input
 
 /-- Parse a user input string to an input command. -/
 def parse (query : String) : IO Input := do
@@ -121,13 +121,13 @@ where loop : M Unit := do
     return ()
   if query.startsWith "#" || query.startsWith "--" then loop else
   IO.println <| toString <| ← match ← parse query with
-  | .command r => return toJson (← runCommand r)
-  | .file r => return toJson (← processFile r)
-  | .proofStep r => return toJson (← runProofStep r)
-  | .pickleEnvironment r => return toJson (← pickleCommandSnapshot r)
-  | .unpickleEnvironment r => return toJson (← unpickleCommandSnapshot r)
-  | .pickleProofSnapshot r => return toJson (← pickleProofSnapshot r)
-  | .unpickleProofSnapshot r => return toJson (← unpickleProofSnapshot r)
+    | .command r => return toJson (← runCommand r)
+    | .file r => return toJson (← processFile r)
+    | .proofStep r => return toJson (← runProofStep r)
+    | .pickleEnvironment r => return toJson (← pickleCommandSnapshot (m := M) r)
+    | .unpickleEnvironment r => return toJson (← unpickleCommandSnapshot r)
+    | .pickleProofSnapshot r => return toJson (← pickleProofSnapshot (m := M) r)
+    | .unpickleProofSnapshot r => return toJson (← unpickleProofSnapshot r)
   printFlush "\n" -- easier to parse the output if there are blank lines
   loop
 

From 6475fa66901a37d71e5bd9077d2d32f36318382f Mon Sep 17 00:00:00 2001
From: Qiu233 <lopi258@outlook.com>
Date: Tue, 15 Jul 2025 17:02:50 +0800
Subject: [PATCH 4/4] chore: decouple request message, request handler and
 dispatch

---
 REPL/Actions/Command.lean |   7 ++-
 REPL/Actions/File.lean    |   2 +
 REPL/Actions/Tactic.lean  |   8 ++-
 REPL/Basic.lean           |  74 +++++++++++++++++++++++++
 REPL/Main.lean            | 113 ++++++++++++++++++++++++++------------
 5 files changed, 167 insertions(+), 37 deletions(-)

diff --git a/REPL/Actions/Command.lean b/REPL/Actions/Command.lean
index 59dbca04..ec6a9195 100644
--- a/REPL/Actions/Command.lean
+++ b/REPL/Actions/Command.lean
@@ -20,6 +20,7 @@ structure CommandOptions where
 If `env = none`, starts a new session (in which you can use `import`).
 If `env = some n`, builds on the existing environment `n`.
 -/
+@[repl_request]
 structure Command extends CommandOptions where
   env : Option Nat
   cmd : String
@@ -46,11 +47,13 @@ instance : ToJson CommandResponse where
     match r.infotree with | some j => [("infotree", j)] | none => []
   ]
 
+@[repl_request]
 structure PickleEnvironment where
   env : Nat
   pickleTo : System.FilePath
 deriving ToJson, FromJson
 
+@[repl_request]
 structure UnpickleEnvironment where
   unpickleEnvFrom : System.FilePath
 deriving ToJson, FromJson
@@ -65,6 +68,7 @@ def recordCommandSnapshot (state : CommandSnapshot) : m Nat := do
 /--
 Run a command, returning the id of the new environment, and any messages and sorries.
 -/
+@[repl_request_handler Command]
 def runCommand (s : Command) : ResultT M CommandResponse := do
   let cmdSnapshot? ← s.env.mapM fun i => do
     let some env := (← get).cmdStates[i]? | throw ⟨"Unknown environment."⟩
@@ -112,13 +116,14 @@ def runCommand (s : Command) : ResultT M CommandResponse := do
       infotree }
 
 /-- Pickle a `CommandSnapshot`, generating a JSON response. -/
-@[specialize]
+@[specialize, repl_request_handler PickleEnvironment]
 def pickleCommandSnapshot (n : PickleEnvironment) : ResultT m CommandResponse := do
   let some env := (← getCmdSnaps)[n.env]? | throw ⟨"Unknown environment."⟩
   discard <| env.pickle n.pickleTo
   return { env := n.env }
 
 /-- Unpickle a `CommandSnapshot`, generating a JSON response. -/
+@[repl_request_handler UnpickleEnvironment]
 def unpickleCommandSnapshot (n : UnpickleEnvironment) : M CommandResponse := do
   let (env, _) ← CommandSnapshot.unpickle n.unpickleEnvFrom
   let env ← recordCommandSnapshot env
diff --git a/REPL/Actions/File.lean b/REPL/Actions/File.lean
index db97df38..53839e68 100644
--- a/REPL/Actions/File.lean
+++ b/REPL/Actions/File.lean
@@ -7,11 +7,13 @@ open Lean Elab InfoTree
 namespace REPL.Actions
 
 /-- Process a Lean file in a fresh environment if `env` is not provided. -/
+@[repl_request]
 structure File extends CommandOptions where
   env : Option Nat
   path : System.FilePath
 deriving ToJson, FromJson
 
+@[repl_request_handler File]
 def processFile (s : File) : ResultT M CommandResponse := do
   try
     let cmd ← IO.FS.readFile s.path
diff --git a/REPL/Actions/Tactic.lean b/REPL/Actions/Tactic.lean
index 6f7e2de4..7d9a84ff 100644
--- a/REPL/Actions/Tactic.lean
+++ b/REPL/Actions/Tactic.lean
@@ -26,6 +26,7 @@ def Tactic.of (goals tactic : String) (pos endPos : Lean.Position) (proofState :
 /--
 Run a tactic in a proof state.
 -/
+@[repl_request]
 structure ProofStep where
   proofState : Nat
   tactic : String
@@ -54,11 +55,13 @@ instance : ToJson ProofStepResponse where
     [("proofStatus", r.proofStatus)]
   ]
 
+@[repl_request pickleProofSnapshot]
 structure PickleProofState where
   proofState : Nat
   pickleTo : System.FilePath
 deriving ToJson, FromJson
 
+@[repl_request unpickleProofSnapshot]
 structure UnpickleProofState where
   unpickleProofStateFrom : System.FilePath
   env : Option Nat
@@ -209,13 +212,14 @@ def createProofStepReponse (proofState : ProofSnapshot) (old? : Option ProofSnap
 
 /-- Pickle a `ProofSnapshot`, generating a JSON response. -/
 -- This generates a new identifier, which perhaps is not what we want?
-@[specialize]
+@[specialize, repl_request_handler PickleProofState]
 def pickleProofSnapshot (n : PickleProofState) : ResultT m ProofStepResponse := do
   let some proofState := (← getProofSnaps)[n.proofState]? | throw ⟨"Unknown proof State."⟩
   discard <| proofState.pickle n.pickleTo
   createProofStepReponse proofState
 
 /-- Unpickle a `ProofSnapshot`, generating a JSON response. -/
+@[repl_request_handler UnpickleProofState]
 def unpickleProofSnapshot (n : UnpickleProofState) : ResultT M ProofStepResponse := do
   let cmdSnapshot? ← n.env.mapM fun i => do
     let some env := (← get).cmdStates[i]? | throw ⟨"Unknown environment."⟩
@@ -227,7 +231,7 @@ def unpickleProofSnapshot (n : UnpickleProofState) : ResultT M ProofStepResponse
 Run a single tactic, returning the id of the new proof statement, and the new goals.
 -/
 -- TODO detect sorries?
-@[specialize]
+@[specialize, repl_request_handler ProofStep]
 def runProofStep (s : ProofStep) : ResultT M ProofStepResponse := do
   let some proofState := (← get).proofStates[s.proofState]? | throw ⟨"Unknown proof state."⟩
   try
diff --git a/REPL/Basic.lean b/REPL/Basic.lean
index 433d8db8..007b0ee0 100644
--- a/REPL/Basic.lean
+++ b/REPL/Basic.lean
@@ -65,3 +65,77 @@ instance {m n} [MonadLift m n] [MonadREPL m] : MonadREPL n where
   getProofSnaps := MonadLift.monadLift (m := m) (n := n) getProofSnaps
   setProofSnaps snaps := MonadLift.monadLift (m := m) (n := n) <| setProofSnaps snaps
   modifyGetProofSnaps f := MonadLift.monadLift (m := m) (n := n) <| modifyGetProofSnaps f
+
+namespace Attr
+
+/-- Specify that a type is an repl request, and optionally its corresponding constructor name. -/
+syntax (name := repl_request) &"repl_request " (ident)? : attr
+
+/-- Specify a handler for a request. -/
+syntax (name := repl_request_handler) &"repl_request_handler " ident : attr
+
+private def validateReplRequestAttr : Name → Syntax → AttrM Name := fun name stx => do
+  let `(attr| repl_request $[$cname?]?) := stx | throwUnsupportedSyntax
+  if !name.isStr then
+    throwError "type must have simple name"
+  let .inductInfo info ← getConstInfo name | throwError "can only apply to an inductive/structure type"
+  if info.numCtors = 0 then
+    throwError "type must have at least one constructor"
+  if info.numParams != 0 then
+    throwError "type must have no type parameters"
+  if info.numIndices != 0 then
+    throwError "type must have no indices"
+  let s := name.getString!.data
+  let s := String.mk <| s[0]!.toLower :: s.tail
+  let cname := cname?.map (fun x => x.getId) |>.getD (Name.mkStr1 s)
+  return cname
+
+initialize replRequestAttr : ParametricAttribute Name ←
+  registerParametricAttribute {
+      name := `repl_request
+      descr := "the corresponding constructor name, optional"
+      getParam := validateReplRequestAttr
+  }
+
+def getAllReplRequests [Monad m] [MonadEnv m] : m (Array (Name × Name)) := do
+  let env ← getEnv
+  let ext := replRequestAttr.ext
+  let s := ext.toEnvExtension.getState env
+  let imported := s.importedEntries
+  let current := s.state.toArray
+  return imported.flatMap id ++ current
+
+open Meta in
+private def validateReplRequestHandlerAttr : Name → Syntax → AttrM (Name × Name × Term) := fun name stx => do
+  let `(attr| repl_request_handler $reqNameStx) := stx | throwUnsupportedSyntax
+  let reqName ← resolveGlobalConstNoOverload reqNameStx
+  addConstInfo reqNameStx reqName
+  let env ← getEnv
+  let some reqCtor := replRequestAttr.getParam? env reqName |
+    throwErrorAt reqNameStx "type {MessageData.ofConstName reqName} has no attribute `repl_request`"
+  let .defnInfo info ← getConstInfo name |
+    throwError "this attribute can only be applied to a function"
+  let fnType := info.type.consumeMData
+  let arity := fnType.getForallArity
+  if arity = 0 then
+    throwError "type of {MessageData.ofConstName name} is not a function"
+  if arity > 1 then
+    let s ← `($(mkIdent name) ($(mkIdent `m) := $(mkIdent ``M)))
+    return (reqName, reqCtor, s)
+  let s ← `($(mkIdent name))
+  return (reqName, reqCtor, s)
+
+initialize replRequestHandlerAttr : ParametricAttribute (Name × Name × Term) ←
+  registerParametricAttribute {
+      name := `repl_request_handler
+      descr := "the request type name"
+      getParam := validateReplRequestHandlerAttr
+  }
+
+def getAllReplRequestHandlers [Monad m] [MonadEnv m] : m (Array (Name × Name × Name × Term)) := do
+  let env ← getEnv
+  let ext := replRequestHandlerAttr.ext
+  let s := ext.toEnvExtension.getState env
+  let imported := s.importedEntries
+  let current := s.state.toArray
+  return imported.flatMap id ++ current
diff --git a/REPL/Main.lean b/REPL/Main.lean
index 1b871e77..5c697344 100644
--- a/REPL/Main.lean
+++ b/REPL/Main.lean
@@ -4,7 +4,6 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Scott Morrison
 -/
 import REPL.Basic
--- import REPL.JSON
 import REPL.Frontend
 import REPL.Util.Path
 import REPL.Lean.ContextInfo
@@ -73,38 +72,89 @@ local instance [ToJson ε] [ToJson α] : ToJson (Except ε α) where
   | .error e => toJson e
   | .ok a => toJson a
 
-/-- Commands accepted by the REPL. -/
-inductive Input
-  | command : Actions.Command → Input
-  | file : Actions.File → Input
-  | proofStep : Actions.ProofStep → Input
-  | pickleEnvironment : Actions.PickleEnvironment → Input
-  | unpickleEnvironment : Actions.UnpickleEnvironment → Input
-  | pickleProofSnapshot : Actions.PickleProofState → Input
-  | unpickleProofSnapshot : Actions.UnpickleProofState → Input
+def makeFromJson_Input (input_name : Ident) : Command.CommandElabM Unit := do
+  let .inductInfo info ← getConstInfo input_name.getId | throwError "type is not an inductive type"
+  if info.numCtors = 0 then
+    throwError "type must have at least one constructor"
+  if info.numParams != 0 then
+    throwError "type must have no type parameters"
+  if info.numIndices != 0 then
+    throwError "type must have no indices"
+  let ctors ← info.ctors.mapM getConstInfo
+  let ts ← ctors.mapM fun c => do
+    let type := c.type.consumeMData
+    if h : type.isForall then
+      let arity := type.getForallArity
+      if arity != 1 then
+        throwError "constructor {MessageData.ofConstName c.name} has non-singular arity, got {arity}"
+      return type.forallDomain h
+    else
+      throwError "constructor {MessageData.ofConstName c.name} has non-singular arity, got 0"
+  let br ← (ctors.zip ts).mapM fun (x, y) => do
+    let t ← Command.runTermElabM fun _ => PrettyPrinter.delab y
+    `($(mkIdent x.name):ident <$> tryThe $t)
+  let br := br.reverse
+  let ctors ← match br with
+    | [] => unreachable!
+    | [x] => pure x
+    | x :: tail => tail.foldlM (init := x) fun r l => `($l <|> $r)
+  let code ← `(command|
+    instance : FromJson $input_name where
+      fromJson? json := do
+        let rec @[specialize, always_inline] tryThe (α : Type) [FromJson α] : Except String α := fromJson? (α := α) json
+        $ctors:term)
+  trace[REPL.mkInput] "{code}"
+  Command.elabCommand code
+
+local elab "#mkInput " colGt input_name:ident : command => do
+  let names ← Attr.getAllReplRequests
+  let cs ← names.mapM fun (t, c) =>
+    `(Parser.Command.ctor| | $(mkIdent c):ident : $(mkIdent t):ident → $input_name)
+  let code ← `(command|
+    /-- Commands accepted by the REPL. -/
+    inductive $input_name where
+    $cs*)
+  trace[REPL.mkInput] "{code}"
+  Command.elabCommand code
+  makeFromJson_Input input_name
+
+#mkInput Input
+
+local syntax (name := makeReplDispatch) "make_repl_dispatch% " colGt term : term
+
+@[local term_elab makeReplDispatch]
+def elabMakeReplDispatch : Term.TermElab := fun stx type? => do
+  let `(make_repl_dispatch% $input) := stx | throwUnsupportedSyntax
+  let requests ← Attr.getAllReplRequests
+  let rs := RBTree.fromArray requests.unzip.fst Name.cmp
+  let handlers ← Attr.getAllReplRequestHandlers
+  let (r, cf) := handlers.unzip.snd.unzip
+  let r := RBTree.fromArray r Name.cmp
+  for d in RBTree.diff rs r do -- short-circuit by the first throw
+    throwError "no handler found for request type {MessageData.ofConstName d}"
+  let (c, f) := cf.unzip
+  let c := c.map mkIdent
+  let toJson := mkIdent ``toJson
+  let toJsons := Array.replicate c.size toJson
+  let s ← `(
+    match $input:term with
+    $[| .$c:ident r => $toJsons <$> ($f:term r)]*
+    )
+  trace[REPL.make_repl_dispatch] "{s}"
+  Term.withMacroExpansion stx s do
+    Term.elabTerm s type?
 
 /-- Parse a user input string to an input command. -/
 def parse (query : String) : IO Input := do
-  let json := Json.parse query
-  match json with
+  let json? := Json.parse query
+  match json? with
   | .error e => throw <| IO.userError <| toString <| toJson <|
       (⟨"Could not parse JSON:\n" ++ e⟩ : Actions.Error)
-  | .ok j => match fromJson? j with
-    | .ok (r : Actions.ProofStep) => return .proofStep r
-    | .error _ => match fromJson? j with
-    | .ok (r : Actions.PickleEnvironment) => return .pickleEnvironment r
-    | .error _ => match fromJson? j with
-    | .ok (r : Actions.UnpickleEnvironment) => return .unpickleEnvironment r
-    | .error _ => match fromJson? j with
-    | .ok (r : Actions.PickleProofState) => return .pickleProofSnapshot r
-    | .error _ => match fromJson? j with
-    | .ok (r : Actions.UnpickleProofState) => return .unpickleProofSnapshot r
-    | .error _ => match fromJson? j with
-    | .ok (r : Actions.Command) => return .command r
-    | .error _ => match fromJson? j with
-    | .ok (r : Actions.File) => return .file r
+  | .ok j =>
+    match fromJson? (α := Input) j with
     | .error e => throw <| IO.userError <| toString <| toJson <|
         (⟨"Could not parse as a valid JSON command:\n" ++ e⟩ : Actions.Error)
+    | .ok r => return r
 
 /-- Avoid buffering the output. -/
 def printFlush [ToString α] (s : α) : IO Unit := do
@@ -120,14 +170,9 @@ where loop : M Unit := do
   if query = "" then
     return ()
   if query.startsWith "#" || query.startsWith "--" then loop else
-  IO.println <| toString <| ← match ← parse query with
-    | .command r => return toJson (← runCommand r)
-    | .file r => return toJson (← processFile r)
-    | .proofStep r => return toJson (← runProofStep r)
-    | .pickleEnvironment r => return toJson (← pickleCommandSnapshot (m := M) r)
-    | .unpickleEnvironment r => return toJson (← unpickleCommandSnapshot r)
-    | .pickleProofSnapshot r => return toJson (← pickleProofSnapshot (m := M) r)
-    | .unpickleProofSnapshot r => return toJson (← unpickleProofSnapshot r)
+  let input ← parse query
+  let response ← make_repl_dispatch% input
+  IO.println <| toString response
   printFlush "\n" -- easier to parse the output if there are blank lines
   loop