feat: eta-expand local function declarations that are not being inlined

src/Lean/Compiler/LCNF/Basic.lean

@@ -255,6 +255,10 @@ def Code.isDecl : Code → Bool
   | .let .. | .fun .. | .jp .. => true
   | _ => false
 
+def Code.isFun : Code → Bool
+  | .fun .. => true
+  | _ => false
+
 def Code.isReturnOf : Code → FVarId → Bool
   | .return fvarId, fvarId' => fvarId == fvarId'
   | _, _ => false

src/Lean/Compiler/LCNF/Bind.lean

@@ -46,7 +46,7 @@ where
       return c
     | .unreach .. => return c
 
-def FunDecl.etaExpand (decl : FunDecl) : CompilerM FunDecl := do
+def FunDeclCore.etaExpand (decl : FunDecl) : CompilerM FunDecl := do
   let typeArity := getArrowArity decl.type
   let valueArity := decl.getArity
   if typeArity <= valueArity then

src/Lean/Compiler/LCNF/Simp.lean

@@ -524,6 +524,8 @@ partial def simp (code : Code) : SimpM Code := do
       Note that functions in `decl` will be marked as used even if `decl` is not actually used.
       They will only be deleted in the next pass.
       -/
+      if code.isFun then
+        decl ← decl.etaExpand
       decl ← simpFunDecl decl
     let k ← simp k
     if (← isUsed decl.fvarId) then
@@ -629,23 +631,6 @@ private def simpUsingEtaReduction (e : Expr) : Expr :=
   | .letE n t v b d => .letE n t v (simpUsingEtaReduction b) d
   | _ => e
 
-private def etaExpand (type : Expr) (value : Expr) : SimpM Expr := do
-  let typeArity := getArrowArity type
-  let valueArity := getLambdaArity value
-  if typeArity <= valueArity then
-    -- It can be < because of the "any" type
-    return value
-  else
-    withNewScope do
-      let (xs, _) ← visitArrow type
-      let value := getLambdaBody value
-      let value := value.instantiateRev xs[:valueArity]
-      let valueType ← inferType value
-      let f ← mkLocalDecl (← mkFreshUserName `_f) valueType
-      let k ← mkLambda #[f] (mkAppN f xs[valueArity:])
-      let value ← attachJp value k
-      mkLambda xs value
-
 
 /--
 Try "cases on cases" simplification.