Refactor floating of bindings (fiBind)

This is just a local refactoring.

I originally planned to try floating top-level bindings inwards,
but I backed off from that leaving only this (harmless) refactoring,
which has no behavioural effect.

I also make FloatIn into a ModGuts -> ModGuts function; again not
necessary now, but no harm either.

My attempt also used the new function CoreFVs.freeVarsBind; but
that too is a plausible refactorig of freeVars, so I left it in too.

Reviewers: austin, bgamari

Subscribers: thomie

Differential Revision: https://phabricator.haskell.org/D3180

3 files changed, 126 insertions, 106 deletions

diff --git a/compiler/coreSyn/CoreFVs.hs b/compiler/coreSyn/CoreFVs.hs
index 3a90ea0f03..511ffc1c9f 100644
--- a/compiler/coreSyn/CoreFVs.hs
+++ b/compiler/coreSyn/CoreFVs.hs
@@ -53,6 +53,7 @@ module CoreFVs (
         CoreBindWithFVs,        -- = AnnBind Id FVAnn
         CoreAltWithFVs,         -- = AnnAlt Id FVAnn
         freeVars,               -- CoreExpr -> CoreExprWithFVs
+        freeVarsBind,           -- CoreBind -> DVarSet -> (DVarSet, CoreBindWithFVs)
         freeVarsOf,             -- CoreExprWithFVs -> DIdSet
         freeVarsOfType,         -- CoreExprWithFVs -> TyCoVarSet
         freeVarsOfAnn, freeVarsOfTypeAnn,
@@ -701,6 +702,29 @@ stableUnfoldingFVs unf
 ************************************************************************
 -}
 
+freeVarsBind :: CoreBind
+             -> DVarSet                     -- Free vars of scope of binding
+             -> (CoreBindWithFVs, DVarSet)  -- Return free vars of binding + scope
+freeVarsBind (NonRec binder rhs) body_fvs
+  = ( AnnNonRec binder rhs2
+    , freeVarsOf rhs2 `unionFVs` body_fvs2
+                      `unionFVs` fvDVarSet (bndrRuleAndUnfoldingFVs binder) )
+    where
+      rhs2      = freeVars rhs
+      body_fvs2 = binder `delBinderFV` body_fvs
+
+freeVarsBind (Rec binds) body_fvs
+  = ( AnnRec (binders `zip` rhss2)
+    , delBindersFV binders all_fvs )
+  where
+    (binders, rhss) = unzip binds
+    rhss2        = map freeVars rhss
+    rhs_body_fvs = foldr (unionFVs . freeVarsOf) body_fvs rhss2
+    binders_fvs  = fvDVarSet $ mapUnionFV idRuleAndUnfoldingFVs binders
+    all_fvs      = rhs_body_fvs `unionFVs` binders_fvs
+            -- The "delBinderFV" happens after adding the idSpecVars,
+            -- since the latter may add some of the binders as fvs
+
 freeVars :: CoreExpr -> CoreExprWithFVs
 -- ^ Annotate a 'CoreExpr' with its (non-global) free type and value variables at every tree node
 freeVars = go
@@ -761,37 +785,14 @@ freeVars = go
                               where
                                  rhs2 = go rhs
 
