From 6c5475725eeff4ffefaf3ab959cf48263ac16424 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Sun, 9 Oct 2022 17:45:15 -0700
Subject: [PATCH] feat: `(lcCast _ _ g) a_1 ... a_n` => `g a_1 ... a_n` if type
 correct

---
 src/Lean/Compiler/LCNF/Simp/SimpValue.lean | 47 +++++++++++++++++++++-
 1 file changed, 45 insertions(+), 2 deletions(-)

diff --git a/src/Lean/Compiler/LCNF/Simp/SimpValue.lean b/src/Lean/Compiler/LCNF/Simp/SimpValue.lean
index f3d16d13c6..cc632d1fb5 100644
--- a/src/Lean/Compiler/LCNF/Simp/SimpValue.lean
+++ b/src/Lean/Compiler/LCNF/Simp/SimpValue.lean
@@ -17,6 +17,34 @@ def simpProj? (e : Expr) : OptionT SimpM Expr := do
   let some (ctorVal, args) := s.constructorApp? (← getEnv) | failure
   return args[ctorVal.numParams + i]!
 
+/--
+Return `some type` if `f args` is type correct, and has type `type`.
+-/
+def checkApp? (f : Expr) (args : Array Expr) : CompilerM (Option Expr) := do
+  let mut fType ← inferType f
+  let mut j := 0
+  for i in [:args.size] do
+    let arg := args[i]!
+    if fType.isErased then
+      return fType
+    fType := fType.headBeta
+    let (d, b) ←
+      match fType with
+      | .forallE _ d b _ => pure (d, b)
+      | _ =>
+        fType := fType.instantiateRevRange j i args |>.headBeta
+        match fType with
+        | .forallE _ d b _ => j := i; pure (d, b)
+        | _ =>
+          if fType.isErased then return fType
+          return none
+    let argType ← inferType arg
+    let expectedType := d.instantiateRevRange j i args
+    unless (← compatibleTypes argType expectedType) do
+      return none
+    fType := b
+  return fType
+
 /--
 Application over application.
 ```
@@ -31,8 +59,23 @@ def simpAppApp? (e : Expr) : OptionT SimpM Expr := do
   guard f.isFVar
   let f ← findExpr f
   guard <| f.isApp || f.isConst
-  guard <| !f.isAppOf ``lcCast
-  return mkAppN f e.getAppArgs
+  if f.isAppOf ``lcCast then
+    /-
+    Given
+    ```
+    let _x.i := lcCast _ _ g
+    let _x.j := _x.i a_1 ... a_n
+    ```
+    Try to eliminate cast when `g a_1 ... a_n` is also type correct
+    -/
+    guard <| f.getAppNumArgs == 3
+    let g := f.appArg!
+    let args := e.getAppArgs
+    let some type ← checkApp? g args | failure
+    guard (← compatibleTypes (← inferType e) type)
+    return mkAppN g args
+  else
+    return mkAppN f e.getAppArgs
 
 def simpCtorDiscr? (e : Expr) : OptionT SimpM Expr := do
   let some v ← simpCtorDiscrCore? e | failure