ghc/compiler/simplCore/SATMonad.lhs

   1 %
   2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1995
   3 %
   4 %************************************************************************
   5 %*                                                                      *
   6 \section[SATMonad]{The Static Argument Transformation pass Monad}
   7 %*                                                                      *
   8 %************************************************************************
   9
  10 \begin{code}
  11 #include "HsVersions.h"
  12
  13 module SATMonad (
  14         SATInfo(..), updSAEnv,
  15         SatM(..), initSAT, emptyEnvSAT,
  16         returnSAT, thenSAT, thenSAT_, mapSAT, getSATInfo, newSATName,
  17         getArgLists, Arg(..), insSAEnv, saTransform,
  18
  19         SATEnv(..), isStatic, dropStatics
  20     ) where
  21
  22 import Type             ( mkTyVarTy, mkSigmaTy, TyVarTemplate,
  23                           extractTyVarsFromTy, splitSigmaTy, splitTyArgs,
  24                           glueTyArgs, instantiateTy, TauType(..),
  25                           Class, ThetaType(..), SigmaType(..),
  26                           InstTyEnv(..)
  27                         )
  28 import Id               ( mkSysLocal, idType )
  29 import Maybes           ( Maybe(..) )
  30 import SrcLoc           ( SrcLoc, mkUnknownSrcLoc )
  31 import UniqSupply
  32 import Util
  33
  34 infixr 9 `thenSAT`, `thenSAT_`
  35 \end{code}
  36
  37 %************************************************************************
  38 %*                                                                      *
  39 \subsection{Static Argument Transformation Environment}
  40 %*                                                                      *
  41 %************************************************************************
  42
  43 \begin{code}
  44 type SATEnv = IdEnv SATInfo
  45
  46 type SATInfo = ([Arg Type],[Arg Id])
  47
  48 data Arg a = Static a | NotStatic
  49     deriving Eq
  50
  51 delOneFromSAEnv v us env
  52   = ((), delOneFromIdEnv env v)
  53
  54 updSAEnv :: Maybe (Id,SATInfo) -> SatM ()
  55 updSAEnv Nothing
  56   = returnSAT ()
  57 updSAEnv (Just (b,(tyargs,args)))
  58   = getSATInfo b      `thenSAT` (\ r ->
  59     case r of
  60       Nothing              -> returnSAT ()
  61       Just (tyargs',args') -> delOneFromSAEnv b `thenSAT_`
  62                               insSAEnv b (checkArgs tyargs tyargs',
  63                                           checkArgs args args')
  64     )
  65
  66 checkArgs as [] = notStatics (length as)
  67 checkArgs [] as = notStatics (length as)
  68 checkArgs (a:as) (a':as') | a == a' = a:checkArgs as as'
  69 checkArgs (_:as) (_:as') = NotStatic:checkArgs as as'
  70
  71 notStatics :: Int -> [Arg a]
  72 notStatics n = nOfThem n NotStatic
  73
  74 insSAEnv :: Id -> SATInfo -> SatM ()
  75 insSAEnv b info us env
  76   = ((), addOneToIdEnv env b info)
  77 \end{code}
  78
  79 %************************************************************************
  80 %*                                                                      *
  81 \subsection{Static Argument Transformation Monad}
  82 %*                                                                      *
  83 %************************************************************************
  84
  85 Two items of state to thread around: a UniqueSupply and a SATEnv.
  86
  87 \begin{code}
  88 type SatM result
  89   =  UniqSupply -> SATEnv -> (result, SATEnv)
  90
  91 initSAT :: SatM a -> UniqSupply -> a
  92
  93 initSAT f us = fst (f us nullIdEnv)
  94
  95 thenSAT m k us env
  96   = case splitUniqSupply us     of { (s1, s2) ->
  97     case m s1 env               of { (m_result, menv) ->
  98     k m_result s2 menv }}
  99
 100 thenSAT_ m k us env
 101   = case splitUniqSupply us     of { (s1, s2) ->
 102     case m s1 env               of { (_, menv) ->
 103     k s2 menv }}
 104
 105 emptyEnvSAT :: SatM ()
 106 emptyEnvSAT us _ = ((), nullIdEnv)
 107
 108 returnSAT v us env = (v, env)
 109
 110 mapSAT f []     = returnSAT []
 111 mapSAT f (x:xs)
 112   = f x         `thenSAT` \ x'  ->
 113     mapSAT f xs `thenSAT` \ xs' ->
 114     returnSAT (x':xs')
 115 \end{code}
 116
 117 %************************************************************************
 118 %*                                                                      *
 119 \subsection{Utility Functions}
 120 %*                                                                      *
 121 %************************************************************************
 122
 123 \begin{code}
 124 getSATInfo :: Id -> SatM (Maybe SATInfo)
 125 getSATInfo var us env
 126   = (lookupIdEnv env var, env)
 127
 128 newSATName :: Id -> Type -> SatM Id
 129 newSATName id ty us env
 130   = case (getUnique us) of { unique ->
 131     (mkSysLocal new_str unique ty mkUnknownSrcLoc, env) }
 132   where
 133     new_str = getOccurrenceName id _APPEND_ SLIT("_sat")
 134
 135 getArgLists :: CoreExpr -> ([Arg Type],[Arg Id])
 136 getArgLists expr
 137   = let
 138         (uvs, tvs, lambda_bounds, body) = digForLambdas expr
 139     in
 140     ([ Static (mkTyVarTy tv) | tv <- tvs ],
 141      [ Static v              | v <- lambda_bounds ])
 142
 143 dropArgs :: CoreExpr -> CoreExpr
 144 dropArgs (Lam   _ e)    = dropArgs e
 145 dropArgs (CoTyLam _ e)  = dropArgs e
 146 dropArgs e              = e
 147 \end{code}
 148
 149 We implement saTransform using shadowing of binders, that is
 150 we transform
 151 map = \f as -> case as of
 152                  [] -> []
 153                  (a':as') -> let x = f a'
 154                                  y = map f as'
 155                              in x:y
 156 to
 157 map = \f as -> let map = \f as -> map' as
 158                in let rec map' = \as -> case as of
 159                                           [] -> []
 160                                           (a':as') -> let x = f a'
 161                                                           y = map f as'
 162                                                       in x:y
 163                   in map' as
 164
 165 the inner map should get inlined and eliminated.
 166 \begin{code}
 167 saTransform :: Id -> CoreExpr -> SatM CoreBinding
 168 saTransform binder rhs
 169   = getSATInfo binder `thenSAT` \ r ->
 170     case r of
 171       -- [Andre] test: do it only if we have more than one static argument.
 172       --Just (tyargs,args) | any isStatic args
 173       Just (tyargs,args) | length (filter isStatic args) > 1
 174         -> newSATName binder (new_ty tyargs args)  `thenSAT` \ binder' ->
 175            mkNewRhs binder binder' tyargs args rhs `thenSAT` \ new_rhs ->
 176            trace ("SAT "++ show (length (filter isStatic args))) (
 177            returnSAT (NonRec binder new_rhs)
 178            )
 179       _ -> returnSAT (Rec [(binder, rhs)])
 180   where
 181     mkNewRhs binder binder' tyargs args rhs
 182       = let
 183             non_static_args :: [Id]
 184             non_static_args
 185                = get_nsa args (snd (getArgLists rhs))
 186                where
 187                  get_nsa :: [Arg a] -> [Arg a] -> [a]
 188                  get_nsa [] _ = []
 189                  get_nsa _ [] = []
 190                  get_nsa (NotStatic:args) (Static v:as) = v:get_nsa args as
 191                  get_nsa (_:args)         (_:as)        =   get_nsa args as
 192
 193             local_body = foldl App (Var binder')
 194                                 [VarArg a | a <- non_static_args]
 195
 196             nonrec_rhs = origLams local_body
 197
 198             -- HACK! The following is a fake SysLocal binder with
 199             -- *the same* unique as binder.
 200             -- the reason for this is the following:
 201             -- this binder *will* get inlined but if it happen to be
 202             -- a top level binder it is never removed as dead code,
 203             -- therefore we have to remove that information (of it being
 204             -- top-level or exported somehow.
 205             -- A better fix is to use binder directly but with the TopLevel
 206             -- tag (or Exported tag) modified.
 207             fake_binder = mkSysLocal
 208                             (getOccurrenceName binder _APPEND_ SLIT("_fsat"))
 209                             (getItsUnique binder)
 210                             (idType binder)
 211                             mkUnknownSrcLoc
 212             rec_body = mkValLam non_static_args
 213                                ( Let (NonRec fake_binder nonrec_rhs)
 214                                  {-in-} (dropArgs rhs))
 215         in
 216         returnSAT (
 217             origLams (Let (Rec [(binder',rec_body)]) {-in-} local_body)
 218         )
 219       where
 220         origLams = origLams' rhs
 221                  where
 222                    origLams' (Lam v e)     e' = Lam   v  (origLams' e e')
 223                    origLams' (CoTyLam ty e)  e' = CoTyLam ty (origLams' e e')
 224                    origLams' _               e' = e'
 225
 226     new_ty tyargs args
 227       = instantiateTy (mk_inst_tyenv tyargs tv_tmpl)
 228                       (mkSigmaTy tv_tmpl' dict_tys' tau_ty')
 229       where
 230         -- get type info for the local function:
 231         (tv_tmpl, dict_tys, tau_ty) = (splitSigmaTy . idType) binder
 232         (reg_arg_tys, res_type)     = splitTyArgs tau_ty
 233
 234         -- now, we drop the ones that are
 235         -- static, that is, the ones we will not pass to the local function
 236         l            = length dict_tys
 237         tv_tmpl'     = dropStatics tyargs tv_tmpl
 238         dict_tys'    = dropStatics (take l args) dict_tys
 239         reg_arg_tys' = dropStatics (drop l args) reg_arg_tys
 240         tau_ty'      = glueTyArgs reg_arg_tys' res_type
 241
 242         mk_inst_tyenv []                    _ = []
 243         mk_inst_tyenv (Static s:args) (t:ts)  = (t,s) : mk_inst_tyenv args ts
 244         mk_inst_tyenv (_:args)      (_:ts)    = mk_inst_tyenv args ts
 245
 246 dropStatics [] t = t
 247 dropStatics (Static _:args) (t:ts) = dropStatics args ts
 248 dropStatics (_:args)        (t:ts) = t:dropStatics args ts
 249
 250 isStatic :: Arg a -> Bool
 251 isStatic NotStatic = False
 252 isStatic _         = True
 253 \end{code}