Eliminate per-call closure for InlineIfLambda partial applications

docs/release-notes/.FSharp.Compiler.Service/11.0.100.md

@@ -194,6 +194,7 @@
 * Make Entity's adhoc members list lazy ([PR #20286](https://github.com/dotnet/fsharp/pull/20286/changes))
 * Constraint solver: `TryD` is now `inline` with `[<InlineIfLambda>]` on its always-run continuation, so the argument closures are no longer allocated at the (very hot) constraint-solver call sites; `IgnoreFailedMemberConstraintResolution` is `inline` so its forwarded continuation stays a literal. ([PR #20367](https://github.com/dotnet/fsharp/pull/20367))
 * `DelayedILModuleReader` no longer boxes its cached `ILModuleReader` on every read: the field is typed `ILModuleReader | null` and matched directly. ([PR #20413](https://github.com/dotnet/fsharp/pull/20413))
+* Optimizer: passing a partial application of a module-level function to an `[<InlineIfLambda>]` parameter (e.g. `xs |> Option.map (f a b)`) no longer allocates a per-call `FSharpFunc` closure when a captured argument is non-trivial (a field read, a call). Under optimization the argument is eta-expanded to a lambda with its captured evaluations floated above the binding, so the parameter's uses beta-reduce and the closure is eliminated. Captured arguments are still evaluated exactly once, in their original left-to-right order, and the binding keeps its sequence point. ([PR #20487](https://github.com/dotnet/fsharp/pull/20487))
 
 ### Changed
 * The `--warnaserror` option now ignores unrecognized diagnostic identifiers in warning lists while still applying recognized F# warning codes. ([PR #20246](https://github.com/dotnet/fsharp/pull/20246))

src/Compiler/Optimize/LowerCalls.fs

@@ -19,17 +19,13 @@ let InterceptExpr g cont expr =
 
     // App (Val v, tys, args)
     | Expr.App (Expr.Val (vref, flags, _) as f0, f0ty, tyargsl, argsl, m) ->
-        // Only transform if necessary, i.e. there are not enough arguments
         match vref.ValReprInfo with
-        | Some(valReprInfo) ->
-            let argsl = List.map cont argsl
-            let f0 =
-                if valReprInfo.AritiesOfArgs.Length > argsl.Length
-                then fst(AdjustValForExpectedValReprInfo g m vref flags valReprInfo)
-                else f0
-
-            Some (MakeApplicationAndBetaReduce g (f0, f0ty, [tyargsl], argsl, m))
         | None -> None
+        | Some _ ->
+            let argsl = List.map cont argsl
+            match TryEtaExpandUnderAppliedValApp g m vref flags tyargsl f0ty argsl with
+            | Some e -> Some e
+            | None -> Some (MakeApplicationAndBetaReduce g (f0, f0ty, [tyargsl], argsl, m))
 
     | Expr.App (f0, f0ty, tyargsl, argsl, m) ->
         Some (MakeApplicationAndBetaReduce g (f0, f0ty, [tyargsl], argsl, m) )

src/Compiler/Optimize/Optimizer.fs

@@ -2514,6 +2514,37 @@ let shouldForceInlineInDebug cenv env (vref: ValRef) : bool =
 
     HasFrameLocalBody cenv env vref
 
+/// `let p = f a b`, p an [<InlineIfLambda>] parameter binding whose right-hand side is an under-applied
+/// call to a known-arity value.
+[<return: Struct>]
+let private (|EtaFloatableValLet|_|) g expr =
+    match expr with
+    | Expr.Let(bind, body, m, _) when bind.Var.InlineIfLambda ->
+        match stripExpr bind.Expr with
+        | Expr.App(Expr.Val(vf, flags, _), f0ty, tyargs, args, mApp) ->
+            match TryEtaExpandUnderAppliedValApp g mApp vf flags tyargs f0ty args with
+            | Some etaExpanded -> ValueSome(bind, body, m, etaExpanded)
+            | None -> ValueNone
+        | _ -> ValueNone
+    | _ -> ValueNone
+
+/// `let p = f a`  ~>  `let p = fun x -> f a0 x`, with `let a0 = a` floated above the binding.
+let private floatEtaCaptures (bind: Binding) body m etaExpanded =
+    let rec rebindP e =
+        match e with
+        | Expr.Let(capture, inner, mLet, _) -> mkLetBind mLet capture (rebindP inner)
+        | rhs -> mkLet bind.DebugPoint m bind.Var rhs body
+
+    rebindP etaExpanded
+
+/// Float a let-bound partial application so its right-hand side is a bare lambda (`CurriedLambdaValue`) the
+/// optimizer can inline away the closure. LowerCalls does the same eta-expansion later but nests the
+/// captures, so the binding keeps `UnknownValue`.
+let EtaExpandUnderAppliedValBinding g expr =
+    match expr with
+    | EtaFloatableValLet g (bind, body, m, etaExpanded) -> floatEtaCaptures bind body m etaExpanded
+    | _ -> expr
+
 /// Optimize/analyze an expression
 let rec OptimizeExpr cenv (env: IncrementalOptimizationEnv) expr =
     cenv.stackGuard.Guard(fun () ->
@@ -3034,6 +3065,7 @@ and OptimizeLinearExpr cenv env expr contf =
     // complete inference types.
     let expr = DetectAndOptimizeForEachExpression g OptimizeAllForExpressions expr
     let expr = if cenv.settings.ExpandStructuralValues() then ExpandStructuralBinding cenv expr else expr
+    let expr = if cenv.settings.alwaysInline then EtaExpandUnderAppliedValBinding g expr else expr
     let expr = stripExpr expr
 
     // Matching on 'match __resumableEntry() with ...` is really a first-class language construct which we

src/Compiler/TypedTree/TypedTreeOps.ExprOps.fs

@@ -1916,6 +1916,13 @@ module internal ExprTransforms =
         // Build a type-lambda expression for the toplevel value if needed...
         mkTypeLambda m tpsR (tauexpr, tauty), tpsR +-> tauty
 
+    let TryEtaExpandUnderAppliedValApp g m (vref: ValRef) flags tyargs fty args =
+        match vref.ValReprInfo with
+        | Some valReprInfo when valReprInfo.NumCurriedArgs > List.length args ->
+            let etaExpr = fst (AdjustValForExpectedValReprInfo g m vref flags valReprInfo)
+            Some(MakeApplicationAndBetaReduce g (etaExpr, fty, [ tyargs ], args, m))
+        | _ -> None
+
     let stripTupledFunTy g ty =
         let argTys, retTy = stripFunTy g ty
         let curriedArgTys = argTys |> List.map (tryDestRefTupleTy g)

src/Compiler/TypedTree/TypedTreeOps.ExprOps.fsi

@@ -524,6 +524,11 @@ module internal ExprTransforms =
 
     val AdjustValForExpectedValReprInfo: TcGlobals -> range -> ValRef -> ValUseFlag -> ValReprInfo -> Expr * TType
 
+    /// Eta-expand an under-applied application of a known-arity value, binding the supplied arguments;
+    /// None when the value is not under-applied (or has no known arity).
+    val TryEtaExpandUnderAppliedValApp:
+        TcGlobals -> range -> ValRef -> ValUseFlag -> tyargs: TypeInst -> fty: TType -> args: Exprs -> Expr option
+
     val AdjustValToHaveValReprInfo: Val -> ParentRef -> ValReprInfo -> unit
 
     val stripTupledFunTy: TcGlobals -> TType -> TType list list * TType

tests/FSharp.Compiler.ComponentTests/EmittedIL/InlineIfLambdaEtaFloat.fs

@@ -0,0 +1,55 @@
+// Copyright (c) Microsoft Corporation.  All Rights Reserved.  See License.txt in the project root for license information.
+
+namespace EmittedIL
+
+open Xunit
+open FSharp.Test.Compiler
+
+module InlineIfLambdaEtaFloat =
+
+    // No-closure facts assert the closure *type* is absent; the negative controls assert it is invoked.
+    let private intIntFunc = "FSharpFunc`2<int32,int32>"
+    let private intIntInvoke = intIntFunc + "::Invoke"
+
+    let private compileOpt source =
+        FSharp source |> withOptimize |> asLibrary |> compile |> shouldSucceed
+
+    [<Fact>]
+    let ``Module-level function capture allocates no closure`` () =
+        compileOpt """
+module Test
+type Box(v: int) = member _.Value = v
+let f4 (a:int) (b:int) (c:int) (x:int) = a + b + c + x
+let shapeB (b: Box) (o: int option) = o |> Option.map (f4 1 b.Value 3)
+"""
+        |> verifyILNotPresent [ intIntFunc ]
+
+    [<Fact>]
+    let ``Generic module-level function capture allocates no closure`` () =
+        compileOpt """
+module Test
+type Box(v: int) = member _.Value = v
+let gpick (a:'T) (b:'T) (c:'T) (x:'T) : 'T = a
+let shapeG (b: Box) (o: int option) = o |> Option.map (gpick 1 b.Value 3)
+"""
+        |> verifyILNotPresent [ intIntFunc ]
+
+    // A first-class function argument has no known arity, so there is nothing to eta-expand.
+    [<Fact>]
+    let ``First-class function argument keeps its closure`` () =
+        compileOpt """
+module Test
+let shapeD (g: int -> int) (o: int option) = o |> Option.map g
+"""
+        |> verifyILPresent [ intIntInvoke ]
+
+    // A static member lacks the module-level known-arity shape the transform keys on.
+    [<Fact>]
+    let ``Static-member callee keeps its closure`` () =
+        compileOpt """
+module Test
+type Box(v: int) = member _.Value = v
+type H = static member SF (a:int) (b:int) (c:int) (x:int) = a + b + c + x
+let shapeM (b: Box) (o: int option) = o |> Option.map (H.SF 1 b.Value 3)
+"""
+        |> verifyILPresent [ intIntInvoke ]

tests/FSharp.Compiler.ComponentTests/FSharp.Compiler.ComponentTests.fsproj

@@ -265,6 +265,7 @@
     <Compile Include="EmittedIL\TailCalls.fs" />
     <Compile Include="EmittedIL\TupleElimination.fs" />
     <Compile Include="EmittedIL\UncheckedDefaultofOptimization.fs" />
+    <Compile Include="EmittedIL\InlineIfLambdaEtaFloat.fs" />
     <Compile Include="EmittedIL\TypeTestsInPatternMatching.fs" />
     <Compile Include="EmittedIL\WhileLoops.fs" />
     <Compile Include="EmittedIL\ArgumentNames.fs" />
@@ -294,6 +295,7 @@
     <Compile Include="EmittedIL\Misc\Misc.fs" />
     <Compile Include="Misc.fs" />
     <Compile Include="Optimizations\TaskCEUnitPropertyAccess.fs" />
+    <Compile Include="Optimizations\InlineIfLambdaEtaFloat.fs" />
     <Compile Include="EmittedIL\operators\Operators.fs" />
     <Compile Include="EmittedIL\Platform\Platform.fs" />
     <Compile Include="EmittedIL\QueryExpressionStepping\QueryExpressionStepping.fs" />

tests/FSharp.Compiler.ComponentTests/Optimizations/InlineIfLambdaEtaFloat.fs

@@ -0,0 +1,129 @@
+// Copyright (c) Microsoft Corporation.  All Rights Reserved.  See License.txt in the project root for license information.
+
+namespace Optimizations
+
+open Xunit
+open FSharp.Test.Compiler
+
+module InlineIfLambdaEtaFloat =
+
+    // Each snippet runs under both --optimize+ (transform on) and --optimize- (off, as a reference oracle),
+    // in-process, so a failure signals with failwith, not exit.
+    let private run (optimize: bool) (source: string) =
+        Fsx source |> withOptimization optimize |> compileExeAndRun |> shouldSucceed |> ignore
+
+    [<Theory; InlineData(true); InlineData(false)>]
+    let ``Captured argument is read once, before the parameter body is copied`` (optimize: bool) =
+        run optimize """
+let mutable state = 10
+let f4 (a:int) (b:int) (c:int) (x:int) = a + b + c + x
+let inline twiceMut ([<InlineIfLambda>] f: int -> int) (x: int) =
+    let r1 = f x
+    state <- 99
+    let r2 = f x
+    r1 + r2
+if twiceMut (f4 1 state 3) 5 <> 38 then failwith "capture was not read exactly once at 10"
+"""
+
+    [<Theory; InlineData(true); InlineData(false)>]
+    let ``Captured effect runs exactly once for a multi-use parameter`` (optimize: bool) =
+        run optimize """
+let mutable reads = 0
+type Box() = member _.Value = reads <- reads + 1; 10
+let f4 (a:int) (b:int) (c:int) (x:int) = a + b + c + x
+let inline twice ([<InlineIfLambda>] f: int -> int) (x: int) = f x + f x
+let r = twice (f4 1 (Box().Value) 3) 5
+if r <> 38 then failwithf "Expected 38 but got %d" r
+if reads <> 1 then failwithf "Captured getter evaluated %d times, expected 1" reads
+"""
+
+    // Floating must not sink the capture into the used branch.
+    [<Theory; InlineData(true); InlineData(false)>]
+    let ``Captured effect runs even when the parameter is unused`` (optimize: bool) =
+        run optimize """
+let mutable eff = 0
+type Box() = member _.Value = eff <- eff + 1; 10
+let f4 (a:int) (b:int) (c:int) (x:int) = a + b + c + x
+None   |> Option.map (f4 1 (Box().Value) 3) |> ignore
+Some 5 |> Option.map (f4 1 (Box().Value) 3) |> ignore
+if eff <> 2 then failwithf "Expected 2 captures but got %d" eff
+"""
+
+    [<Theory; InlineData(true); InlineData(false)>]
+    let ``Captured arguments preserve left-to-right evaluation order`` (optimize: bool) =
+        run optimize """
+let log = System.Collections.Generic.List<string>()
+let tap (name: string) (v: int) = log.Add name; v
+let f4 (a:int) (b:int) (c:int) (x:int) = a + b + c + x
+Some 5 |> Option.map (f4 (tap "a" 1) (tap "b" 2) (tap "c" 3)) |> ignore
+if String.concat "," (List.ofSeq log) <> "a,b,c" then failwithf "Wrong order: %A" (List.ofSeq log)
+"""
+
+    // The transform re-applies to the second parameter's binding, nested under the first.
+    [<Theory; InlineData(true); InlineData(false)>]
+    let ``Two InlineIfLambda parameters each capture once, left-to-right`` (optimize: bool) =
+        run optimize """
+let log = System.Collections.Generic.List<string>()
+let eff (n: string) (v: int) = log.Add n; v
+let f4 (a:int) (b:int) (c:int) (x:int) = a + b + c + x
+let inline combine ([<InlineIfLambda>] f: int -> int) ([<InlineIfLambda>] g: int -> int) (x: int) =
+    f x + g x + f x + g x
+let r = combine (f4 1 (eff "f" 10) 3) (f4 2 (eff "g" 20) 4) 5
+if r <> 100 then failwithf "Expected 100 but got %d" r
+if String.concat "," (List.ofSeq log) <> "f,g" then failwithf "Wrong capture order/count: %A" (List.ofSeq log)
+"""
+
+    // The "body" marker stays out of the log, proving the throw happens during capture, before the body.
+    [<Theory; InlineData(true); InlineData(false)>]
+    let ``A throwing capture is raised eagerly, in order, before the body`` (optimize: bool) =
+        run optimize """
+let log = System.Collections.Generic.List<string>()
+let tapOk (name: string) (v: int) = log.Add name; v
+let tapThrow (name: string) : int = log.Add name; failwith ("throw-" + name)
+let f4 (a:int) (b:int) (c:int) (x:int) = a + b + c + x
+let inline twice ([<InlineIfLambda>] f: int -> int) (x: int) = log.Add "body"; f x + f x
+let mutable msg = "NOTHROW"
+(try twice (f4 (tapOk "a" 1) (tapThrow "b") (tapOk "c" 3)) 5 |> ignore
+ with e -> msg <- e.Message)
+if String.concat "," (List.ofSeq log) <> "a,b" then failwithf "Wrong pre-throw order: %A" (List.ofSeq log)
+if msg <> "throw-b" then failwithf "Wrong exception surfaced: %s" msg
+"""
+
+    [<Theory; InlineData(true); InlineData(false)>]
+    let ``Capture is evaluated once when the parameter escapes in a returned closure`` (optimize: bool) =
+        run optimize """
+let mutable reads = 0
+type Box() = member _.Value = reads <- reads + 1; 10
+let f4 (a:int) (b:int) (c:int) (x:int) = a + b + c + x
+let inline makeAdder ([<InlineIfLambda>] f: int -> int) = fun y -> f y + f y
+let g = makeAdder (f4 1 (Box().Value) 3)
+let r1 = g 5
+let r2 = g 6
+if r1 <> 38 || r2 <> 40 then failwithf "Expected 38/40 but got %d/%d" r1 r2
+if reads <> 1 then failwithf "Captured getter evaluated %d times, expected 1" reads
+"""
+
+    [<Theory; InlineData(true); InlineData(false)>]
+    let ``A compound-expression capture is evaluated once`` (optimize: bool) =
+        run optimize """
+let mutable reads = 0
+let bump () = reads <- reads + 1; reads
+let f4 (a:int) (b:int) (c:int) (x:int) = a + b + c + x
+let inline twice ([<InlineIfLambda>] f: int -> int) (x: int) = f x + f x
+let r = twice (f4 1 (let n = bump () in n * 10 + 5) 3) 5
+if r <> 48 then failwithf "Expected 48 but got %d" r
+if reads <> 1 then failwithf "Compound capture evaluated %d times, expected 1" reads
+"""
+
+    // The transform does not fire here: the inline SRTP body collapses first.
+    [<Theory; InlineData(true); InlineData(false)>]
+    let ``SRTP partial application stays correct`` (optimize: bool) =
+        run optimize """
+let mutable reads = 0
+type Box() = member _.Value = reads <- reads + 1; 10
+let inline addThree (a: ^T) (b: ^T) (c: ^T) (x: ^T) = a + b + c + x
+let inline mapTwice ([<InlineIfLambda>] f: ^U -> ^U) (x: ^U) = f (f x)
+let r = mapTwice (addThree 1 (Box()).Value 3) 5
+if r <> 33 then failwithf "Expected 33 but got %d" r
+if reads <> 1 then failwithf "Captured getter evaluated %d times, expected 1" reads
+"""