Improve static compilation of state machines

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

@@ -14,6 +14,7 @@
 * Fix FS0229 B-stream misalignment when reading metadata from assemblies compiled with LangVersion < 9.0, introduced by [#17706](https://github.com/dotnet/fsharp/pull/17706). ([PR #19260](https://github.com/dotnet/fsharp/pull/19260))
 * Fix FS3356 false positive for instance extension members with same name on different types, introduced by [#18821](https://github.com/dotnet/fsharp/pull/18821). ([PR #19260](https://github.com/dotnet/fsharp/pull/19260))
 * F# Scripts: Fix default reference paths resolving when an SDK directory is specified. ([PR #19270](https://github.com/dotnet/fsharp/pull/19270))
+* Improve static compilation of state machines. ([PR #19297](https://github.com/dotnet/fsharp/pull/19297))
 
 ### Added
 * FSharpType: add ImportILType ([PR #19300](https://github.com/dotnet/fsharp/pull/19300))

src/Compiler/Optimize/LowerStateMachines.fs

@@ -189,6 +189,11 @@ type LowerStateMachine(g: TcGlobals) =
             if sm_verbose then printfn "eliminating 'if __useResumableCode...'"
             BindResumableCodeDefinitions env thenExpr
 
+         // Look through debug points to find resumable code bindings inside
+         | Expr.DebugPoint (_, innerExpr) ->
+            let envR, _ = BindResumableCodeDefinitions env innerExpr
+            (envR, expr)
+
          | _ ->
             (env, expr)
 
@@ -235,7 +240,16 @@ type LowerStateMachine(g: TcGlobals) =
                     TryReduceApp env expandedExpr laterArgs
 
         | Expr.Let (bind, bodyExpr, m, _) -> 
-            match TryReduceApp env bodyExpr args with 
+            // If the binding returns resumable code, add it to the env so that
+            // references to it in the body can be resolved during reduction.
+            // This handles patterns like 'let cont = (fun () -> ...; Zero()) in cont()'
+            // generated by the optimizer for CE if-then branches.
+            let envR =
+                if isExpandVar g bind.Var then
+                    { env with ResumableCodeDefns = env.ResumableCodeDefns.Add bind.Var bind.Expr }
+                else
+                    env
+            match TryReduceApp envR bodyExpr args with 
             | Some bodyExpr2 -> Some (mkLetBind m bind bodyExpr2)
             | None -> None
 
@@ -308,6 +322,14 @@ type LowerStateMachine(g: TcGlobals) =
             | Some innerExpr2 -> Some (Expr.DebugPoint (dp, innerExpr2))
             | None -> None
 
+        // Resolve variables known to the env, e.g. locally-bound resumable code continuations
+        | Expr.Val (vref, _, _) when env.ResumableCodeDefns.ContainsVal vref.Deref ->
+            TryReduceApp env env.ResumableCodeDefns[vref.Deref] args
+
+        // Push through function applications by combining the arg lists
+        | Expr.App (f, _fty, _tyargs, fArgs, _m) ->
+            TryReduceApp env f (fArgs @ args)
+
         | _ -> 
             None
 
@@ -349,7 +371,19 @@ type LowerStateMachine(g: TcGlobals) =
 
     // Repeated top-down rewrite
     let makeRewriteEnv (env: env) = 
-        { PreIntercept = Some (fun cont e -> match TryReduceExpr env e [] id with Some e2 -> Some (cont e2) | None -> None)
+        { PreIntercept = Some (fun cont e ->
+            match e with
+            // Don't recurse into nested state machine expressions - they will be
+            // processed by their own LowerStateMachineExpr during codegen.
+            // This prevents modification of the nested machine's internal
+            // 'if __useResumableCode' patterns which select its dynamic fallback.
+            | _ when Option.isSome (IsStateMachineExpr g e) -> Some e
+            // Eliminate 'if __useResumableCode' - nested state machines are already
+            // guarded above, so any remaining occurrences at this level are from
+            // beta-reduced inline helpers and should take the static branch.
+            | IfUseResumableStateMachinesExpr g (thenExpr, _) -> Some (cont thenExpr)
+            | _ ->
+            match TryReduceExpr env e [] id with Some e2 -> Some (cont e2) | None -> None)
           PostTransform = (fun _ -> None)
           PreInterceptBinding = None
           RewriteQuotations=true