compiler: Fix performance regression in backport of "Make FloatIn robust to shadowing" (6206cb9287f3f6e70c669660a646a65274870d2b)

In 9.4, we have

noFloatIntoArg :: CoreExprWithFVs' -> Type -> Bool
noFloatIntoArg expr expr_ty = ...

But in master when 6206cb92 landed, after "Drop the app invariant" (dcf30da8) we had

noFloatIntoArg :: CoreExprWithFVs' -> Bool
noFloatIntoArg expr = ...

When deciding whether to float things into the argument of a function, in 9.4 we
must know the type of the argument. This was previously done by extracting the
type of the argument from the function type, computed as we walked through all
the arguments.

However, this backport regressed compile time performance due to allocations by
`exprType` particularly in T16577 and T5642, where it turns out that computing
the type of the arguments to a function is quite expensive.

Instead, we can compute the type of the argument by looking at the argument
term directly, which turns out to be much faster and eliminates the performance
regression.

1 files changed, 8 insertions, 10 deletions

diff --git a/compiler/GHC/Core/Opt/FloatIn.hs b/compiler/GHC/Core/Opt/FloatIn.hs
index a6a504e664..227f921c3c 100644
--- a/compiler/GHC/Core/Opt/FloatIn.hs
+++ b/compiler/GHC/Core/Opt/FloatIn.hs
@@ -177,7 +177,6 @@ fiExpr platform to_drop ann_expr@(_,AnnApp {})
            -- length ann_args = length arg_fvs = length arg_drops
   where
     (ann_fun, ann_args, ticks) = collectAnnArgsTicks tickishFloatable ann_expr
-    fun_ty  = exprType (deAnnotate ann_fun)
     fun_fvs = freeVarsOf ann_fun
 
     (drop_here, fun_drop : arg_drops)
@@ -189,7 +188,7 @@ fiExpr platform to_drop ann_expr@(_,AnnApp {})
          -- lists without evaluating extra_fvs, and hence without
          -- peering into each argument
 
-    ((_,here_fvs), arg_fvs) = mapAccumL add_arg (fun_ty, here_fvs0) ann_args
+    (here_fvs, arg_fvs) = mapAccumL add_arg here_fvs0 ann_args
     here_fvs0 = case ann_fun of
                    (_, AnnVar _) -> fun_fvs
                    _             -> emptyDVarSet
@@ -198,15 +197,14 @@ fiExpr platform to_drop ann_expr@(_,AnnApp {})
           -- join point, floating it in isn't especially harmful but it's
           -- useless since the simplifier will immediately float it back out.)
 
-    add_arg :: (Type,FreeVarSet) -> CoreExprWithFVs -> ((Type,FreeVarSet),FreeVarSet)
-    add_arg (fun_ty, here_fvs) (arg_fvs, AnnType ty)
-      = ((piResultTy fun_ty ty, here_fvs), arg_fvs)
-    -- We can't float into some arguments, so put them into the here_fvs
-    add_arg (fun_ty, here_fvs) (arg_fvs, arg)
-      | noFloatIntoArg arg arg_ty = ((res_ty,here_fvs `unionDVarSet` arg_fvs), emptyDVarSet)
-      | otherwise          = ((res_ty,here_fvs), arg_fvs)
+    add_arg :: FreeVarSet -> CoreExprWithFVs -> (FreeVarSet,FreeVarSet)
+    add_arg here_fvs (arg_fvs, AnnType _)
+      = (here_fvs, arg_fvs)
+    add_arg here_fvs (arg_fvs, arg)
+      | noFloatIntoArg arg arg_ty = (here_fvs `unionDVarSet` arg_fvs, emptyDVarSet)
+      | otherwise          = (here_fvs, arg_fvs)
       where
-       (_, arg_ty, res_ty) = splitFunTy fun_ty
+       arg_ty = exprType $ deAnnotate' arg
 
 {- Note [Dead bindings]
 ~~~~~~~~~~~~~~~~~~~~~~~