Improve static compilation of state machines

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

@@ -20,6 +20,7 @@
 * 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))
 * Fix graph-based type checking incorrectly resolving dependencies when the same module name is defined across multiple files in the same namespace. ([PR #19280](https://github.com/dotnet/fsharp/pull/19280))
 * 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))

docs/release-notes/.VisualStudio/18.vNext.md

@@ -3,5 +3,6 @@
 * Fixed Rename incorrectly renaming `get` and `set` keywords for properties with explicit accessors. ([Issue #18270](https://github.com/dotnet/fsharp/issues/18270), [PR #19252](https://github.com/dotnet/fsharp/pull/19252))
 * Fixed Find All References crash when F# project contains non-F# files like `.cshtml`. ([Issue #16394](https://github.com/dotnet/fsharp/issues/16394), [PR #19252](https://github.com/dotnet/fsharp/pull/19252))
 * Find All References for external DLL symbols now only searches projects that reference the specific assembly. ([Issue #10227](https://github.com/dotnet/fsharp/issues/10227), [PR #19252](https://github.com/dotnet/fsharp/pull/19252))
+* Improve static compilation of state machines. ([PR #19297](https://github.com/dotnet/fsharp/pull/19297))
 
 ### Changed

src/Compiler/CodeGen/IlxGen.fs

@@ -1261,6 +1261,11 @@ and IlxGenEnv =
         intraAssemblyInfo: IlxGenIntraAssemblyInfo
 
         realsig: bool
+
+        /// Expression definitions of variables returning resumable code from outer scopes.
+        /// Used by state machine lowering to resolve otherwise-free expand variables
+        /// when the state machine is inside a lambda whose outer let-binding provides the definition.
+        resumableCodeDefinitions: ValMap<Expr>
     }
 
     override _.ToString() = "<IlxGenEnv>"
@@ -3004,7 +3009,9 @@ and GenExprPreSteps (cenv: cenv) (cgbuf: CodeGenBuffer) eenv expr sequel =
                 true
             | None ->
 
-                match LowerStateMachineExpr cenv.g expr with
+                let smResult = LowerStateMachineExpr cenv.g eenv.resumableCodeDefinitions expr
+
+                match smResult with
                 | LoweredStateMachineResult.Lowered res ->
                     let eenv = RemoveTemplateReplacement eenv
                     checkLanguageFeatureError cenv.g.langVersion LanguageFeature.ResumableStateMachines expr.Range
@@ -3499,6 +3506,16 @@ and GenLinearExpr cenv cgbuf eenv expr sequel preSteps (contf: FakeUnit -> FakeU
             GenDebugPointForBind cenv cgbuf bind
             GenBindingAfterDebugPoint cenv cgbuf eenv bind false (Some startMark)
 
+            // Track expand-var (resumable code) definitions so state machine lowering
+            // inside nested lambdas can resolve otherwise-free expand variables.
+            let eenv =
+                if isReturnsResumableCodeTy cenv.g bind.Var.TauType then
+                    { eenv with
+                        resumableCodeDefinitions = eenv.resumableCodeDefinitions.Add bind.Var bind.Expr
+                    }
+                else
+                    eenv
+
             // Generate the body
             GenLinearExpr cenv cgbuf eenv body (EndLocalScope(sequel, endMark)) true contf
 
@@ -12048,6 +12065,7 @@ let GetEmptyIlxGenEnv (g: TcGlobals) ccu =
         intraAssemblyInfo = IlxGenIntraAssemblyInfo.Create()
         realsig = g.realsig
         initClassFieldSpec = None
+        resumableCodeDefinitions = ValMap<_>.Empty
     }
 
 type IlxGenResults =

src/Compiler/Optimize/LowerStateMachines.fs

@@ -167,7 +167,7 @@ type LoweredStateMachineResult =
     | NotAStateMachine
 
 /// Used to scope the action of lowering a state machine expression
-type LowerStateMachine(g: TcGlobals) =
+type LowerStateMachine(g: TcGlobals, outerResumableCodeDefns: ValMap<Expr>) =
 
     let mutable pcCount = 0
     let genPC() =
@@ -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 
@@ -375,7 +409,9 @@ type LowerStateMachine(g: TcGlobals) =
     [<return: Struct>]
     let (|ExpandedStateMachineInContext|_|) inputExpr = 
         // All expanded resumable code state machines e.g. 'task { .. }' begin with a bind of @builder or 'defn'
-        let env, expr = BindResumableCodeDefinitions env.Empty inputExpr 
+        // Seed the env with any expand-var definitions from outer scopes (e.g. across lambda boundaries)
+        let initialEnv = { env.Empty with ResumableCodeDefns = outerResumableCodeDefns }
+        let env, expr = BindResumableCodeDefinitions initialEnv inputExpr 
         match expr with
         | StructStateMachineExpr g 
                (dataTy, 
@@ -858,8 +894,8 @@ type LowerStateMachine(g: TcGlobals) =
             let env, codeExprR = RepeatBindAndApplyOuterDefinitions env codeExpr
             let frees = (freeInExpr CollectLocals overallExpr).FreeLocals
 
-            if frees |> Zset.exists (isExpandVar g) then 
-                let nonfree = frees |> Zset.elements |> List.filter (isExpandVar g) |> List.map (fun v -> v.DisplayName) |> String.concat ","
+            if frees |> Zset.exists (fun v -> isExpandVar g v && not (env.ResumableCodeDefns.ContainsVal v)) then 
+                let nonfree = frees |> Zset.elements |> List.filter (fun v -> isExpandVar g v && not (env.ResumableCodeDefns.ContainsVal v)) |> List.map (fun v -> v.DisplayName) |> String.concat ","
                 let msg = FSComp.SR.reprResumableCodeValueHasNoDefinition(nonfree)
                 fallback msg
             else
@@ -913,12 +949,12 @@ type LowerStateMachine(g: TcGlobals) =
             let msg = FSComp.SR.reprStateMachineInvalidForm()
             fallback msg
 
-let LowerStateMachineExpr g (overallExpr: Expr) : LoweredStateMachineResult =
+let LowerStateMachineExpr g (outerResumableCodeDefns: ValMap<Expr>) (overallExpr: Expr) : LoweredStateMachineResult =
     // Detect a state machine and convert it
     let stateMachine = IsStateMachineExpr g overallExpr
 
     match stateMachine with 
     | None -> LoweredStateMachineResult.NotAStateMachine
     | Some altExprOpt ->
 
-    LowerStateMachine(g).Apply(overallExpr, altExprOpt)
+    LowerStateMachine(g, outerResumableCodeDefns).Apply(overallExpr, altExprOpt)

src/Compiler/Optimize/LowerStateMachines.fsi

@@ -3,6 +3,7 @@
 module internal FSharp.Compiler.LowerStateMachines
 
 open FSharp.Compiler.TypedTree
+open FSharp.Compiler.TypedTreeOps
 open FSharp.Compiler.TcGlobals
 
 type LoweredStateMachine =
@@ -30,4 +31,5 @@ type LoweredStateMachineResult =
 
 /// Analyze a TAST expression to detect the elaborated form of a state machine expression, a special kind
 /// of object expression that uses special code generation constructs.
-val LowerStateMachineExpr: g: TcGlobals -> overallExpr: Expr -> LoweredStateMachineResult
+val LowerStateMachineExpr:
+    g: TcGlobals -> outerResumableCodeDefns: ValMap<Expr> -> overallExpr: Expr -> LoweredStateMachineResult

src/FSharp.Core/FSharp.Core.fsproj

@@ -18,8 +18,6 @@
     <OtherFlags>$(OtherFlags) --warnon:3520</OtherFlags>
     <!-- Turn off 57: Use of construct with Experimental attribute -->
     <OtherFlags>$(OtherFlags) --nowarn:57</OtherFlags>
-    <!-- Turn off 3511: state machine not compilable - expected for inlined functions defining state machine generators -->
-    <OtherFlags>$(OtherFlags) --nowarn:3511</OtherFlags>
     <!-- Turn off 3513: resumable code invocation' - expected for resumable code combinators -->
     <OtherFlags>$(OtherFlags) --nowarn:3513</OtherFlags>
     <OtherFlags>$(OtherFlags) --compiling-fslib --compiling-fslib-40 --maxerrors:100 --extraoptimizationloops:1</OtherFlags>