fix: unify types in calc

src/Lean/Elab/Calc.lean

@@ -43,24 +43,62 @@ def mkCalcTrans (result resultType step stepType : Expr) : MetaM (Expr × Expr)
     return (result, resultType)
   | _ => throwError "invalid 'calc' step, failed to synthesize `Trans` instance{indentExpr selfType}"
 
+/--
+Adds a type annotation to a hole that occurs immediately at the beginning of the term.
+This is so that coercions can trigger when elaborating the term.
+See https://github.com/leanprover/lean4/issues/2040 for futher rationale.
+
+- `_ < 3` is annotated
+- `(_) < 3` is not, because it occurs after an atom
+- in `_ < _` only the first one is annotated
+- `_ + 2 < 3` is annotated (not the best heuristic, ideally we'd like to annotate `_ + 2`)
+- `lt _ 3` is not, because it occurs after an identifier
+-/
+partial def annotateFirstHoleWithType (t : Syntax) (type : Expr) : TermElabM Syntax :=
+  -- The state is true if we should annotate the immediately next hole with the type.
+  StateT.run' (go t) true
+where
+  go (t : Syntax) := do
+    unless ← get do return t
+    match t with
+    | .node _ ``Lean.Parser.Term.hole _ =>
+      set false
+      `(($(⟨t⟩) : $(← exprToSyntax type)))
+    | .node i k as => return .node i k (← as.mapM go)
+    | _ => set false; return t
+
 /--
   Elaborate `calc`-steps
 -/
-def elabCalcSteps (steps : Array Syntax) : TermElabM Expr := do
+def elabCalcSteps (steps : Array Syntax) (goalType? : Option Expr)
+    (enforceLastRhs := true) : TermElabM Expr := do
+  let (goalLhs?, goalRhs?) ← id do
+    if let some goalType := goalType? then
+      unless goalType.getAppFn.isMVar do
+        if let some (_, goalLhs, goalRhs) ← getCalcRelation? goalType then
+          return (some goalLhs, some goalRhs)
+    return (none, none)
   let mut proofs := #[]
   let mut types  := #[]
+  let mut prevRhs? := goalLhs?
   for step in steps do
-    let type  ← elabType step[0]
-    let some (_, lhs, _) ← getCalcRelation? type |
+    let mut pred := step[0]
+    if let some prevRhs := prevRhs? then
+      pred ← annotateFirstHoleWithType pred (← inferType prevRhs)
+    let type ← elabType pred
+    let some (_, lhs, rhs) ← getCalcRelation? type |
       throwErrorAt step[0] "invalid 'calc' step, relation expected{indentExpr type}"
-    if types.size > 0 then
-      let some (_, _, prevRhs) ← getCalcRelation? types.back | unreachable!
+    if let some prevRhs := prevRhs? then
       unless (← isDefEqGuarded lhs prevRhs) do
-        throwErrorAt step[0] "invalid 'calc' step, left-hand-side is {indentD m!"{lhs} : {← inferType lhs}"}\nprevious right-hand-side is{indentD m!"{prevRhs} : {← inferType prevRhs}"}"
+        throwErrorAt step[0] "invalid 'calc' step, left-hand-side is{indentD m!"{lhs} : {← inferType lhs}"}\nexpected{indentD m!"{prevRhs} : {← inferType prevRhs}"}"
     types := types.push type
     let proof ← elabTermEnsuringType step[2] type
     synthesizeSyntheticMVars
     proofs := proofs.push proof
+    prevRhs? := rhs
+  if let (some prevRhs, some goalRhs, some prevStep) := (prevRhs?, goalRhs?, steps.back?) then
+    if enforceLastRhs && !(← isDefEqGuarded prevRhs goalRhs) then
+      throwErrorAt prevStep "invalid 'calc' step, right-hand-side is {indentD m!"{prevRhs} : {← inferType prevRhs}"}\nexpected{indentD m!"{goalRhs} : {← inferType goalRhs}"}"
   let mut result := proofs[0]!
   let mut resultType := types[0]!
   for i in [1:proofs.size] do
@@ -71,5 +109,7 @@ def elabCalcSteps (steps : Array Syntax) : TermElabM Expr := do
 @[builtin_term_elab «calc»]
 def elabCalc : TermElab :=  fun stx expectedType? => do
   let steps := #[stx[1]] ++ stx[2].getArgs
-  let result ← elabCalcSteps steps
+  if let some expectedType := expectedType? then
+    tryPostponeIfMVar expectedType
+  let result ← elabCalcSteps steps expectedType?
   ensureHasType expectedType? result

src/Lean/Elab/Tactic/Calc.lean

@@ -9,7 +9,7 @@ import Lean.Elab.Tactic.ElabTerm
 namespace Lean.Elab.Tactic
 open Meta
 
-def elabCalcSteps (steps : Array Syntax) : TacticM Expr := do
+def elabCalcSteps (steps : Array Syntax) (target : Expr) : TacticM Expr := do
   /- If error recovery is disabled, we disable `Term.withoutErrToSorry` -/
   if (← read).recover then
     go
@@ -17,7 +17,7 @@ def elabCalcSteps (steps : Array Syntax) : TacticM Expr := do
     Term.withoutErrToSorry go
 where
   go : TermElabM Expr := do
-    let e ← Term.elabCalcSteps steps
+    let e ← Term.elabCalcSteps steps target (enforceLastRhs := false)
     Term.synthesizeSyntheticMVars (mayPostpone := false)
     instantiateMVars e
 
@@ -27,9 +27,9 @@ def evalCalc : Tactic := fun stx => do
   withMainContext do
     let steps := #[stx[1]] ++ stx[2].getArgs
     let (val, mvarIds) ← withCollectingNewGoalsFrom (tagSuffix := `calc) do
