From 5720c725154cb54b3db6b8a6e6a6a1d6e0082192 Mon Sep 17 00:00:00 2001
From: Leni Aniva <v@leni.sh>
Date: Tue, 30 Jan 2024 17:22:20 -0800
Subject: [PATCH] feat: Prevent crash during rootExpr call

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

diff --git a/Pantograph/Goal.lean b/Pantograph/Goal.lean
index bd28944..630637d 100644
--- a/Pantograph/Goal.lean
+++ b/Pantograph/Goal.lean
@@ -210,8 +210,8 @@ protected def GoalState.continue (target: GoalState) (branch: GoalState): Except
   else
     target.resume (goals := branch.goals)
 
-protected def GoalState.rootExpr? (goalState: GoalState): Option Expr :=
-  let expr := goalState.mctx.eAssignment.find! goalState.root
+protected def GoalState.rootExpr? (goalState: GoalState): Option Expr := do
+  let expr ← goalState.mctx.eAssignment.find? goalState.root
   let (expr, _) := instantiateMVarsCore (mctx := goalState.mctx) (e := expr)
   if expr.hasMVar then
     -- Must not assert that the goal state is empty here. We could be in a branch goal.
@@ -219,7 +219,7 @@ protected def GoalState.rootExpr? (goalState: GoalState): Option Expr :=
     .none
   else
     assert! goalState.goals.isEmpty
-    .some expr
+    return expr
 protected def GoalState.parentExpr? (goalState: GoalState): Option Expr := do
   let parent ← goalState.parentMVar
   let expr := goalState.mctx.eAssignment.find! parent
diff --git a/Test/Proofs.lean b/Test/Proofs.lean
index 224bb22..07e4cea 100644
--- a/Test/Proofs.lean
+++ b/Test/Proofs.lean
@@ -210,6 +210,8 @@ def proof_or_comm: TestM Unit := do
     | .none => do
       addTest $ assertUnreachable "Goal could not parse"
       return ()
+  addTest $ LSpec.check "(0 parent)" state0.parentExpr?.isNone
+  addTest $ LSpec.check "(0 root)" state0.rootExpr?.isNone
 
   let state1 ← match ← state0.execute (goalId := 0) (tactic := "intro p q h") with
     | .success state => pure state
@@ -218,6 +220,8 @@ def proof_or_comm: TestM Unit := do
       return ()
   addTest $ LSpec.check "intro n m" ((← state1.serializeGoals (options := ← read)).map (·.devolatilize) =
     #[buildGoal [("p", "Prop"), ("q", "Prop"), ("h", "p ∨ q")] "q ∨ p"])
+  addTest $ LSpec.check "(1 parent)" state1.parentExpr?.isSome
+  addTest $ LSpec.check "(1 root)" state1.rootExpr?.isNone
   let state2 ← match ← state1.execute (goalId := 0) (tactic := "cases h") with
     | .success state => pure state
     | other => do
@@ -225,6 +229,8 @@ def proof_or_comm: TestM Unit := do
       return ()
   addTest $ LSpec.check "cases h" ((← state2.serializeGoals (options := ← read)).map (·.devolatilize) =
     #[branchGoal "inl" "p", branchGoal "inr" "q"])
+  addTest $ LSpec.check "(2 parent)" state2.parentExpr?.isSome
+  addTest $ LSpec.check "(2 root)" state2.rootExpr?.isNone
 
   let state3_1 ← match ← state2.execute (goalId := 0) (tactic := "apply Or.inr") with
     | .success state => pure state