{-
(c) The GRASP/AQUA Project, Glasgow University, 1992-1998

************************************************************************
*                                                                      *
\section[FloatIn]{Floating Inwards pass}
*                                                                      *
************************************************************************

The main purpose of @floatInwards@ is floating into branches of a
case, so that we don't allocate things, save them on the stack, and
then discover that they aren't needed in the chosen branch.
-}

{-# LANGUAGE CPP #-}
{-# OPTIONS_GHC -fprof-auto #-}

module FloatIn ( floatInwards ) where

#include "HsVersions.h"

import GhcPrelude

import CoreSyn
import MkCore hiding    ( wrapFloats )
import HscTypes         ( ModGuts(..) )
import CoreUtils
import CoreFVs
import CoreMonad        ( CoreM )
import Id               ( isOneShotBndr, idType, isJoinId, isJoinId_maybe )
import Var
import Type
import VarSet
import Util
import DynFlags
import Outputable
-- import Data.List        ( mapAccumL )
import BasicTypes       ( RecFlag(..), isRec )

{-
Top-level interface function, @floatInwards@.  Note that we do not
actually float any bindings downwards from the top-level.
-}

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 dflags (NonRec binder rhs)
      = NonRec binder (fiExpr dflags [] (freeVars rhs))
    fi_top_bind dflags (Rec pairs)
      = Rec [ (b, fiExpr dflags [] (freeVars rhs)) | (b, rhs) <- pairs ]

{-
************************************************************************
*                                                                      *
\subsection{Mail from Andr\'e [edited]}
*                                                                      *
************************************************************************

{\em Will wrote: What??? I thought the idea was to float as far
inwards as possible, no matter what.  This is dropping all bindings
every time it sees a lambda of any kind.  Help! }

You are assuming we DO DO full laziness AFTER floating inwards!  We
have to [not float inside lambdas] if we don't.

If we indeed do full laziness after the floating inwards (we could
check the compilation flags for that) then I agree we could be more
aggressive and do float inwards past lambdas.

Actually we are not doing a proper full laziness (see below), which
was another reason for not floating inwards past a lambda.

This can easily be fixed.  The problem is that we float lets outwards,
but there are a few expressions which are not let bound, like case
scrutinees and case alternatives.  After floating inwards the
simplifier could decide to inline the let and the laziness would be
lost, e.g.

\begin{verbatim}
let a = expensive             ==> \b -> case expensive of ...
in \ b -> case a of ...
\end{verbatim}
The fix is
\begin{enumerate}
\item
to let bind the algebraic case scrutinees (done, I think) and
the case alternatives (except the ones with an
unboxed type)(not done, I think). This is best done in the
SetLevels.hs module, which tags things with their level numbers.
\item
do the full laziness pass (floating lets outwards).
\item
simplify. The simplifier inlines the (trivial) lets that were
 created but were not floated outwards.
\end{enumerate}

With the fix I think Will's suggestion that we can gain even more from
strictness by floating inwards past lambdas makes sense.

We still gain even without going past lambdas, as things may be
strict in the (new) context of a branch (where it was floated to) or
of a let rhs, e.g.
\begin{verbatim}
let a = something            case x of
in case x of                   alt1 -> case something of a -> a + a
     alt1 -> a + a      ==>    alt2 -> b
     alt2 -> b

let a = something           let b = case something of a -> a + a
in let b = a + a        ==> in (b,b)
in (b,b)
\end{verbatim}
Also, even if a is not found to be strict in the new context and is
still left as a let, if the branch is not taken (or b is not entered)
the closure for a is not built.

************************************************************************
*                                                                      *
\subsection{Main floating-inwards code}
*                                                                      *
************************************************************************
-}

type FreeVarSet  = DIdSet
type BoundVarSet = DIdSet

data FloatInBind = FB BoundVarSet FreeVarSet FloatBind
        -- The FreeVarSet is the free variables of the binding.  In the case
        -- of recursive bindings, the set doesn't include the bound
        -- variables.

type FloatInBinds = [FloatInBind]
        -- In reverse dependency order (innermost binder first)

fiExpr :: DynFlags
       -> FloatInBinds      -- Binds we're trying to drop
                            -- as far "inwards" as possible
       -> CoreExprWithFVs   -- Input expr
       -> CoreExpr          -- Result

fiExpr _ to_drop (_, AnnLit lit)     = wrapFloats to_drop (Lit lit)
                                       -- See Note [Dead bindings]
fiExpr _ to_drop (_, AnnType ty)     = ASSERT( null to_drop ) Type ty
fiExpr _ to_drop (_, AnnVar v)       = wrapFloats to_drop (Var v)
fiExpr _ to_drop (_, AnnCoercion co) = wrapFloats to_drop (Coercion co)
fiExpr dflags to_drop (_, AnnCast expr (co_ann, co))
  = wrapFloats (drop_here ++ co_drop) $
    Cast (fiExpr dflags e_drop expr) co
  where
    [drop_here, e_drop, co_drop]
      = sepBindsByDropPoint dflags False
          [freeVarsOf expr, freeVarsOfAnn co_ann]
          to_drop

{-
Applications: we do float inside applications, mainly because we
need to get at all the arguments.  The next simplifier run will
pull out any silly ones.
-}

fiExpr dflags to_drop ann_expr@(_,AnnApp {})
  = wrapFloats drop_here $ wrapFloats extra_drop $
    mkTicks ticks $
    mkApps (fiExpr dflags fun_drop ann_fun)
           (zipWith (fiExpr dflags) arg_drops ann_args)
  where
    (ann_fun, ann_args, ticks) = collectAnnArgsTicks tickishFloatable ann_expr
    fun_ty  = exprType (deAnnotate ann_fun)
    fun_fvs = freeVarsOf ann_fun
    arg_fvs = map freeVarsOf ann_args

    (drop_here : extra_drop : fun_drop : arg_drops)
       = sepBindsByDropPoint dflags False
                             (extra_fvs : fun_fvs : arg_fvs)
                             to_drop
         -- Shortcut behaviour: if to_drop is empty,
         -- sepBindsByDropPoint returns a suitable bunch of empty
         -- lists without evaluating extra_fvs, and hence without
         -- peering into each argument

    (_, extra_fvs) = foldl' add_arg (fun_ty, extra_fvs0) ann_args
    extra_fvs0 = case ann_fun of
                   (_, AnnVar _) -> fun_fvs
                   _             -> emptyDVarSet
          -- Don't float the binding for f into f x y z; see Note [Join points]
          -- for why we *can't* do it when f is a join point. (If f isn't a
          -- 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)
    add_arg (fun_ty, extra_fvs) (_, AnnType ty)
      = (piResultTy fun_ty ty, extra_fvs)

    add_arg (fun_ty, extra_fvs) (arg_fvs, arg)
      | noFloatIntoArg arg arg_ty
      = (res_ty, extra_fvs `unionDVarSet` arg_fvs)
      | otherwise
      = (res_ty, extra_fvs)
      where
       (arg_ty, res_ty) = splitFunTy fun_ty

{- Note [Dead bindings]
~~~~~~~~~~~~~~~~~~~~~~~
At a literal we won't usually have any floated bindings; the
only way that can happen is if the binding wrapped the literal
/in the original input program/.  e.g.
   case x of { DEFAULT -> 1# }
But, while this may be unusual it is not actually wrong, and it did
once happen (Trac #15696).

Note [Do not destroy the let/app invariant]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Watch out for
   f (x +# y)
We don't want to float bindings into here
   f (case ... of { x -> x +# y })
because that might destroy the let/app invariant, which requires
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.

We do need to be careful about jumps, however:

  joinrec j x y z = ... in
  jump j a b c

Previous versions often floated the definition of a recursive function into its
only non-recursive occurrence. But for a join point, this is a disaster:

  (joinrec j x y z = ... in
  jump j) a b c -- wrong!

Every jump must be exact, so the jump to j must have three arguments. Hence
we're careful not to float into the target of a jump (though we can float into
the arguments just fine).

Note [Floating in past a lambda group]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
* We must be careful about floating inside a value lambda.
  That risks losing laziness.
  The float-out pass might rescue us, but then again it might not.

* We must be careful about type lambdas too.  At one time we did, and
  there is no risk of duplicating work thereby, but we do need to be
  careful.  In particular, here is a bad case (it happened in the
  cichelli benchmark:
        let v = ...
        in let f = /\t -> \a -> ...
           ==>
        let f = /\t -> let v = ... in \a -> ...
  This is bad as now f is an updatable closure (update PAP)
  and has arity 0.

* Hack alert!  We only float in through one-shot lambdas,
  not (as you might guess) through lone big lambdas.
  Reason: we float *out* past big lambdas (see the test in the Lam
  case of FloatOut.floatExpr) and we don't want to float straight
  back in again.

  It *is* important to float into one-shot lambdas, however;
  see the remarks with noFloatIntoRhs.

So we treat lambda in groups, using the following rule:

 Float in if (a) there is at least one Id,
         and (b) there are no non-one-shot Ids

 Otherwise drop all the bindings outside the group.

This is what the 'go' function in the AnnLam case is doing.

(Join points are handled similarly: a join point is considered one-shot iff
it's non-recursive, so we float only into non-recursive join points.)

Urk! if all are tyvars, and we don't float in, we may miss an
      opportunity to float inside a nested case branch

Note [Floating coercions]
~~~~~~~~~~~~~~~~~~~~~~~~~
We could, in principle, have a coercion binding like
   case f x of co { DEFAULT -> e1 e2 }
It's not common to have a function that returns a coercion, but nothing
in Core prohibits it.  If so, 'co' might be mentioned in e1 or e2
/only in a type/.  E.g. suppose e1 was
  let (x :: Int |> co) = blah in blah2

But, with coercions appearing in types, there is a complication: we
might be floating in a "strict let" -- that is, a case. Case expressions
mention their return type. We absolutely can't float a coercion binding
inward to the point that the type of the expression it's about to wrap
mentions the coercion. So we include the union of the sets of free variables
of the types of all the drop points involved. If any of the floaters
bind a coercion variable mentioned in any of the types, that binder must
be dropped right away.

-}

fiExpr dflags to_drop lam@(_, AnnLam _ _)
  | noFloatIntoLam bndrs       -- Dump it all here
     -- NB: Must line up with noFloatIntoRhs (AnnLam...); see Trac #7088
  = wrapFloats to_drop (mkLams bndrs (fiExpr dflags [] body))

  | otherwise           -- Float inside
  = mkLams bndrs (fiExpr dflags to_drop body)

  where
    (bndrs, body) = collectAnnBndrs lam

{-
We don't float lets inwards past an SCC.
        ToDo: keep info on current cc, and when passing
        one, if it is not the same, annotate all lets in binds with current
        cc, change current cc to the new one and float binds into expr.
-}

fiExpr dflags to_drop (_, AnnTick tickish expr)
  | tickish `tickishScopesLike` SoftScope
  = Tick tickish (fiExpr dflags to_drop expr)

  | otherwise -- Wimp out for now - we could push values in
  = wrapFloats to_drop (Tick tickish (fiExpr dflags [] expr))

{-
For @Lets@, the possible ``drop points'' for the \tr{to_drop}
bindings are: (a)~in the body, (b1)~in the RHS of a NonRec binding,
or~(b2), in each of the RHSs of the pairs of a @Rec@.

Note that we do {\em weird things} with this let's binding.  Consider:
\begin{verbatim}
let
    w = ...
in {
    let v = ... w ...
    in ... v .. w ...
}
\end{verbatim}
Look at the inner \tr{let}.  As \tr{w} is used in both the bind and
body of the inner let, we could panic and leave \tr{w}'s binding where
it is.  But \tr{v} is floatable further into the body of the inner let, and
{\em then} \tr{w} will also be only in the body of that inner let.

So: rather than drop \tr{w}'s binding here, we add it onto the list of
things to drop in the outer let's body, and let nature take its
course.

Note [extra_fvs (1): avoid floating into RHS]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Consider let x=\y....t... in body.  We do not necessarily want to float
a binding for t into the RHS, because it'll immediately be floated out
again.  (It won't go inside the lambda else we risk losing work.)
In letrec, we need to be more careful still. We don't want to transform
        let x# = y# +# 1#
        in
        letrec f = \z. ...x#...f...
        in ...
into
        letrec f = let x# = y# +# 1# in \z. ...x#...f... in ...
because now we can't float the let out again, because a letrec
can't have unboxed bindings.

So we make "extra_fvs" which is the rhs_fvs of such bindings, and
arrange to dump bindings that bind extra_fvs before the entire let.

Note [extra_fvs (2): free variables of rules]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Consider
  let x{rule mentioning y} = rhs in body
Here y is not free in rhs or body; but we still want to dump bindings
that bind y outside the let.  So we augment extra_fvs with the
idRuleAndUnfoldingVars of x.  No need for type variables, hence not using
idFreeVars.
-}

fiExpr dflags to_drop (_,AnnLet bind body)
  = fiExpr dflags (after ++ new_float : before) body
           -- to_drop is in reverse dependency order
  where
    (before, new_float, after) = fiBind dflags to_drop bind body_fvs
    body_fvs    = freeVarsOf body

{- Note [Floating primops]
~~~~~~~~~~~~~~~~~~~~~~~~~~
We try to float-in a case expression over an unlifted type.  The
motivating example was Trac #5658: in particular, this change allows
array indexing operations, which have a single DEFAULT alternative
without any binders, to be floated inward.

SIMD primops for unpacking SIMD vectors into an unboxed tuple of unboxed
scalars also need to be floated inward, but unpacks have a single non-DEFAULT
alternative that binds the elements of the tuple. We now therefore also support
floating in cases with a single alternative that may bind values.

But there are wrinkles

* Which unlifted cases do we float? See PrimOp.hs
  Note [PrimOp can_fail and has_side_effects] which explains:
   - We can float-in can_fail primops, but we can't float them out.
   - But we can float a has_side_effects primop, but NOT inside a lambda,
     so for now we don't float them at all.
  Hence exprOkForSideEffects

* Because we can float can-fail primops (array indexing, division) inwards
  but not outwards, we must be careful not to transform
     case a /# b of r -> f (F# r)
  ===>
    f (case a /# b of r -> F# r)
  because that creates a new thunk that wasn't there before.  And
  because it can't be floated out (can_fail), the thunk will stay
  there.  Disaster!  (This happened in nofib 'simple' and 'scs'.)

  Solution: only float cases into the branches of other cases, and
  not into the arguments of an application, or the RHS of a let. This
  is somewhat conservative, but it's simple.  And it still hits the
  cases like Trac #5658.   This is implemented in sepBindsByJoinPoint;
  if is_case is False we dump all floating cases right here.

* Trac #14511 is another example of why we want to restrict float-in
  of case-expressions.  Consider
     case indexArray# a n of (# r #) -> writeArray# ma i (f r)
  Now, floating that indexing operation into the (f r) thunk will
  not create any new thunks, but it will keep the array 'a' alive
  for much longer than the programmer expected.

  So again, not floating a case into a let or argument seems like
  the Right Thing

For @Case@, the possible drop points for the 'to_drop'
bindings are:
  (a) inside the scrutinee
  (b) inside one of the alternatives/default (default FVs always /first/!).

-}

fiExpr dflags to_drop (_, AnnCase scrut case_bndr _ [(con,alt_bndrs,rhs)])
  | isUnliftedType (idType case_bndr)
  , exprOkForSideEffects (deAnnotate scrut)
      -- See Note [Floating primops]
  = wrapFloats shared_binds $
    fiExpr dflags (case_float : rhs_binds) rhs
  where
    case_float = FB (mkDVarSet (case_bndr : alt_bndrs)) scrut_fvs
                    (FloatCase scrut' case_bndr con alt_bndrs)
    scrut'     = fiExpr dflags scrut_binds scrut
    rhs_fvs    = freeVarsOf rhs `delDVarSetList` (case_bndr : alt_bndrs)
    scrut_fvs  = freeVarsOf scrut

    [shared_binds, scrut_binds, rhs_binds]
       = sepBindsByDropPoint dflags False
           [scrut_fvs, rhs_fvs]
           to_drop

fiExpr dflags to_drop (_, AnnCase scrut case_bndr ty alts)
  = wrapFloats drop_here1 $
    wrapFloats drop_here2 $
    Case (fiExpr dflags scrut_drops scrut) case_bndr ty
         (zipWith fi_alt alts_drops_s alts)
  where
        -- Float into the scrut and alts-considered-together just like App
    [drop_here1, scrut_drops, alts_drops]
       = sepBindsByDropPoint dflags False
           [scrut_fvs, all_alts_fvs]
           to_drop

        -- Float into the alts with the is_case flag set
    (drop_here2 : alts_drops_s)
      | [ _ ] <- alts = [] : [alts_drops]
      | otherwise     = sepBindsByDropPoint dflags True alts_fvs alts_drops

    scrut_fvs    = freeVarsOf scrut
    alts_fvs     = map alt_fvs alts
    all_alts_fvs = unionDVarSets alts_fvs
    alt_fvs (_con, args, rhs)
      = foldl' delDVarSet (freeVarsOf rhs) (case_bndr:args)
           -- Delete case_bndr and args from free vars of rhs
           -- to get free vars of alt

    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
       -> ( FloatInBinds    -- Land these before
          , FloatInBind     -- The binding itself
          , FloatInBinds)   -- Land these after

fiBind dflags to_drop (AnnNonRec id ann_rhs@(rhs_fvs, rhs)) body_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

    rule_fvs = bndrRuleAndUnfoldingVarsDSet id        -- See Note [extra_fvs (2): free variables of rules]
    extra_fvs | noFloatIntoRhs NonRecursive id rhs
              = rule_fvs `unionDVarSet` rhs_fvs
              | 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]
            to_drop

        -- Push rhs_binds into the right hand side of the binding
    rhs'     = fiRhs dflags rhs_binds id ann_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
  = ( 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 bndrRuleAndUnfoldingVarsDSet ids
    extra_fvs = rule_fvs `unionDVarSet`
                unionDVarSets [ rhs_fvs | (bndr, (rhs_fvs, rhs)) <- bindings
                              , noFloatIntoRhs Recursive bndr rhs ]

    (shared_binds:extra_binds:body_binds:rhss_binds)
        = sepBindsByDropPoint dflags False
            (extra_fvs:body_fvs:rhss_fvs)
            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
  , let (bndrs, body) = collectNAnnBndrs join_arity rhs
  = mkLams bndrs (fiExpr dflags to_drop body)
  | otherwise
  = fiExpr dflags to_drop rhs

------------------
noFloatIntoLam :: [Var] -> Bool
noFloatIntoLam bndrs = any bad bndrs
  where
    bad b = isId b && not (isOneShotBndr b)
    -- Don't float inside a non-one-shot lambda

noFloatIntoRhs :: RecFlag -> Id -> CoreExprWithFVs' -> Bool
-- ^ True if it's a bad idea to float bindings into this RHS
noFloatIntoRhs is_rec bndr rhs
  | isJoinId bndr
  = isRec is_rec -- Joins are one-shot iff non-recursive

  | otherwise
  = noFloatIntoArg rhs (idType bndr)

noFloatIntoArg :: CoreExprWithFVs' -> Type -> Bool
noFloatIntoArg expr expr_ty
  | isUnliftedType expr_ty
  = True  -- See Note [Do not destroy the let/app invariant]

   | AnnLam bndr e <- expr
   , (bndrs, _) <- collectAnnBndrs e
   =  noFloatIntoLam (bndr:bndrs)  -- Wrinkle 1 (a)
   || all isTyVar (bndr:bndrs)     -- Wrinkle 1 (b)
      -- See Note [noFloatInto considerations] wrinkle 2

  | otherwise  -- Note [noFloatInto considerations] wrinkle 2
  = exprIsTrivial deann_expr || exprIsHNF deann_expr
  where
    deann_expr = deAnnotate' expr

{- Note [noFloatInto considerations]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
When do we want to float bindings into
   - noFloatIntoRHs: the RHS of a let-binding
   - noFloatIntoArg: the argument of a function application

Definitely don't float in if it has unlifted type; that
would destroy the let/app invariant.

* Wrinkle 1: do not float in if
     (a) any non-one-shot value lambdas
  or (b) all type lambdas
  In both cases we'll float straight back out again
  NB: Must line up with fiExpr (AnnLam...); see Trac #7088

  (a) is important: we /must/ float into a one-shot lambda group
  (which includes join points). This makes a big difference
  for things like
     f x# = let x = I# x#
            in let j = \() -> ...x...
               in if <condition> then normal-path else j ()
  If x is used only in the error case join point, j, we must float the
  boxing constructor into it, else we box it every time which is very
  bad news indeed.

* Wrinkle 2: for RHSs, do not float into a HNF; we'll just float right
  back out again... not tragic, but a waste of time.

  For function arguments we will still end up with this
  in-then-out stuff; consider
    letrec x = e in f x
  Here x is not a HNF, so we'll produce
    f (letrec x = e in x)
  which is OK... it's not that common, and we'll end up
  floating out again, in CorePrep if not earlier.
  Still, we use exprIsTrivial to catch this case (sigh)

************************************************************************
*                                                                      *
\subsection{@sepBindsByDropPoint@}
*                                                                      *
************************************************************************

This is the crucial function.  The idea is: We have a wad of bindings
that we'd like to distribute inside a collection of {\em drop points};
insides the alternatives of a \tr{case} would be one example of some
drop points; the RHS and body of a non-recursive \tr{let} binding
would be another (2-element) collection.

So: We're given a list of sets-of-free-variables, one per drop point,
and a list of floating-inwards bindings.  If a binding can go into
only one drop point (without suddenly making something out-of-scope),
in it goes.  If a binding is used inside {\em multiple} drop points,
then it has to go in a you-must-drop-it-above-all-these-drop-points
point.

We have to maintain the order on these drop-point-related lists.
-}

-- pprFIB :: FloatInBinds -> SDoc
-- pprFIB fibs = text "FIB:" <+> ppr [b | FB _ _ b <- fibs]

sepBindsByDropPoint
    :: DynFlags
    -> Bool                -- True <=> is case expression
    -> [FreeVarSet]        -- One set of FVs per drop point
                           -- Always at least two long!
    -> FloatInBinds        -- Candidate floaters
    -> [FloatInBinds]      -- FIRST one is bindings which must not be floated
                           -- inside any drop point; the rest correspond
                           -- one-to-one with the input list of FV sets

-- Every input floater is returned somewhere in the result;
-- none are dropped, not even ones which don't seem to be
-- free in *any* of the drop-point fvs.  Why?  Because, for example,
-- a binding (let x = E in B) might have a specialised version of
-- x (say x') stored inside x, but x' isn't free in E or B.

type DropBox = (FreeVarSet, FloatInBinds)

sepBindsByDropPoint dflags is_case drop_pts floaters
  | null floaters  -- Shortcut common case
  = [] : [[] | _ <- drop_pts]

  | otherwise
  = ASSERT( drop_pts `lengthAtLeast` 2 )
    go floaters (map (\fvs -> (fvs, [])) (emptyDVarSet : drop_pts))
  where
    n_alts = length drop_pts

    go :: FloatInBinds -> [DropBox] -> [FloatInBinds]
        -- The *first* one in the argument list is the drop_here set
        -- The FloatInBinds in the lists are in the reverse of
        -- the normal FloatInBinds order; that is, they are the right way round!

    go [] drop_boxes = map (reverse . snd) drop_boxes

    go (bind_w_fvs@(FB bndrs bind_fvs bind) : binds) drop_boxes@(here_box : fork_boxes)
        = go binds new_boxes
        where
          -- "here" means the group of bindings dropped at the top of the fork

          (used_here : used_in_flags) = [ fvs `intersectsDVarSet` bndrs
                                        | (fvs, _) <- drop_boxes]

          drop_here = used_here || cant_push

          n_used_alts = count id used_in_flags -- returns number of Trues in list.

          cant_push
            | is_case   = n_used_alts == n_alts   -- Used in all, don't push
                                                  -- Remember n_alts > 1
                          || (n_used_alts > 1 && not (floatIsDupable dflags bind))
                             -- floatIsDupable: see Note [Duplicating floats]

            | otherwise = floatIsCase bind || n_used_alts > 1
                             -- floatIsCase: see Note [Floating primops]

          new_boxes | drop_here = (insert here_box : fork_boxes)
                    | otherwise = (here_box : new_fork_boxes)

          new_fork_boxes = zipWithEqual "FloatIn.sepBinds" insert_maybe
                                        fork_boxes used_in_flags

          insert :: DropBox -> DropBox
          insert (fvs,drops) = (fvs `unionDVarSet` bind_fvs, bind_w_fvs:drops)

          insert_maybe box True  = insert box
          insert_maybe box False = box

    go _ _ = panic "sepBindsByDropPoint/go"

{- Note [Duplicating floats]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~

For case expressions we duplicate the binding if it is reasonably
small, and if it is not used in all the RHSs This is good for
situations like
     let x = I# y in
     case e of
       C -> error x
       D -> error x
       E -> ...not mentioning x...

If the thing is used in all RHSs there is nothing gained,
so we don't duplicate then.
-}

floatedBindsFVs :: FloatInBinds -> FreeVarSet
floatedBindsFVs binds = mapUnionDVarSet fbFVs binds

fbFVs :: FloatInBind -> DVarSet
fbFVs (FB _ fvs _) = fvs

wrapFloats :: FloatInBinds -> CoreExpr -> CoreExpr
-- Remember FloatInBinds is in *reverse* dependency order
wrapFloats []               e = e
wrapFloats (FB _ _ fl : bs) e = wrapFloats bs (wrapFloat fl e)

floatIsDupable :: DynFlags -> FloatBind -> Bool
floatIsDupable dflags (FloatCase scrut _ _ _) = exprIsDupable dflags scrut
floatIsDupable dflags (FloatLet (Rec prs))    = all (exprIsDupable dflags . snd) prs
floatIsDupable dflags (FloatLet (NonRec _ r)) = exprIsDupable dflags r

floatIsCase :: FloatBind -> Bool
floatIsCase (FloatCase {}) = True
floatIsCase (FloatLet {})  = False


{-
(c) The GRASP/AQUA Project, Glasgow University, 1992-1998

\section[FloatOut]{Float bindings outwards (towards the top level)}

``Long-distance'' floating of bindings towards the top level.
-}

{-# LANGUAGE CPP #-}

module FloatOut ( floatOutwards ) where

import GhcPrelude

import CoreSyn
import CoreUtils
import MkCore
import CoreArity        ( etaExpand )
import CoreMonad        ( FloatOutSwitches(..) )

import DynFlags
import ErrUtils         ( dumpIfSet_dyn )
import Id               ( Id, idArity, idType, isBottomingId,
                          isJoinId, isJoinId_maybe )
import SetLevels
import UniqSupply       ( UniqSupply )
import Bag
import Util
import Maybes
import Outputable
import Type
import qualified Data.IntMap as M

import Data.List        ( partition )

#include "HsVersions.h"

{-
        -----------------
        Overall game plan
        -----------------

The Big Main Idea is:

        To float out sub-expressions that can thereby get outside
        a non-one-shot value lambda, and hence may be shared.

To achieve this we may need to do two things:

   a) Let-bind the sub-expression:

        f (g x)  ==>  let lvl = f (g x) in lvl

      Now we can float the binding for 'lvl'.

   b) More than that, we may need to abstract wrt a type variable

        \x -> ... /\a -> let v = ...a... in ....

      Here the binding for v mentions 'a' but not 'x'.  So we
      abstract wrt 'a', to give this binding for 'v':

            vp = /\a -> ...a...
            v  = vp a

      Now the binding for vp can float out unimpeded.
      I can't remember why this case seemed important enough to
      deal with, but I certainly found cases where important floats
      didn't happen if we did not abstract wrt tyvars.

With this in mind we can also achieve another goal: lambda lifting.
We can make an arbitrary (function) binding float to top level by
abstracting wrt *all* local variables, not just type variables, leaving
a binding that can be floated right to top level.  Whether or not this
happens is controlled by a flag.

Random comments
~~~~~~~~~~~~~~~

At the moment we never float a binding out to between two adjacent
lambdas.  For example:

@
        \x y -> let t = x+x in ...
===>
        \x -> let t = x+x in \y -> ...
@
Reason: this is less efficient in the case where the original lambda
is never partially applied.

But there's a case I've seen where this might not be true.  Consider:
@
elEm2 x ys
  = elem' x ys
  where
    elem' _ []  = False
    elem' x (y:ys)      = x==y || elem' x ys
@
It turns out that this generates a subexpression of the form
@
        \deq x ys -> let eq = eqFromEqDict deq in ...
@
which might usefully be separated to
@
        \deq -> let eq = eqFromEqDict deq in \xy -> ...
@
Well, maybe.  We don't do this at the moment.

Note [Join points]
~~~~~~~~~~~~~~~~~~
Every occurrence of a join point must be a tail call (see Note [Invariants on
join points] in CoreSyn), so we must be careful with how far we float them. The
mechanism for doing so is the *join ceiling*, detailed in Note [Join ceiling]
in SetLevels. For us, the significance is that a binder might be marked to be
dropped at the nearest boundary between tail calls and non-tail calls. For
example:

  (< join j = ... in
     let x = < ... > in
     case < ... > of
       A -> ...
       B -> ...
   >) < ... > < ... >

Here the join ceilings are marked with angle brackets. Either side of an
application is a join ceiling, as is the scrutinee position of a case
expression or the RHS of a let binding (but not a join point).

Why do we *want* do float join points at all? After all, they're never
allocated, so there's no sharing to be gained by floating them. However, the
other benefit of floating is making RHSes small, and this can have a significant
impact. In particular, stream fusion has been known to produce nested loops like
this:

  joinrec j1 x1 =
    joinrec j2 x2 =
      joinrec j3 x3 = ... jump j1 (x3 + 1) ... jump j2 (x3 + 1) ...
      in jump j3 x2
    in jump j2 x1
  in jump j1 x

(Assume x1 and x2 do *not* occur free in j3.)

Here j1 and j2 are wholly superfluous---each of them merely forwards its
argument to j3. Since j3 only refers to x3, we can float j2 and j3 to make
everything one big mutual recursion:

  joinrec j1 x1 = jump j2 x1
          j2 x2 = jump j3 x2
          j3 x3 = ... jump j1 (x3 + 1) ... jump j2 (x3 + 1) ...
  in jump j1 x

Now the simplifier will happily inline the trivial j1 and j2, leaving only j3.
Without floating, we're stuck with three loops instead of one.

************************************************************************
*                                                                      *
\subsection[floatOutwards]{@floatOutwards@: let-floating interface function}
*                                                                      *
************************************************************************
-}

floatOutwards :: FloatOutSwitches
              -> DynFlags
              -> UniqSupply
              -> CoreProgram -> IO CoreProgram

floatOutwards float_sws dflags us pgm
  = do {
        let { annotated_w_levels = setLevels float_sws pgm us ;
              (fss, binds_s')    = unzip (map floatTopBind annotated_w_levels)
            } ;

        dumpIfSet_dyn dflags Opt_D_verbose_core2core "Levels added:"
                  (vcat (map ppr annotated_w_levels));

        let { (tlets, ntlets, lams) = get_stats (sum_stats fss) };

        dumpIfSet_dyn dflags Opt_D_dump_simpl_stats "FloatOut stats:"
                (hcat [ int tlets,  text " Lets floated to top level; ",
                        int ntlets, text " Lets floated elsewhere; from ",
                        int lams,   text " Lambda groups"]);

        return (bagToList (unionManyBags binds_s'))
    }

floatTopBind :: LevelledBind -> (FloatStats, Bag CoreBind)
floatTopBind bind
  = case (floatBind bind) of { (fs, floats, bind') ->
    let float_bag = flattenTopFloats floats
    in case bind' of
      -- bind' can't have unlifted values or join points, so can only be one
      -- value bind, rec or non-rec (see comment on floatBind)
      [Rec prs]    -> (fs, unitBag (Rec (addTopFloatPairs float_bag prs)))
      [NonRec b e] -> (fs, float_bag `snocBag` NonRec b e)
      _            -> pprPanic "floatTopBind" (ppr bind') }

{-
************************************************************************
*                                                                      *
\subsection[FloatOut-Bind]{Floating in a binding (the business end)}
*                                                                      *
************************************************************************
-}

floatBind :: LevelledBind -> (FloatStats, FloatBinds, [CoreBind])
  -- Returns a list with either
  --   * A single non-recursive binding (value or join point), or
  --   * The following, in order:
  --     * Zero or more non-rec unlifted bindings
  --     * One or both of:
  --       * A recursive group of join binds
  --       * A recursive group of value binds
  -- See Note [Floating out of Rec rhss] for why things get arranged this way.
floatBind (NonRec (TB var _) rhs)
  = case (floatRhs var rhs) of { (fs, rhs_floats, rhs') ->

        -- A tiresome hack:
        -- see Note [Bottoming floats: eta expansion] in SetLevels
    let rhs'' | isBottomingId var = etaExpand (idArity var) rhs'
              | otherwise         = rhs'

    in (fs, rhs_floats, [NonRec var rhs'']) }

floatBind (Rec pairs)
  = case floatList do_pair pairs of { (fs, rhs_floats, new_pairs) ->
    let (new_ul_pairss, new_other_pairss) = unzip new_pairs
        (new_join_pairs, new_l_pairs)     = partition (isJoinId . fst)
                                                      (concat new_other_pairss)
        -- Can't put the join points and the values in the same rec group
        new_rec_binds | null new_join_pairs = [ Rec new_l_pairs    ]
                      | null new_l_pairs    = [ Rec new_join_pairs ]
                      | otherwise           = [ Rec new_l_pairs
                                              , Rec new_join_pairs ]
        new_non_rec_binds = [ NonRec b e | (b, e) <- concat new_ul_pairss ]
    in
    (fs, rhs_floats, new_non_rec_binds ++ new_rec_binds) }
  where
    do_pair :: (LevelledBndr, LevelledExpr)
            -> (FloatStats, FloatBinds,
                ([(Id,CoreExpr)],  -- Non-recursive unlifted value bindings
                 [(Id,CoreExpr)])) -- Join points and lifted value bindings
    do_pair (TB name spec, rhs)
      | isTopLvl dest_lvl  -- See Note [floatBind for top level]
      = case (floatRhs name rhs) of { (fs, rhs_floats, rhs') ->
        (fs, emptyFloats, ([], addTopFloatPairs (flattenTopFloats rhs_floats)
                                                [(name, rhs')]))}
      | otherwise         -- Note [Floating out of Rec rhss]
      = case (floatRhs name rhs) of { (fs, rhs_floats, rhs') ->
        case (partitionByLevel dest_lvl rhs_floats) of { (rhs_floats', heres) ->
        case (splitRecFloats heres) of { (ul_pairs, pairs, case_heres) ->
        let pairs' = (name, installUnderLambdas case_heres rhs') : pairs in
        (fs, rhs_floats', (ul_pairs, pairs')) }}}
      where
        dest_lvl = floatSpecLevel spec

splitRecFloats :: Bag FloatBind
               -> ([(Id,CoreExpr)], -- Non-recursive unlifted value bindings
                   [(Id,CoreExpr)], -- Join points and lifted value bindings
                   Bag FloatBind)   -- A tail of further bindings
-- The "tail" begins with a case
-- See Note [Floating out of Rec rhss]
splitRecFloats fs
  = go [] [] (bagToList fs)
  where
    go ul_prs prs (FloatLet (NonRec b r) : fs) | isUnliftedType (idType b)
                                               , not (isJoinId b)
                                               = go ((b,r):ul_prs) prs fs
                                               | otherwise
                                               = go ul_prs ((b,r):prs) fs
    go ul_prs prs (FloatLet (Rec prs')   : fs) = go ul_prs (prs' ++ prs) fs
    go ul_prs prs fs                           = (reverse ul_prs, prs,
                                                  listToBag fs)
                                                   -- Order only matters for
                                                   -- non-rec

installUnderLambdas :: Bag FloatBind -> CoreExpr -> CoreExpr
-- Note [Floating out of Rec rhss]
installUnderLambdas floats e
  | isEmptyBag floats = e
  | otherwise         = go e
  where
    go (Lam b e)                 = Lam b (go e)
    go e                         = install floats e

---------------
floatList :: (a -> (FloatStats, FloatBinds, b)) -> [a] -> (FloatStats, FloatBinds, [b])
floatList _ [] = (zeroStats, emptyFloats, [])
floatList f (a:as) = case f a            of { (fs_a,  binds_a,  b)  ->
                     case floatList f as of { (fs_as, binds_as, bs) ->
                     (fs_a `add_stats` fs_as, binds_a `plusFloats`  binds_as, b:bs) }}

{-
Note [Floating out of Rec rhss]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Consider   Rec { f<1,0> = \xy. body }
From the body we may get some floats. The ones with level <1,0> must
stay here, since they may mention f.  Ideally we'd like to make them
part of the Rec block pairs -- but we can't if there are any
FloatCases involved.

Nor is it a good idea to dump them in the rhs, but outside the lambda
    f = case x of I# y -> \xy. body
because now f's arity might get worse, which is Not Good. (And if
there's an SCC around the RHS it might not get better again.
See Trac #5342.)

So, gruesomely, we split the floats into
 * the outer FloatLets, which can join the Rec, and
 * an inner batch starting in a FloatCase, which are then
   pushed *inside* the lambdas.
This loses full-laziness the rare situation where there is a
FloatCase and a Rec interacting.

If there are unlifted FloatLets (that *aren't* join points) among the floats,
we can't add them to the recursive group without angering Core Lint, but since
they must be ok-for-speculation, they can't actually be making any recursive
calls, so we can safely pull them out and keep them non-recursive.

(Why is something getting floated to <1,0> that doesn't make a recursive call?
The case that came up in testing was that f *and* the unlifted binding were
getting floated *to the same place*:

  \x<2,0> ->
    ... <3,0>
    letrec { f<F<2,0>> =
      ... let x'<F<2,0>> = x +# 1# in ...
    } in ...

Everything gets labeled "float to <2,0>" because it all depends on x, but this
makes f and x' look mutually recursive when they're not.

The test was shootout/k-nucleotide, as compiled using commit 47d5dd68 on the
wip/join-points branch.

TODO: This can probably be solved somehow in SetLevels. The difference between
"this *is at* level <2,0>" and "this *depends on* level <2,0>" is very
important.)

Note [floatBind for top level]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
We may have a *nested* binding whose destination level is (FloatMe tOP_LEVEL), thus
         letrec { foo <0,0> = .... (let bar<0,0> = .. in ..) .... }
The binding for bar will be in the "tops" part of the floating binds,
and thus not partioned by floatBody.

We could perhaps get rid of the 'tops' component of the floating binds,
but this case works just as well.

************************************************************************

\subsection[FloatOut-Expr]{Floating in expressions}
*                                                                      *
************************************************************************
-}

floatBody :: Level
          -> LevelledExpr
          -> (FloatStats, FloatBinds, CoreExpr)

floatBody lvl arg       -- Used rec rhss, and case-alternative rhss
  = case (floatExpr arg) of { (fsa, floats, arg') ->
    case (partitionByLevel lvl floats) of { (floats', heres) ->
        -- Dump bindings are bound here
    (fsa, floats', install heres arg') }}

-----------------

{- Note [Floating past breakpoints]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

We used to disallow floating out of breakpoint ticks (see #10052). However, I
think this is too restrictive.

Consider the case of an expression scoped over by a breakpoint tick,

  tick<...> (let x = ... in f x)

In this case it is completely legal to float out x, despite the fact that
breakpoint ticks are scoped,

  let x = ... in (tick<...>  f x)

The reason here is that we know that the breakpoint will still be hit when the
expression is entered since the tick still scopes over the RHS.

-}

floatExpr :: LevelledExpr
          -> (FloatStats, FloatBinds, CoreExpr)
floatExpr (Var v)   = (zeroStats, emptyFloats, Var v)
floatExpr (Type ty) = (zeroStats, emptyFloats, Type ty)
floatExpr (Coercion co) = (zeroStats, emptyFloats, Coercion co)
floatExpr (Lit lit) = (zeroStats, emptyFloats, Lit lit)

floatExpr (App e a)
  = case (atJoinCeiling $ floatExpr  e) of { (fse, floats_e, e') ->
    case (atJoinCeiling $ floatExpr  a) of { (fsa, floats_a, a') ->
    (fse `add_stats` fsa, floats_e `plusFloats` floats_a, App e' a') }}

floatExpr lam@(Lam (TB _ lam_spec) _)
  = let (bndrs_w_lvls, body) = collectBinders lam
        bndrs                = [b | TB b _ <- bndrs_w_lvls]
        bndr_lvl             = asJoinCeilLvl (floatSpecLevel lam_spec)
        -- All the binders have the same level
        -- See SetLevels.lvlLamBndrs
        -- Use asJoinCeilLvl to make this the join ceiling
    in
    case (floatBody bndr_lvl body) of { (fs, floats, body') ->
    (add_to_stats fs floats, floats, mkLams bndrs body') }

floatExpr (Tick tickish expr)
  | tickish `tickishScopesLike` SoftScope -- not scoped, can just float
  = case (atJoinCeiling $ floatExpr expr)    of { (fs, floating_defns, expr') ->
    (fs, floating_defns, Tick tickish expr') }

  | not (tickishCounts tickish) || tickishCanSplit tickish
  = case (atJoinCeiling $ floatExpr expr)    of { (fs, floating_defns, expr') ->
    let -- Annotate bindings floated outwards past an scc expression
        -- with the cc.  We mark that cc as "duplicated", though.
        annotated_defns = wrapTick (mkNoCount tickish) floating_defns
    in
    (fs, annotated_defns, Tick tickish expr') }

  -- Note [Floating past breakpoints]
  | Breakpoint{} <- tickish
  = case (floatExpr expr)    of { (fs, floating_defns, expr') ->
    (fs, floating_defns, Tick tickish expr') }

  | otherwise
  = pprPanic "floatExpr tick" (ppr tickish)

floatExpr (Cast expr co)
  = case (atJoinCeiling $ floatExpr expr) of { (fs, floating_defns, expr') ->
    (fs, floating_defns, Cast expr' co) }

floatExpr (Let bind body)
  = case bind_spec of
      FloatMe dest_lvl
        -> case (floatBind bind) of { (fsb, bind_floats, binds') ->
           case (floatExpr body) of { (fse, body_floats, body') ->
           let new_bind_floats = foldr plusFloats emptyFloats
                                   (map (unitLetFloat dest_lvl) binds') in
           ( add_stats fsb fse
           , bind_floats `plusFloats` new_bind_floats
                         `plusFloats` body_floats
           , body') }}

      StayPut bind_lvl  -- See Note [Avoiding unnecessary floating]
        -> case (floatBind bind)          of { (fsb, bind_floats, binds') ->
           case (floatBody bind_lvl body) of { (fse, body_floats, body') ->
           ( add_stats fsb fse
           , bind_floats `plusFloats` body_floats
           , foldr Let body' binds' ) }}
  where
    bind_spec = case bind of
                 NonRec (TB _ s) _     -> s
                 Rec ((TB _ s, _) : _) -> s
                 Rec []                -> panic "floatExpr:rec"

floatExpr (Case scrut (TB case_bndr case_spec) ty alts)
  = case case_spec of
      FloatMe dest_lvl  -- Case expression moves
        | [(con@(DataAlt {}), bndrs, rhs)] <- alts
        -> case atJoinCeiling $ floatExpr scrut of { (fse, fde, scrut') ->
           case                 floatExpr rhs   of { (fsb, fdb, rhs') ->
           let
             float = unitCaseFloat dest_lvl scrut'
                          case_bndr con [b | TB b _ <- bndrs]
           in
           (add_stats fse fsb, fde `plusFloats` float `plusFloats` fdb, rhs') }}
        | otherwise
        -> pprPanic "Floating multi-case" (ppr alts)

      StayPut bind_lvl  -- Case expression stays put
        -> case atJoinCeiling $ floatExpr scrut of { (fse, fde, scrut') ->
           case floatList (float_alt bind_lvl) alts of { (fsa, fda, alts')  ->
           (add_stats fse fsa, fda `plusFloats` fde, Case scrut' case_bndr ty alts')
           }}
  where
    float_alt bind_lvl (con, bs, rhs)
        = case (floatBody bind_lvl rhs) of { (fs, rhs_floats, rhs') ->
          (fs, rhs_floats, (con, [b | TB b _ <- bs], rhs')) }

floatRhs :: CoreBndr
         -> LevelledExpr
         -> (FloatStats, FloatBinds, CoreExpr)
floatRhs bndr rhs
  | Just join_arity <- isJoinId_maybe bndr
  , Just (bndrs, body) <- try_collect join_arity rhs []
  = case bndrs of
      []                -> floatExpr rhs
      (TB _ lam_spec):_ ->
        let lvl = floatSpecLevel lam_spec in
        case floatBody lvl body of { (fs, floats, body') ->
        (fs, floats, mkLams [b | TB b _ <- bndrs] body') }
  | otherwise
  = atJoinCeiling $ floatExpr rhs
  where
    try_collect 0 expr      acc = Just (reverse acc, expr)
    try_collect n (Lam b e) acc = try_collect (n-1) e (b:acc)
    try_collect _ _         _   = Nothing

{-
Note [Avoiding unnecessary floating]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
In general we want to avoid floating a let unnecessarily, because
it might worsen strictness:
    let
       x = ...(let y = e in y+y)....
Here y is demanded.  If we float it outside the lazy 'x=..' then
we'd have to zap its demand info, and it may never be restored.

So at a 'let' we leave the binding right where the are unless
the binding will escape a value lambda, e.g.

(\x -> let y = fac 100 in y)

That's what the partitionByMajorLevel does in the floatExpr (Let ...)
case.

Notice, though, that we must take care to drop any bindings
from the body of the let that depend on the staying-put bindings.

We used instead to do the partitionByMajorLevel on the RHS of an '=',
in floatRhs.  But that was quite tiresome.  We needed to test for
values or trival rhss, because (in particular) we don't want to insert
new bindings between the "=" and the "\".  E.g.
        f = \x -> let <bind> in <body>
We do not want
        f = let <bind> in \x -> <body>
(a) The simplifier will immediately float it further out, so we may
        as well do so right now; in general, keeping rhss as manifest
        values is good
(b) If a float-in pass follows immediately, it might add yet more
        bindings just after the '='.  And some of them might (correctly)
        be strict even though the 'let f' is lazy, because f, being a value,
        gets its demand-info zapped by the simplifier.
And even all that turned out to be very fragile, and broke
altogether when profiling got in the way.

So now we do the partition right at the (Let..) itself.

************************************************************************
*                                                                      *
\subsection{Utility bits for floating stats}
*                                                                      *
************************************************************************

I didn't implement this with unboxed numbers.  I don't want to be too
strict in this stuff, as it is rarely turned on.  (WDP 95/09)
-}

data FloatStats
  = FlS Int  -- Number of top-floats * lambda groups they've been past
        Int  -- Number of non-top-floats * lambda groups they've been past
        Int  -- Number of lambda (groups) seen

get_stats :: FloatStats -> (Int, Int, Int)
get_stats (FlS a b c) = (a, b, c)

zeroStats :: FloatStats
zeroStats = FlS 0 0 0

sum_stats :: [FloatStats] -> FloatStats
sum_stats xs = foldr add_stats zeroStats xs

add_stats :: FloatStats -> FloatStats -> FloatStats
add_stats (FlS a1 b1 c1) (FlS a2 b2 c2)
  = FlS (a1 + a2) (b1 + b2) (c1 + c2)

add_to_stats :: FloatStats -> FloatBinds -> FloatStats
add_to_stats (FlS a b c) (FB tops ceils others)
  = FlS (a + lengthBag tops)
        (b + lengthBag ceils + lengthBag (flattenMajor others))
        (c + 1)

{-
************************************************************************
*                                                                      *
\subsection{Utility bits for floating}
*                                                                      *
************************************************************************

Note [Representation of FloatBinds]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The FloatBinds types is somewhat important.  We can get very large numbers
of floating bindings, often all destined for the top level.  A typical example
is     x = [4,2,5,2,5, .... ]
Then we get lots of small expressions like (fromInteger 4), which all get
lifted to top level.

The trouble is that
  (a) we partition these floating bindings *at every binding site*
  (b) SetLevels introduces a new bindings site for every float
So we had better not look at each binding at each binding site!

That is why MajorEnv is represented as a finite map.

We keep the bindings destined for the *top* level separate, because
we float them out even if they don't escape a *value* lambda; see
partitionByMajorLevel.
-}

type FloatLet = CoreBind        -- INVARIANT: a FloatLet is always lifted
type MajorEnv = M.IntMap MinorEnv         -- Keyed by major level
type MinorEnv = M.IntMap (Bag FloatBind)  -- Keyed by minor level

data FloatBinds  = FB !(Bag FloatLet)           -- Destined for top level
                      !(Bag FloatBind)          -- Destined for join ceiling
                      !MajorEnv                 -- Other levels
     -- See Note [Representation of FloatBinds]

instance Outputable FloatBinds where
  ppr (FB fbs ceils defs)
      = text "FB" <+> (braces $ vcat
           [ text "tops ="     <+> ppr fbs
           , text "ceils ="    <+> ppr ceils
           , text "non-tops =" <+> ppr defs ])

flattenTopFloats :: FloatBinds -> Bag CoreBind
flattenTopFloats (FB tops ceils defs)
  = ASSERT2( isEmptyBag (flattenMajor defs), ppr defs )
    ASSERT2( isEmptyBag ceils, ppr ceils )
    tops

addTopFloatPairs :: Bag CoreBind -> [(Id,CoreExpr)] -> [(Id,CoreExpr)]
addTopFloatPairs float_bag prs
  = foldrBag add prs float_bag
  where
    add (NonRec b r) prs  = (b,r):prs
    add (Rec prs1)   prs2 = prs1 ++ prs2

flattenMajor :: MajorEnv -> Bag FloatBind
flattenMajor = M.foldr (unionBags . flattenMinor) emptyBag

flattenMinor :: MinorEnv -> Bag FloatBind
flattenMinor = M.foldr unionBags emptyBag

emptyFloats :: FloatBinds
emptyFloats = FB emptyBag emptyBag M.empty

unitCaseFloat :: Level -> CoreExpr -> Id -> AltCon -> [Var] -> FloatBinds
unitCaseFloat (Level major minor t) e b con bs
  | t == JoinCeilLvl
  = FB emptyBag floats M.empty
  | otherwise
  = FB emptyBag emptyBag (M.singleton major (M.singleton minor floats))
  where
    floats = unitBag (FloatCase e b con bs)

unitLetFloat :: Level -> FloatLet -> FloatBinds
unitLetFloat lvl@(Level major minor t) b
  | isTopLvl lvl     = FB (unitBag b) emptyBag M.empty
  | t == JoinCeilLvl = FB emptyBag floats M.empty
  | otherwise        = FB emptyBag emptyBag (M.singleton major
                                              (M.singleton minor floats))
  where
    floats = unitBag (FloatLet b)

plusFloats :: FloatBinds -> FloatBinds -> FloatBinds
plusFloats (FB t1 c1 l1) (FB t2 c2 l2)
  = FB (t1 `unionBags` t2) (c1 `unionBags` c2) (l1 `plusMajor` l2)

plusMajor :: MajorEnv -> MajorEnv -> MajorEnv
plusMajor = M.unionWith plusMinor

plusMinor :: MinorEnv -> MinorEnv -> MinorEnv
plusMinor = M.unionWith unionBags

install :: Bag FloatBind -> CoreExpr -> CoreExpr
install defn_groups expr
  = foldrBag wrapFloat expr defn_groups

partitionByLevel
        :: Level                -- Partitioning level
        -> FloatBinds           -- Defns to be divided into 2 piles...
        -> (FloatBinds,         -- Defns  with level strictly < partition level,
            Bag FloatBind)      -- The rest

{-
--       ---- partitionByMajorLevel ----
-- Float it if we escape a value lambda,
--     *or* if we get to the top level
--     *or* if it's a case-float and its minor level is < current
--
-- If we can get to the top level, say "yes" anyway. This means that
--      x = f e
-- transforms to
--    lvl = e
--    x = f lvl
-- which is as it should be

partitionByMajorLevel (Level major _) (FB tops defns)
  = (FB tops outer, heres `unionBags` flattenMajor inner)
  where
    (outer, mb_heres, inner) = M.splitLookup major defns
    heres = case mb_heres of
               Nothing -> emptyBag
               Just h  -> flattenMinor h
-}

partitionByLevel (Level major minor typ) (FB tops ceils defns)
  = (FB tops ceils' (outer_maj `plusMajor` M.singleton major outer_min),
     here_min `unionBags` here_ceil
              `unionBags` flattenMinor inner_min
              `unionBags` flattenMajor inner_maj)

  where
    (outer_maj, mb_here_maj, inner_maj) = M.splitLookup major defns
    (outer_min, mb_here_min, inner_min) = case mb_here_maj of
                                            Nothing -> (M.empty, Nothing, M.empty)
                                            Just min_defns -> M.splitLookup minor min_defns
    here_min = mb_here_min `orElse` emptyBag
    (here_ceil, ceils') | typ == JoinCeilLvl = (ceils, emptyBag)
                        | otherwise          = (emptyBag, ceils)

-- Like partitionByLevel, but instead split out the bindings that are marked
-- to float to the nearest join ceiling (see Note [Join points])
partitionAtJoinCeiling :: FloatBinds -> (FloatBinds, Bag FloatBind)
partitionAtJoinCeiling (FB tops ceils defs)
  = (FB tops emptyBag defs, ceils)

-- Perform some action at a join ceiling, i.e., don't let join points float out
-- (see Note [Join points])
atJoinCeiling :: (FloatStats, FloatBinds, CoreExpr)
              -> (FloatStats, FloatBinds, CoreExpr)
atJoinCeiling (fs, floats, expr')
  = (fs, floats', install ceils expr')
  where
    (floats', ceils) = partitionAtJoinCeiling floats

wrapTick :: Tickish Id -> FloatBinds -> FloatBinds
wrapTick t (FB tops ceils defns)
  = FB (mapBag wrap_bind tops) (wrap_defns ceils)
       (M.map (M.map wrap_defns) defns)
  where
    wrap_defns = mapBag wrap_one

    wrap_bind (NonRec binder rhs) = NonRec binder (maybe_tick rhs)
    wrap_bind (Rec pairs)         = Rec (mapSnd maybe_tick pairs)

    wrap_one (FloatLet bind)      = FloatLet (wrap_bind bind)
    wrap_one (FloatCase e b c bs) = FloatCase (maybe_tick e) b c bs

    maybe_tick e | exprIsHNF e = tickHNFArgs t e
                 | otherwise   = mkTick t e
      -- we don't need to wrap a tick around an HNF when we float it
      -- outside a tick: that is an invariant of the tick semantics
      -- Conversely, inlining of HNFs inside an SCC is allowed, and
      -- indeed the HNF we're floating here might well be inlined back
      -- again, and we don't want to end up with duplicate ticks.


{-
(c) The AQUA Project, Glasgow University, 1994-1998

\section[LiberateCase]{Unroll recursion to allow evals to be lifted from a loop}
-}

{-# LANGUAGE CPP #-}
module LiberateCase ( liberateCase ) where

#include "HsVersions.h"

import GhcPrelude

import DynFlags
import CoreSyn
import CoreUnfold       ( couldBeSmallEnoughToInline )
import TysWiredIn       ( unitDataConId )
import Id
import VarEnv
import Util             ( notNull )

{-
The liberate-case transformation
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
This module walks over @Core@, and looks for @case@ on free variables.
The criterion is:
        if there is case on a free on the route to the recursive call,
        then the recursive call is replaced with an unfolding.

Example

   f = \ t -> case v of
                 V a b -> a : f t

=> the inner f is replaced.

   f = \ t -> case v of
                 V a b -> a : (letrec
                                f =  \ t -> case v of
                                               V a b -> a : f t
                               in f) t
(note the NEED for shadowing)

=> Simplify

  f = \ t -> case v of
                 V a b -> a : (letrec
                                f = \ t -> a : f t
                               in f t)

Better code, because 'a' is  free inside the inner letrec, rather
than needing projection from v.

Note that this deals with *free variables*.  SpecConstr deals with
*arguments* that are of known form.  E.g.

        last []     = error
        last (x:[]) = x
        last (x:xs) = last xs

Note [Scrutinee with cast]
~~~~~~~~~~~~~~~~~~~~~~~~~~
Consider this:
    f = \ t -> case (v `cast` co) of
                 V a b -> a : f t

Exactly the same optimisation (unrolling one call to f) will work here,
despite the cast.  See mk_alt_env in the Case branch of libCase.

To think about (Apr 94)
~~~~~~~~~~~~~~
Main worry: duplicating code excessively.  At the moment we duplicate
the entire binding group once at each recursive call.  But there may
be a group of recursive calls which share a common set of evaluated
free variables, in which case the duplication is a plain waste.

Another thing we could consider adding is some unfold-threshold thing,
so that we'll only duplicate if the size of the group rhss isn't too
big.

Data types
~~~~~~~~~~
The ``level'' of a binder tells how many
recursive defns lexically enclose the binding
A recursive defn "encloses" its RHS, not its
scope.  For example:
\begin{verbatim}
        letrec f = let g = ... in ...
        in
        let h = ...
        in ...
\end{verbatim}
Here, the level of @f@ is zero, the level of @g@ is one,
and the level of @h@ is zero (NB not one).

************************************************************************
*                                                                      *
         Top-level code
*                                                                      *
************************************************************************
-}

liberateCase :: DynFlags -> CoreProgram -> CoreProgram
liberateCase dflags binds = do_prog (initEnv dflags) binds
  where
    do_prog _   [] = []
    do_prog env (bind:binds) = bind' : do_prog env' binds
                             where
                               (env', bind') = libCaseBind env bind

{-
************************************************************************
*                                                                      *
         Main payload
*                                                                      *
************************************************************************

Bindings
~~~~~~~~
-}

libCaseBind :: LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)

libCaseBind env (NonRec binder rhs)
  = (addBinders env [binder], NonRec binder (libCase env rhs))

libCaseBind env (Rec pairs)
  = (env_body, Rec pairs')
  where
    binders = map fst pairs

    env_body = addBinders env binders

    pairs' = [(binder, libCase env_rhs rhs) | (binder,rhs) <- pairs]

        -- We extend the rec-env by binding each Id to its rhs, first
        -- processing the rhs with an *un-extended* environment, so
        -- that the same process doesn't occur for ever!
    env_rhs | is_dupable_bind = addRecBinds env dup_pairs
            | otherwise       = env

    dup_pairs = [ (localiseId binder, libCase env_body rhs)
                | (binder, rhs) <- pairs ]
        -- localiseID : see Note [Need to localiseId in libCaseBind]

    is_dupable_bind = small_enough && all ok_pair pairs

    -- Size: we are going to duplicate dup_pairs; to find their
    --       size, build a fake binding (let { dup_pairs } in (),
    --       and find the size of that
    -- See Note [Small enough]
    small_enough = case bombOutSize env of
                      Nothing   -> True   -- Infinity
                      Just size -> couldBeSmallEnoughToInline (lc_dflags env) size $
                                   Let (Rec dup_pairs) (Var unitDataConId)

    ok_pair (id,_)
        =  idArity id > 0          -- Note [Only functions!]
        && not (isBottomingId id)  -- Note [Not bottoming ids]

{- Note [Not bottoming Ids]
~~~~~~~~~~~~~~~~~~~~~~~~~~~
Do not specialise error-functions (this is unusual, but I once saw it,
(acually in Data.Typable.Internal)

Note [Only functions!]
~~~~~~~~~~~~~~~~~~~~~~
Consider the following code

       f = g (case v of V a b -> a : t f)

where g is expensive. If we aren't careful, liberate case will turn this into

       f = g (case v of
               V a b -> a : t (letrec f = g (case v of V a b -> a : f t)
                                in f)
             )

Yikes! We evaluate g twice. This leads to a O(2^n) explosion
if g calls back to the same code recursively.

Solution: make sure that we only do the liberate-case thing on *functions*

Note [Small enough]
~~~~~~~~~~~~~~~~~~~
Consider
  \fv. letrec
         f = \x. BIG...(case fv of { (a,b) -> ...g.. })...
         g = \y. SMALL...f...

Then we *can* in principle do liberate-case on 'g' (small RHS) but not
for 'f' (too big).  But doing so is not profitable, because duplicating
'g' at its call site in 'f' doesn't get rid of any cases.  So we just
ask for the whole group to be small enough.

Note [Need to localiseId in libCaseBind]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The call to localiseId is needed for two subtle reasons
(a)  Reset the export flags on the binders so
        that we don't get name clashes on exported things if the
        local binding floats out to top level.  This is most unlikely
        to happen, since the whole point concerns free variables.
        But resetting the export flag is right regardless.

(b)  Make the name an Internal one.  External Names should never be
        nested; if it were floated to the top level, we'd get a name
        clash at code generation time.

Expressions
~~~~~~~~~~~
-}

libCase :: LibCaseEnv
        -> CoreExpr
        -> CoreExpr

libCase env (Var v)             = libCaseApp env v []
libCase _   (Lit lit)           = Lit lit
libCase _   (Type ty)           = Type ty
libCase _   (Coercion co)       = Coercion co
libCase env e@(App {})          | let (fun, args) = collectArgs e
                                , Var v <- fun
                                = libCaseApp env v args
libCase env (App fun arg)       = App (libCase env fun) (libCase env arg)
libCase env (Tick tickish body) = Tick tickish (libCase env body)
libCase env (Cast e co)         = Cast (libCase env e) co

libCase env (Lam binder body)
  = Lam binder (libCase (addBinders env [binder]) body)

libCase env (Let bind body)
  = Let bind' (libCase env_body body)
  where
    (env_body, bind') = libCaseBind env bind

libCase env (Case scrut bndr ty alts)
  = Case (libCase env scrut) bndr ty (map (libCaseAlt env_alts) alts)
  where
    env_alts = addBinders (mk_alt_env scrut) [bndr]
    mk_alt_env (Var scrut_var) = addScrutedVar env scrut_var
    mk_alt_env (Cast scrut _)  = mk_alt_env scrut       -- Note [Scrutinee with cast]
    mk_alt_env _               = env

libCaseAlt :: LibCaseEnv -> (AltCon, [CoreBndr], CoreExpr)
                         -> (AltCon, [CoreBndr], CoreExpr)
libCaseAlt env (con,args,rhs) = (con, args, libCase (addBinders env args) rhs)

{-
Ids
~~~

To unfold, we can't just wrap the id itself in its binding if it's a join point:

  jump j a b c  =>  (joinrec j x y z = ... in jump j) a b c -- wrong!!!

Every jump must provide all arguments, so we have to be careful to wrap the
whole jump instead:

  jump j a b c  =>  joinrec j x y z = ... in jump j a b c -- right

-}

libCaseApp :: LibCaseEnv -> Id -> [CoreExpr] -> CoreExpr
libCaseApp env v args
  | Just the_bind <- lookupRecId env v  -- It's a use of a recursive thing
  , notNull free_scruts                 -- with free vars scrutinised in RHS
  = Let the_bind expr'

  | otherwise
  = expr'

  where
    rec_id_level = lookupLevel env v
    free_scruts  = freeScruts env rec_id_level
    expr'        = mkApps (Var v) (map (libCase env) args)

freeScruts :: LibCaseEnv
           -> LibCaseLevel      -- Level of the recursive Id
           -> [Id]              -- Ids that are scrutinised between the binding
                                -- of the recursive Id and here
freeScruts env rec_bind_lvl
  = [v | (v, scrut_bind_lvl, scrut_at_lvl) <- lc_scruts env
       , scrut_bind_lvl <= rec_bind_lvl
       , scrut_at_lvl > rec_bind_lvl]
        -- Note [When to specialise]
        -- Note [Avoiding fruitless liberate-case]

{-
Note [When to specialise]
~~~~~~~~~~~~~~~~~~~~~~~~~
Consider
  f = \x. letrec g = \y. case x of
                           True  -> ... (f a) ...
                           False -> ... (g b) ...

We get the following levels
          f  0
          x  1
          g  1
          y  2

Then 'x' is being scrutinised at a deeper level than its binding, so
it's added to lc_sruts:  [(x,1)]

We do *not* want to specialise the call to 'f', because 'x' is not free
in 'f'.  So here the bind-level of 'x' (=1) is not <= the bind-level of 'f' (=0).

We *do* want to specialise the call to 'g', because 'x' is free in g.
Here the bind-level of 'x' (=1) is <= the bind-level of 'g' (=1).

Note [Avoiding fruitless liberate-case]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Consider also:
  f = \x. case top_lvl_thing of
                I# _ -> let g = \y. ... g ...
                        in ...

Here, top_lvl_thing is scrutinised at a level (1) deeper than its
binding site (0).  Nevertheless, we do NOT want to specialise the call
to 'g' because all the structure in its free variables is already
visible at the definition site for g.  Hence, when considering specialising
an occurrence of 'g', we want to check that there's a scruted-var v st

   a) v's binding site is *outside* g
   b) v's scrutinisation site is *inside* g

************************************************************************
*                                                                      *
        Utility functions
*                                                                      *
************************************************************************
-}

addBinders :: LibCaseEnv -> [CoreBndr] -> LibCaseEnv
addBinders env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env }) binders
  = env { lc_lvl_env = lvl_env' }
  where
    lvl_env' = extendVarEnvList lvl_env (binders `zip` repeat lvl)

addRecBinds :: LibCaseEnv -> [(Id,CoreExpr)] -> LibCaseEnv
addRecBinds env@(LibCaseEnv {lc_lvl = lvl, lc_lvl_env = lvl_env,
                             lc_rec_env = rec_env}) pairs
  = env { lc_lvl = lvl', lc_lvl_env = lvl_env', lc_rec_env = rec_env' }
  where
    lvl'     = lvl + 1
    lvl_env' = extendVarEnvList lvl_env [(binder,lvl) | (binder,_) <- pairs]
    rec_env' = extendVarEnvList rec_env [(binder, Rec pairs) | (binder,_) <- pairs]

addScrutedVar :: LibCaseEnv
              -> Id             -- This Id is being scrutinised by a case expression
              -> LibCaseEnv

addScrutedVar env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env,
                                lc_scruts = scruts }) scrut_var
  | bind_lvl < lvl
  = env { lc_scruts = scruts' }
        -- Add to scruts iff the scrut_var is being scrutinised at
        -- a deeper level than its defn

  | otherwise = env
  where
    scruts'  = (scrut_var, bind_lvl, lvl) : scruts
    bind_lvl = case lookupVarEnv lvl_env scrut_var of
                 Just lvl -> lvl
                 Nothing  -> topLevel

lookupRecId :: LibCaseEnv -> Id -> Maybe CoreBind
lookupRecId env id = lookupVarEnv (lc_rec_env env) id

lookupLevel :: LibCaseEnv -> Id -> LibCaseLevel
lookupLevel env id
  = case lookupVarEnv (lc_lvl_env env) id of
      Just lvl -> lvl
      Nothing  -> topLevel

{-
************************************************************************
*                                                                      *
         The environment
*                                                                      *
************************************************************************
-}

type LibCaseLevel = Int

topLevel :: LibCaseLevel
topLevel = 0

data LibCaseEnv
  = LibCaseEnv {
        lc_dflags :: DynFlags,

        lc_lvl :: LibCaseLevel, -- Current level
                -- The level is incremented when (and only when) going
                -- inside the RHS of a (sufficiently small) recursive
                -- function.

        lc_lvl_env :: IdEnv LibCaseLevel,
                -- Binds all non-top-level in-scope Ids (top-level and
                -- imported things have a level of zero)

        lc_rec_env :: IdEnv CoreBind,
                -- Binds *only* recursively defined ids, to their own
                -- binding group, and *only* in their own RHSs

        lc_scruts :: [(Id, LibCaseLevel, LibCaseLevel)]
                -- Each of these Ids was scrutinised by an enclosing
                -- case expression, at a level deeper than its binding
                -- level.
                --
                -- The first LibCaseLevel is the *binding level* of
                --   the scrutinised Id,
                -- The second is the level *at which it was scrutinised*.
                --   (see Note [Avoiding fruitless liberate-case])
                -- The former is a bit redundant, since you could always
                -- look it up in lc_lvl_env, but it's just cached here
                --
                -- The order is insignificant; it's a bag really
                --
                -- There's one element per scrutinisation;
                --    in principle the same Id may appear multiple times,
                --    although that'd be unusual:
                --       case x of { (a,b) -> ....(case x of ...) .. }
        }

initEnv :: DynFlags -> LibCaseEnv
initEnv dflags
  = LibCaseEnv { lc_dflags = dflags,
                 lc_lvl = 0,
                 lc_lvl_env = emptyVarEnv,
                 lc_rec_env = emptyVarEnv,
                 lc_scruts = [] }

-- Bomb-out size for deciding if
-- potential liberatees are too big.
-- (passed in from cmd-line args)
bombOutSize :: LibCaseEnv -> Maybe Int
bombOutSize = liberateCaseThreshold . lc_dflags