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


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

               Static Argument Transformation pass

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

May be seen as removing invariants from loops:
Arguments of recursive functions that do not change in recursive
calls are removed from the recursion, which is done locally
and only passes the arguments which effectively change.

Example:
map = /\ ab -> \f -> \xs -> case xs of
                 []       -> []
                 (a:b) -> f a : map f b

as map is recursively called with the same argument f (unmodified)
we transform it to

map = /\ ab -> \f -> \xs -> let map' ys = case ys of
                       []     -> []
                       (a:b) -> f a : map' b
                in map' xs

Notice that for a compiler that uses lambda lifting this is
useless as map' will be transformed back to what map was.

We could possibly do the same for big lambdas, but we don't as
they will eventually be removed in later stages of the compiler,
therefore there is no penalty in keeping them.

We only apply the SAT when the number of static args is > 2. This
produces few bad cases.  See
                should_transform
in saTransform.

Here are the headline nofib results:
                  Size    Allocs   Runtime
Min             +0.0%    -13.7%    -21.4%
Max             +0.1%     +0.0%     +5.4%
Geometric Mean  +0.0%     -0.2%     -6.9%

The previous patch, to fix polymorphic floatout demand signatures, is
essential to make this work well!
-}

{-# LANGUAGE CPP #-}
module SAT ( doStaticArgs ) where

import GhcPrelude

import Var
import CoreSyn
import CoreUtils
import Type
import Coercion
import Id
import Name
import VarEnv
import UniqSupply
import Util
import UniqFM
import VarSet
import Unique
import UniqSet
import Outputable

import Data.List
import FastString

#include "HsVersions.h"

doStaticArgs :: UniqSupply -> CoreProgram -> CoreProgram
doStaticArgs :: UniqSupply -> CoreProgram -> CoreProgram
doStaticArgs us :: UniqSupply
us binds :: CoreProgram
binds = (UniqSupply, CoreProgram) -> CoreProgram
forall a b. (a, b) -> b
snd ((UniqSupply, CoreProgram) -> CoreProgram)
-> (UniqSupply, CoreProgram) -> CoreProgram
forall a b. (a -> b) -> a -> b
$ (UniqSupply -> CoreBind -> (UniqSupply, CoreBind))
-> UniqSupply -> CoreProgram -> (UniqSupply, CoreProgram)
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
mapAccumL UniqSupply -> CoreBind -> (UniqSupply, CoreBind)
sat_bind_threaded_us UniqSupply
us CoreProgram
binds
  where
    sat_bind_threaded_us :: UniqSupply -> CoreBind -> (UniqSupply, CoreBind)
sat_bind_threaded_us us :: UniqSupply
us bind :: CoreBind
bind =
        let (us1 :: UniqSupply
us1, us2 :: UniqSupply
us2) = UniqSupply -> (UniqSupply, UniqSupply)
splitUniqSupply UniqSupply
us
        in (UniqSupply
us1, (CoreBind, IdSATInfo) -> CoreBind
forall a b. (a, b) -> a
fst ((CoreBind, IdSATInfo) -> CoreBind)
-> (CoreBind, IdSATInfo) -> CoreBind
forall a b. (a -> b) -> a -> b
$ UniqSupply -> SatM (CoreBind, IdSATInfo) -> (CoreBind, IdSATInfo)
forall a. UniqSupply -> SatM a -> a
runSAT UniqSupply
us2 (CoreBind -> IdSet -> SatM (CoreBind, IdSATInfo)
satBind CoreBind
bind IdSet
forall a. UniqSet a
emptyUniqSet))

-- We don't bother to SAT recursive groups since it can lead
-- to massive code expansion: see Andre Santos' thesis for details.
-- This means we only apply the actual SAT to Rec groups of one element,
-- but we want to recurse into the others anyway to discover other binds
satBind :: CoreBind -> IdSet -> SatM (CoreBind, IdSATInfo)
satBind :: CoreBind -> IdSet -> SatM (CoreBind, IdSATInfo)
satBind (NonRec binder :: CoreBndr
binder expr :: Expr CoreBndr
expr) interesting_ids :: IdSet
interesting_ids = do
    (expr' :: Expr CoreBndr
expr', sat_info_expr :: IdSATInfo
sat_info_expr, expr_app :: Maybe IdAppInfo
expr_app) <- Expr CoreBndr
-> IdSet -> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
satExpr Expr CoreBndr
expr IdSet
interesting_ids
    (CoreBind, IdSATInfo) -> SatM (CoreBind, IdSATInfo)
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreBndr -> Expr CoreBndr -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec CoreBndr
binder Expr CoreBndr
expr', Maybe IdAppInfo -> IdSATInfo -> IdSATInfo
finalizeApp Maybe IdAppInfo
expr_app IdSATInfo
sat_info_expr)
satBind (Rec [(binder :: CoreBndr
binder, rhs :: Expr CoreBndr
rhs)]) interesting_ids :: IdSet
interesting_ids = do
    let interesting_ids' :: IdSet
interesting_ids' = IdSet
interesting_ids IdSet -> CoreBndr -> IdSet
forall a. Uniquable a => UniqSet a -> a -> UniqSet a
`addOneToUniqSet` CoreBndr
binder
        (rhs_binders :: [CoreBndr]
rhs_binders, rhs_body :: Expr CoreBndr
rhs_body) = Expr CoreBndr -> ([CoreBndr], Expr CoreBndr)
forall b. Expr b -> ([b], Expr b)
collectBinders Expr CoreBndr
rhs
    (rhs_body' :: Expr CoreBndr
rhs_body', sat_info_rhs_body :: IdSATInfo
sat_info_rhs_body) <- Expr CoreBndr -> IdSet -> SatM (Expr CoreBndr, IdSATInfo)
satTopLevelExpr Expr CoreBndr
rhs_body IdSet
interesting_ids'
    let sat_info_rhs_from_args :: IdSATInfo
