fix: erase propositon formers, add isErasedCompatible, remove approx. from compatibleTypes

src/Lean/Compiler/LCNF/Check.lean

@@ -51,11 +51,11 @@ def checkAppArgs (f : Expr) (args : Array Expr) : CheckM Unit := do
           throwError "function expected at{indentExpr (mkAppN f args)}\narrow type expected{indentExpr fType}"
     let argType ← inferType arg
     let expectedType := d.instantiateRevRange j i args
-    unless compatibleTypes argType expectedType do
+    unless (← compatibleTypes argType expectedType) do
       throwError "type mismatch at LCNF application{indentExpr (mkAppN f args)}\nargument {arg} has type{indentExpr argType}\nbut is expected to have type{indentExpr expectedType}"
-    unless maybeTypeFormerType expectedType || expectedType.isErased do
+    unless (← pure (maybeTypeFormerType expectedType) <||> isErasedCompatible expectedType) do
       unless arg.isFVar do
-        throwError "invalid LCNF application{indentExpr (mkAppN f args)}\nargument{indentExpr arg}\nmust be a free variable"
+        throwError "invalid LCNF application{indentExpr (mkAppN f args)}\nargument{indentExpr arg}\nhas type{indentExpr expectedType}\nmust be a free variable"
       checkFVar arg.fvarId!
     fType := b
 
@@ -100,7 +100,7 @@ def checkParams (params : Array Param) : CheckM Unit :=
 def checkLetDecl (letDecl : LetDecl) : CheckM Unit := do
   checkExpr letDecl.value
   let valueType ← inferType letDecl.value
-  unless compatibleTypes letDecl.type valueType do
+  unless (← compatibleTypes letDecl.type valueType) do
     throwError "type mismatch at `{letDecl.binderName}`, value has type{indentExpr valueType}\nbut is expected to have type{indentExpr letDecl.type}"
   unless letDecl == (← getLetDecl letDecl.fvarId) do
     throwError "LCNF let declaration mismatch at `{letDecl.binderName}`, does not match value in local context"