-    go (Let (NonRec binder rhs) body)
-      = ( FVAnn { fva_fvs    = freeVarsOf rhs2
-                               `unionFVs` body_fvs
-                               `unionFVs` fvDVarSet
-                                            (bndrRuleAndUnfoldingFVs binder)
-                               -- Remember any rules; cf rhs_fvs above
-                , fva_ty_fvs = freeVarsOfType body2
-                , fva_ty     = exprTypeFV body2 }
-        , AnnLet (AnnNonRec binder rhs2) body2 )
-      where
-        rhs2     = go rhs
-        body2    = go body
-        body_fvs = binder `delBinderFV` freeVarsOf body2
-
-    go (Let (Rec binds) body)
-      = ( FVAnn { fva_fvs    = delBindersFV binders all_fvs
+    go (Let bind body)
+      = ( FVAnn { fva_fvs    = bind_fvs
                 , fva_ty_fvs = freeVarsOfType body2
                 , fva_ty     = exprTypeFV body2 }
-        , AnnLet (AnnRec (binders `zip` rhss2)) body2 )
+        , AnnLet bind2 body2 )
       where
-        (binders, rhss) = unzip binds
-
-        rhss2        = map go rhss
-        rhs_body_fvs = foldr (unionFVs . freeVarsOf) body_fvs rhss2
-        binders_fvs  = fvDVarSet $ mapUnionFV idRuleAndUnfoldingFVs binders
-        all_fvs      = rhs_body_fvs `unionFVs` binders_fvs
-            -- The "delBinderFV" happens after adding the idSpecVars,
-            -- since the latter may add some of the binders as fvs
-
-        body2    = go body
-        body_fvs = freeVarsOf body2
+        (bind2, bind_fvs) = freeVarsBind bind (freeVarsOf body2)
+        body2             = go body
 
     go (Cast expr co)
       = ( FVAnn (freeVarsOf expr2 `unionFVs` cfvs) (tyCoVarsOfTypeDSet to_ty) to_ty
diff --git a/compiler/simplCore/FloatIn.hs b/compiler/simplCore/FloatIn.hs
index 1fd969e638..cabdc3b430 100644
--- a/compiler/simplCore/FloatIn.hs
+++ b/compiler/simplCore/FloatIn.hs
@@ -20,17 +20,19 @@ module FloatIn ( floatInwards ) where
 
 import CoreSyn
 import MkCore
+import HscTypes         ( ModGuts(..) )
 import CoreUtils        ( exprIsDupable, exprIsExpandable,
                           exprOkForSideEffects, mkTicks )
 import CoreFVs
-import Id               ( isJoinId, isJoinId_maybe, isOneShotBndr, idType )
+import CoreMonad        ( CoreM )
+import Id               ( isOneShotBndr, idType )
 import Var
 import Type             ( isUnliftedType )
 import VarSet
 import Util
 import DynFlags
 import Outputable
-import Data.List( mapAccumL )
+import Data.List        ( mapAccumL )
 import BasicTypes       ( RecFlag(..), isRec )
 
 {-
@@ -38,14 +40,17 @@ Top-level interface function, @floatInwards@.  Note that we do not
 actually float any bindings downwards from the top-level.
 -}
 
-floatInwards :: DynFlags -> CoreProgram -> CoreProgram
-floatInwards dflags = map fi_top_bind
+floatInwards :: ModGuts -> CoreM ModGuts
+floatInwards pgm@(ModGuts { mg_binds = binds })
+  = do { dflags <- getDynFlags
+       ; return (pgm { mg_binds = map (fi_top_bind dflags) binds }) }
   where
-    fi_top_bind (NonRec binder rhs)
+    fi_top_bind dflags (NonRec binder rhs)
       = NonRec binder (fiExpr dflags [] (freeVars rhs))
-    fi_top_bind (Rec pairs)
+    fi_top_bind dflags (Rec pairs)
       = Rec [ (b, fiExpr dflags [] (freeVars rhs)) | (b, rhs) <- pairs ]
 
+
 {-
 ************************************************************************
 *                                                                      *
@@ -196,7 +201,6 @@ unlifted function arguments to be ok-for-speculation.
 
 Note [Join points]
 ~~~~~~~~~~~~~~~~~~
-
 Generally, we don't need to worry about join points - there are places we're
 not allowed to float them, but since they can't have occurrences in those
 places, we're not tempted.
@@ -334,77 +338,13 @@ idRuleAndUnfoldingVars of x.  No need for type variables, hence not using
 idFreeVars.
 -}
 
-fiExpr dflags to_drop (_,AnnLet (AnnNonRec id rhs) body)
-  = fiExpr dflags new_to_drop body
-  where
-    body_fvs = freeVarsOf body `delDVarSet` id
-    rhs_fvs  = freeVarsOf rhs
-
-    rule_fvs = idRuleAndUnfoldingVarsDSet id        -- See Note [extra_fvs (2): free variables of rules]
-    extra_fvs | noFloatIntoRhs (isJoinId id) NonRecursive rhs
-              = rule_fvs `unionDVarSet` freeVarsOf rhs
-              | otherwise
-              = rule_fvs
-        -- See Note [extra_fvs (1): avoid floating into RHS]
-        -- No point in floating in only to float straight out again
-        -- We *can't* float into ok-for-speculation unlifted RHSs
-        -- But do float into join points
-
-    [shared_binds, extra_binds, rhs_binds, body_binds]
-        = sepBindsByDropPoint dflags False
-            [extra_fvs, rhs_fvs, body_fvs]
-            (freeVarsOfType rhs `unionDVarSet` freeVarsOfType body)
-            to_drop
-
-    new_to_drop = body_binds ++                         -- the bindings used only in the body
-                  [FB (unitDVarSet id) rhs_fvs'
-                      (FloatLet (NonRec id rhs'))] ++   -- the new binding itself
-                  extra_binds ++                        -- bindings from extra_fvs
-                  shared_binds                          -- the bindings used both in rhs and body
-
-        -- Push rhs_binds into the right hand side of the binding
-    rhs'     = fiRhs dflags rhs_binds id rhs
-    rhs_fvs' = rhs_fvs `unionDVarSet` floatedBindsFVs rhs_binds `unionDVarSet` rule_fvs
-                        -- Don't forget the rule_fvs; the binding mentions them!
-
-fiExpr dflags to_drop (_,AnnLet (AnnRec bindings) body)
-  = fiExpr dflags new_to_drop body
+fiExpr dflags to_drop (_,AnnLet bind body)
+  = fiExpr dflags (after ++ new_float : before) body
+           -- to_drop is in reverse dependency order
   where
-    (ids, rhss) = unzip bindings
-    rhss_fvs = map freeVarsOf rhss
-    body_fvs = freeVarsOf body
-
-        -- See Note [extra_fvs (1,2)]
-    rule_fvs = mapUnionDVarSet idRuleAndUnfoldingVarsDSet ids
-    extra_fvs = rule_fvs `unionDVarSet`
-                unionDVarSets [ freeVarsOf rhs | (bndr, rhs) <- bindings
-                              , noFloatIntoRhs (isJoinId bndr) Recursive rhs ]
-
-    (shared_binds:extra_binds:body_binds:rhss_binds)
-        = sepBindsByDropPoint dflags False
-            (extra_fvs:body_fvs:rhss_fvs)
-            (freeVarsOfType body `unionDVarSet` mapUnionDVarSet freeVarsOfType rhss)
-            to_drop
-
-    new_to_drop = body_binds ++         -- the bindings used only in the body
-                  [FB (mkDVarSet ids) rhs_fvs'
-                      (FloatLet (Rec (fi_bind rhss_binds bindings)))] ++
-                                        -- The new binding itself
-                  extra_binds ++        -- Note [extra_fvs (1,2)]
-                  shared_binds          -- Used in more than one place
-
-    rhs_fvs' = unionDVarSets rhss_fvs `unionDVarSet`
-               unionDVarSets (map floatedBindsFVs rhss_binds) `unionDVarSet`
-               rule_fvs         -- Don't forget the rule variables!
-
-    -- Push rhs_binds into the right hand side of the binding
-    fi_bind :: [FloatInBinds]       -- one per "drop pt" conjured w/ fvs_of_rhss
-            -> [(Id, CoreExprWithFVs)]
-            -> [(Id, CoreExpr)]
-
-    fi_bind to_drops pairs
-      = [ (binder, fiRhs dflags to_drop binder rhs)
-        | ((binder, rhs), to_drop) <- zipEqual "fi_bind" pairs to_drops ]
+    (before, new_float, after) = fiBind dflags to_drop bind body_fvs body_ty_fvs
+    body_fvs    = freeVarsOf body
+    body_ty_fvs = freeVarsOfType body
 
 {-
 For @Case@, the possible ``drop points'' for the \tr{to_drop}
@@ -471,6 +411,84 @@ fiExpr dflags to_drop (_, AnnCase scrut case_bndr ty alts)
 
     fi_alt to_drop (con, args, rhs) = (con, args, fiExpr dflags to_drop rhs)
 
+------------------
+fiBind :: DynFlags
+       -> FloatInBinds      -- Binds we're trying to drop
+                            -- as far "inwards" as possible
+       -> CoreBindWithFVs   -- Input binding
+       -> DVarSet           -- Free in scope of binding
+       -> DVarSet           -- Free in type of body of binding
+       -> ( FloatInBinds    -- Land these before
+          , FloatInBind     -- The binding itself
+          , FloatInBinds)   -- Land these after
+
+fiBind dflags to_drop (AnnNonRec id rhs) body_fvs body_ty_fvs
+  = ( extra_binds ++ shared_binds          -- Land these before
+                                           -- See Note [extra_fvs (1,2)]
+    , FB (unitDVarSet id) rhs_fvs'         -- The new binding itself
+          (FloatLet (NonRec id rhs'))
+    , body_binds )                         -- Land these after
+
+  where
+    body_fvs2 = body_fvs `delDVarSet` id
+    rhs_fvs   = freeVarsOf rhs
+
+    rule_fvs = idRuleAndUnfoldingVarsDSet id        -- See Note [extra_fvs (2): free variables of rules]
+    extra_fvs | noFloatIntoRhs (isJoinId id) NonRecursive rhs
+              = rule_fvs `unionDVarSet` freeVarsOf rhs
+              | otherwise
+              = rule_fvs
+        -- See Note [extra_fvs (1): avoid floating into RHS]
+        -- No point in floating in only to float straight out again
+        -- We *can't* float into ok-for-speculation unlifted RHSs
+        -- But do float into join points
+
+    [shared_binds, extra_binds, rhs_binds, body_binds]
+        = sepBindsByDropPoint dflags False
+            [extra_fvs, rhs_fvs, body_fvs2]
+            (freeVarsOfType rhs `unionDVarSet` body_ty_fvs)
+            to_drop
+
+        -- Push rhs_binds into the right hand side of the binding
+    rhs'     = fiRhs dflags rhs_binds id rhs
+    rhs_fvs' = rhs_fvs `unionDVarSet` floatedBindsFVs rhs_binds `unionDVarSet` rule_fvs
+                        -- Don't forget the rule_fvs; the binding mentions them!
+
+fiBind dflags to_drop (AnnRec bindings) body_fvs body_ty_fvs
+  = ( extra_binds ++ shared_binds
+    , FB (mkDVarSet ids) rhs_fvs'
+         (FloatLet (Rec (fi_bind rhss_binds bindings)))
+    , body_binds )
+  where
+    (ids, rhss) = unzip bindings
+    rhss_fvs = map freeVarsOf rhss
+
+        -- See Note [extra_fvs (1,2)]
+    rule_fvs = mapUnionDVarSet idRuleAndUnfoldingVarsDSet ids
+    extra_fvs = rule_fvs `unionDVarSet`
+                unionDVarSets [ freeVarsOf rhs | (bndr, rhs) <- bindings
+                              , noFloatIntoRhs (isJoinId bndr) Recursive rhs ]
+
+    (shared_binds:extra_binds:body_binds:rhss_binds)
+        = sepBindsByDropPoint dflags False
+            (extra_fvs:body_fvs:rhss_fvs)
+            (body_ty_fvs `unionDVarSet` mapUnionDVarSet freeVarsOfType rhss)
+            to_drop
+
+    rhs_fvs' = unionDVarSets rhss_fvs `unionDVarSet`
+               unionDVarSets (map floatedBindsFVs rhss_binds) `unionDVarSet`
+               rule_fvs         -- Don't forget the rule variables!
+
+    -- Push rhs_binds into the right hand side of the binding
+    fi_bind :: [FloatInBinds]       -- one per "drop pt" conjured w/ fvs_of_rhss
+            -> [(Id, CoreExprWithFVs)]
+            -> [(Id, CoreExpr)]
+
+    fi_bind to_drops pairs
+      = [ (binder, fiRhs dflags to_drop binder rhs)
+        | ((binder, rhs), to_drop) <- zipEqual "fi_bind" pairs to_drops ]
+
+------------------
 fiRhs :: DynFlags -> FloatInBinds -> CoreBndr -> CoreExprWithFVs -> CoreExpr
 fiRhs dflags to_drop bndr rhs
   | Just join_arity <- isJoinId_maybe bndr
@@ -479,6 +497,7 @@ fiRhs dflags to_drop bndr rhs
   | otherwise
   = fiExpr dflags to_drop rhs
 
+------------------
 okToFloatInside :: [Var] -> Bool
 okToFloatInside bndrs = all ok bndrs
   where
diff --git a/compiler/simplCore/SimplCore.hs b/compiler/simplCore/SimplCore.hs
index 3c6c3115bc..7c89dc9099 100644
--- a/compiler/simplCore/SimplCore.hs
+++ b/compiler/simplCore/SimplCore.hs
@@ -432,7 +432,7 @@ doCorePass CoreLiberateCase          = {-# SCC "LiberateCase" #-}
                                        doPassD liberateCase
 
 doCorePass CoreDoFloatInwards        = {-# SCC "FloatInwards" #-}
-                                       doPassD floatInwards
+                                       floatInwards
 
 doCorePass (CoreDoFloatOutwards f)   = {-# SCC "FloatOutwards" #-}
                                        doPassDUM (floatOutwards f)