From dd7c6c36c8e40d29c8bbeea3b51072fab6643746 Mon Sep 17 00:00:00 2001
From: Leni Aniva <v@leni.sh>
Date: Mon, 30 Jun 2025 15:00:30 -0700
Subject: [PATCH] fix(goal): Allow parent expr to be fragments

---
 Pantograph/Goal.lean | 13 ++++++++++---
 Repl.lean            |  9 +++++++--
 Test/Serial.lean     |  2 +-
 3 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/Pantograph/Goal.lean b/Pantograph/Goal.lean
index 9e4bbba..e646448 100644
--- a/Pantograph/Goal.lean
+++ b/Pantograph/Goal.lean
@@ -64,7 +64,11 @@ structure GoalState where
   -- The root goal which is the search target
   root: MVarId
 
-  -- Parent goals assigned to produce this state
+  /--
+  Parent goals which became assigned or fragmented to produce this state.
+  Note that due to the existence of tactic fragments, parent goals do not
+  necessarily have an expression assignment.
+  -/
   parentMVars : List MVarId := []
 
   -- Any goal associated with a fragment has a partial tactic which has not
@@ -202,8 +206,11 @@ protected def GoalState.getMVarEAssignment (goalState: GoalState) (mvarId: MVarI
   let (expr, _) := instantiateMVarsCore (mctx := goalState.mctx) (e := expr)
   return expr
 @[export pantograph_goal_state_parent_exprs]
-protected def GoalState.parentExprs (state : GoalState) : List Expr :=
-  state.parentMVars.map λ goal => state.getMVarEAssignment goal |>.get!
+protected def GoalState.parentExprs (state : GoalState) : List (Except Fragment Expr) :=
+  state.parentMVars.map λ goal => match state.getMVarEAssignment goal with
+    | .some e => .ok e
+    -- A parent goal which is not assigned must have a fragment
+    | .none => .error state.fragments[goal]!
 @[always_inline]
 protected def GoalState.hasUniqueParent (state : GoalState) : Bool :=
   state.parentMVars.length == 1
diff --git a/Repl.lean b/Repl.lean
index 3811f96..963503b 100644
--- a/Repl.lean
+++ b/Repl.lean
@@ -147,10 +147,15 @@ def goal_tactic (args: Protocol.GoalTactic): EMainM Protocol.GoalTacticResult :=
   match nextGoalState? with
   | .error error => Protocol.throw error
   | .ok (.success nextGoalState messages) => do
+    let env ← getEnv
     let nextStateId ← newGoalState nextGoalState
     let parentExprs := nextGoalState.parentExprs
-    let hasSorry := parentExprs.any (·.hasSorry)
-    let hasUnsafe := parentExprs.any ((← getEnv).hasUnsafe ·)
+    let hasSorry := parentExprs.any λ
+      | .ok e => e.hasSorry
+      | .error _ => false
+    let hasUnsafe := parentExprs.any λ
+      | .ok e => env.hasUnsafe e
+      | .error _ => false
     let goals ← runCoreM $ nextGoalState.serializeGoals (options := state.options) |>.run'
     return {
       nextStateId? := .some nextStateId,
diff --git a/Test/Serial.lean b/Test/Serial.lean
index cca7dbd..05dc4af 100644
--- a/Test/Serial.lean
+++ b/Test/Serial.lean
@@ -84,7 +84,7 @@ def test_pickling_env_extensions : TestM Unit := do
     let goal := state.goals[0]!
     let type ← goal.withContext do
       let .ok type ← elabTerm (← `(term|(2: Nat) ≤ 3)) (.some $ .sort 0) | unreachable!
-      pure type
+      instantiateMVars type
     let .success state1 _ ← state.tryTacticM goal (Tactic.assignWithAuxLemma type) | unreachable!
     let parentExpr := state1.parentExpr!
     checkTrue "src has aux lemma" $ parentExpr.getUsedConstants.any (·.isAuxLemma)