-      let val ← elabCalcSteps steps
-      let valType ← inferType val
       let target ← getMainTarget
+      let val ← elabCalcSteps steps target
+      let valType ← inferType val
       if (← isDefEq valType target) then
         return val
       else

tests/lean/2040.lean

@@ -0,0 +1,36 @@
+example (n : Nat) (a : Int) : a = 22 :=
+  calc
+    a = 2 ^ n := sorry -- error
+    _ = (22 : Int) := sorry
+
+example (n : Nat) (a : Int) : a = 22 :=
+  calc
+    a = (37 : Int) := sorry
+    _ = 2 ^ n := sorry -- should be same error as above
+    _ = (22 : Int) := sorry
+
+example (n : Nat) (a : Int) : a = (2 : Int) ^ n :=
+  calc
+    a = (37 : Int) := sorry
+    _ = 2 ^ n := sorry -- could be an error, but here unification figures out that (2 : Int) from the goal
+
+example (n : Nat) (h : n = 42) : 42 = (n : Int) :=
+  calc
+    (42 : Int) = 42 := rfl
+    _ = n := h ▸ rfl
+      --^ coe needs to be inserted here
+
+example (n : Nat) (h : n = 42) : 42 = (n : Int) :=
+  calc
+    42 = 42 := rfl -- type of 42 should match goal (i.e., `Int`)
+    _ = n := h ▸ rfl
+      --^ coe needs to be inserted here
+
+example (n : Nat) (h : n = 42) : 42 = (n : Int) :=
+  calc
+    _ = 42 := rfl -- type of 42 should match goal (i.e., `Int`)
+    _ = n := h ▸ rfl
+      --^ coe needs to be inserted here
+
+example := calc 1 = 1 := rfl
+example := by calc 1 = 1 := rfl

tests/lean/2040.lean.expected.out

@@ -0,0 +1,5 @@
+2040.lean:3:8-3:13: error: failed to synthesize instance
+  HPow Nat Nat Int
+2040.lean:9:8-9:13: error: failed to synthesize instance
+  HPow Nat Nat Int
+2040.lean:12:0-12:7: warning: declaration uses 'sorry'

tests/lean/calcErrors.lean

@@ -33,6 +33,6 @@ infix:50 "===" => HEqRel
 instance {α β γ} : Trans (HEqRel : α → β → Prop) (HEqRel : β → γ → Prop) (HEqRel : α → γ → Prop) where
   trans h₁ h₂ := HEq.trans h₁ h₂
 
-theorem ex7 {a : α} {b : β} {c : γ} (h₁ : a ≅ b) (h₂ : b ≅ c) : a ≅ c :=
+theorem ex7 {a : α} {b : β} {c : γ} (h₁ : a ≅ b) (h₂ : b ≅ c) : a === c :=
   calc a === b := h₁
        _ === c := h₂

tests/lean/calcErrors.lean.expected.out

@@ -4,15 +4,15 @@ has type
   b + b = b + b : Prop
 but is expected to have type
   b + b = 0 + c + b : Prop
-calcErrors.lean:7:8-7:29: error: invalid 'calc' step, left-hand-side is 
+calcErrors.lean:7:8-7:29: error: invalid 'calc' step, left-hand-side is
   0 + c + b : Nat
-previous right-hand-side is
+expected
   b + b : Nat
 calcErrors.lean:15:8-15:11: error: invalid 'calc' step, relation expected
   p a
 calcErrors.lean:20:8-20:19: error: invalid 'calc' step, failed to synthesize `Trans` instance
   Trans p p ?m
-calcErrors.lean:27:7-27:12: error: invalid 'calc' step, left-hand-side is 
+calcErrors.lean:26:7-26:12: error: invalid 'calc' step, left-hand-side is
+  β : Sort u_1
+expected
   γ : Sort u_1
-previous right-hand-side is
-  b : β