From 9e5a818de5a476a74445f2a506c648699c366fbb Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Wed, 14 Sep 2022 15:23:34 -0700
Subject: [PATCH] fix: bug at LCNF `toDecl`

---
 src/Lean/Compiler/LCNF/Bind.lean   | 26 +++++++++++++++++---------
 src/Lean/Compiler/LCNF/ToDecl.lean | 10 ++++++----
 tests/lean/run/toDeclEtaBug.lean   | 19 +++++++++++++++++++
 3 files changed, 42 insertions(+), 13 deletions(-)
 create mode 100644 tests/lean/run/toDeclEtaBug.lean

diff --git a/src/Lean/Compiler/LCNF/Bind.lean b/src/Lean/Compiler/LCNF/Bind.lean
index 3e1169d423..cd0d96c0a7 100644
--- a/src/Lean/Compiler/LCNF/Bind.lean
+++ b/src/Lean/Compiler/LCNF/Bind.lean
@@ -66,20 +66,28 @@ where
       else
         return ps
 
-def FunDeclCore.etaExpand (decl : FunDecl) : CompilerM FunDecl := do
-  let typeArity := getArrowArity decl.type
-  let valueArity := decl.getArity
+def etaExpandCore? (type : Expr) (params : Array Param) (value : Code) : CompilerM (Option (Array Param × Code)) := do
+  let typeArity := getArrowArity type
+  let valueArity := params.size
   if typeArity <= valueArity then
     -- It can be < because of the "any" type
-    return decl
+    return none
   else
-    let valueType ← instantiateForall decl.type decl.params
+    let valueType ← instantiateForall type params
     let psNew ← mkNewParams valueType
-    let params := decl.params ++ psNew
+    let params := params ++ psNew
     let xs := psNew.map fun p => Expr.fvar p.fvarId
-    let value ← decl.value.bind fun fvarId => do
+    let value ← value.bind fun fvarId => do
       let auxDecl ← mkAuxLetDecl (mkAppN (.fvar fvarId) xs)
       return .let auxDecl (.return auxDecl.fvarId)
-    decl.update decl.type params value
+    return (params, value)
 
-end Lean.Compiler.LCNF
\ No newline at end of file
+def FunDeclCore.etaExpand (decl : FunDecl) : CompilerM FunDecl := do
+  let some (params, value) ← etaExpandCore? decl.type decl.params decl.value | return decl
+  decl.update decl.type params value
+
+def Decl.etaExpand (decl : Decl) : CompilerM Decl := do
+  let some (params, value) ← etaExpandCore? decl.type decl.params decl.value | return decl
+  return { decl with params, value }
+
+end Lean.Compiler.LCNF
diff --git a/src/Lean/Compiler/LCNF/ToDecl.lean b/src/Lean/Compiler/LCNF/ToDecl.lean
index ef3d8470cf..ca6597facc 100644
--- a/src/Lean/Compiler/LCNF/ToDecl.lean
+++ b/src/Lean/Compiler/LCNF/ToDecl.lean
@@ -100,10 +100,12 @@ def toDecl (declName : Name) : CompilerM Decl := do
     -- let value ← applyCasesOnImplementedBy value
     return (type, value)
   let value ← toLCNF value
-  if let .fun decl (.return _) := value then
+  let decl ← if let .fun decl (.return _) := value then
     eraseFVar decl.fvarId (recursive := false)
-    return { name := declName, params := decl.params, type, value := decl.value, levelParams := info.levelParams }
+    pure { name := declName, params := decl.params, type, value := decl.value, levelParams := info.levelParams : Decl }
   else
-    return { name := declName, params := #[], type, value, levelParams := info.levelParams }
+    pure { name := declName, params := #[], type, value, levelParams := info.levelParams }
+  /- `toLCNF` may eta-reduce simple declarations. -/
+  decl.etaExpand
 
-end Lean.Compiler.LCNF
\ No newline at end of file
+end Lean.Compiler.LCNF
diff --git a/tests/lean/run/toDeclEtaBug.lean b/tests/lean/run/toDeclEtaBug.lean
new file mode 100644
index 0000000000..a0a0688dbd
--- /dev/null
+++ b/tests/lean/run/toDeclEtaBug.lean
@@ -0,0 +1,19 @@
+import Lean
+
+open Lean Compiler LCNF Testing
+
+@[inline] def f (c b : Bool) (x : Nat) : Nat :=
+  if c && b then
+    x + 1
+  else
+    x + 2
+
+def g (c : Nat) : Nat :=
+  f true false c
+
+#eval Lean.Compiler.compile #[``g]
+
+@[cpass] def cseSizeTest : PassInstaller := Testing.assertNoFun |>.install `specialize `noFun
+
+set_option trace.Compiler.result true
+#eval Lean.Compiler.compile #[``g]