From 77907fd0609eb5db081b25dabca4fbf4a5f98e00 Mon Sep 17 00:00:00 2001
From: Leni Aniva <v@leni.sh>
Date: Fri, 12 Apr 2024 21:30:56 -0700
Subject: [PATCH] feat: `goalLet` function

---
 Pantograph/Goal.lean    | 46 +++++++++++++++++++++++++++++++++++++++++
 Pantograph/Library.lean |  3 +++
 Test/Proofs.lean        | 10 ++++++---
 3 files changed, 56 insertions(+), 3 deletions(-)

diff --git a/Pantograph/Goal.lean b/Pantograph/Goal.lean
index f1c2503..484ff51 100644
--- a/Pantograph/Goal.lean
+++ b/Pantograph/Goal.lean
@@ -177,6 +177,7 @@ protected def GoalState.assign (state: GoalState) (goal: MVarId) (expr: Expr):
       },
       newMVars,
       parentMVar? := .some goal,
+      calcPrevRhs? := .none
     }
   catch exception =>
     return .failure #[← exception.toMessageData.toString]
@@ -249,6 +250,49 @@ protected def GoalState.tryHave (state: GoalState) (goalId: Nat) (binderName: St
       },
       newMVars := nextGoals.toSSet,
       parentMVar? := .some goal,
+      calcPrevRhs? := .none
+    }
+  catch exception =>
+    return .failure #[← exception.toMessageData.toString]
+protected def GoalState.tryLet (state: GoalState) (goalId: Nat) (binderName: String) (type: String):
+      Elab.TermElabM TacticResult := do
+  state.restoreElabM
+  let goal ← match state.savedState.tactic.goals.get? goalId with
+    | .some goal => pure goal
+    | .none => return .indexError goalId
+  let type ← match Parser.runParserCategory
+    (env := state.env)
+    (catName := `term)
+    (input := type)
+    (fileName := filename) with
+    | .ok syn => pure syn
+    | .error error => return .parseError error
+  let binderName := binderName.toName
+  try
+    -- Implemented similarly to the intro tactic
+    let nextGoals: List MVarId ← goal.withContext do
+      let type ← Elab.Term.elabType (stx := type)
+      let lctx ← MonadLCtx.getLCtx
+
+      -- The branch goal inherits the same context, but with a different type
+      let mvarBranch ← Meta.mkFreshExprMVarAt lctx (← Meta.getLocalInstances) type
+
+      let upstreamType := .letE binderName type mvarBranch (← goal.getType) false
+      let mvarUpstream ← Meta.mkFreshExprMVarAt (← getLCtx) (← Meta.getLocalInstances)
+        upstreamType (kind := MetavarKind.synthetic) (userName := (← goal.getTag))
+
+      goal.assign mvarUpstream
+
+      pure [mvarBranch.mvarId!, mvarUpstream.mvarId!]
+    return .success {
+      root := state.root,
+      savedState := {
+        term := ← MonadBacktrack.saveState,
+        tactic := { goals := nextGoals }
+      },
+      newMVars := nextGoals.toSSet,
+      parentMVar? := .some goal,
+      calcPrevRhs? := .none
     }
   catch exception =>
     return .failure #[← exception.toMessageData.toString]
@@ -280,6 +324,7 @@ protected def GoalState.conv (state: GoalState) (goalId: Nat):
       newMVars := newMVarSet prevMCtx nextMCtx,
       parentMVar? := .some goal,
       convMVar? := .some (convRhs, goal),
+      calcPrevRhs? := .none
     }
   catch exception =>
     return .failure #[← exception.toMessageData.toString]
@@ -318,6 +363,7 @@ protected def GoalState.convExit (state: GoalState):
       newMVars := newMVarSet prevMCtx nextMCtx,
       parentMVar? := .some convGoal,
       convMVar? := .none
+      calcPrevRhs? := .none
     }
   catch exception =>
     return .failure #[← exception.toMessageData.toString]
diff --git a/Pantograph/Library.lean b/Pantograph/Library.lean
index ff365b2..17925cd 100644
--- a/Pantograph/Library.lean
+++ b/Pantograph/Library.lean
@@ -169,6 +169,9 @@ def goalAssign (state: GoalState) (goalId: Nat) (expr: String): Lean.CoreM Tacti
 @[export pantograph_goal_have_m]
 def goalHave (state: GoalState) (goalId: Nat) (binderName: String) (type: String): Lean.CoreM TacticResult :=
   runTermElabM <| state.tryHave goalId binderName type
+@[export pantograph_goal_let_m]
+def goalLet (state: GoalState) (goalId: Nat) (binderName: String) (type: String): Lean.CoreM TacticResult :=
+  runTermElabM <| state.tryLet goalId binderName type
 
 @[export pantograph_goal_conv_m]
 def goalConv (state: GoalState) (goalId: Nat): Lean.CoreM TacticResult :=
diff --git a/Test/Proofs.lean b/Test/Proofs.lean
index 9837e04..4f81085 100644
--- a/Test/Proofs.lean
+++ b/Test/Proofs.lean
@@ -548,7 +548,7 @@ def test_calc: TestM Unit := do
         ("h1", "a + b = b + c"), ("h2", "b + c = c + d")] ++ free
       buildGoal free target userName?
 
-def test_let: TestM Unit := do
+def test_let (specialized: Bool): TestM Unit := do
   let state? ← startProof (.expr "∀ (a: Nat) (p: Prop), p → p ∨ ¬p")
   let state0 ← match state? with
     | .some state => pure state
@@ -565,7 +565,10 @@ def test_let: TestM Unit := do
     #[interiorGoal [] "p ∨ ¬p"])
 
   let expr := "let b: Nat := _; _"
-  let state2 ← match ← state1.tryAssign (goalId := 0) (expr := expr) with
+  let result2 ← match specialized with
+    | true => state1.tryLet (goalId := 0) (binderName := "b") (type := "Nat")
+    | false => state1.tryAssign (goalId := 0) (expr := expr)
+  let state2 ← match result2 with
     | .success state => pure state
     | other => do
       addTest $ assertUnreachable $ other.toString
@@ -627,7 +630,8 @@ def suite (env: Environment): List (String × IO LSpec.TestSeq) :=
     ("have", test_have),
     ("conv", test_conv),
     ("calc", test_calc),
-    ("let", test_let),
+    ("let via assign", test_let false),
+    ("let via tryLet", test_let true),
   ]
   tests.map (fun (name, test) => (name, proofRunner env test))