Optimize simple mappings with binds and/or sequentials before single yield (#17419)

Author: brianrourkeboll

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

@@ -2,6 +2,7 @@
 
 * Compiler hangs when compiling inline recursive invocation ([Issue #17376](https://github.com/dotnet/fsharp/issues/17376), [PR #17394](https://github.com/dotnet/fsharp/pull/17394))
 * Fix reporting IsFromComputationExpression only for CE builder type constructors and let bindings. ([PR #17375](https://github.com/dotnet/fsharp/pull/17375))
+* Optimize simple mappings in comprehensions when the body of the mapping has `let`-bindings and/or sequential expressions before a single yield. ([PR #17419](https://github.com/dotnet/fsharp/pull/17419))
 
 ### Added
 

src/Compiler/Optimize/LowerComputedCollections.fs

@@ -524,22 +524,46 @@ module Array =
                     )
             )
 
-/// f (); …; Seq.singleton x
+/// Matches Seq.singleton and returns the body expression.
+[<return: Struct>]
+let (|SeqSingleton|_|) g expr : Expr voption =
+    match expr with
+    | ValApp g g.seq_singleton_vref (_, [body], _) -> ValueSome body
+    | _ -> ValueNone
+
+/// Matches the compiled representation of the mapping in
+///
+///     for … in … do f (); …; yield …
+///     for … in … do let … = … in yield …
+///     for … in … do f (); …; …
+///     for … in … do let … = … in …
 ///
-/// E.g., in [for x in … do f (); …; yield x]
+/// i.e.,
+///
+///     f (); …; Seq.singleton …
+///     let … = … in Seq.singleton …
 [<return: Struct>]
-let (|SimpleSequential|_|) g expr : Expr voption =
+let (|SingleYield|_|) g expr : Expr voption =
     let rec loop expr cont =
         match expr with
-        | Expr.Sequential (expr1, DebugPoints (ValApp g g.seq_singleton_vref (_, [body], _), debug), kind, m) ->
-            ValueSome (cont (expr1, debug body, kind, m))
+        | Expr.Let (binding, DebugPoints (SeqSingleton g body, debug), m, frees) ->
+            ValueSome (cont (Expr.Let (binding, debug body, m, frees)))
+
+        | Expr.Let (binding, DebugPoints (body, debug), m, frees) ->
+            loop body (cont << fun body -> Expr.Let (binding, debug body, m, frees))
+
+        | Expr.Sequential (expr1, DebugPoints (SeqSingleton g body, debug), kind, m) ->
+            ValueSome (cont (Expr.Sequential (expr1, debug body, kind, m)))
 
         | Expr.Sequential (expr1, DebugPoints (body, debug), kind, m) ->
-            loop body (cont >> fun body -> Expr.Sequential (expr1, debug body, kind, m))
+            loop body (cont << fun body -> Expr.Sequential (expr1, debug body, kind, m))
+
+        | SeqSingleton g body ->
+            ValueSome (cont body)
 
         | _ -> ValueNone
 
-    loop expr Expr.Sequential
+    loop expr id
 
 /// Extracts any let-bindings or sequential
 /// expressions that directly precede the specified mapping application, e.g.,
@@ -573,11 +597,9 @@ let gatherPrelude ((|App|_|) : _ -> _ voption) expr =
 
 /// The representation used for
 ///
-/// for … in … -> …
-///
-/// and
-///
-/// for … in … do yield …
+///     for … in … -> …
+///     for … in … do yield …
+///     for … in … do …
 [<return: Struct>]
 let (|SeqMap|_|) g =
     gatherPrelude (function
@@ -592,30 +614,41 @@ let (|SeqMap|_|) g =
 
 /// The representation used for
 ///
-/// for … in … do f (); …; yield …
+///     for … in … do f (); …; yield …
+///     for … in … do let … = … in yield …
+///     for … in … do f (); …; …
+///     for … in … do let … = … in …
 [<return: Struct>]
 let (|SeqCollectSingle|_|) g =
     gatherPrelude (function
-        | ValApp g g.seq_collect_vref ([ty1; _; ty2], [Expr.Lambda (valParams = [loopVal]; bodyExpr = SimpleSequential g body; range = mIn) as mapping; input], mFor) ->
+        | ValApp g g.seq_collect_vref ([ty1; _; ty2], [Expr.Lambda (valParams = [loopVal]; bodyExpr = DebugPoints (SingleYield g body, debug); range = mIn) as mapping; input], mFor) ->
             let spIn = match mIn.NotedSourceConstruct with NotedSourceConstruct.InOrTo -> DebugPointAtInOrTo.Yes mIn | _ -> DebugPointAtInOrTo.No
             let spFor = DebugPointAtBinding.Yes mFor
             let spInWhile = match spIn with DebugPointAtInOrTo.Yes m -> DebugPointAtWhile.Yes m | DebugPointAtInOrTo.No -> DebugPointAtWhile.No
             let ranges = body.Range, spFor, spIn, mFor, mIn, spInWhile
-            ValueSome (ty1, ty2, input, mapping, loopVal, body, ranges)
+            ValueSome (ty1, ty2, input, mapping, loopVal, debug body, ranges)
 
         | _ -> ValueNone)
 
 /// for … in … -> …
 /// for … in … do yield …
+/// for … in … do …
 /// for … in … do f (); …; yield …
+/// for … in … do let … = … in yield …
+/// for … in … do f (); …; …
+/// for … in … do let … = … in …
 [<return: Struct>]
 let (|SimpleMapping|_|) g expr =
     match expr with
     // for … in … -> …
     // for … in … do yield …
+    // for … in … do …
     | ValApp g g.seq_delay_vref (_, [Expr.Lambda (bodyExpr = DebugPoints (SeqMap g (cont, (ty1, ty2, input, mapping, loopVal, body, ranges)), debug))], _)
 
     // for … in … do f (); …; yield …
+    // for … in … do let … = … in yield …
+    // for … in … do f (); …; …
+    // for … in … do let … = … in …
     | ValApp g g.seq_delay_vref (_, [Expr.Lambda (bodyExpr = DebugPoints (SeqCollectSingle g (cont, (ty1, ty2, input, mapping, loopVal, body, ranges)), debug))], _) ->
         ValueSome (debug >> cont, (ty1, ty2, input, mapping, loopVal, body, ranges))
 

src/Compiler/TypedTree/TypedTreeOps.fs

@@ -1271,11 +1271,14 @@ let rec stripDebugPoints expr =
     | Expr.DebugPoint (_, innerExpr) -> stripDebugPoints innerExpr
     | expr -> expr
 
-// Strip debug points and remember how to recrete them
+// Strip debug points and remember how to recreate them
 let (|DebugPoints|) expr =
-    match stripExpr expr with
-    | Expr.DebugPoint (dp, innerExpr) -> innerExpr, (fun e -> Expr.DebugPoint(dp, e))
-    | expr -> expr, id
+    let rec loop expr debug =
+        match stripExpr expr with
+        | Expr.DebugPoint (dp, innerExpr) -> loop innerExpr (debug << fun e -> Expr.DebugPoint (dp, e))
+        | expr -> expr, debug
+
+    loop expr id
 
 let mkCase (a, b) = TCase(a, b)
 

tests/FSharp.Compiler.ComponentTests/EmittedIL/ComputedCollections/ForNInRangeArrays.fs

@@ -2,6 +2,13 @@
 let f00 f g = [|for n in 1..10 do f (); g (); yield n|]
 let f000 f = [|for n in 1..10 do f (); yield n; yield n + 1|]
 let f0000 () = [|for n in 1..10 do yield n|]
+let f00000 () = [|for n in 1..10 do n|]
+let f000000 () = [|for n in 1..10 do let n = n in n|]
+let f0000000 () = [|for n in 1..10 do let n = n in yield n|]
+let f00000000 () = [|for n in 1..10 do let n = n in let n = n in yield n|]
+let f000000000 x y = [|for n in 1..10 do let foo = n + x in let bar = n + y in yield n + foo + bar|]
+let f0000000000 f g = [|for n in 1..10 do f (); g (); n|]
+let f00000000000 (f : unit -> int) (g : unit -> int) = [|for n in 1..10 do f (); g (); n|]
 let f1 () = [|for n in 1..10 -> n|]
 let f2 () = [|for n in 10..1 -> n|]
 let f3 () = [|for n in 1..1..10 -> n|]