refactor(goal): Allow multiple fragments

2025-06-24 13:31:55 -07:00 · 2025-06-24 13:31:55 -07:00 · 9d4b1ae755
parent 58ae791da3
commit 9d4b1ae755
6 changed files with 74 additions and 74 deletions
--- a/Pantograph/Goal.lean
+++ b/Pantograph/Goal.lean
@ -22,8 +22,9 @@ structure GoalState where
  -- Parent state metavariable source
  parentMVar?: Option MVarId := .none

-  -- If this is not `.none` there is a partial tactic being executed
-  fragment? : Option TacticFragment := .none
+  -- Any goal associated with a fragment has a partial tactic which has not
+  -- finished executing.
+  fragments : FragmentMap := .empty

@[export pantograph_goal_state_create_m]
 protected def GoalState.create (expr: Expr): Elab.TermElabM GoalState := do
@ -95,16 +96,12 @@ private def GoalState.restoreTacticM (state: GoalState) (goal: MVarId): Elab.Tac

@[export pantograph_goal_state_focus]
 protected def GoalState.focus (state : GoalState) (goal : MVarId): Option GoalState := do
-  let fragment? := match state.fragment? with
-    | .some { goal := goal', .. } => if goal' == goal then state.fragment? else .none
-    | .none => .none
  return {
    state with
    savedState := {
      state.savedState with
      tactic := { goals := [goal] },
    },
-    fragment?,
  }

 /-- Immediately bring all parent goals back into scope. Used in automatic mode -/
@ -125,10 +122,10 @@ protected def GoalState.immediateResume (state: GoalState) (parent: GoalState):
  }

 /--
-Brings into scope a list of goals. User must ensure `goals` is distinct.
+Brings into scope a list of goals. User must ensure `goals` are distinct.
 -/
@[export pantograph_goal_state_resume]
-protected def GoalState.resume (state: GoalState) (goals: List MVarId): Except String GoalState := do
+protected def GoalState.resume (state : GoalState) (goals : List MVarId) : Except String GoalState := do
  if ¬ (goals.all (λ goal => state.mvars.contains goal)) then
    let invalid_goals := goals.filter (λ goal => ¬ state.mvars.contains goal) |>.map (·.name.toString)
    .error s!"Goals {invalid_goals} are not in scope"
@ -147,7 +144,7 @@ protected def GoalState.resume (state: GoalState) (goals: List MVarId): Except S
 Brings into scope all goals from `branch`
 -/
@[export pantograph_goal_state_continue]
-protected def GoalState.continue (target: GoalState) (branch: GoalState): Except String GoalState :=
+protected def GoalState.continue (target : GoalState) (branch : GoalState) : Except String GoalState :=
  if !target.goals.isEmpty then
    .error s!"Target state has unresolved goals"
  else if target.root != branch.root then
@ -156,7 +153,7 @@ protected def GoalState.continue (target: GoalState) (branch: GoalState): Except
    target.resume (goals := branch.goals)

@[export pantograph_goal_state_root_expr]
-protected def GoalState.rootExpr? (goalState : GoalState): Option Expr := do
+protected def GoalState.rootExpr? (goalState : GoalState) : Option Expr := do
  if goalState.root.name == .anonymous then
    .none
  let expr ← goalState.mctx.eAssignment.find? goalState.root
@ -452,14 +449,21 @@ protected def GoalState.tryTacticM
 protected def GoalState.tryTactic (state: GoalState) (goal: MVarId) (tactic: String):
    Elab.TermElabM TacticResult := do
  state.restoreElabM