sat_info_rhs_from_args = CoreBndr -> SATInfo -> IdSATInfo
forall a. CoreBndr -> a -> VarEnv a
unitVarEnv CoreBndr
binder ([CoreBndr] -> SATInfo
bindersToSATInfo [CoreBndr]
rhs_binders)
        sat_info_rhs' :: IdSATInfo
sat_info_rhs' = IdSATInfo -> IdSATInfo -> IdSATInfo
mergeIdSATInfo IdSATInfo
sat_info_rhs_from_args IdSATInfo
sat_info_rhs_body

        shadowing :: Bool
shadowing = CoreBndr
binder CoreBndr -> IdSet -> Bool
forall a. Uniquable a => a -> UniqSet a -> Bool
`elementOfUniqSet` IdSet
interesting_ids
        sat_info_rhs'' :: IdSATInfo
sat_info_rhs'' = if Bool
shadowing
                        then IdSATInfo
sat_info_rhs' IdSATInfo -> CoreBndr -> IdSATInfo
forall key elt. Uniquable key => UniqFM elt -> key -> UniqFM elt
`delFromUFM` CoreBndr
binder -- For safety
                        else IdSATInfo
sat_info_rhs'

    CoreBind
bind' <- CoreBndr
-> Maybe SATInfo -> [CoreBndr] -> Expr CoreBndr -> SatM CoreBind
saTransformMaybe CoreBndr
binder (IdSATInfo -> CoreBndr -> Maybe SATInfo
forall key elt. Uniquable key => UniqFM elt -> key -> Maybe elt
lookupUFM IdSATInfo
sat_info_rhs' CoreBndr
binder)
                              [CoreBndr]
rhs_binders Expr CoreBndr
rhs_body'
    (CoreBind, IdSATInfo) -> SatM (CoreBind, IdSATInfo)
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreBind
bind', IdSATInfo
sat_info_rhs'')
satBind (Rec pairs :: [(CoreBndr, Expr CoreBndr)]
pairs) interesting_ids :: IdSet
interesting_ids = do
    let (binders :: [CoreBndr]
binders, rhss :: [Expr CoreBndr]
rhss) = [(CoreBndr, Expr CoreBndr)] -> ([CoreBndr], [Expr CoreBndr])
forall a b. [(a, b)] -> ([a], [b])
unzip [(CoreBndr, Expr CoreBndr)]
pairs
    [(Expr CoreBndr, IdSATInfo)]
