From 036fab0ad6a5e137059c634e9931ea7e1724c060 Mon Sep 17 00:00:00 2001
From: Leni Aniva <v@leni.sh>
Date: Thu, 11 Apr 2024 16:11:06 -0700
Subject: [PATCH] fix: Prevent incorrect inheritance of calc rhs

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

diff --git a/Pantograph/Goal.lean b/Pantograph/Goal.lean
index fbb850a..75dc3f3 100644
--- a/Pantograph/Goal.lean
+++ b/Pantograph/Goal.lean
@@ -325,6 +325,11 @@ protected def GoalState.convExit (state: GoalState):
   catch exception =>
     return .failure #[← exception.toMessageData.toString]
 
+protected def GoalState.calcPrevRhsOf? (state: GoalState) (goalId: Nat) :=
+  if goalId == 1 then
+    state.calcPrevRhs?
+  else
+    .none
 protected def GoalState.tryCalc (state: GoalState) (goalId: Nat) (pred: String):
       Elab.TermElabM TacticResult := do
   state.restoreElabM
@@ -340,20 +345,21 @@ protected def GoalState.tryCalc (state: GoalState) (goalId: Nat) (pred: String):
     (fileName := filename) with
     | .ok syn => pure syn
     | .error error => return .parseError error
+  let calcPrevRhs? := state.calcPrevRhsOf? goalId
   try
     goal.withContext do
     let target ← instantiateMVars (← goal.getDecl).type
     let tag := (← goal.getDecl).userName
 
     let mut step ← Elab.Term.elabType <| ← do
-      if let some prevRhs := state.calcPrevRhs? then
+      if let some prevRhs := calcPrevRhs? then
         Elab.Term.annotateFirstHoleWithType pred (← Meta.inferType prevRhs)
       else
         pure pred
 
     let some (_, lhs, rhs) ← Elab.Term.getCalcRelation? step |
       throwErrorAt pred "invalid 'calc' step, relation expected{indentExpr step}"
-    if let some prevRhs := state.calcPrevRhs? then
+    if let some prevRhs := calcPrevRhs? then
       unless (← Meta.isDefEqGuarded lhs prevRhs) do
         throwErrorAt pred "invalid 'calc' step, left-hand-side is{indentD m!"{lhs} : {← Meta.inferType lhs}"}\nprevious right-hand-side is{indentD m!"{prevRhs} : {← Meta.inferType prevRhs}"}" -- "
 
diff --git a/Test/Proofs.lean b/Test/Proofs.lean
index 9ede63e..1b1b9de 100644
--- a/Test/Proofs.lean
+++ b/Test/Proofs.lean
@@ -510,6 +510,8 @@ def test_calc: TestM Unit := do
       interiorGoal [] "a + b = b + c" (.some "calc"),
       interiorGoal [] "b + c = c + d"
     ])
+  addTest $ LSpec.test "(2.0 prev rhs)" (state2.calcPrevRhsOf? 0 |>.isNone)
+  addTest $ LSpec.test "(2.1 prev rhs)" (state2.calcPrevRhsOf? 1 |>.isSome)
 
   let tactic := "apply h1"
   let state2m ← match ← state2.tryTactic (goalId := 0) (tactic := tactic) with
@@ -532,6 +534,7 @@ def test_calc: TestM Unit := do
     #[
       interiorGoal [] "b + c = c + d" (.some "calc")
     ])
+  addTest $ LSpec.test "(4.0 prev rhs)" (state4.calcPrevRhsOf? 0 |>.isNone)
   let tactic := "apply h2"
   let state4m ← match ← state4.tryTactic (goalId := 0) (tactic := tactic) with
     | .success state => pure state