-  if let .some { goal := goal', fragment } := state.fragment? then
-    if goal == goal' then
-      return ← withCapturingError do
-        let (fragment?', state') ← state.step' goal (fragment.step goal tactic)
-        return { state' with fragment? := fragment?' }
+  /-
+  if let .some fragment := state.fragments[goal]? then
+    return ← withCapturingError do
+      let (moreFragments, state') ← state.step' goal (fragment.step goal tactic)
+      let fragments := moreFragments.fold (init := state.fragments.erase goal) λ acc mvarId f =>
+        acc.insert mvarId f
+      return { state' with fragments } -/
+
+  let catName := match isLHSGoal? (← goal.getType) with
+    | .some _ => `conv
+    | .none => `tactic
+  -- Normal tactic without fragment
  let tactic ← match Parser.runParserCategory
    (env := ← MonadEnv.getEnv)
-    (catName := `tactic)
+    (catName := catName)
    (input := tactic)
    (fileName := ← getFileName) with
    | .ok stx => pure $ stx
@ -468,6 +472,8 @@ protected def GoalState.tryTactic (state: GoalState) (goal: MVarId) (tactic: Str
  withCapturingError do
    state.step goal tacticM (guardMVarErrors := true)

+-- Specialized Tactics
+
 protected def GoalState.tryAssign (state: GoalState) (goal: MVarId) (expr: String):
      Elab.TermElabM TacticResult := do
  state.restoreElabM
@ -480,8 +486,6 @@ protected def GoalState.tryAssign (state: GoalState) (goal: MVarId) (expr: Strin
    | .error error => return .parseError error
  state.tryTacticM goal $ Tactic.evalAssign expr

-- Specialized Tactics
-
 protected def GoalState.tryLet (state: GoalState) (goal: MVarId) (binderName: String) (type: String):
      Elab.TermElabM TacticResult := do
  state.restoreElabM
@ -495,55 +499,49 @@ protected def GoalState.tryLet (state: GoalState) (goal: MVarId) (binderName: St
  state.tryTacticM goal $ Tactic.evalLet binderName.toName type

 /-- Enter conv tactic mode -/
-protected def GoalState.conv (state: GoalState) (goal: MVarId):
+@[export pantograph_goal_state_conv_enter_m]
+protected def GoalState.conv (state : GoalState) (goal : MVarId) :
      Elab.TermElabM TacticResult := do
-  if state.fragment? matches .some { fragment := .conv .., .. } then
+  if let .some (.conv ..) := state.fragments[goal]? then
    return .invalidAction "Already in conv state"
  goal.checkNotAssigned `GoalState.conv
  withCapturingError do
    let (fragment, state') ← state.step' goal Fragment.enterConv
-    let fragment? := .some {
-      goal := state'.goals[0]!,
-      fragment,
-    }
    return {
      state' with
-      fragment?
+      fragments := state'.fragments.insert goal fragment
    }

 /-- Exit from `conv` mode. Resumes all goals before the mode starts and applys the conv -/
@[export pantograph_goal_state_conv_exit_m]
-protected def GoalState.convExit (state: GoalState):
+protected def GoalState.convExit (state : GoalState) (goal : MVarId):
      Elab.TermElabM TacticResult := do
-  let .some { goal, fragment } := state.fragment? |
-    return .invalidAction "Not in conv state"
-  unless fragment matches .conv .. do
+  let .some fragment@(.conv ..) := state.fragments[goal]? |
    return .invalidAction "Not in conv state"
  withCapturingError do
-    let state' ← state.step goal (fragment.exit goal)
+    let (fragments, state') ← state.step' goal (fragment.exit goal state.fragments)
    return {
      state' with
-      fragment? := .none,
+      fragments,
    }

-protected def GoalState.calcPrevRhsOf? (state: GoalState) (goal: MVarId): Option Expr := do
-  match state.fragment? with
-  | .some { goal := goal', fragment := .calc prevRhs? } =>
-    if goal == goal' then prevRhs? else .none
-  | .some _ => unreachable!
-  | .none => .none
+protected def GoalState.calcPrevRhsOf? (state : GoalState) (goal : MVarId) : Option Expr := do
+  let .some (.calc prevRhs?) := state.fragments[goal]? | .none
+  prevRhs?

@[export pantograph_goal_state_try_calc_m]
-protected def GoalState.tryCalc (state: GoalState) (goal: MVarId) (pred: String):
-      Elab.TermElabM TacticResult := do
+protected def GoalState.tryCalc (state : GoalState) (goal : MVarId) (pred : String)
+  : Elab.TermElabM TacticResult := do
  let prevRhs? := state.calcPrevRhsOf? goal
  withCapturingError do
-    let (fragment?, state') ← state.step' goal do
+    let (moreFragments, state') ← state.step' goal do
      let fragment := Fragment.calc prevRhs?
      fragment.step goal pred
+    let fragments := moreFragments.fold (init := state.fragments.erase goal) λ acc mvarId f =>
+      acc.insert mvarId f
    return {
      state' with
-      fragment?,
+      fragments,
    }

 end Pantograph
--- a/Pantograph/Serial.lean
+++ b/Pantograph/Serial.lean
@ -104,7 +104,7 @@ def goalStatePickle (goalState : GoalState) (path : System.FilePath) : IO Unit :
    }
    root,
    parentMVar?,
-    fragment?,
+    fragments,
  } := goalState
  Pantograph.pickle path (
    env.constants.map₂,
@ -116,7 +116,7 @@ def goalStatePickle (goalState : GoalState) (path : System.FilePath) : IO Unit :

    root,
    parentMVar?,
-    fragment?,
+    fragments,
  )

@[export pantograph_goal_state_unpickle_m]
@ -132,7 +132,7 @@ def goalStateUnpickle (path : System.FilePath) (env : Environment)

    root,
    parentMVar?,
-    fragment?,
+    fragments,
  ), region) ← Pantograph.unpickle (
    PHashMap Name ConstantInfo ×

@ -143,7 +143,7 @@ def goalStateUnpickle (path : System.FilePath) (env : Environment)

    MVarId ×
    Option MVarId ×
-    Option TacticFragment
+    FragmentMap
  ) path
  let env ← env.replay (Std.HashMap.ofList map₂.toList)
  let goalState := {
@ -163,7 +163,7 @@ def goalStateUnpickle (path : System.FilePath) (env : Environment)
    },
    root,
    parentMVar?,
-    fragment?,
+    fragments,
  }
  return (goalState, region)

--- a/Pantograph/Tactic/Fragment.lean
+++ b/Pantograph/Tactic/Fragment.lean
@ -15,42 +15,50 @@ namespace Pantograph
 inductive Fragment where
  | calc (prevRhs? : Option Expr)
  | conv (rhs : MVarId) (dormant : List MVarId)
+  deriving BEq
+abbrev FragmentMap := Std.HashMap MVarId Fragment
+def FragmentMap.empty : FragmentMap := Std.HashMap.emptyWithCapacity 2

 protected def Fragment.enterCalc : Elab.Tactic.TacticM Fragment := do
  return .calc .none
 protected def Fragment.enterConv : Elab.Tactic.TacticM Fragment := do
  let goal ← Elab.Tactic.getMainGoal
-  let convGoal ← goal.withContext do
+  let rhs ← goal.withContext do
    let (rhs, newGoal) ← Elab.Tactic.Conv.mkConvGoalFor (← Elab.Tactic.getMainTarget)
    Elab.Tactic.replaceMainGoal [newGoal.mvarId!]
    pure rhs.mvarId!
  let otherGoals := (← Elab.Tactic.getGoals).filter (· != goal)
-  return .conv convGoal otherGoals
+  return conv rhs otherGoals

-protected def Fragment.exit (fragment : Fragment) (goal: MVarId) : Elab.Tactic.TacticM Unit :=
+protected def Fragment.exit (fragment : Fragment) (goal : MVarId) (fragments : FragmentMap)
+  : Elab.Tactic.TacticM FragmentMap :=
  match fragment with
-  | .calc _prevRhs? => Elab.Tactic.setGoals [goal]
+  | .calc .. => do
+    Elab.Tactic.setGoals [goal]
+    return fragments.erase goal
  | .conv rhs otherGoals => do
+    -- FIXME: Only process the goals that are descendants of `goal`
+    let goals := (← Elab.Tactic.getGoals).filter λ goal => true
+      --match fragments[goal]? with
+      --| .some f => fragment == f
+      --| .none => false -- Not a conv goal from this
    -- Close all existing goals with `refl`
-    for mvarId in (← Elab.Tactic.getGoals) do
+    for mvarId in goals do
      liftM <| mvarId.refl <|> mvarId.inferInstance <|> pure ()
    Elab.Tactic.pruneSolvedGoals
-    unless (← Elab.Tactic.getGoals).isEmpty do
-      throwError "convert tactic failed, there are unsolved goals\n{Elab.goalsToMessageData (← Elab.Tactic.getGoals)}"
+    unless (← goals.filterM (·.isAssignedOrDelayedAssigned)).isEmpty do
+      throwError "convert tactic failed, there are unsolved goals\n{Elab.goalsToMessageData (goals)}"

    Elab.Tactic.setGoals $ [goal] ++ otherGoals
    let targetNew ← instantiateMVars (.mvar rhs)
    let proof ← instantiateMVars (.mvar goal)

    Elab.Tactic.liftMetaTactic1 (·.replaceTargetEq targetNew proof)
-
-structure TacticFragment where
-  -- The goal which the fragment acts on
-  goal : MVarId
-  fragment : Fragment
+    return fragments.erase goal

 protected def Fragment.step (fragment : Fragment) (goal : MVarId) (s : String)
-  : Elab.Tactic.TacticM (Option TacticFragment) := goal.withContext do
+  : Elab.Tactic.TacticM FragmentMap := goal.withContext do
+  assert! ¬ (← goal.isAssigned)
  match fragment with
  | .calc prevRhs? => do
    let .ok stx := Parser.runParserCategory
@ -103,16 +111,9 @@ protected def Fragment.step (fragment : Fragment) (goal : MVarId) (s : String)
    let goals := [ mvarBranch ] ++ remainder?.toList
    Elab.Tactic.setGoals goals
    match remainder? with
-    | .some goal => return .some { goal, fragment := .calc (.some rhs) }
-    | .none => return .none
-  | fragment@(.conv _ _) => do
-    let .ok tactic := Parser.runParserCategory
-      (env := ← MonadEnv.getEnv)
-      (catName := `conv)
-      (input := s)
-      (fileName := ← getFileName) | throwError "Could not parse `conv tactic {s}"
-    Elab.Tactic.evalTactic tactic
-    let goal ← Elab.Tactic.getMainGoal
-    return .some { goal, fragment }
+    | .some goal => return FragmentMap.empty.insert goal $ .calc (.some rhs)
+    | .none => return .empty
+  | .conv .. => do
+    throwError "Direct operation on conversion tactic parent goal is not allowed"

 end Pantograph
--- a/Repl.lean
+++ b/Repl.lean
@ -303,6 +303,7 @@ def execute (command: Protocol.Command): MainM Json := do
    let state ← getMainState
    let .some goalState := state.goalStates[args.stateId]? |
      Protocol.throw $ errorIndex s!"Invalid state index {args.stateId}"
+    -- FIXME: Allow proper conversion tactic exit even in automatic mode
    let .some goal := goalState.goals[args.goalId]? |
      Protocol.throw $ errorIndex s!"Invalid goal index {args.goalId}"
    let nextGoalState?: Except _ TacticResult ← liftTermElabM do
@ -323,7 +324,7 @@ def execute (command: Protocol.Command): MainM Json := do
      | .none, .none, .none, .none, .none, .some true, .none => do
        pure <| Except.ok <| ← goalState.conv goal
      | .none, .none, .none, .none, .none, .some false, .none => do
-        pure <| Except.ok <| ← goalState.convExit
+        pure <| Except.ok <| ← goalState.convExit goal
      | .none, .none, .none, .none, .none, .none, .some draft => do
        pure <| Except.ok <| ← goalState.tryDraft goal draft
      | _, _, _, _, _, _, _ =>
--- a/Test/Proofs.lean
+++ b/Test/Proofs.lean
@ -362,6 +362,7 @@ def test_conv: TestM Unit := do
  addTest $ LSpec.check tactic ((← state1.serializeGoals (options := ← read)).map (·.devolatilize) =
    #[interiorGoal [] "a + b + c1 = b + a + c2"])

+  let goalConv := state1.goals[0]!
  let state2 ← match ← state1.conv (state1.get! 0) with
    | .success state _ => pure state
    | other => do
@ -433,7 +434,7 @@ def test_conv: TestM Unit := do
  addTest $ LSpec.check s!"  · {convTactic}" ((← state6.serializeGoals (options := ← read)).map (·.devolatilize) =
    #[])

-  let state1_1 ← match ← state6.convExit with
+  let state1_1 ← match ← state6.convExit goalConv with
    | .success state _ => pure state
    | other => do
      addTest $ assertUnreachable $ other.toString
--- a/doc/repl.md
+++ b/doc/repl.md
@ -36,8 +36,7 @@ See `Pantograph/Protocol.lean` for a description of the parameters and return va
  - `{ "calc": <expr> }`: Execute one step of a `calc` tactic. Each step must
    be of the form `lhs op rhs`. An `lhs` of `_` indicates that it should be set
    to the previous `rhs`.
-  - `{ "conv": <bool> }`: Enter or exit conversion tactic mode. In the case of
-    exit, the goal id is ignored.
+  - `{ "conv": <bool> }`: Enter or exit conversion tactic mode.
  - `{ "draft": <expr> }`: Draft an expression with `sorry`s, turning them into goals. Coupling is not allowed.
  If the `goals` field does not exist, the tactic execution has failed. Read
  `messages` to find the reason.