WIP: don't float out lets between lambdaswip/T15606

This is incomplete work in prograss on Trac #T15606

4 files changed, 49 insertions, 15 deletions

diff --git a/compiler/main/DynFlags.hs b/compiler/main/DynFlags.hs
index 9f0ba57bf5..f281aba7be 100644
--- a/compiler/main/DynFlags.hs
+++ b/compiler/main/DynFlags.hs
@@ -451,6 +451,7 @@ data GeneralFlag
    | Opt_KillOneShot
    | Opt_FullLaziness
    | Opt_FloatIn
+   | Opt_FloatBetweenLambdas
    | Opt_LateSpecialise
    | Opt_Specialise
    | Opt_SpecialiseAggressively
@@ -3930,6 +3931,7 @@ fFlagsDeps = [
   flagSpec "external-interpreter"             Opt_ExternalInterpreter,
   flagSpec "flat-cache"                       Opt_FlatCache,
   flagSpec "float-in"                         Opt_FloatIn,
+  flagSpec "float-between-lambdas"            Opt_FloatBetweenLambdas,
   flagSpec "force-recomp"                     Opt_ForceRecomp,
   flagSpec "ignore-optim-changes"             Opt_IgnoreOptimChanges,
   flagSpec "ignore-hpc-changes"               Opt_IgnoreHpcChanges,
diff --git a/compiler/simplCore/CoreMonad.hs b/compiler/simplCore/CoreMonad.hs
index 0c5d8d9fd2..155976ca0d 100644
--- a/compiler/simplCore/CoreMonad.hs
+++ b/compiler/simplCore/CoreMonad.hs
@@ -195,7 +195,12 @@ data FloatOutSwitches = FloatOutSwitches {
                              --            based on arity information.
                              -- See Note [Floating over-saturated applications]
                              -- in SetLevels
-  floatToTopLevelOnly :: Bool      -- ^ Allow floating to the top level only.
+  floatToTopLevelOnly :: Bool,     -- ^ Allow floating to the top level only.
+
+  floatBetweenLambdas :: Bool -- True <=> \x. let ... in \y... (x+1)...
+                              --          Float out the (x+1)
+                              -- False <=> do not do so
+                              -- See Trac #15606
   }
 instance Outputable FloatOutSwitches where
     ppr = pprFloatOutSwitches
@@ -206,6 +211,7 @@ pprFloatOutSwitches sw
      sep $ punctuate comma $
      [ text "Lam ="    <+> ppr (floatOutLambdas sw)
      , text "Consts =" <+> ppr (floatOutConstants sw)
+     , text "Between =" <+> ppr (floatBetweenLambdas sw)
      , text "OverSatApps ="   <+> ppr (floatOutOverSatApps sw) ])
 
 -- The core-to-core pass ordering is derived from the DynFlags:
diff --git a/compiler/simplCore/SetLevels.hs b/compiler/simplCore/SetLevels.hs
index b8212c72f3..71afb22a8e 100644
--- a/compiler/simplCore/SetLevels.hs
+++ b/compiler/simplCore/SetLevels.hs
@@ -401,7 +401,7 @@ lvlApp env orig_expr ((_,AnnVar fn), args)
   , arity > 0
   , arity < n_val_args
   , Nothing <- isClassOpId_maybe fn
