From e0c1afd12d4fc6b0e520774959aed06bf122aba9 Mon Sep 17 00:00:00 2001
From: Arthur Adjedj <arthur.adjedj@ens-paris-saclay.fr>
Date: Mon, 6 May 2024 08:50:08 +0200
Subject: [PATCH] fix: occurence check in `mkInjectiveTheoremTypeCore?` (#3398)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Closes #3386

Currently, when generating the signature of an injectivity lemma for a
certain constructor `c : forall xs, Foo a_1 ... a_n`,
`mkInjectiveTheoremTypeCore?` will differentiate between variables which
are bound to stay the same between the two equal values (i.e inductive
indices), and non-fixed ones. To do that, the function currently checks
whether a variable `x ∈ xs` appears in the final co-domain `Foo a_1 ...
a_n` of the constructor. This condition isn't enough however. As shown
in the linked issue, the codomain may also depend on variables which
appears in the type of free vars contained in `Foo a_1 ... a_n`, but not
in the term itself. This PR fixes the issue by also checking the types
of any free variable occuring in the final codomain, so as to ensure
injectivity lemmas are well-typed.
---
 src/Lean/Meta/Injective.lean | 31 ++++++++++++++++++++++++++++++-
 tests/lean/run/3386.lean     |  9 +++++++++
 2 files changed, 39 insertions(+), 1 deletion(-)
 create mode 100644 tests/lean/run/3386.lean

diff --git a/src/Lean/Meta/Injective.lean b/src/Lean/Meta/Injective.lean
index 95e4dad4b747..7fefdd21234f 100644
--- a/src/Lean/Meta/Injective.lean
+++ b/src/Lean/Meta/Injective.lean
@@ -31,6 +31,35 @@ def elimOptParam (type : Expr) : CoreM Expr := do
     else
       return .continue
 
+
+/-- Returns true if `e` occurs either in `t`, or in the type of a sub-expression of `t`.
+  Consider the following example:
+  ```lean
+  inductive Tyₛ : Type (u+1)
+  | SPi : (T : Type u) -> (T -> Tyₛ) -> Tyₛ
+
+  inductive Tmₛ.{u} :  Tyₛ.{u} -> Type (u+1)
+  | app : Tmₛ (.SPi T A) -> (arg : T) -> Tmₛ (A arg)```
+  ```
+  When looking for fixed arguments in `Tmₛ.app`, if we only consider occurences in the term `Tmₛ (A arg)`,
+  `T` is considered non-fixed despite the fact that `A : T -> Tyₛ`.
+  This leads to an ill-typed injectivity theorem signature:
+  ```lean
+  theorem Tmₛ.app.inj {T : Type u} {A : T → Tyₛ} {a : Tmₛ (Tyₛ.SPi T A)} {arg : T} {T_1 : Type u} {a_1 : Tmₛ (Tyₛ.SPi T_1 A)} :
+  Tmₛ.app a arg = Tmₛ.app a_1 arg →
+    T = T_1 ∧ HEq a a_1 := fun x => Tmₛ.noConfusion x fun T_eq A_eq a_eq arg_eq => eq_of_heq a_eq
+  ```
+  Instead of checking the type of every subterm, we only need to check the type of free variables, since free variables introduced in
+  the constructor may only appear in the type of other free variables introduced after them.
+-/
+def occursOrInType (e : Expr) (t : Expr) : MetaM Bool := do
+  let_fun f (s : Expr) := do
+    if !s.isFVar then
+      return s == e
+    let ty ← inferType s
+    return s == e || e.occurs ty
+  return (← t.findM? f).isSome
+
 private partial def mkInjectiveTheoremTypeCore? (ctorVal : ConstructorVal) (useEq : Bool) : MetaM (Option Expr) := do
   let us := ctorVal.levelParams.map mkLevelParam
   let type ← elimOptParam ctorVal.type
@@ -58,7 +87,7 @@ private partial def mkInjectiveTheoremTypeCore? (ctorVal : ConstructorVal) (useE
         match (← whnf type) with
         | Expr.forallE n d b _ =>
           let arg1 := args1.get ⟨i, h⟩
-          if arg1.occurs resultType then
+          if ← occursOrInType arg1 resultType then
             mkArgs2 (i + 1) (b.instantiate1 arg1) (args2.push arg1) args2New
           else
             withLocalDecl n (if useEq then BinderInfo.default else BinderInfo.implicit) d fun arg2 =>
diff --git a/tests/lean/run/3386.lean b/tests/lean/run/3386.lean
new file mode 100644
index 000000000000..f470fdb72993
--- /dev/null
+++ b/tests/lean/run/3386.lean
@@ -0,0 +1,9 @@
+/- Verify that injectivity lemmas are constructed with the right level of generality
+   in order to avoid type errors.
+-/
+
+inductive Tyₛ : Type (u+1)
+| SPi : (T : Type u) -> (T -> Tyₛ) -> Tyₛ
+
+inductive Tmₛ.{u} :  Tyₛ.{u} -> Type (u+1)
+| app : Tmₛ (.SPi T A) -> (arg : T) -> Tmₛ (A arg)