From 7837ddff4f17ffd9e32a93e8439cb903d1607ad6 Mon Sep 17 00:00:00 2001
From: Leni Aniva <v@leni.sh>
Date: Tue, 24 Jun 2025 13:51:33 -0700
Subject: [PATCH] feat(goal): Unify fragments in replay

---
 Pantograph/Goal.lean            |  9 ++++++++-
 Pantograph/Tactic/Fragment.lean | 10 ++++++++++
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/Pantograph/Goal.lean b/Pantograph/Goal.lean
index 88a7f0e..a02e875 100644
--- a/Pantograph/Goal.lean
+++ b/Pantograph/Goal.lean
@@ -218,6 +218,7 @@ protected def GoalState.replay (dst : GoalState) (src src' : GoalState) : CoreM
     | .mvar { name } => .mvar ⟨mapId name⟩
     | l => l
   let mapExpr (e : Expr) : CoreM Expr := Core.transform e λ
+    | .sort level => pure $ .done $ .sort (mapLevel level)
     | .mvar { name } => pure $ .done $ .mvar ⟨mapId name⟩
     | _ => pure .continue
   let mapDelayedAssignment (d : DelayedMetavarAssignment) : CoreM DelayedMetavarAssignment := do
@@ -335,9 +336,14 @@ protected def GoalState.replay (dst : GoalState) (src src' : GoalState) : CoreM
         throwError "Conflicting assignment of expr metavariable (d != d) {mvarId.name}"
 
     Meta.saveState
-  -- FIXME: Handle fragments
   let goals :=dst.savedState.tactic.goals ++
     src'.savedState.tactic.goals.map (⟨mapId ·.name⟩)
+  let fragments ← src'.fragments.foldM (init := dst.fragments) λ acc mvarId' fragment' => do
+    let mvarId := ⟨mapId mvarId'.name⟩
+    let fragment ← fragment'.map mapExpr
+    if let .some _fragment0 := acc[mvarId]? then
+      throwError "Conflicting fragments on {mvarId.name}"
+    return acc.insert mvarId fragment
   return {
     dst with
     savedState := {
@@ -349,6 +355,7 @@ protected def GoalState.replay (dst : GoalState) (src src' : GoalState) : CoreM
         meta := ← m.run',
       },
     },
+    fragments,
   }
 
 --- Tactic execution functions ---
diff --git a/Pantograph/Tactic/Fragment.lean b/Pantograph/Tactic/Fragment.lean
index 0fbbb03..23ec332 100644
--- a/Pantograph/Tactic/Fragment.lean
+++ b/Pantograph/Tactic/Fragment.lean
@@ -19,6 +19,16 @@ inductive Fragment where
 abbrev FragmentMap := Std.HashMap MVarId Fragment
 def FragmentMap.empty : FragmentMap := Std.HashMap.emptyWithCapacity 2
 
+protected def Fragment.map (fragment : Fragment) (mapExpr : Expr → CoreM Expr) : CoreM Fragment :=
+  let mapMVarId (mvarId : MVarId) : CoreM MVarId :=
+    return (← mapExpr (.mvar mvarId)) |>.mvarId!
+  match fragment with
+  | .calc prevRhs? => return .calc (← prevRhs?.mapM mapExpr)
+  | .conv rhs dormant => do
+    let rhs' ← mapMVarId rhs
+    let dormant' ← dormant.mapM mapMVarId
+    return .conv rhs' dormant'
+
 protected def Fragment.enterCalc : Elab.Tactic.TacticM Fragment := do
   return .calc .none
 protected def Fragment.enterConv : Elab.Tactic.TacticM Fragment := do