fix: calc indentation and allow underscore in first relation

src/Init/NotationExtra.lean

@@ -87,9 +87,11 @@ macro:35 xs:bracketedExplicitBinders " × " b:term:35  : term => expandBrackedBi
 macro:35 xs:bracketedExplicitBinders " ×' " b:term:35 : term => expandBrackedBinders ``PSigma xs b
 end
 
+-- first step of a `calc` block
+syntax calcFirstStep := ppIndent(colGe term (" := " term)?)
 -- enforce indentation of calc steps so we know when to stop parsing them
-syntax calcStep := ppIndent(colGe term " := " withPosition(term))
-syntax calcSteps := ppLine withPosition(calcStep) ppLine withPosition((calcStep ppLine)*)
+syntax calcStep := ppIndent(colGe term " := " term)
+syntax calcSteps := ppLine withPosition(calcFirstStep) ppLine withPosition((calcStep ppLine)*)
 
 /-- Step-wise reasoning over transitive relations.
 ```
@@ -102,6 +104,23 @@ calc
 proves `a = z` from the given step-wise proofs. `=` can be replaced with any
 relation implementing the typeclass `Trans`. Instead of repeating the right-
 hand sides, subsequent left-hand sides can be replaced with `_`.
+```
+calc
+  a = b := pab
+  _ = c := pbc
+  ...
+  _ = z := pyz
+```
+It is also possible to write the *first* relation as `<lhs>\n  _ = <rhs> :=
+<proof>`. This is useful for aligning relation symbols, especially on longer:
+identifiers:
+```
+calc abc
+  _ = bce := pabce
+  _ = cef := pbcef
+  ...
+  _ = xyz := pwxyz
+```
 
 `calc` has term mode and tactic mode variants. This is the term mode variant.
 
@@ -122,6 +141,22 @@ calc
 proves `a = z` from the given step-wise proofs. `=` can be replaced with any
 relation implementing the typeclass `Trans`. Instead of repeating the right-
 hand sides, subsequent left-hand sides can be replaced with `_`.
+```
+calc
+  a = b := pab
+  _ = c := pbc
+  ...
+  _ = z := pyz
+```
+It is also possible to write the *first* relation as `<lhs>\n  _ = <rhs> :=
+<proof>`. This is useful for aligning relation symbols:
+```
+calc abc
+  _ = bce := pabce
+  _ = cef := pbcef
+  ...
+  _ = xyz := pwxyz
+```
 
 `calc` has term mode and tactic mode variants. This is the tactic mode variant,
 which supports an additional feature: it works even if the goal is `a = z'`

src/Lean/Elab/Calc.lean

@@ -67,15 +67,25 @@ where
     | .node i k as => return .node i k (← as.mapM go)
     | _ => set false; return t
 
-def getCalcSteps (steps : TSyntax ``calcSteps) : Array (TSyntax ``calcStep) :=
+def getCalcFirstStep (step0 : TSyntax ``calcFirstStep) : TermElabM (TSyntax ``calcStep) :=
+  match step0  with
+  | `(calcFirstStep| $term:term) =>
+    `(calcStep| $term = _ := rfl)
+  | `(calcFirstStep| $term := $proof) =>
+    `(calcStep| $term := $proof)
+  | _ => throwUnsupportedSyntax
+
+def getCalcSteps (steps : TSyntax ``calcSteps) : TermElabM (Array (TSyntax ``calcStep)) :=
   match steps with
