feat: eagerly inline simple join points

src/Lean/Compiler/TerminalCases.lean

@@ -52,7 +52,11 @@ partial def visitLet (e : Expr) (fvars : Array Expr) : M Expr := do
         let body ← mkLetUsingScope body
         let bodyAbst := body.abstract #[x]
         return bodyAbst
-      let jp ← mkJpDecl (.lam binderName type bodyAbst .default)
+      let jp ← if (← isSimpleLCNF bodyAbst) then
+        -- Join point is too simple, we eagerly inline it.
+        pure <| .lam binderName type bodyAbst .default
+      else
+        mkJpDecl (.lam binderName type bodyAbst .default)
       withReader (fun _ => { jp? := some jp }) do
         visitCases casesInfo value
     else
@@ -72,18 +76,19 @@ partial def visitLet (e : Expr) (fvars : Array Expr) : M Expr := do
           return e
       | some jp =>
         let .forallE _ d b _ ← inferType jp | unreachable!
+        let mkJpApp (x : Expr) := mkApp jp x |>.headBeta
         if isLcUnreachable e then
           mkLcUnreachable b
         else if compatibleTypes (← inferType e) d then
           let x ← mkAuxLetDecl e
-          return mkApp jp x
+          return mkJpApp x
         else if let some x := isLcCast? e then
-          let e ← mkAuxLetDecl (← mkLcCast x d)
-          return mkApp jp e
+          let x ← mkAuxLetDecl (← mkLcCast x d)
+          return mkJpApp x
         else
           let x ← mkAuxLetDecl e
           let x ← mkAuxLetDecl (← mkLcCast x d)
-          return mkApp jp x
+          return mkJpApp x
 
 end
 

src/Lean/Compiler/Util.lean

@@ -109,6 +109,17 @@ def getCtorArity? (declName : Name) : CoreM (Option Nat) := do
   let .ctorInfo val ← getConstInfo declName | return none
   return val.numParams + val.numFields
 
+/--
+Return `true` if the `value` if not a `lambda`, `let` or cases-like expression.
+-/
+def isSimpleLCNF (value : Expr) : CoreM Bool := do
+  if value.isLet || value.isLambda then
+    return false
+  else if let some _ ← isCasesApp? value then
+    return false
+  else
+    return true
+
 /--
 List of types that have builtin runtime support
 -/

tests/lean/run/lcnf1.lean

@@ -41,8 +41,6 @@ def Vec.head : Vec α (n+1) → α
 
 #eval Compiler.compile #[``Vec.zip]
 
-#eval Compiler.compile #[``Vec.zip.match_1]
-
 #eval Compiler.compile #[``Lean.Elab.Term.reportStuckSyntheticMVar]
 
 #eval Compiler.compile #[``Lean.Elab.Term.synthesizeSyntheticMVars]
@@ -61,3 +59,5 @@ def foo (a b : Nat) :=
   Nat.add d e
 
 #eval Compiler.compile #[``foo]
+
+#eval Compiler.compile #[``Vec.zip.match_1]