From 82acc2b39c26779196e92501f55202dc2a57c7eb Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Mon, 22 Aug 2022 21:07:23 -0700
Subject: [PATCH] feat: improve `Code.bind`

Add `joinTypes`.
---
 src/Lean/Compiler/LCNF/Basic.lean     |  5 ++++
 src/Lean/Compiler/LCNF/Bind.lean      | 11 +++++++--
 src/Lean/Compiler/LCNF/InferType.lean | 10 ++++++++
 src/Lean/Compiler/LCNF/Types.lean     | 35 +++++++++++++++++++++++++++
 4 files changed, 59 insertions(+), 2 deletions(-)

diff --git a/src/Lean/Compiler/LCNF/Basic.lean b/src/Lean/Compiler/LCNF/Basic.lean
index 50be1dba9c..780c317b12 100644
--- a/src/Lean/Compiler/LCNF/Basic.lean
+++ b/src/Lean/Compiler/LCNF/Basic.lean
@@ -16,12 +16,14 @@ structure Param where
 inductive AltCore (Code : Type) where
   | alt (ctorName : Name) (params : Array Param) (code : Code)
   | default (code : Code)
+  deriving Inhabited
 
 structure LetDecl where
   fvarId : FVarId
   binderName : Name
   type : Expr
   value : Expr
+  deriving Inhabited
 
 structure FunDeclCore (Code : Type) where
   fvarId : FVarId
@@ -29,12 +31,14 @@ structure FunDeclCore (Code : Type) where
   params : Array Param
   type : Expr
   value : Code
+  deriving Inhabited
 
 structure CasesCore (Code : Type) where
   typeName : Name
   resultType : Expr
   discr : FVarId
   alts : Array (AltCore Code)
+  deriving Inhabited
 
 inductive Code where
   | let (decl : LetDecl) (k : Code)
@@ -44,6 +48,7 @@ inductive Code where
   | cases (cases : CasesCore Code)
   | return (fvarId : FVarId)
   | unreach (type : Expr)
+  deriving Inhabited
 
 abbrev Alt := AltCore Code
 abbrev FunDecl := FunDeclCore Code
diff --git a/src/Lean/Compiler/LCNF/Bind.lean b/src/Lean/Compiler/LCNF/Bind.lean
index da8a7fbdc1..e26f8da378 100644
--- a/src/Lean/Compiler/LCNF/Bind.lean
+++ b/src/Lean/Compiler/LCNF/Bind.lean
@@ -24,14 +24,21 @@ where
     | .let decl k => return .let decl (← go k)
     | .fun decl k => return .fun decl (← go k)
     | .jp decl k =>
-      let decl := { decl with value := (← go decl.value) }
+      let value ← go decl.value
+      let type ← value.inferParamType decl.params
+      let decl := { decl with value, type }
       withReader (fun s => s.insert decl.fvarId) do
         return .jp decl (← go k)
     | .cases c =>
       let alts ← c.alts.mapM fun
         | .alt ctorName params k => return .alt ctorName params (← go k)
         | .default k => return .default (← go k)
-      return .cases { c with alts }
+      if alts.isEmpty then
+        throwError "`Code.bind` failed, empty `cases` found"
+      let mut resultType ← alts[0]!.inferType
+      for alt in alts[1:] do
+        resultType := joinTypes resultType (← alt.inferType)
+      return .cases { c with alts, resultType }
     | .return fvarId => f fvarId
     | .jmp fvarId .. =>
       unless (← read).contains fvarId do
diff --git a/src/Lean/Compiler/LCNF/InferType.lean b/src/Lean/Compiler/LCNF/InferType.lean
index 8dd13b5f54..6c8ad93e59 100644
--- a/src/Lean/Compiler/LCNF/InferType.lean
+++ b/src/Lean/Compiler/LCNF/InferType.lean
@@ -170,4 +170,14 @@ def Code.inferType (code : Code) : CompilerM Expr := do
   | .unreach type => return type
   | .cases c => return c.resultType
 
+def Code.inferParamType (params : Array Param) (code : Code) : CompilerM Expr := do
+  let type ← code.inferType
+  let xs := params.map fun p => .fvar p.fvarId
+  InferType.mkForallFVars xs type |>.run {}
+
+def AltCore.inferType (alt : Alt) : CompilerM Expr := do
+  match alt with
+  | .default k => k.inferType
+  | .alt _ params k => k.inferParamType params
+
 end Lean.Compiler.LCNF
diff --git a/src/Lean/Compiler/LCNF/Types.lean b/src/Lean/Compiler/LCNF/Types.lean
index 2b6391fb8a..117276897d 100644
--- a/src/Lean/Compiler/LCNF/Types.lean
+++ b/src/Lean/Compiler/LCNF/Types.lean
@@ -183,6 +183,41 @@ partial def compatibleTypes (a b : Expr) : Bool :=
       | .const n us, .const m vs => n == m && List.isEqv us vs Level.isEquiv
       | _, _ => false
 
+mutual
+
+partial def joinTypes (a b : Expr) : Expr :=
+  joinTypes? a b |>.getD anyTypeExpr
+
+partial def joinTypes? (a b : Expr) : Option Expr := do
+  if a.isAnyType then return a
+  else if b.isAnyType then return b
+  else if a == b then return a
+  else if a.erased || b.erased then failure
+  else
+    let a' := a.headBeta
+    let b' := b.headBeta
+    if a != a' || b != b' then
+      joinTypes? a' b'
+    else
+      match a, b with
+      | .mdata _ a, b => joinTypes? a b
+      | a, .mdata _ b => joinTypes? a b
+      | .app f a, .app g b =>
+        (do return .app (← joinTypes? f g) (← joinTypes? a b))
+         <|>
+        return anyTypeExpr
+      | .forallE n d₁ b₁ _, .forallE _ d₂ b₂ _ =>
+        (do return .forallE n (← joinTypes? d₁ d₂) (joinTypes b₁ b₂) .default)
+        <|>
+        return anyTypeExpr
+      | .lam n d₁ b₁ _, .lam _ d₂ b₂ _ =>
+        (do return .lam n (← joinTypes? d₁ d₂) (joinTypes b₁ b₂) .default)
+        <|>
+        return anyTypeExpr
+      | _, _ => return anyTypeExpr
+
+end
+
 /--
 Return `true` if `type` is a LCNF type former type.
 -/