From ab0ef229acaec008f94509d8aa69d1b78428c539 Mon Sep 17 00:00:00 2001
From: Daniel Fabian <daniel.fabian@integral-it.ch>
Date: Tue, 18 May 2021 05:08:28 +0000
Subject: [PATCH] feat: add `getBelowIndices`.

---
 src/Lean/Meta/IndPredBelow.lean    | 25 +++++++++++++++++++++++++
 tests/lean/run/inductive_pred.lean | 19 +++++++++++++++++++
 2 files changed, 44 insertions(+)

diff --git a/src/Lean/Meta/IndPredBelow.lean b/src/Lean/Meta/IndPredBelow.lean
index 5a466055fb..d66bb57a2a 100644
--- a/src/Lean/Meta/IndPredBelow.lean
+++ b/src/Lean/Meta/IndPredBelow.lean
@@ -348,6 +348,31 @@ where
         mkForallFVars ys (←mkArrow premise conclusion)
     (←name, mkDomain)
 
+/-- Given a constructor name, find the indices of the corresponding `below` version thereof. -/
+partial def getBelowIndices (ctorName : Name) : MetaM $ Array Nat := do
+  let ctorInfo ← getConstInfoCtor ctorName
+  let belowCtorInfo ← getConstInfoCtor (ctorName.updatePrefix $ ctorInfo.induct ++ `below)
+  let belowInductInfo ← getConstInfoInduct belowCtorInfo.induct
+  forallTelescope ctorInfo.type fun xs t => do
+  loop xs belowCtorInfo.type #[] 0 0 
+
+where 
+  loop 
+      (xs : Array Expr) 
+      (rest : Expr) 
+      (belowIndices : Array Nat) 
+      (xIdx yIdx : Nat) : MetaM $ Array Nat := do 
+    if xIdx ≥ xs.size then belowIndices else
+    let x := xs[xIdx]
+    let xTy ← inferType x
+    let yTy := rest.bindingDomain!
+    if ←isDefEq xTy yTy then 
+      let rest ← instantiateForall rest #[x]
+      loop xs rest (belowIndices.push yIdx) (xIdx + 1) (yIdx + 1)
+    else 
+      forallBoundedTelescope rest (some 1) fun ys rest =>
+      loop xs rest belowIndices xIdx (yIdx + 1)
+
 def mkBelow (declName : Name) : MetaM Unit := do
   if (←isInductivePredicate declName) then
     let x ← getConstInfoInduct declName
diff --git a/tests/lean/run/inductive_pred.lean b/tests/lean/run/inductive_pred.lean
index 8bb7fd05c1..2b6ed21a09 100644
--- a/tests/lean/run/inductive_pred.lean
+++ b/tests/lean/run/inductive_pred.lean
@@ -1,7 +1,17 @@
+import Lean 
+open Lean
+
+def checkGetBelowIndices (ctorName : Name) (indices : Array Nat) : MetaM Unit := do
+  let actualIndices ← Meta.IndPredBelow.getBelowIndices ctorName
+  if actualIndices != indices then
+    throwError "wrong indices for {ctorName}: {actualIndices} ≟ {indices}"
+
 namespace Ex
 inductive LE : Nat → Nat → Prop
   | refl : LE n n
   | succ : LE n m → LE n m.succ
+#eval checkGetBelowIndices ``LE.refl #[1]
+#eval checkGetBelowIndices ``LE.succ #[1, 2, 3]
 
 def typeOf {α : Sort u} (a : α) := α
 
@@ -22,6 +32,8 @@ theorem LE.trans' : LE m n → LE n o → LE m o
 inductive Even : Nat → Prop
   | zero : Even 0
   | ss   : Even n → Even n.succ.succ
+#eval checkGetBelowIndices ``Even.zero #[]
+#eval checkGetBelowIndices ``Even.ss #[1, 2]
 
 theorem Even_brecOn : typeOf @Even.brecOn = ∀ {motive : (a : Nat) → Even a → Prop} {a : Nat} (x : Even a),
   (∀ (a : Nat) (x : Even a), @Even.below motive a x → motive a x) → motive a x := rfl
@@ -42,6 +54,8 @@ theorem mul_left_comm (n m o : Nat) : n * (m * o) = m * (n * o) := by
 inductive Power2 : Nat → Prop
   | base : Power2 1
   | ind  : Power2 n → Power2 (2*n) -- Note that index here is not a constructor
+#eval checkGetBelowIndices ``Power2.base #[]
+#eval checkGetBelowIndices ``Power2.ind #[1, 2]
 
 theorem Power2_brecOn : typeOf @Power2.brecOn = ∀ {motive : (a : Nat) → Power2 a → Prop} {a : Nat} (x : Power2 a),
   (∀ (a : Nat) (x : Power2 a), @Power2.below motive a x → motive a x) → motive a x := rfl
@@ -75,6 +89,9 @@ inductive step : tm → tm → Prop :=
   | ST_Plus2 : ∀ n1 t2 t2',
       t2 ==> t2' →
       P (C n1) t2 ==> P (C n1) t2'
+#eval checkGetBelowIndices ``step.ST_PlusConstConst #[1, 2]
+#eval checkGetBelowIndices ``step.ST_Plus1 #[1, 2, 3, 4]
+#eval checkGetBelowIndices ``step.ST_Plus2 #[1, 2, 3, 4]
 
 def deterministic {X : Type} (R : X → X → Prop) :=
   ∀ x y1 y2 : X, R x y1 → R x y2 → y1 = y2
@@ -96,6 +113,8 @@ axiom f : Nat → Nat
 inductive is_nat : Nat -> Prop
 | Z : is_nat 0
 | S {n} : is_nat n → is_nat (f n)
+#eval checkGetBelowIndices ``is_nat.Z #[]
+#eval checkGetBelowIndices ``is_nat.S #[1, 2]
 
 axiom P : Nat → Prop
 axiom F0 : P 0