Simple mapping optimization fixes (#17711)

dotnet · Sep 13, 2024 · 91d5a0a · 91d5a0a
1 parent 180d339
commit 91d5a0a
Show file tree

Hide file tree

Showing 6 changed files with 2,104 additions and 768 deletions.
diff --git a/docs/release-notes/.FSharp.Compiler.Service/9.0.100.md b/docs/release-notes/.FSharp.Compiler.Service/9.0.100.md
@@ -18,6 +18,7 @@
 * Fix IsUnionCaseTester throwing for non-methods/properties [#17301](/~https://github.com/dotnet/fsharp/pull/17634)
 * Consider `open type` used when the type is an enum and any of the enum cases is used unqualified. ([PR #17628](/~https://github.com/dotnet/fsharp/pull/17628))
 * Guard for possible StackOverflowException when typechecking non recursive modules and namespaces ([PR #17654](/~https://github.com/dotnet/fsharp/pull/17654))
+* Fixes for the optimization of simple mappings in array and list comprehensions. ([Issue #17708](/~https://github.com/dotnet/fsharp/issues/17708), [PR #17711](/~https://github.com/dotnet/fsharp/pull/17711))
 
 ### Added
 

diff --git a/src/Compiler/Optimize/LowerComputedCollections.fs b/src/Compiler/Optimize/LowerComputedCollections.fs
@@ -373,60 +373,70 @@ module Array =
         else any ilTy
 
     /// Makes the equivalent of an inlined call to Array.map.
-    let mkMap g m (mBody, _spFor, _spIn, mFor, mIn, spInWhile) srcArray srcIlTy destIlTy overallElemTy loopVal body =
-        let len = mkLdlen g mIn srcArray
-        let arrayTy = mkArrayType g overallElemTy
-
-        /// (# "newarr !0" type ('T) count : 'T array #)
-        let array =
-            mkAsmExpr
-                (
-                    [I_newarr (ILArrayShape.SingleDimensional, destIlTy)],
-                    [],
-                    [len],
-                    [arrayTy],
-                    m
-                )
-
-        let ldelem = mkIlInstr g I_ldelem (fun ilTy -> I_ldelem_any (ILArrayShape.SingleDimensional, ilTy)) srcIlTy
-        let stelem = mkIlInstr g I_stelem (fun ilTy -> I_stelem_any (ILArrayShape.SingleDimensional, ilTy)) destIlTy
-
-        let mapping =
-            mkCompGenLetIn m (nameof array) arrayTy array (fun (_, array) ->
-                mkCompGenLetMutableIn mFor "i" g.int32_ty (mkTypedZero g mIn g.int32_ty) (fun (iVal, i) ->
-                    let body =
-                        // Rebind the loop val to pull directly from the source array.
-                        let body = mkInvisibleLet mBody loopVal (mkAsmExpr ([ldelem], [], [srcArray; i], [loopVal.val_type], mBody)) body
-
-                        // destArray[i] <- body srcArray[i]
-                        let setArrSubI = mkAsmExpr ([stelem], [], [array; i; body], [], mIn)
-
-                        // i <- i + 1
-                        let incrI = mkValSet mIn (mkLocalValRef iVal) (mkAsmExpr ([AI_add], [], [i; mkTypedOne g mIn g.int32_ty], [g.int32_ty], mIn))
-
-                        mkSequential mIn setArrSubI incrI
-
-                    let guard = mkILAsmClt g mFor i (mkLdlen g mFor array)
+    let mkMap g m (mBody, _spFor, _spIn, mFor, mIn, spInWhile) srcArray srcIlTy destIlTy overallElemTy (loopVal: Val) body =
+        mkCompGenLetIn m (nameof srcArray) (tyOfExpr g srcArray) srcArray (fun (_, srcArray) ->
+            let len = mkLdlen g mIn srcArray
+            let arrayTy = mkArrayType g overallElemTy
+
+            /// (# "newarr !0" type ('T) count : 'T array #)
+            let array =
+                mkAsmExpr
+                    (
+                        [I_newarr (ILArrayShape.SingleDimensional, destIlTy)],
+                        [],
+                        [len],
+                        [arrayTy],
+                        m
+                    )
 
-                    let loop =
-                        mkWhile
-                            g
-                            (
-                                spInWhile,
-                                NoSpecialWhileLoopMarker,
-                                guard,
-                                body,
-                                mIn
-                            )
+            let ldelem = mkIlInstr g I_ldelem (fun ilTy -> I_ldelem_any (ILArrayShape.SingleDimensional, ilTy)) srcIlTy
+            let stelem = mkIlInstr g I_stelem (fun ilTy -> I_stelem_any (ILArrayShape.SingleDimensional, ilTy)) destIlTy
 
-                    // while i < array.Length do <body> done
-                    // array
-                    mkSequential m loop array
+            let mapping =
+                mkCompGenLetIn m (nameof array) arrayTy array (fun (_, array) ->
+                    mkCompGenLetMutableIn mFor "i" g.int32_ty (mkTypedZero g mIn g.int32_ty) (fun (iVal, i) ->
+                        let body =
+                            // If the loop val is used in the loop body,
+                            // rebind it to pull directly from the source array.
+                            // Otherwise, don't bother reading from the source array at all.
+                            let body =
+                                let freeLocals = (freeInExpr CollectLocals body).FreeLocals
+
+                                if freeLocals.Contains loopVal then
+                                    mkInvisibleLet mBody loopVal (mkAsmExpr ([ldelem], [], [srcArray; i], [loopVal.val_type], mBody)) body
+                                else
+                                    body
+
+                            // destArray[i] <- body srcArray[i]
+                            let setArrSubI = mkAsmExpr ([stelem], [], [array; i; body], [], mIn)
+
+                            // i <- i + 1
+                            let incrI = mkValSet mIn (mkLocalValRef iVal) (mkAsmExpr ([AI_add], [], [i; mkTypedOne g mIn g.int32_ty], [g.int32_ty], mIn))
+
+                            mkSequential mIn setArrSubI incrI
+
+                        let guard = mkILAsmClt g mFor i (mkLdlen g mFor array)
+
+                        let loop =
+                            mkWhile
+                                g
+                                (
+                                    spInWhile,
+                                    NoSpecialWhileLoopMarker,
+                                    guard,
+                                    body,
+                                    mIn
+                                )
+
+                        // while i < array.Length do <body> done
+                        // array
+                        mkSequential m loop array
+                    )
                 )
-            )
 
-        // Add a debug point at the `for`, before anything gets evaluated.
-        Expr.DebugPoint (DebugPointAtLeafExpr.Yes mFor, mapping)
+            // Add a debug point at the `for`, before anything gets evaluated.
+            Expr.DebugPoint (DebugPointAtLeafExpr.Yes mFor, mapping)
+        )
 
     /// Whether to check for overflow when converting a value to a native int.
     [<NoEquality; NoComparison>]
@@ -558,6 +568,9 @@ let (|SingleYield|_|) g expr : Expr voption =
         | Expr.Sequential (expr1, DebugPoints (body, debug), kind, m) ->
             loop body (cont << fun body -> Expr.Sequential (expr1, debug body, kind, m))
 
+        | Expr.Match (debugPoint, mInput, decision, [|TTarget (boundVals, DebugPoints (SeqSingleton g body, debug), isStateVarFlags)|], mFull, exprType) ->
+            ValueSome (cont (Expr.Match (debugPoint, mInput, decision, [|TTarget (boundVals, debug body, isStateVarFlags)|], mFull, exprType)))
+
         | SeqSingleton g body ->
             ValueSome (cont body)
 

diff --git a/tests/FSharp.Compiler.ComponentTests/EmittedIL/ComputedCollections/ForXInArray_ToArray.fs b/tests/FSharp.Compiler.ComponentTests/EmittedIL/ComputedCollections/ForXInArray_ToArray.fs
@@ -21,3 +21,17 @@ let f8 f g (array: int array) = [|let y = f () in let z = g () in for x in array
 let f9 f g (array: int array) = [|let y = f () in g (); for x in array -> x + y|]
 let f10 f g (array: int array) = [|f (); g (); for x in array -> x|]
 let f11 f g (array: int array) = [|f (); let y = g () in for x in array -> x + y|]
+let f12 (f: unit -> int array) y = [|for x in f () -> x + y|]
+
+// /~https://github.com/dotnet/fsharp/issues/17708
+// Don't read or rebind the loop variable when it is not in scope in the body.
+let ``for _ in Array.groupBy id [||] do ...`` () = [|for _ in Array.groupBy id [||] do 0|]
+let ``for _ | _ in Array.groupBy id [||] do ...`` () = [|for _ | _ in Array.groupBy id [||] do 0|]
+let ``for _ & _ in Array.groupBy id [||] do ...`` () = [|for _ & _ in Array.groupBy id [||] do 0|]
+let ``for _, _group in Array.groupBy id [||] do ...`` () = [|for _, _group in Array.groupBy id [||] do 0|]
+let ``for _, group in Array.groupBy id [||] do ...`` () = [|for _, group in Array.groupBy id [||] do group.Length|]
+let ``for 1 | 2 | _ in ...`` () = [|for 1 | 2 | _ in [||] do 0|]
+let ``for Failure _ | _ in ...`` () = [|for Failure _ | _ in [||] do 0|]
+let ``for true | false in ...`` () = [|for true | false in [||] do 0|]
+let ``for true | _ in ...`` () = [|for true | _ in [||] do 0|]
+let ``for _ | true in ...`` () = [|for _ | true in [||] do 0|]