-  =  do { rargs' <- mapM (lvlNonTailMFE env False) rargs
+  =  do { rargs' <- mapM (lvlNonTailMFE env_arg False) rargs
         ; lapp'  <- lvlNonTailMFE env False lapp
         ; return (foldl' App lapp' rargs') }
 
@@ -411,6 +411,8 @@ lvlApp env orig_expr ((_,AnnVar fn), args)
             -- Note [Floating to the top]
        ; return (foldl' App (lookupVar env fn) args') }
   where
+    env_arg    = switchBumpingOn env
+
     n_val_args = count (isValArg . deAnnotate) args
     arity      = idArity fn
 
@@ -439,10 +441,10 @@ lvlApp env orig_expr ((_,AnnVar fn), args)
     lvl_arg :: [Demand] -> CoreExprWithFVs -> LvlM ([Demand], LevelledExpr)
     lvl_arg strs arg | (str1 : strs') <- strs
                      , is_val_arg arg
-                     = do { arg' <- lvlMFE env (isStrictDmd str1) arg
+                     = do { arg' <- lvlMFE env_arg (isStrictDmd str1) arg
                           ; return (strs', arg') }
                      | otherwise
-                     = do { arg' <- lvlMFE env False arg
+                     = do { arg' <- lvlMFE env_arg False arg
                           ; return (strs, arg') }
 
 lvlApp env _ (fun, args)
@@ -1196,13 +1198,14 @@ lvlFloatRhs abs_vars dest_lvl env rec is_bot mb_join_arity rhs
                       = collectNAnnBndrs join_arity rhs
                       | otherwise
                       = collectAnnBndrs rhs
-    (env1, bndrs1)    = substBndrsSL NonRecursive env bndrs
+    env1              = switchBumpingOn env
+    (env2, bndrs1)    = substBndrsSL NonRecursive env1 bndrs
     all_bndrs         = abs_vars ++ bndrs1
     (body_env, bndrs') | Just _ <- mb_join_arity
-                      = lvlJoinBndrs env1 dest_lvl rec all_bndrs
-                      | otherwise
-                      = case lvlLamBndrs env1 dest_lvl all_bndrs of
-                          (env2, bndrs') -> (placeJoinCeiling env2, bndrs')
+                       = lvlJoinBndrs env2 dest_lvl rec all_bndrs
+                       | otherwise
+                       = case lvlLamBndrs env2 dest_lvl all_bndrs of
+                           (env3, bndrs') -> (placeJoinCeiling env3, bndrs')
         -- The important thing here is that we call lvlLamBndrs on
         -- all these binders at once (abs_vars and bndrs), so they
         -- all get the same major level.  Otherwise we create stupid
@@ -1279,12 +1282,20 @@ substBndrsSL is_rec env@(LE { le_subst = subst, le_env = id_env }) bndrs
 
 lvlLamBndrs :: LevelEnv -> Level -> [OutVar] -> (LevelEnv, [LevelledBndr])
 -- Compute the levels for the binders of a lambda group
+-- Bump a major level if
+--      Any "major binder"
+-- and  le_bump is True
+--
+-- If we bump a major level,
+--   then set le_bump to False if floatBetweenLambdas is False
 lvlLamBndrs env lvl bndrs
-  = lvlBndrs env new_lvl bndrs
-  where
-    new_lvl | any is_major bndrs = incMajorLvl lvl
-            | otherwise          = incMinorLvl lvl
+  | le_bump env
+  , any is_major bndrs
+  = lvlBndrs (switchBumpingOff env) (incMajorLvl lvl) bndrs
 
+  | otherwise
+  = lvlBndrs env (incMinorLvl lvl) bndrs
+  where
     is_major bndr = isId bndr && not (isProbablyOneShotLambda bndr)
        -- The "probably" part says "don't float things out of a
        -- probable one-shot lambda"
@@ -1423,6 +1434,9 @@ countFreeIds = nonDetFoldUDFM add 0
 data LevelEnv
   = LE { le_switches :: FloatOutSwitches
        , le_ctxt_lvl :: Level           -- The current level
+       , le_bump     :: Bool            -- True <=> bump major level when you meet
+                                        --          a value lambda
+                                        -- False <=> do not bump
        , le_lvl_env  :: VarEnv Level    -- Domain is *post-cloned* TyVars and Ids
        , le_join_ceil:: Level           -- Highest level to which joins float
                                         -- Invariant: always >= le_ctxt_lvl
@@ -1435,6 +1449,13 @@ data LevelEnv
        , le_env      :: IdEnv ([OutVar], LevelledExpr)  -- Domain is pre-cloned Ids
     }
 
+switchBumpingOff :: LevelEnv -> LevelEnv
+switchBumpingOff env@(LE { le_switches = sw })
+  = env { le_bump = floatBetweenLambdas sw }
+
+switchBumpingOn :: LevelEnv -> LevelEnv
+switchBumpingOn env = env { le_bump = True }
+
 {- Note [le_subst and le_env]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 We clone let- and case-bound variables so that they are still distinct
@@ -1470,6 +1491,7 @@ initialEnv :: FloatOutSwitches -> LevelEnv
 initialEnv float_lams
   = LE { le_switches = float_lams
        , le_ctxt_lvl = tOP_LEVEL
+       , le_bump     = True
        , le_join_ceil = panic "initialEnv"
        , le_lvl_env = emptyVarEnv
        , le_subst = emptySubst
diff --git a/compiler/simplCore/SimplCore.hs b/compiler/simplCore/SimplCore.hs
index 168ece971c..a7be25b0fc 100644
--- a/compiler/simplCore/SimplCore.hs
+++ b/compiler/simplCore/SimplCore.hs
@@ -136,6 +136,7 @@ getCoreToDo dflags
     eta_expand_on = gopt Opt_DoLambdaEtaExpansion         dflags
     ww_on         = gopt Opt_WorkerWrapper                dflags
     static_ptrs   = xopt LangExt.StaticPointers           dflags
+    float_between = gopt Opt_FloatBetweenLambdas          dflags
 
     maybe_rule_check phase = runMaybe rule_check (CoreDoRuleCheck phase)
 
@@ -196,6 +197,7 @@ getCoreToDo dflags
           , floatOutConstants = True
           , floatOutOverSatApps = False
           , floatToTopLevelOnly = True
+          , floatBetweenLambdas = False
           }
         ]
 
@@ -224,8 +226,9 @@ getCoreToDo dflags
 
         if full_laziness then
            CoreDoFloatOutwards FloatOutSwitches {
-                                 floatOutLambdas   = Just 0,
-                                 floatOutConstants = True,
+                                 floatOutLambdas     = Just 0,
+                                 floatBetweenLambdas = float_between,
+                                 floatOutConstants   = True,
                                  floatOutOverSatApps = False,
                                  floatToTopLevelOnly = False }
                 -- Was: gentleFloatOutSwitches
@@ -285,6 +288,7 @@ getCoreToDo dflags
         runWhen full_laziness $
            CoreDoFloatOutwards FloatOutSwitches {
                                  floatOutLambdas     = floatLamArgs dflags,
+                                 floatBetweenLambdas = float_between,
                                  floatOutConstants   = True,
                                  floatOutOverSatApps = True,
                                  floatToTopLevelOnly = False },