-  | `(calcSteps| $step0 $rest*) => #[step0] ++ rest
+  | `(calcSteps| $step0:calcFirstStep $rest*) => do
+    let step0 ← getCalcFirstStep step0
+    pure (#[step0] ++ rest)
   | _ => unreachable!
 
 def elabCalcSteps (steps : TSyntax ``calcSteps) : TermElabM Expr := do
   let mut result? := none
   let mut prevRhs? := none
-  for step in getCalcSteps steps do
+  for step in ← getCalcSteps steps do
     let `(calcStep| $pred := $proofTerm) := step | unreachable!
     let type ← elabType <| ← do
       if let some prevRhs := prevRhs? then

tests/lean/run/1813.lean

@@ -1,4 +1,4 @@
 example [Add α] [Neg α] [Mul α] [∀ n, OfNat α n] {a b : α} :=
   calc
     4 + 5 * b = -6 + 5 * (b + 2) := sorry
-            _ = -6 + 5 * 3 := sorry
+    _         = -6 + 5 * 3       := sorry

tests/lean/run/860.lean

@@ -9,7 +9,7 @@ private theorem pack_loop_terminates : (n : Nat) → n / 2 < n.succ
     · rw [Nat.add_sub_self_right]
       have := pack_loop_terminates n
       calc n/2 + 1 < Nat.succ n + 1   := Nat.add_le_add_right this 1
-             _     < Nat.succ (n + 2) := Nat.succ_lt_succ (Nat.succ_lt_succ (Nat.lt_succ_self _))
+           _       < Nat.succ (n + 2) := Nat.succ_lt_succ (Nat.succ_lt_succ (Nat.lt_succ_self _))
     · apply Nat.zero_lt_succ
 
 def pack (n: Nat) : List Nat :=

tests/lean/run/calc.lean

@@ -4,8 +4,8 @@ variable (pf12 : t1 = t2) (pf23 : t2 = t3) (pf34 : t3 = t4)
 theorem foo : t1 = t4 :=
   calc
     t1 = t2 := pf12
-     _ = t3 := pf23
-     _ = t4 := pf34
+    _  = t3 := pf23
+    _  = t4 := pf34
 
 variable (t5 : Nat)
 variable (pf23' : t2 < t3) (pf45' : t4 < t5)
@@ -13,11 +13,67 @@ variable (pf23' : t2 < t3) (pf45' : t4 < t5)
 instance [LT α] : Trans (α := α) (· < ·) (· < ·) (· < ·) where
   trans := sorry
 
-theorem foo' : t1 < t5 :=
+theorem foo₁ : t1 < t5 :=
   let p := calc
-   t1 = t2 := pf12
+    t1 = t2 := pf12
+    _  < t3 := pf23'
+    _  = t4 := pf34
+    _  < t5 := pf45'
+  -- dedent terminates the block
+  p
+
+-- same-line `calc <first relation>` with normal indent afterwards
+theorem foo₂ : t1 < t5 :=
+  calc t1 = t2 := pf12
+    _  < t3 := pf23'
+    _  = t4 := pf34
+    _  < t5 := pf45'
+
+-- `calc <first relation LHS>\n<indent><relation and relation RHS>`
+theorem foo₃ : t1 < t5 :=
+  calc t1
+      = t2 := pf12
+    _ < t3 := pf23'
+    _ = t4 := pf34
+    _ < t5 := pf45'
+
+-- `calc <first relation LHS>\n<indent><relation and relation RHS>`
+theorem foo₄ : t1 < t5 :=
+  calc t1 = t2 := pf12
+       _  < t3 := pf23'
+       _  = t4 := pf34
+       _  < t5 := pf45'
+
+-- `by` with indented sequence of tactics in `calc`-item RHS
+theorem foo₅ : t1 = t4 :=
+  calc
+    t1 = t2 := pf12
+    _  = t3 := by
+      skip
+      skip
+      exact pf23
+    _  = t4 := pf34
+
+-- function application with indented argument in `calc`-item RHS
+theorem foo₆ : t1 = t4 :=
+  calc
+    t1 = t2 := pf12
+    _  = t3 := id
+      pf23
+    _  = t4 := pf34
+
+-- `calc <first relation LHS>\n<indent>_ <rel> <rhs> := <proof>` (term)
+theorem foo₇ : t1 < t5 :=
+  calc t1
+    _ = t2 := pf12
+    _ < t3 := pf23'
+    _ = t4 := pf34
+    _ < t5 := pf45'
+
+-- `calc <first relation LHS>\n<indent>_ <rel> <rhs> := <proof>` (tactic)
+theorem foo₈ : t1 < t5 := by
+  calc t1
+    _ = t2 := pf12
     _ < t3 := pf23'
     _ = t4 := pf34
     _ < t5 := pf45'
-  -- dedent terminates the block
-  p

tests/lean/run/calcBug.lean

@@ -9,7 +9,7 @@ theorem zero_add (n : Nat) : 0 + n = n :=
    (fun (n : Nat) (ih : 0 + n = n) =>
     show 0 + succ n = succ n from
     calc
-       0 + succ n = succ (0 + n) := rfl
-                _ = succ n       := by rw [ih])
+      0 + succ n = succ (0 + n) := rfl
+      _          = succ n       := by rw [ih])
 
 end Hidden