Handle binds & sequentials in mapping body

src/Compiler/Optimize/LowerComputedCollections.fs

@@ -524,25 +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))
 
-        | ValApp g g.seq_singleton_vref (_, [body], _) ->
-            ValueSome body
+        | 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.,
@@ -576,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
@@ -595,11 +614,14 @@ 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 = DebugPoints (SimpleSequential g body, debug); 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
@@ -610,15 +632,23 @@ let (|SeqCollectSingle|_|) g =
 
 /// 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

@@ -1236,11 +1236,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

@@ -4,6 +4,11 @@ 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|]