@@ -129,7 +129,7 @@ partial def checkFunDeclCore (declName : Name) (type : Expr) (params : Array Par
   checkParams params
   let valueType ← withParams params do
     mkForallParams params (← check value)
-  unless compatibleTypes type valueType do
+  unless (← compatibleTypes type valueType) do
     throwError "type mismatch at `{declName}`, value has type{indentExpr valueType}\nbut is expected to have type{indentExpr type}"
 
 partial def checkFunDecl (funDecl : FunDecl) : CheckM Unit := do
@@ -167,7 +167,7 @@ partial def checkCases (c : Cases) : CheckM Expr := do
           throwError "invalid LCNF `cases`, `{ctorName}` has # {val.numFields} fields, but alternative has # {params.size} alternatives"
         -- TODO: check whether the ctor field types as parameter types match.
         withParams params do check k
-    unless compatibleTypes type c.resultType do
+    unless (← compatibleTypes type c.resultType) do
       throwError "type mismatch at LCNF `cases` alternative\nhas type{indentExpr type}\nbut is expected to have type{indentExpr c.resultType}"
   return c.resultType
 

src/Lean/Compiler/LCNF/InferType.lean

@@ -7,6 +7,95 @@ import Lean.Compiler.LCNF.CompilerM
 import Lean.Compiler.LCNF.Types
 
 namespace Lean.Compiler.LCNF
+/--
+Return `true` if `type` is compatible with `lcErased`.
+After instantiating universe polymorphic code, we may have
+some types may propositions which are compatible with `lcErased`.
+
+For example, suppose we have
+```
+def f (α : Sort u) (x : α → α → Sort v) (a b : α) (h : x a b) ...
+```
+The LCNF type for this universe polymorphic declaration is
+```
+def f (α : Sort u) (x : α → α → Sort v) (a b : α) (h : x ◾ ◾) ...
+```
+Now, if we instantiate with `v` with `0`, we have that `x ◾ ◾` is also a proposition
+and should be equivalent to `◾`.
+
+Remark: `predVars` is a bitmask that indicates whether de-bruijn variables are predicates or not.
+That is, `#i` is a predicate if `predVars[predVars.size - i - 1] = true`
+-/
+partial def isErasedCompatible (type : Expr) (predVars : Array Bool := #[]): CompilerM Bool :=
+  go type predVars
+where
+  go (type : Expr) (predVars : Array Bool) : CompilerM Bool := do
+    let type := type.headBeta
+    match type with
+    | .const ..        => return type.isErased
+    | .sort ..         => return false
+    | .mdata _ e       => go e predVars
+    | .forallE _ t b _ => go b (predVars.push <| isPredicateType t)
+    | .app f _         => go f predVars
+    | .bvar idx        => return predVars[predVars.size - idx - 1]!
+    | .fvar fvarId     => return isPredicateType (← getType fvarId)
+    | .proj .. | .mvar .. | .lam .. | .letE .. | .lit .. => unreachable!
+
+/--
+Return true if the LCNF types `a` and `b` are compatible.
+
+Remark: `a` and `b` can be type formers (e.g., `List`, or `fun (α : Type) => Nat → Nat × α`)
+
+Remark: LCNFs types are eagerly eta reduced.
+
+Remark: see comment at `isErasedCompatible`.
+
+The below checks do not appear exhaustive, but are
+in fact exhaustive due to LCNF constraints:
+  - bvar: handled by ==.
+  - fvar: handled by ==.
+  - mvar: should be resolved by the time we get to LCNF.
+  - sort: matched.
+  - const: matched.
+  - app: handled by β reduction + match.
+  - lam: matched.
+  - forallE: matched.
+  - letE: LCNF does not contain let at the type level.
+  - lit: We don't have data in LCNF, so we don't need to handle it.
+         Erased is handled by `const`.
+  - mdata: matched.
+  - proj: Becomes Any/Erased depending on what it should become.
+          type inside structure becomes Any, value inside structure becomes Erased.
+-/
+partial def compatibleTypes (a b : Expr) : CompilerM Bool := do
+  go a b #[]
+where
+  go (a b : Expr) (predVars : Array Bool) : CompilerM Bool := do
+    if a.isAnyType || b.isAnyType then
+      return true
+    else if a.isErased then
+      isErasedCompatible b
+    else if b.isErased then
+      isErasedCompatible a
+    else
+      let a' := a.headBeta
+      let b' := b.headBeta
+      if a != a' || b != b' then
+        go a' b' predVars
+      else if a == b then
+        return true
+      else
+        match a, b with
+        | .mdata _ a, b => go a b predVars
+        | a, .mdata _ b => go a b predVars
+        -- Note that even after reducing to headβ, we can still have `.app` terms. For example,
+        -- an inductive constructor application such as `List Int`
+        | .app f a, .app g b => go f g predVars <&&> go a b predVars
+        | .forallE _ d₁ b₁ _, .forallE _ d₂ b₂ _ => go d₁ d₂ predVars <&&> go b₁ b₂ (predVars.push (isPredicateType d₁))
+        | .lam _ d₁ b₁ _, .lam _ d₂ b₂ _ => go d₁ d₂ predVars <&&> go b₁ b₂ (predVars.push (isPredicateType d₂))
+        | .sort u, .sort v => return Level.isEquiv u v
+        | .const n us, .const m vs => return n == m && List.isEqv us vs Level.isEquiv
+        | _, _ => return false
 
 namespace InferType
 

src/Lean/Compiler/LCNF/ToLCNF.lean

@@ -378,11 +378,11 @@ where
 
   visit (e : Expr) : M Expr := withIncRecDepth do
     if isLCProof e then
-      return mkConst ``lcErased
+      return erasedExpr
     let type ← liftMetaM <| Meta.inferType e
     if (← liftMetaM <| Meta.isProp type) then
       /- We erase proofs. -/
-      return mkConst ``lcErased
+      return erasedExpr
     if (← isTypeFormerType type) then
       /-
       We erase type formers unless they occur as application arguments.
@@ -391,25 +391,29 @@ where
 
       See `visitAppArg`
       -/
-      return mkConst ``lcErased
+      return erasedExpr
     visitCore e
 
   visitAppArg (e : Expr) : M Expr := do
     if isLCProof e then
-      return mkConst ``lcErased
+      return erasedExpr
     let type ← liftMetaM <| Meta.inferType e
     if (← liftMetaM <| Meta.isProp type) then
       /- We erase proofs. -/
-      return mkConst ``lcErased
+      return erasedExpr
     if (← isTypeFormerType type) then
-      /- Types and Type formers are not put into A-normal form -/
-      toLCNFType e
+      /- Predicates are erased (e.g., `Eq`) -/
+      if isPredicateType (← toLCNFType type) then
+        return erasedExpr
+      else
+        /- Types and Type formers are not put into A-normal form -/
+        toLCNFType e
     else
       visitCore e
 
   /-- Visit args, and return `f args` -/
   visitAppDefault (f : Expr) (args : Array Expr) : M Expr := do
-    if f.isConstOf ``lcErased then
+    if f.isErased then
       return f
     else
       let args ← args.mapM visitAppArg
@@ -466,7 +470,7 @@ where
       let .inductInfo indVal ← getConstInfo typeName | unreachable!
       for i in casesInfo.altsRange, numParams in casesInfo.altNumParams, ctorName in indVal.ctors do
         let (altType, alt) ← visitAlt ctorName numParams args[i]!
-        unless compatibleTypes altType resultType do
+        unless (← compatibleTypes altType resultType) do
           resultType := anyTypeExpr
         alts := alts.push alt
       let cases : Cases := { typeName, discr := discr.fvarId!, resultType, alts }
@@ -504,7 +508,7 @@ where
       let minor := if e.isAppOf ``Eq.rec || e.isAppOf ``Eq.ndrec then args[3]! else args[5]!
       let minor ← visit minor
       let minorType ← inferType minor
-      let cast ← if compatibleTypes minorType recType then
+      let cast ← if (← compatibleTypes minorType recType) then
         -- Recall that many types become compatible after LCNF conversion
         -- Example: `Fin 10` and `Fin n`
         pure minor

src/Lean/Compiler/LCNF/Types.lean

@@ -24,6 +24,36 @@ def _root_.Lean.Expr.isAnyType (e : Expr) :=
 def _root_.Lean.Expr.isErased (e : Expr) :=
   e.isConstOf ``lcErased
 
+def isPropFormerTypeQuick : Expr → Bool
+  | .forallE _ _ b _ => isPropFormerTypeQuick b
+  | .sort .zero => true
+  | _ => false
+
+/--
+Return true iff `type` is `Prop` or `As → Prop`.
+-/
+partial def isPropFormerType (type : Expr) : MetaM Bool := do
+  match isPropFormerTypeQuick type with
+  | true => return true
+  | false => go type #[]
+where
+  go (type : Expr) (xs : Array Expr) : MetaM Bool := do
+    match type with
+    | .sort .zero => return true
+    | .forallE n d b c => Meta.withLocalDecl n c (d.instantiateRev xs) fun x => go b (xs.push x)
+    | _ =>
+      let type ← Meta.whnfD (type.instantiateRev xs)
+      match type with
+      | .sort .zero => return true
+      | .forallE .. => go type #[]
+      | _ => return false
+
+/--
+Return true iff `e : Prop` or `e : As → Prop`.
+-/
+def isPropFormer (e : Expr) : MetaM Bool := do
+  isPropFormerType (← Meta.inferType e)
+
 /-!
 The code generator uses a format based on A-normal form.
 This normal form uses many let-expressions and it is very convenient for
@@ -139,6 +169,8 @@ where
     for arg in args do
       if (← isProp arg) then
         result := mkApp result erasedExpr
+      else if (← isPropFormer arg) then
+        result := mkApp result erasedExpr
       else if (← isTypeFormer arg) then
         result := mkApp result (← toLCNFType arg)
       else
@@ -180,69 +212,6 @@ def instantiateLCNFTypeLevelParams (declName : Name) (us : List Level) : CoreM E
     modifyEnv fun env => lcnfTypeExt.modifyState env fun s => { s with instLevelType := s.instLevelType.insert declName (us, r) }
     return r
 
-/--
-Return true if the LCNF types `a` and `b` are compatible.
-
-Remark: `a` and `b` can be type formers (e.g., `List`, or `fun (α : Type) => Nat → Nat × α`)
-
-Remark: LCNFs types are eagerly eta reduced.
-
-The below checks do not appear exhaustive, but are
-in fact exhaustive due to LCNF constraints:
-  - bvar: handled by ==.
-  - fvar: handled by ==.
-  - mvar: should be resolved by the time we get to LCNF.
-  - sort: matched.
-  - const: matched.
-  - app: handled by β reduction + match.
-  - lam: matched.
-  - forallE: matched.
-  - letE: LCNF does not contain let at the type level.
-  - lit: We don't have data in LCNF, so we don't need to handle it.
-         Erased is handled by `const`.
-  - mdata: matched.
-  - proj: Becomes Any/Erased depending on what it should become.
-          type inside structure becomes Any, value inside structure becomes Erased.
--/
-partial def compatibleTypes (a b : Expr) : Bool :=
-  if a.isAnyType || b.isAnyType then
-    true
-  else if a.isErased || b.isErased then
-    /-
-    This is an approximation. This can happen when processing universe polymorphic code.
-    After instantiating universes, we may have propositions. For example, suppose we have
-    ```
-    def f (α : Sort u) (x : α → α → Sort v) (a b : α) (h : x a b) ...
-    ```
-    The LCNF type for this universe polymorphic declaration is
-    ```
-    def f (α : Sort u) (x : α → α → Sort v) (a b : α) (h : x ◾ ◾) ...
-    ```
-    Now, if we instantiate with `v = 0`, we have that `x ◾ ◾` is also a proposition
-    and should be equivalent to `◾`. However, to perform this check we need the environment
-    and local context. We may consider doing at `Check.lean`
-    -/
-    true
-  else
-    let a' := a.headBeta
-    let b' := b.headBeta
-    if a != a' || b != b' then
-      compatibleTypes a' b'
-    else if a == b then
-      true
-    else
-      match a, b with
-      | .mdata _ a, b => compatibleTypes a b
-      | a, .mdata _ b => compatibleTypes a b
-      -- Note that even after reducing to headβ, we can still have `.app` terms. For example,
-      -- an inductive constructor application such as `List Int`
-      | .app f a, .app g b => compatibleTypes f g && compatibleTypes a b
-      | .forallE _ d₁ b₁ _, .forallE _ d₂ b₂ _ => compatibleTypes d₁ d₂ && compatibleTypes b₁ b₂
-      | .lam _ d₁ b₁ _, .lam _ d₂ b₂ _ => compatibleTypes d₁ d₂ && compatibleTypes b₁ b₂
-      | .sort u, .sort v => Level.isEquiv u v
-      | .const n us, .const m vs => n == m && List.isEqv us vs Level.isEquiv
-      | _, _ => false
-
 mutual
 
 partial def joinTypes (a b : Expr) : Expr :=
@@ -291,6 +260,16 @@ partial def isTypeFormerType (type : Expr) : Bool :=
   | .forallE _ _ b _ => isTypeFormerType b
   | _ => false
 
+/--
+Return `true` if `type` is a predicate.
+Examples: `Nat → Prop`, `Prop`, `Int → Bool → Prop`.
+-/
+partial def isPredicateType (type : Expr) : Bool :=
+  match type.headBeta with
+  | .sort .zero => true
+  | .forallE _ _ b _ => isPredicateType b
+  | _ => false
+
 /--
 Return `true` if `type` is a LCNF type former type or it is an "any" type.
 This function is similar to `isTypeFormerType`, but more liberal.

tests/lean/lcnfTypes.lean

@@ -113,7 +113,7 @@ def weird1 (c : Bool) : (cond c List Array) Nat :=
 def compatible (declName₁ declName₂ : Name) : MetaM Unit := do
   let type₁ ← getDeclLCNFType declName₁
   let type₂ ← getDeclLCNFType declName₂
-  unless compatibleTypes type₁ type₂ do
+  unless (← compatibleTypes type₁ type₂ |>.run' {}) do
     throwError "{declName₁} : {← ppExpr type₁}\ntype is not compatible with\n{declName₂} : {← ppExpr type₂}"
 
 axiom monadList₁.{u} : Monad List.{u}

tests/lean/lcnfTypes.lean.expected.out

@@ -5,7 +5,7 @@ Vec.cons : {α : Type u} → {n : Nat} → α → Vec α ◾ → Vec α ◾
 Eq.rec : {α : Sort u_1} → {a : α} → {motive : α → ◾ → Sort u} → motive ◾ ◾ → {a : α} → ◾ → motive ◾ ◾
 GetElem.getElem : {cont : Type u} →
   {idx : Type v} →
-    {elem : Type w} → {dom : cont → idx → Prop} → [self : GetElem cont idx elem dom] → cont → idx → ◾ → elem
+    {elem : Type w} → {dom : cont → idx → Prop} → [self : GetElem cont idx elem ◾] → cont → idx → ◾ → elem
 Term.constFold : {ctx : List Ty} → {ty : Ty} → Term ◾ ◾ → Term ◾ ◾
 Term.denote : {ctx : List Ty} → {ty : Ty} → Term ◾ ◾ → HList ⊤ ◾ → ⊤
 HList.get : {α : Type u_1} → {β : α → Type u_2} → {is : List α} → {i : α} → HList β ◾ → Member ◾ ◾ → β ◾