rhss_SATed <- (Expr CoreBndr -> SatM (Expr CoreBndr, IdSATInfo))
-> [Expr CoreBndr] -> UniqSM [(Expr CoreBndr, IdSATInfo)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\e :: Expr CoreBndr
e -> Expr CoreBndr -> IdSet -> SatM (Expr CoreBndr, IdSATInfo)
satTopLevelExpr Expr CoreBndr
e IdSet
interesting_ids) [Expr CoreBndr]
rhss
    let (rhss' :: [Expr CoreBndr]
rhss', sat_info_rhss' :: [IdSATInfo]
sat_info_rhss') = [(Expr CoreBndr, IdSATInfo)] -> ([Expr CoreBndr], [IdSATInfo])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Expr CoreBndr, IdSATInfo)]
rhss_SATed
    (CoreBind, IdSATInfo) -> SatM (CoreBind, IdSATInfo)
forall (m :: * -> *) a. Monad m => a -> m a
return ([(CoreBndr, Expr CoreBndr)] -> CoreBind
forall b. [(b, Expr b)] -> Bind b
Rec (String
-> [CoreBndr] -> [Expr CoreBndr] -> [(CoreBndr, Expr CoreBndr)]
forall a b. String -> [a] -> [b] -> [(a, b)]
zipEqual "satBind" [CoreBndr]
binders [Expr CoreBndr]
rhss'), [IdSATInfo] -> IdSATInfo
mergeIdSATInfos [IdSATInfo]
sat_info_rhss')

data App = VarApp Id | TypeApp Type | CoApp Coercion
data Staticness a = Static a | NotStatic

type IdAppInfo = (Id, SATInfo)

type SATInfo = [Staticness App]
type IdSATInfo = IdEnv SATInfo
emptyIdSATInfo :: IdSATInfo
emptyIdSATInfo :: IdSATInfo
emptyIdSATInfo = IdSATInfo
forall elt. UniqFM elt
emptyUFM

{-
pprIdSATInfo id_sat_info = vcat (map pprIdAndSATInfo (Map.toList id_sat_info))
  where pprIdAndSATInfo (v, sat_info) = hang (ppr v <> colon) 4 (pprSATInfo sat_info)
-}

pprSATInfo :: SATInfo -> SDoc
pprSATInfo :: SATInfo -> SDoc
pprSATInfo staticness :: SATInfo
staticness = [SDoc] -> SDoc
hcat ([SDoc] -> SDoc) -> [SDoc] -> SDoc
forall a b. (a -> b) -> a -> b
$ (Staticness App -> SDoc) -> SATInfo -> [SDoc]
forall a b. (a -> b) -> [a] -> [b]
map Staticness App -> SDoc
pprStaticness SATInfo
staticness

pprStaticness :: Staticness App -> SDoc
pprStaticness :: Staticness App -> SDoc
pprStaticness (Static (VarApp _))  = String -> SDoc
text "SV"
pprStaticness (Static (TypeApp _)) = String -> SDoc
text "ST"
pprStaticness (Static (CoApp _))   = String -> SDoc
text "SC"
pprStaticness NotStatic            = String -> SDoc
text "NS"


mergeSATInfo :: SATInfo -> SATInfo -> SATInfo
mergeSATInfo :: SATInfo -> SATInfo -> SATInfo
mergeSATInfo l :: SATInfo
l r :: SATInfo
r = (Staticness App -> Staticness App -> Staticness App)
-> SATInfo -> SATInfo -> SATInfo
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Staticness App -> Staticness App -> Staticness App
mergeSA SATInfo
l SATInfo
r
  where
    mergeSA :: Staticness App -> Staticness App -> Staticness App
mergeSA NotStatic _ = Staticness App
forall a. Staticness a
NotStatic
    mergeSA _ NotStatic = Staticness App
forall a. Staticness a
NotStatic
    mergeSA (Static (VarApp v :: CoreBndr
v)) (Static (VarApp v' :: CoreBndr
v'))
      | CoreBndr
v CoreBndr -> CoreBndr -> Bool
forall a. Eq a => a -> a -> Bool
== CoreBndr
v'   = App -> Staticness App
forall a. a -> Staticness a
Static (CoreBndr -> App
VarApp CoreBndr
v)
      | Bool
otherwise = Staticness App
forall a. Staticness a
NotStatic
    mergeSA (Static (TypeApp t :: Type
t)) (Static (TypeApp t' :: Type
t'))
      | Type
t Type -> Type -> Bool
`eqType` Type
t' = App -> Staticness App
forall a. a -> Staticness a
Static (Type -> App
TypeApp Type
t)
      | Bool
otherwise     = Staticness App
forall a. Staticness a
NotStatic
    mergeSA (Static (CoApp c :: Coercion
c)) (Static (CoApp c' :: Coercion
c'))
      | Coercion
c Coercion -> Coercion -> Bool
`eqCoercion` Coercion
c' = App -> Staticness App
forall a. a -> Staticness a
Static (Coercion -> App
CoApp Coercion
c)
      | Bool
otherwise             = Staticness App
forall a. Staticness a
NotStatic
    mergeSA _ _  = String -> SDoc -> Staticness App
forall a. HasCallStack => String -> SDoc -> a
pprPanic "mergeSATInfo" (SDoc -> Staticness App) -> SDoc -> Staticness App
forall a b. (a -> b) -> a -> b
$
                          String -> SDoc
text "Left:"
                       SDoc -> SDoc -> SDoc
<> SATInfo -> SDoc
pprSATInfo SATInfo
l SDoc -> SDoc -> SDoc
<> String -> SDoc
text ", "
                       SDoc -> SDoc -> SDoc
<> String -> SDoc
text "Right:"
                       SDoc -> SDoc -> SDoc
<> SATInfo -> SDoc
pprSATInfo SATInfo
r

mergeIdSATInfo :: IdSATInfo -> IdSATInfo -> IdSATInfo
mergeIdSATInfo :: IdSATInfo -> IdSATInfo -> IdSATInfo
mergeIdSATInfo = (SATInfo -> SATInfo -> SATInfo)
-> IdSATInfo -> IdSATInfo -> IdSATInfo
forall elt.
(elt -> elt -> elt) -> UniqFM elt -> UniqFM elt -> UniqFM elt
plusUFM_C SATInfo -> SATInfo -> SATInfo
mergeSATInfo

mergeIdSATInfos :: [IdSATInfo] -> IdSATInfo
mergeIdSATInfos :: [IdSATInfo] -> IdSATInfo
mergeIdSATInfos = (IdSATInfo -> IdSATInfo -> IdSATInfo)
-> IdSATInfo -> [IdSATInfo] -> IdSATInfo
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' IdSATInfo -> IdSATInfo -> IdSATInfo
mergeIdSATInfo IdSATInfo
emptyIdSATInfo

bindersToSATInfo :: [Id] -> SATInfo
bindersToSATInfo :: [CoreBndr] -> SATInfo
bindersToSATInfo vs :: [CoreBndr]
vs = (CoreBndr -> Staticness App) -> [CoreBndr] -> SATInfo
forall a b. (a -> b) -> [a] -> [b]
map (App -> Staticness App
forall a. a -> Staticness a
Static (App -> Staticness App)
-> (CoreBndr -> App) -> CoreBndr -> Staticness App
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CoreBndr -> App
binderToApp) [CoreBndr]
vs
    where binderToApp :: CoreBndr -> App
binderToApp v :: CoreBndr
v | CoreBndr -> Bool
isId CoreBndr
v    = CoreBndr -> App
VarApp CoreBndr
v
                        | CoreBndr -> Bool
isTyVar CoreBndr
v = Type -> App
TypeApp (Type -> App) -> Type -> App
forall a b. (a -> b) -> a -> b
$ CoreBndr -> Type
mkTyVarTy CoreBndr
v
                        | Bool
otherwise = Coercion -> App
CoApp (Coercion -> App) -> Coercion -> App
forall a b. (a -> b) -> a -> b
$ CoreBndr -> Coercion
mkCoVarCo CoreBndr
v

finalizeApp :: Maybe IdAppInfo -> IdSATInfo -> IdSATInfo
finalizeApp :: Maybe IdAppInfo -> IdSATInfo -> IdSATInfo
finalizeApp Nothing id_sat_info :: IdSATInfo
id_sat_info = IdSATInfo
id_sat_info
finalizeApp (Just (v :: CoreBndr
v, sat_info' :: SATInfo
sat_info')) id_sat_info :: IdSATInfo
id_sat_info =
    let sat_info'' :: SATInfo
sat_info'' = case IdSATInfo -> CoreBndr -> Maybe SATInfo
forall key elt. Uniquable key => UniqFM elt -> key -> Maybe elt
lookupUFM IdSATInfo
id_sat_info CoreBndr
v of
                        Nothing -> SATInfo
sat_info'
                        Just sat_info :: SATInfo
sat_info -> SATInfo -> SATInfo -> SATInfo
mergeSATInfo SATInfo
sat_info SATInfo
sat_info'
    in IdSATInfo -> CoreBndr -> SATInfo -> IdSATInfo
forall a. VarEnv a -> CoreBndr -> a -> VarEnv a
extendVarEnv IdSATInfo
id_sat_info CoreBndr
v SATInfo
sat_info''

satTopLevelExpr :: CoreExpr -> IdSet -> SatM (CoreExpr, IdSATInfo)
satTopLevelExpr :: Expr CoreBndr -> IdSet -> SatM (Expr CoreBndr, IdSATInfo)
satTopLevelExpr expr :: Expr CoreBndr
expr interesting_ids :: IdSet
interesting_ids = do
    (expr' :: Expr CoreBndr
expr', sat_info_expr :: IdSATInfo
sat_info_expr, expr_app :: Maybe IdAppInfo
expr_app) <- Expr CoreBndr
-> IdSet -> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
satExpr Expr CoreBndr
expr IdSet
interesting_ids
    (Expr CoreBndr, IdSATInfo) -> SatM (Expr CoreBndr, IdSATInfo)
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr CoreBndr
expr', Maybe IdAppInfo -> IdSATInfo -> IdSATInfo
finalizeApp Maybe IdAppInfo
expr_app IdSATInfo
sat_info_expr)

satExpr :: CoreExpr -> IdSet -> SatM (CoreExpr, IdSATInfo, Maybe IdAppInfo)
satExpr :: Expr CoreBndr
-> IdSet -> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
satExpr var :: Expr CoreBndr
var@(Var v :: CoreBndr
v) interesting_ids :: IdSet
interesting_ids = do
    let app_info :: Maybe (CoreBndr, [a])
app_info = if CoreBndr
v CoreBndr -> IdSet -> Bool
forall a. Uniquable a => a -> UniqSet a -> Bool
`elementOfUniqSet` IdSet
interesting_ids
                   then (CoreBndr, [a]) -> Maybe (CoreBndr, [a])
forall a. a -> Maybe a
Just (CoreBndr
v, [])
                   else Maybe (CoreBndr, [a])
forall a. Maybe a
Nothing
    (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
-> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr CoreBndr
var, IdSATInfo
emptyIdSATInfo, Maybe IdAppInfo
forall a. Maybe (CoreBndr, [a])
app_info)

satExpr lit :: Expr CoreBndr
lit@(Lit _) _ = do
    (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
-> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr CoreBndr
lit, IdSATInfo
emptyIdSATInfo, Maybe IdAppInfo
forall a. Maybe a
Nothing)

satExpr (Lam binders :: CoreBndr
binders body :: Expr CoreBndr
body) interesting_ids :: IdSet
interesting_ids = do
    (body' :: Expr CoreBndr
body', sat_info :: IdSATInfo
sat_info, this_app :: Maybe IdAppInfo
this_app) <- Expr CoreBndr
-> IdSet -> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
satExpr Expr CoreBndr
body IdSet
interesting_ids
    (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
-> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreBndr -> Expr CoreBndr -> Expr CoreBndr
forall b. b -> Expr b -> Expr b
Lam CoreBndr
binders Expr CoreBndr
body', Maybe IdAppInfo -> IdSATInfo -> IdSATInfo
finalizeApp Maybe IdAppInfo
this_app IdSATInfo
sat_info, Maybe IdAppInfo
forall a. Maybe a
Nothing)

satExpr (App fn :: Expr CoreBndr
fn arg :: Expr CoreBndr
arg) interesting_ids :: IdSet
interesting_ids = do
    (fn' :: Expr CoreBndr
fn', sat_info_fn :: IdSATInfo
sat_info_fn, fn_app :: Maybe IdAppInfo
fn_app) <- Expr CoreBndr
-> IdSet -> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
satExpr Expr CoreBndr
fn IdSet
interesting_ids
    let satRemainder :: Maybe IdAppInfo -> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
satRemainder = Expr CoreBndr
-> IdSATInfo
-> Maybe IdAppInfo
-> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
boring Expr CoreBndr
fn' IdSATInfo
sat_info_fn
    case Maybe IdAppInfo
fn_app of
        Nothing -> Maybe IdAppInfo -> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
satRemainder Maybe IdAppInfo
forall a. Maybe a
Nothing
        Just (fn_id :: CoreBndr
fn_id, fn_app_info :: SATInfo
fn_app_info) ->
            -- TODO: remove this use of append somehow (use a data structure with O(1) append but a left-to-right kind of interface)
            let satRemainderWithStaticness :: Staticness App -> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
satRemainderWithStaticness arg_staticness :: Staticness App
arg_staticness = Maybe IdAppInfo -> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
satRemainder (Maybe IdAppInfo
 -> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo))
-> Maybe IdAppInfo
-> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
forall a b. (a -> b) -> a -> b
$ IdAppInfo -> Maybe IdAppInfo
forall a. a -> Maybe a
Just (CoreBndr
fn_id, SATInfo
fn_app_info SATInfo -> SATInfo -> SATInfo
forall a. [a] -> [a] -> [a]
++ [Staticness App
arg_staticness])
            in case Expr CoreBndr
arg of
                Type t :: Type
t     -> Staticness App -> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
satRemainderWithStaticness (Staticness App
 -> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo))
-> Staticness App
-> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
forall a b. (a -> b) -> a -> b
$ App -> Staticness App
forall a. a -> Staticness a
Static (Type -> App
TypeApp Type
t)
                Coercion c :: Coercion
c -> Staticness App -> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
satRemainderWithStaticness (Staticness App
 -> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo))
-> Staticness App
-> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
forall a b. (a -> b) -> a -> b
$ App -> Staticness App
forall a. a -> Staticness a
Static (Coercion -> App
CoApp Coercion
c)
                Var v :: CoreBndr
v      -> Staticness App -> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
satRemainderWithStaticness (Staticness App
 -> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo))
-> Staticness App
-> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
forall a b. (a -> b) -> a -> b
$ App -> Staticness App
forall a. a -> Staticness a
Static (CoreBndr -> App
VarApp CoreBndr
v)
                _          -> Staticness App -> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
satRemainderWithStaticness (Staticness App
 -> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo))
-> Staticness App
-> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
forall a b. (a -> b) -> a -> b
$ Staticness App
forall a. Staticness a
NotStatic
  where
    boring :: CoreExpr -> IdSATInfo -> Maybe IdAppInfo -> SatM (CoreExpr, IdSATInfo, Maybe IdAppInfo)
    boring :: Expr CoreBndr
-> IdSATInfo
-> Maybe IdAppInfo
-> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
boring fn' :: Expr CoreBndr
fn' sat_info_fn :: IdSATInfo
sat_info_fn app_info :: Maybe IdAppInfo
app_info =
        do (arg' :: Expr CoreBndr
arg', sat_info_arg :: IdSATInfo
sat_info_arg, arg_app :: Maybe IdAppInfo
arg_app) <- Expr CoreBndr
-> IdSet -> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
satExpr Expr CoreBndr
arg IdSet
interesting_ids
           let sat_info_arg' :: IdSATInfo
sat_info_arg' = Maybe IdAppInfo -> IdSATInfo -> IdSATInfo
finalizeApp Maybe IdAppInfo
arg_app IdSATInfo
sat_info_arg
               sat_info :: IdSATInfo
sat_info = IdSATInfo -> IdSATInfo -> IdSATInfo
mergeIdSATInfo IdSATInfo
sat_info_fn IdSATInfo
sat_info_arg'
           (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
-> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr CoreBndr -> Expr CoreBndr -> Expr CoreBndr
forall b. Expr b -> Expr b -> Expr b
App Expr CoreBndr
fn' Expr CoreBndr
arg', IdSATInfo
sat_info, Maybe IdAppInfo
app_info)

satExpr (Case expr :: Expr CoreBndr
expr bndr :: CoreBndr
bndr ty :: Type
ty alts :: [Alt CoreBndr]
alts) interesting_ids :: IdSet
interesting_ids = do
    (expr' :: Expr CoreBndr
expr', sat_info_expr :: IdSATInfo
sat_info_expr, expr_app :: Maybe IdAppInfo
expr_app) <- Expr CoreBndr
-> IdSet -> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
satExpr Expr CoreBndr
expr IdSet
interesting_ids
    let sat_info_expr' :: IdSATInfo
sat_info_expr' = Maybe IdAppInfo -> IdSATInfo -> IdSATInfo
finalizeApp Maybe IdAppInfo
expr_app IdSATInfo
sat_info_expr

    [(Alt CoreBndr, IdSATInfo)]
zipped_alts' <- (Alt CoreBndr -> UniqSM (Alt CoreBndr, IdSATInfo))
-> [Alt CoreBndr] -> UniqSM [(Alt CoreBndr, IdSATInfo)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Alt CoreBndr -> UniqSM (Alt CoreBndr, IdSATInfo)
forall a b.
(a, b, Expr CoreBndr) -> UniqSM ((a, b, Expr CoreBndr), IdSATInfo)
satAlt [Alt CoreBndr]
alts
    let (alts' :: [Alt CoreBndr]
alts', sat_infos_alts :: [IdSATInfo]
sat_infos_alts) = [(Alt CoreBndr, IdSATInfo)] -> ([Alt CoreBndr], [IdSATInfo])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Alt CoreBndr, IdSATInfo)]
zipped_alts'
    (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
-> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr CoreBndr
-> CoreBndr -> Type -> [Alt CoreBndr] -> Expr CoreBndr
forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case Expr CoreBndr
expr' CoreBndr
bndr Type
ty [Alt CoreBndr]
alts', IdSATInfo -> IdSATInfo -> IdSATInfo
mergeIdSATInfo IdSATInfo
sat_info_expr' ([IdSATInfo] -> IdSATInfo
mergeIdSATInfos [IdSATInfo]
sat_infos_alts), Maybe IdAppInfo
forall a. Maybe a
Nothing)
  where
    satAlt :: (a, b, Expr CoreBndr) -> UniqSM ((a, b, Expr CoreBndr), IdSATInfo)
satAlt (con :: a
con, bndrs :: b
bndrs, expr :: Expr CoreBndr
expr) = do
        (expr' :: Expr CoreBndr
expr', sat_info_expr :: IdSATInfo
sat_info_expr) <- Expr CoreBndr -> IdSet -> SatM (Expr CoreBndr, IdSATInfo)
satTopLevelExpr Expr CoreBndr
expr IdSet
interesting_ids
        ((a, b, Expr CoreBndr), IdSATInfo)
-> UniqSM ((a, b, Expr CoreBndr), IdSATInfo)
forall (m :: * -> *) a. Monad m => a -> m a
return ((a
con, b
bndrs, Expr CoreBndr
expr'), IdSATInfo
sat_info_expr)

satExpr (Let bind :: CoreBind
bind body :: Expr CoreBndr
body) interesting_ids :: IdSet
interesting_ids = do
    (body' :: Expr CoreBndr
body', sat_info_body :: IdSATInfo
sat_info_body, body_app :: Maybe IdAppInfo
body_app) <- Expr CoreBndr
-> IdSet -> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
satExpr Expr CoreBndr
body IdSet
interesting_ids
    (bind' :: CoreBind
bind', sat_info_bind :: IdSATInfo
sat_info_bind) <- CoreBind -> IdSet -> SatM (CoreBind, IdSATInfo)
satBind CoreBind
bind IdSet
interesting_ids
    (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
-> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreBind -> Expr CoreBndr -> Expr CoreBndr
forall b. Bind b -> Expr b -> Expr b
Let CoreBind
bind' Expr CoreBndr
body', IdSATInfo -> IdSATInfo -> IdSATInfo
mergeIdSATInfo IdSATInfo
sat_info_body IdSATInfo
sat_info_bind, Maybe IdAppInfo
body_app)

satExpr (Tick tickish :: Tickish CoreBndr
tickish expr :: Expr CoreBndr
expr) interesting_ids :: IdSet
interesting_ids = do
    (expr' :: Expr CoreBndr
expr', sat_info_expr :: IdSATInfo
sat_info_expr, expr_app :: Maybe IdAppInfo
expr_app) <- Expr CoreBndr
-> IdSet -> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
satExpr Expr CoreBndr
expr IdSet
interesting_ids
    (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
-> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
forall (m :: * -> *) a. Monad m => a -> m a
return (Tickish CoreBndr -> Expr CoreBndr -> Expr CoreBndr
forall b. Tickish CoreBndr -> Expr b -> Expr b
Tick Tickish CoreBndr
tickish Expr CoreBndr
expr', IdSATInfo
sat_info_expr, Maybe IdAppInfo
expr_app)

satExpr ty :: Expr CoreBndr
ty@(Type _) _ = do
    (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
-> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr CoreBndr
ty, IdSATInfo
emptyIdSATInfo, Maybe IdAppInfo
forall a. Maybe a
Nothing)

satExpr co :: Expr CoreBndr
co@(Coercion _) _ = do
    (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
-> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr CoreBndr
co, IdSATInfo
emptyIdSATInfo, Maybe IdAppInfo
forall a. Maybe a
Nothing)

satExpr (Cast expr :: Expr CoreBndr
expr coercion :: Coercion
coercion) interesting_ids :: IdSet
interesting_ids = do
    (expr' :: Expr CoreBndr
expr', sat_info_expr :: IdSATInfo
sat_info_expr, expr_app :: Maybe IdAppInfo
expr_app) <- Expr CoreBndr
-> IdSet -> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
satExpr Expr CoreBndr
expr IdSet
interesting_ids
    (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
-> SatM (Expr CoreBndr, IdSATInfo, Maybe IdAppInfo)
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr CoreBndr -> Coercion -> Expr CoreBndr
forall b. Expr b -> Coercion -> Expr b
Cast Expr CoreBndr
expr' Coercion
coercion, IdSATInfo
sat_info_expr, Maybe IdAppInfo
expr_app)

{-
************************************************************************

                Static Argument Transformation Monad

************************************************************************
-}

type SatM result = UniqSM result

runSAT :: UniqSupply -> SatM a -> a
runSAT :: UniqSupply -> SatM a -> a
runSAT = UniqSupply -> SatM a -> a
forall a. UniqSupply -> SatM a -> a
initUs_

newUnique :: SatM Unique
newUnique :: SatM Unique
newUnique = SatM Unique
forall (m :: * -> *). MonadUnique m => m Unique
getUniqueM

{-
************************************************************************

                Static Argument Transformation Monad

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

To do the transformation, the game plan is to:

1. Create a small nonrecursive RHS that takes the
   original arguments to the function but discards
   the ones that are static and makes a call to the
   SATed version with the remainder. We intend that
   this will be inlined later, removing the overhead

2. Bind this nonrecursive RHS over the original body
   WITH THE SAME UNIQUE as the original body so that
   any recursive calls to the original now go via
   the small wrapper

3. Rebind the original function to a new one which contains
   our SATed function and just makes a call to it:
   we call the thing making this call the local body

Example: transform this

    map :: forall a b. (a->b) -> [a] -> [b]
    map = /\ab. \(f:a->b) (as:[a]) -> body[map]
to
    map :: forall a b. (a->b) -> [a] -> [b]
    map = /\ab. \(f:a->b) (as:[a]) ->
         letrec map' :: [a] -> [b]
                    -- The "worker function
                map' = \(as:[a]) ->
                         let map :: forall a' b'. (a -> b) -> [a] -> [b]
                                -- The "shadow function
                             map = /\a'b'. \(f':(a->b) (as:[a]).
                                   map' as
                         in body[map]
         in map' as

Note [Shadow binding]
~~~~~~~~~~~~~~~~~~~~~
The calls to the inner map inside body[map] should get inlined
by the local re-binding of 'map'.  We call this the "shadow binding".

But we can't use the original binder 'map' unchanged, because
it might be exported, in which case the shadow binding won't be
discarded as dead code after it is inlined.

So we use a hack: we make a new SysLocal binder with the *same* unique
as binder.  (Another alternative would be to reset the export flag.)

Note [Binder type capture]
~~~~~~~~~~~~~~~~~~~~~~~~~~
Notice that in the inner map (the "shadow function"), the static arguments
are discarded -- it's as if they were underscores.  Instead, mentions
of these arguments (notably in the types of dynamic arguments) are bound
by the *outer* lambdas of the main function.  So we must make up fresh
names for the static arguments so that they do not capture variables
mentioned in the types of dynamic args.

In the map example, the shadow function must clone the static type
argument a,b, giving a',b', to ensure that in the \(as:[a]), the 'a'
is bound by the outer forall.  We clone f' too for consistency, but
that doesn't matter either way because static Id arguments aren't
mentioned in the shadow binding at all.

If we don't we get something like this:

[Exported]
[Arity 3]
GHC.Base.until =
  \ (@ a_aiK)
    (p_a6T :: a_aiK -> GHC.Types.Bool)
    (f_a6V :: a_aiK -> a_aiK)
    (x_a6X :: a_aiK) ->
    letrec {
      sat_worker_s1aU :: a_aiK -> a_aiK
      []
      sat_worker_s1aU =
        \ (x_a6X :: a_aiK) ->
          let {
            sat_shadow_r17 :: forall a_a3O.
                              (a_a3O -> GHC.Types.Bool) -> (a_a3O -> a_a3O) -> a_a3O -> a_a3O
            []
            sat_shadow_r17 =
              \ (@ a_aiK)
                (p_a6T :: a_aiK -> GHC.Types.Bool)
                (f_a6V :: a_aiK -> a_aiK)
                (x_a6X :: a_aiK) ->
                sat_worker_s1aU x_a6X } in
          case p_a6T x_a6X of wild_X3y [ALWAYS Dead Nothing] {
            GHC.Types.False -> GHC.Base.until @ a_aiK p_a6T f_a6V (f_a6V x_a6X);
            GHC.Types.True -> x_a6X
          }; } in
    sat_worker_s1aU x_a6X

Where sat_shadow has captured the type variables of x_a6X etc as it has a a_aiK
type argument. This is bad because it means the application sat_worker_s1aU x_a6X
is not well typed.
-}

saTransformMaybe :: Id -> Maybe SATInfo -> [Id] -> CoreExpr -> SatM CoreBind
saTransformMaybe :: CoreBndr
-> Maybe SATInfo -> [CoreBndr] -> Expr CoreBndr -> SatM CoreBind
saTransformMaybe binder :: CoreBndr
binder maybe_arg_staticness :: Maybe SATInfo
maybe_arg_staticness rhs_binders :: [CoreBndr]
rhs_binders rhs_body :: Expr CoreBndr
rhs_body
  | Just arg_staticness :: SATInfo
arg_staticness <- Maybe SATInfo
maybe_arg_staticness
  , SATInfo -> Bool
should_transform SATInfo
arg_staticness
  = CoreBndr -> SATInfo -> [CoreBndr] -> Expr CoreBndr -> SatM CoreBind
saTransform CoreBndr
binder SATInfo
arg_staticness [CoreBndr]
rhs_binders Expr CoreBndr
rhs_body
  | Bool
otherwise
  = CoreBind -> SatM CoreBind
forall (m :: * -> *) a. Monad m => a -> m a
return ([(CoreBndr, Expr CoreBndr)] -> CoreBind
forall b. [(b, Expr b)] -> Bind b
Rec [(CoreBndr
binder, [CoreBndr] -> Expr CoreBndr -> Expr CoreBndr
forall b. [b] -> Expr b -> Expr b
mkLams [CoreBndr]
rhs_binders Expr CoreBndr
rhs_body)])
  where
    should_transform :: SATInfo -> Bool
should_transform staticness :: SATInfo
staticness = Int
n_static_args Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> 1 -- THIS IS THE DECISION POINT
      where
        n_static_args :: Int
n_static_args = (Staticness App -> Bool) -> SATInfo -> Int
forall a. (a -> Bool) -> [a] -> Int
count Staticness App -> Bool
isStaticValue SATInfo
staticness

saTransform :: Id -> SATInfo -> [Id] -> CoreExpr -> SatM CoreBind
saTransform :: CoreBndr -> SATInfo -> [CoreBndr] -> Expr CoreBndr -> SatM CoreBind
saTransform binder :: CoreBndr
binder arg_staticness :: SATInfo
arg_staticness rhs_binders :: [CoreBndr]
rhs_binders rhs_body :: Expr CoreBndr
rhs_body
  = do  { [CoreBndr]
shadow_lam_bndrs <- ((CoreBndr, Staticness App) -> UniqSM CoreBndr)
-> [(CoreBndr, Staticness App)] -> UniqSM [CoreBndr]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (CoreBndr, Staticness App) -> UniqSM CoreBndr
forall a. (CoreBndr, Staticness a) -> UniqSM CoreBndr
clone [(CoreBndr, Staticness App)]
binders_w_staticness
        ; Unique
uniq             <- SatM Unique
newUnique
        ; CoreBind -> SatM CoreBind
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreBndr -> Expr CoreBndr -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec CoreBndr
binder (Unique -> [CoreBndr] -> Expr CoreBndr
mk_new_rhs Unique
uniq [CoreBndr]
shadow_lam_bndrs)) }
  where
    -- Running example: foldr
    -- foldr \alpha \beta c n xs = e, for some e
    -- arg_staticness = [Static TypeApp, Static TypeApp, Static VarApp, Static VarApp, NonStatic]
    -- rhs_binders = [\alpha, \beta, c, n, xs]
    -- rhs_body = e

    binders_w_staticness :: [(CoreBndr, Staticness App)]
binders_w_staticness = [CoreBndr]
rhs_binders [CoreBndr] -> SATInfo -> [(CoreBndr, Staticness App)]
forall a b. [a] -> [b] -> [(a, b)]
`zip` (SATInfo
arg_staticness SATInfo -> SATInfo -> SATInfo
forall a. [a] -> [a] -> [a]
++ Staticness App -> SATInfo
forall a. a -> [a]
repeat Staticness App
forall a. Staticness a
NotStatic)
                                        -- Any extra args are assumed NotStatic

    non_static_args :: [Var]
            -- non_static_args = [xs]
            -- rhs_binders_without_type_capture = [\alpha', \beta', c, n, xs]
    non_static_args :: [CoreBndr]
non_static_args = [CoreBndr
v | (v :: CoreBndr
v, NotStatic) <- [(CoreBndr, Staticness App)]
binders_w_staticness]

    clone :: (CoreBndr, Staticness a) -> UniqSM CoreBndr
clone (bndr :: CoreBndr
bndr, NotStatic) = CoreBndr -> UniqSM CoreBndr
forall (m :: * -> *) a. Monad m => a -> m a
return CoreBndr
bndr
    clone (bndr :: CoreBndr
bndr, _        ) = do { Unique
uniq <- SatM Unique
newUnique
                                 ; CoreBndr -> UniqSM CoreBndr
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreBndr -> Unique -> CoreBndr
setVarUnique CoreBndr
bndr Unique
uniq) }

    -- new_rhs = \alpha beta c n xs ->
    --           let sat_worker = \xs -> let sat_shadow = \alpha' beta' c n xs ->
    --                                       sat_worker xs
    --                                   in e
    --           in sat_worker xs
    mk_new_rhs :: Unique -> [CoreBndr] -> Expr CoreBndr
mk_new_rhs uniq :: Unique
uniq shadow_lam_bndrs :: [CoreBndr]
shadow_lam_bndrs
        = [CoreBndr] -> Expr CoreBndr -> Expr CoreBndr
forall b. [b] -> Expr b -> Expr b
mkLams [CoreBndr]
rhs_binders (Expr CoreBndr -> Expr CoreBndr) -> Expr CoreBndr -> Expr CoreBndr
forall a b. (a -> b) -> a -> b
$
          CoreBind -> Expr CoreBndr -> Expr CoreBndr
forall b. Bind b -> Expr b -> Expr b
Let ([(CoreBndr, Expr CoreBndr)] -> CoreBind
forall b. [(b, Expr b)] -> Bind b
Rec [(CoreBndr
rec_body_bndr, Expr CoreBndr
rec_body)])
          Expr CoreBndr
local_body
        where
          local_body :: Expr CoreBndr
local_body = Expr CoreBndr -> [CoreBndr] -> Expr CoreBndr
forall b. Expr b -> [CoreBndr] -> Expr b
mkVarApps (CoreBndr -> Expr CoreBndr
forall b. CoreBndr -> Expr b
Var CoreBndr
rec_body_bndr) [CoreBndr]
non_static_args

          rec_body :: Expr CoreBndr
rec_body = [CoreBndr] -> Expr CoreBndr -> Expr CoreBndr
forall b. [b] -> Expr b -> Expr b
mkLams [CoreBndr]
non_static_args (Expr CoreBndr -> Expr CoreBndr) -> Expr CoreBndr -> Expr CoreBndr
forall a b. (a -> b) -> a -> b
$
                     CoreBind -> Expr CoreBndr -> Expr CoreBndr
forall b. Bind b -> Expr b -> Expr b
Let (CoreBndr -> Expr CoreBndr -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec CoreBndr
shadow_bndr Expr CoreBndr
shadow_rhs) Expr CoreBndr
rhs_body

            -- See Note [Binder type capture]
          shadow_rhs :: Expr CoreBndr
shadow_rhs = [CoreBndr] -> Expr CoreBndr -> Expr CoreBndr
forall b. [b] -> Expr b -> Expr b
mkLams [CoreBndr]
shadow_lam_bndrs Expr CoreBndr
local_body
            -- nonrec_rhs = \alpha' beta' c n xs -> sat_worker xs

          rec_body_bndr :: CoreBndr
rec_body_bndr = FastString -> Unique -> Type -> CoreBndr
mkSysLocal (String -> FastString
fsLit "sat_worker") Unique
uniq (Expr CoreBndr -> Type
exprType Expr CoreBndr
rec_body)
            -- rec_body_bndr = sat_worker

            -- See Note [Shadow binding]; make a SysLocal
          shadow_bndr :: CoreBndr
shadow_bndr = FastString -> Unique -> Type -> CoreBndr
mkSysLocal (OccName -> FastString
occNameFS (CoreBndr -> OccName
forall a. NamedThing a => a -> OccName
getOccName CoreBndr
binder))
                                   (CoreBndr -> Unique
idUnique CoreBndr
binder)
                                   (Expr CoreBndr -> Type
exprType Expr CoreBndr
shadow_rhs)

isStaticValue :: Staticness App -> Bool
isStaticValue :: Staticness App -> Bool
isStaticValue (Static (VarApp _)) = Bool
True
isStaticValue _                   = Bool
False