From 1ea92baa21f4d9612f12736fa2b7ae9f36ffd647 Mon Sep 17 00:00:00 2001
From: Arthur Adjedj <arthur.adjedj@ens-paris-saclay.fr>
Date: Mon, 6 May 2024 07:58:15 +0200
Subject: [PATCH] fix: Incorrect promotion from index to paramater (#3591)

Depends on #3590

Closes #3458
---
 src/Lean/Elab/Inductive.lean        | 17 +++++++++++++++--
 src/Lean/Util/ForEachExprWhere.lean | 10 +++++++---
 tests/lean/run/3458_2.lean          | 17 +++++++++++++++++
 3 files changed, 39 insertions(+), 5 deletions(-)
 create mode 100644 tests/lean/run/3458_2.lean

diff --git a/src/Lean/Elab/Inductive.lean b/src/Lean/Elab/Inductive.lean
index b9d9242c9850..cc7e2d688d88 100644
--- a/src/Lean/Elab/Inductive.lean
+++ b/src/Lean/Elab/Inductive.lean
@@ -686,8 +686,8 @@ private def computeFixedIndexBitMask (numParams : Nat) (indType : InductiveType)
               maskRef.modify fun mask => mask.set! i false
           for x in xs[numParams:] do
             let xType ← inferType x
-            let cond (e : Expr) := indFVars.any (fun indFVar => e.getAppFn == indFVar) && e.getAppNumArgs > numParams
-            xType.forEachWhere cond fun e => do
+            let cond (e : Expr) := indFVars.any (fun indFVar => e.getAppFn == indFVar)
+            xType.forEachWhere (stopWhenVisited := true) cond fun e => do
               let eArgs := e.getAppArgs
               for i in [numParams:eArgs.size] do
                 if i >= typeArgs.size then
@@ -695,6 +695,19 @@ private def computeFixedIndexBitMask (numParams : Nat) (indType : InductiveType)
                 else
                   unless eArgs[i]! == typeArgs[i]! do
                     maskRef.modify (resetMaskAt · i)
+              /-If an index is missing in the arguments of the inductive type, then it must be non-fixed.
+                Consider the following example:
+                ```lean
+                inductive All {I : Type u} (P : I → Type v) : List I → Type (max u v) where
+                  | cons : P x → All P xs → All P (x :: xs)
+
+                inductive Iμ {I : Type u}  : I → Type (max u v) where
+                  | mk : (i : I)  → All Iμ [] → Iμ i
+                ```
+                because `i` doesn't appear in `All Iμ []`, the index shouldn't be fixed.
+              -/
+              for i in [eArgs.size:arity] do
+                maskRef.modify (resetMaskAt · i)
         go ctors
     go indType.ctors
 
diff --git a/src/Lean/Util/ForEachExprWhere.lean b/src/Lean/Util/ForEachExprWhere.lean
index 202bbfb016a6..e38bf0c56bc3 100644
--- a/src/Lean/Util/ForEachExprWhere.lean
+++ b/src/Lean/Util/ForEachExprWhere.lean
@@ -58,7 +58,7 @@ def checked (e : Expr) : ForEachM m Bool := do
     return false
 
 /-- `Expr.forEachWhere` (unsafe) implementation -/
-unsafe def visit (p : Expr → Bool) (f : Expr → m Unit) (e : Expr) : m Unit := do
+unsafe def visit (p : Expr → Bool) (f : Expr → m Unit) (e : Expr) (stopWhenVisited : Bool := false) : m Unit := do
   go e |>.run' initCache
 where
   go (e : Expr) : StateRefT' ω State m Unit := do
@@ -66,6 +66,8 @@ where
       if p e then
         unless (← checked e) do
           f e
+          if stopWhenVisited then
+            return ()
       match e with
       | .forallE _ d b _   => go d; go b
       | .lam _ d b _       => go d; go b
@@ -78,9 +80,11 @@ where
 end ForEachExprWhere
 
 /--
-`e.forEachWhere p f` applies `f` to each subterm that satisfies `p`.
+  `e.forEachWhere p f` applies `f` to each subterm that satisfies `p`.
+  If `stopWhenVisited` is `true`, the function doesn't visit subterms of terms
+  which satisfy `p`.
 -/
 @[implemented_by ForEachExprWhere.visit]
-opaque Expr.forEachWhere {ω : Type} {m : Type → Type} [STWorld ω m] [MonadLiftT (ST ω) m] [Monad m] (p : Expr → Bool) (f : Expr → m Unit) (e : Expr) : m Unit
+opaque Expr.forEachWhere {ω : Type} {m : Type → Type} [STWorld ω m] [MonadLiftT (ST ω) m] [Monad m] (p : Expr → Bool) (f : Expr → m Unit) (e : Expr) (stopWhenVisited : Bool := false) : m Unit
 
 end Lean
diff --git a/tests/lean/run/3458_2.lean b/tests/lean/run/3458_2.lean
new file mode 100644
index 000000000000..6945757c4027
--- /dev/null
+++ b/tests/lean/run/3458_2.lean
@@ -0,0 +1,17 @@
+/- Ensure that non-fixed indices don't get mistakenly promoted to parameter
+   See issue `https://github.com/leanprover/lean4/issues/3458`
+ -/
+
+set_option autoImplicit false
+universe u v w
+
+structure ISignature (I : Type u) : Type (max u v + 1) where
+  symbols : I → Type v
+  indices : {i : I} → symbols i → List I
+
+inductive All {I : Type u} (P : I → Type v) : List I → Type (max u v) where
+  | nil : All P []
+  | cons {x xs}: P x → All P xs → All P (x :: xs)
+
+inductive Iμ {I : Type u} (ζ : ISignature.{u,v} I) : I → Type (max u v) where
+  | mk : (i : I) → (s : ζ.symbols i) → All (Iμ ζ) (ζ.indices s) → Iμ ζ i