From ee6e2036ddef27ddea87d0abb8a83b92bd7c8857 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Thu, 28 Jul 2022 10:13:01 -0700
Subject: [PATCH] feat: allow relations in `Type` at `Trans`

It addresses issue reported at https://leanprover.zulipchat.com/#narrow/stream/270676-lean4/topic/Calc.20mode/near/291214574
---
 src/Init/Prelude.lean          |  6 +++---
 src/Lean/Elab/Calc.lean        | 17 +++++++++++++----
 tests/lean/run/calcInType.lean | 15 +++++++++++++++
 3 files changed, 31 insertions(+), 7 deletions(-)
 create mode 100644 tests/lean/run/calcInType.lean

diff --git a/src/Init/Prelude.lean b/src/Init/Prelude.lean
index 53eafaecda..05051f947f 100644
--- a/src/Init/Prelude.lean
+++ b/src/Init/Prelude.lean
@@ -451,15 +451,15 @@ class LT (α : Type u) where lt : α → α → Prop
   ite (LE.le a b) a b
 
 /-- Transitive chaining of proofs, used e.g. by `calc`. -/
-class Trans (r : α → β → Prop) (s : β → γ → Prop) (t : outParam (α → γ → Prop)) where
+class Trans (r : α → β → Sort u) (s : β → γ → Sort v) (t : outParam (α → γ → Sort w)) where
   trans : r a b → s b c → t a c
 
 export Trans (trans)
 
-instance (r : α → γ → Prop) : Trans Eq r r where
+instance (r : α → γ → Sort u) : Trans Eq r r where
   trans heq h' := heq ▸ h'
 
-instance (r : α → β → Prop) : Trans r Eq r where
+instance (r : α → β → Sort u) : Trans r Eq r where
   trans h' heq := heq ▸ h'
 
 class HAdd (α : Type u) (β : Type v) (γ : outParam (Type w)) where
diff --git a/src/Lean/Elab/Calc.lean b/src/Lean/Elab/Calc.lean
index 758f614104..9d9811deb9 100644
--- a/src/Lean/Elab/Calc.lean
+++ b/src/Lean/Elab/Calc.lean
@@ -18,16 +18,25 @@ def getCalcRelation? (e : Expr) : MetaM (Option (Expr × Expr × Expr)) :=
   else
     return some (e.appFn!.appFn!, e.appFn!.appArg!, e.appArg!)
 
+private def getRelUniv (r : Expr) : MetaM Level := do
+  let rType ← inferType r
+  forallTelescopeReducing rType fun _ sort => do
+    let .sort u ← whnf sort | throwError "unexpected relation type{indentExpr rType}"
+    return u
+
 def mkCalcTrans (result resultType step stepType : Expr) : MetaM (Expr × Expr) := do
   let some (r, a, b) ← getCalcRelation? resultType | unreachable!
   let some (s, _, c) ← getCalcRelation? (← instantiateMVars stepType) | unreachable!
+  let u ← getRelUniv r
+  let v ← getRelUniv s
   let (α, β, γ)       := (← inferType a, ← inferType b, ← inferType c)
   let (u_1, u_2, u_3) := (← getLevel α, ← getLevel β, ← getLevel γ)
-  let t ← mkFreshExprMVar (← mkArrow α (← mkArrow γ (mkSort levelZero)))
-  let selfType := mkAppN (Lean.mkConst ``Trans [u_1, u_2, u_3]) #[α, β, γ, r, s, t]
+  let w ← mkFreshLevelMVar
+  let t ← mkFreshExprMVar (← mkArrow α (← mkArrow γ (mkSort w)))
+  let selfType := mkAppN (Lean.mkConst ``Trans [u, v, w, u_1, u_2, u_3]) #[α, β, γ, r, s, t]
   match (← trySynthInstance selfType) with
-  | LOption.some self =>
-    let result := mkAppN (Lean.mkConst ``Trans.trans [u_1, u_2, u_3]) #[α, β, γ, r, s, t, self, a, b, c, result, step]
+  | .some self =>
+    let result := mkAppN (Lean.mkConst ``Trans.trans [u, v, w, u_1, u_2, u_3]) #[α, β, γ, r, s, t, self, a, b, c, result, step]
     let resultType := (← instantiateMVars (← inferType result)).headBeta
     unless (← getCalcRelation? resultType).isSome do
       throwError "invalid 'calc' step, step result is not a relation{indentExpr resultType}"
diff --git a/tests/lean/run/calcInType.lean b/tests/lean/run/calcInType.lean
new file mode 100644
index 0000000000..2cafbb5a9a
--- /dev/null
+++ b/tests/lean/run/calcInType.lean
@@ -0,0 +1,15 @@
+inductive EQ {α : Type} (a : α) : α → Type where
+  | refl : EQ a a
+
+def EQ.trans (h₁ : EQ a b) (h₂ : EQ b c) : EQ a c := by
+  cases h₁; cases h₂; constructor
+
+instance : Trans (@EQ α) (@EQ α) (@EQ α) where
+  trans := EQ.trans
+
+infix:50 " ≋ " => EQ
+
+example (h₁ : EQ a b) (h₂ : b = c) (h₃ : EQ c d) : EQ a d := by
+  calc a ≋ b := h₁
+       _ = c := h₂
+       _ ≋ d := h₃