From fed7ff27e83173a9e2817002095c6bb55abd5ef4 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Mon, 3 Oct 2022 09:51:32 -0700
Subject: [PATCH] fix: issue reported on Zulip

https://leanprover.zulipchat.com/#narrow/stream/270676-lean4/topic/Annoying.20LCNF.20errors/near/302056742
---
 src/Lean/Compiler/LCNF/Internalize.lean   | 58 ++++++++++++++++++++++-
 tests/lean/run/internalizeCasesIssue.lean | 27 +++++++++++
 2 files changed, 83 insertions(+), 2 deletions(-)
 create mode 100644 tests/lean/run/internalizeCasesIssue.lean

diff --git a/src/Lean/Compiler/LCNF/Internalize.lean b/src/Lean/Compiler/LCNF/Internalize.lean
index 2ca59c62e6..66816f935c 100644
--- a/src/Lean/Compiler/LCNF/Internalize.lean
+++ b/src/Lean/Compiler/LCNF/Internalize.lean
@@ -3,6 +3,8 @@ Copyright (c) 2022 Microsoft Corporation. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
+import Lean.Compiler.LCNF.Types
+import Lean.Compiler.LCNF.Bind
 import Lean.Compiler.LCNF.CompilerM
 
 namespace Lean.Compiler.LCNF
@@ -74,10 +76,62 @@ partial def internalizeCode (code : Code) : InternalizeM Code := do
   | .unreach type => return .unreach (← normExpr type)
   | .cases c =>
     let resultType ← normExpr c.resultType
+    let ensureAny := resultType != c.resultType && (resultType.isAnyType || resultType.isErased)
+    /-
+    Note:
+    If the new result type for the cases is `⊤` or `◾`, we must add a cast to `⊤` (the any type)
+    to every alternative if the resulting type is not `⊤`. This is similar to what we do at `ToLCNF.visitCases`.
+    Here is an example to illustrate this issue.
+    Suppose we have
+    ```
+    inductive Id {A : Type u} : A → A → Type u
+      | refl {a : A} : Id a a
+    def transport {A : Type u} (B : A → Type v) {a b : A} (p : Id a b) : B a → B b :=
+    ```
+    Its LCNF type is
+    ```
+    {A : Type u} (B : A → Type v) {a b : A} (p : Id ◾ ◾) (a.1 : B ◾) : B ◾
+    ```
+    and base phase code is
+    ```
+    cases p : B ◾
+    | Id.refl =>
+      a.1
+    ```
+    Now suppose we define
+    ```
+    def transportconst {A B : Type u} : A = B → A → B :=
+      transport id
+    ```
+    By setting `B` as `id`, and then inlining `transport, we would have the following code for `transportconst` is
+    ```
+    cases p : ◾
+    | Id.refl =>
+      a.1
+    ```
+    Which can be checked by `Check.lean` because it assumes `◾` is compatible with anything and `a.1 : A`.
+    However, if inline `transportconst`, we can hit type error since the continuation for transportconst is
+    expecting a `B` instead of an `A`. We avoid this problem by adding a cast to `⊤`. See `ToLCNF.visitCases` for
+    another place where we use this approach.
+    Thus, the resulting code for `transportconst` is
+    ```
+    def MWE.transportconst (A : Type u) (B : Type u) (p : Id A B) (a.1 : A) :=
+      cases p
+      | Id.refl =>
+        let _x.2 := @lcCast A ⊤ a.1
+        _x.2
+    ```
+    -/
+    let internalizeAltCode (k : Code) : InternalizeM Code := do
+      let k ← internalizeCode k
+      if ensureAny then
+        k.ensureAnyType
+      else
+        return k
     let discr ← normFVar c.discr
     let alts ← c.alts.mapM fun
-      | .alt ctorName params k => return .alt ctorName (← params.mapM internalizeParam) (← internalizeCode k)
-      | .default k => return .default (← internalizeCode k)
+      | .alt ctorName params k => return .alt ctorName (← params.mapM internalizeParam) (← internalizeAltCode k)
+      | .default k => return .default (← internalizeAltCode k)
     return .cases { c with discr, alts, resultType }
 
 end
diff --git a/tests/lean/run/internalizeCasesIssue.lean b/tests/lean/run/internalizeCasesIssue.lean
new file mode 100644
index 0000000000..8233bde6a0
--- /dev/null
+++ b/tests/lean/run/internalizeCasesIssue.lean
@@ -0,0 +1,27 @@
+namespace MWE
+
+inductive Id {A : Type u} : A → A → Type u
+| refl {a : A} : Id a a
+
+attribute [eliminator] Id.casesOn
+
+infix:50 (priority := high) " = " => Id
+
+def symm {A : Type u} {a b : A} (p : a = b) : b = a :=
+by { induction p; exact Id.refl }
+
+def transport {A : Type u} (B : A → Type v) {a b : A} (p : a = b) : B a → B b :=
+by { induction p; exact id }
+
+def transportconst {A B : Type u} : A = B → A → B :=
+transport id
+
+def transportconstInv {A B : Type u} (e : A = B) : B → A :=
+transportconst (symm e)
+
+def transportconstOverInv {A B : Type u} (p : A = B) :
+  ∀ x, transportconst (symm p) x = transportconstInv p x :=
+by { intro x; apply Id.refl }
+
+def transportconstInv' {A B : Type u} : A = B → B → A :=
+  transportconst ∘ symm