compiler/ndpFlatten/FlattenMonad.hs

   1 --  $Id$
   2 --
   3 --  Copyright (c) [2001..2002] Manuel M T Chakravarty & Gabriele Keller
   4 --
   5 --  Monad maintaining parallel contexts and substitutions for flattening.
   6 --
   7 --- DESCRIPTION ---------------------------------------------------------------
   8 --
   9 --  The flattening transformation needs to perform a fair amount of plumbing.
  10 --  It needs to mainatin a set of variables, called the parallel context for
  11 --  lifting, variable substitutions in case alternatives, and so on.
  12 --  Moreover, we need to manage uniques to create new variables.  The monad
  13 --  defined in this module takes care of maintaining this state.
  14 --
  15 --- DOCU ----------------------------------------------------------------------
  16 --
  17 --  Language: Haskell 98
  18 --
  19 --  * a parallel context is a set of variables that get vectorised during a
  20 --    lifting transformations (ie, their type changes from `t' to `[:t:]')
  21 --
  22 --  * all vectorised variables in a parallel context have the same size; we
  23 --    call this also the size of the parallel context
  24 --
  25 --  * we represent contexts by maps that give the lifted version of a variable
  26 --    (remember that in GHC, variables contain type information that changes
  27 --    during lifting)
  28 --
  29 --- TODO ----------------------------------------------------------------------
  30 --
  31 --  * Assumptions currently made that should (if they turn out to be true) be
  32 --    documented in The Commentary:
  33 --
  34 --    - Local bindings can be copied without any need to alpha-rename bound
  35 --      variables (or their uniques).  Such renaming is only necessary when
  36 --      bindings in a recursive group are replicated; implying that this is
  37 --      required in the case of top-level bindings).  (Note: The CoreTidy path
  38 --      generates global uniques before code generation.)
  39 --
  40 --  * One FIXME left to resolve.
  41 --
  42
  43 module FlattenMonad (
  44
  45   -- monad definition
  46   --
  47   Flatten, runFlatten,
  48
  49   -- variable generation
  50   --
  51   newVar, mkBind,
  52
  53   -- context management & query operations
  54   --
  55   extendContext, packContext, liftVar, liftConst, intersectWithContext,
  56
  57   -- construction of prelude functions
  58   --
  59   mk'fst, mk'eq, mk'neq, mk'and, mk'or, mk'lengthP, mk'replicateP, mk'mapP,
  60   mk'bpermuteP, mk'bpermuteDftP, mk'indexOfP
  61 ) where
  62
  63 -- standard
  64 import Monad        (mplus)
  65
  66 -- GHC
  67 import Panic        (panic)
  68 import Outputable   (Outputable(ppr), pprPanic)
  69 import UniqSupply   (UniqSupply, splitUniqSupply, uniqFromSupply)
  70 import Var          (Var, idType)
  71 import Id           (Id, mkSysLocal)
  72 import Name         (Name)
  73 import VarSet       (VarSet, emptyVarSet, extendVarSet, varSetElems )
  74 import VarEnv       (VarEnv, emptyVarEnv, zipVarEnv, plusVarEnv,
  75                      elemVarEnv, lookupVarEnv, lookupVarEnv_NF, delVarEnvList)
  76 import Type         (Type, tyConAppTyCon)
  77 import HscTypes     (HomePackageTable,
  78                      ExternalPackageState(eps_PTE), HscEnv(hsc_HPT),
  79                      TyThing(..), lookupType)
  80 import PrelNames    ( fstName, andName, orName,
  81                      lengthPName, replicatePName, mapPName, bpermutePName,
  82                      bpermuteDftPName, indexOfPName)
  83 import TysPrim      ( charPrimTyCon, intPrimTyCon, floatPrimTyCon, doublePrimTyCon )
  84 import PrimOp       ( PrimOp(..) )
  85 import PrelInfo     ( primOpId )
  86 import CoreSyn      (Expr(..), Bind(..), CoreBndr, CoreExpr, CoreBind, mkApps)
  87 import CoreUtils    (exprType)
  88 import FastString   (FastString)
  89
  90 -- friends
  91 import NDPCoreUtils (parrElemTy)
  92
  93
  94 -- definition of the monad
  95 -- -----------------------
  96
  97 -- state maintained by the flattening monad
  98 --
  99 data FlattenState = FlattenState {
 100
 101                       -- our source for uniques
 102                       --
 103                       us       :: UniqSupply,
 104
 105                       -- environment containing all known names (including all
 106                       -- Prelude functions)
 107                       --
 108                       env      :: Name -> Id,
 109
 110                       -- this variable determines the parallel context; if
 111                       -- `Nothing', we are in pure vectorisation mode, no
 112                       -- lifting going on
 113                       --
 114                       ctxtVar  :: Maybe Var,
 115
 116                       -- environment that maps each variable that is
 117                       -- vectorised in the current parallel context to the
 118                       -- vectorised version of that variable
 119                       --
 120                       ctxtEnv :: VarEnv Var,
 121
 122                       -- those variables from the *domain* of `ctxtEnv' that
 123                       -- have been used since the last context restriction (cf.
 124                       -- `restrictContext')
 125                       --
 126                       usedVars :: VarSet
 127                     }
 128
 129 -- initial value of the flattening state
 130 --
 131 initialFlattenState :: ExternalPackageState
 132                     -> HomePackageTable
 133                     -> UniqSupply
 134                     -> FlattenState
 135 initialFlattenState eps hpt us =
 136   FlattenState {
 137     us       = us,
 138     env      = lookup,
 139     ctxtVar  = Nothing,
 140     ctxtEnv  = emptyVarEnv,
 141     usedVars = emptyVarSet
 142   }
 143   where
 144     lookup n =
 145       case lookupType hpt (eps_PTE eps) n of
 146         Just (AnId v) -> v
 147         _             -> pprPanic "FlattenMonad: unknown name:" (ppr n)
 148
 149 -- the monad representation (EXPORTED ABSTRACTLY)
 150 --
 151 newtype Flatten a = Flatten {
 152                       unFlatten :: (FlattenState -> (a, FlattenState))
 153                     }
 154
 155 instance Monad Flatten where
 156   return x = Flatten $ \s -> (x, s)
 157   m >>= n  = Flatten $ \s -> let
 158                                (r, s') = unFlatten m s
 159                              in
 160                              unFlatten (n r) s'
 161
 162 -- execute the given flattening computation (EXPORTED)
 163 --
 164 runFlatten :: HscEnv
 165            -> ExternalPackageState
 166            -> UniqSupply
 167            -> Flatten a
 168            -> a
 169 runFlatten hsc_env eps us m
 170   = fst $ unFlatten m (initialFlattenState eps (hsc_HPT hsc_env) us)
 171
 172
 173 -- variable generation
 174 -- -------------------
 175
 176 -- generate a new local variable whose name is based on the given lexeme and
 177 -- whose type is as specified in the second argument (EXPORTED)
 178 --
 179 newVar           :: FastString -> Type -> Flatten Var
 180 newVar lexeme ty  = Flatten $ \state ->
 181   let
 182     (us1, us2) = splitUniqSupply (us state)
 183     state'     = state {us = us2}
 184   in
 185   (mkSysLocal lexeme (uniqFromSupply us1) ty, state')
 186
 187 -- generate a non-recursive binding using a new binder whose name is derived
 188 -- from the given lexeme (EXPORTED)
 189 --
 190 mkBind          :: FastString -> CoreExpr -> Flatten (CoreBndr, CoreBind)
 191 mkBind lexeme e  =
 192   do
 193     v <- newVar lexeme (exprType e)
 194     return (v, NonRec v e)
 195
 196
 197 -- context management
 198 -- ------------------
 199
 200 -- extend the parallel context by the given set of variables (EXPORTED)
 201 --
 202 --  * if there is no parallel context at the moment, the first element of the
 203 --   variable list will be used to determine the new parallel context
 204 --
 205 --  * the second argument is executed in the current context extended with the
 206 --   given variables
 207 --
 208 --  * the variables must already have been lifted by transforming their type,
 209 --   but they *must* have retained their original name (or, at least, their
 210 --   unique); this is needed so that they match the original variable in
 211 --   variable environments
 212 --
 213 --  * any trace of the given set of variables has to be removed from the state
 214 --   at the end of this operation
 215 --
 216 extendContext      :: [Var] -> Flatten a -> Flatten a
 217 extendContext [] m  = m
 218 extendContext vs m  = Flatten $ \state ->
 219   let
 220     extState       = state {
 221                        ctxtVar = ctxtVar state `mplus` Just (head vs),
 222                        ctxtEnv = ctxtEnv state `plusVarEnv` zipVarEnv vs vs
 223                      }
 224     (r, extState') = unFlatten m extState
 225     resState       = extState' { -- remove `vs' from the result state
 226                        ctxtVar  = ctxtVar state,
 227                        ctxtEnv  = ctxtEnv state,
 228                        usedVars = usedVars extState' `delVarEnvList` vs
 229                      }
 230   in
 231   (r, resState)
 232
 233 -- execute the second argument in a restricted context (EXPORTED)
 234 --
 235 --  * all variables in the current parallel context are packed according to
 236 --   the permutation vector associated with the variable passed as the first
 237 --   argument (ie, all elements of vectorised context variables that are
 238 --   invalid in the restricted context are dropped)
 239 --
 240 --  * the returned list of core binders contains the operations that perform
 241 --   the restriction on all variables in the parallel context that *do* occur
 242 --   during the execution of the second argument (ie, `liftVar' is executed at
 243 --   least once on any such variable)
 244 --
 245 packContext        :: Var -> Flatten a -> Flatten (a, [CoreBind])
 246 packContext perm m  = Flatten $ \state ->
 247   let
 248     -- FIXME: To set the packed environment to the unpacked on is a hack of
 249     --   which I am not sure yet (a) whether it works and (b) whether it's
 250     --   really worth it.  The one advantages is that, we can use a var set,
 251     --   after all, instead of a var environment.
 252     --
 253     --   The idea is the following: If we have to pack a variable `x', we
 254     --   generate `let{-NonRec-} x = bpermuteP perm x in ...'.  As this is a
 255     --   non-recursive binding, the lhs `x' overshadows the rhs `x' in the
 256     --   body of the let.
 257     --
 258     --   NB: If we leave it like this, `mkCoreBind' can be simplified.
 259     packedCtxtEnv     = ctxtEnv state
 260     packedState       = state {
 261                           ctxtVar  = fmap
 262                                        (lookupVarEnv_NF packedCtxtEnv)
 263                                        (ctxtVar state),
 264                           ctxtEnv  = packedCtxtEnv,
 265                           usedVars = emptyVarSet
 266                         }
 267     (r, packedState') = unFlatten m packedState
 268     resState          = state {    -- revert to the unpacked context
 269                           ctxtVar  = ctxtVar state,
 270                           ctxtEnv  = ctxtEnv state
 271                         }
 272     bndrs             = map mkCoreBind . varSetElems . usedVars $ packedState'
 273
 274     -- generate a binding for the packed variant of a context variable
 275     --
 276     mkCoreBind var    = let
 277                           rhs = fst $ unFlatten (mk'bpermuteP (idType var)
 278                                                               (Var perm)
 279                                                               (Var var)
 280                                                 ) state
 281                         in
 282                         NonRec (lookupVarEnv_NF packedCtxtEnv var) $ rhs
 283
 284   in
 285   ((r, bndrs), resState)
 286
 287 -- lift a single variable in the current context (EXPORTED)
 288 --
 289 --  * if the variable does not occur in the context, it's value is vectorised to
 290 --   match the size of the current context
 291 --
 292 --  * otherwise, the variable is replaced by whatever the context environment
 293 --   maps it to (this may either be simply the lifted version of the original
 294 --   variable or a packed variant of that variable)
 295 --
 296 --  * the monad keeps track of all lifted variables that occur in the parallel
 297 --   context, so that `packContext' can determine the correct set of core
 298 --   bindings
 299 --
 300 liftVar     :: Var -> Flatten CoreExpr
 301 liftVar var  = Flatten $ \s ->
 302   let
 303     v          = ctxtVarErr s
 304     v'elemType = parrElemTy . idType $ v
 305     len        = fst $ unFlatten (mk'lengthP v'elemType (Var v)) s
 306     replicated = fst $ unFlatten (mk'replicateP (idType var) len (Var var)) s
 307   in case lookupVarEnv (ctxtEnv s) var of
 308     Just liftedVar -> (Var liftedVar,
 309                        s {usedVars = usedVars s `extendVarSet` var})
 310     Nothing        -> (replicated, s)
 311
 312 -- lift a constant expression in the current context (EXPORTED)
 313 --
 314 --  * the value of the constant expression is vectorised to match the current
 315 --   parallel context
 316 --
 317 liftConst   :: CoreExpr -> Flatten CoreExpr
 318 liftConst e  = Flatten $ \s ->
 319   let
 320      v          = ctxtVarErr s
 321      v'elemType = parrElemTy . idType $ v
 322      len        = fst $ unFlatten (mk'lengthP v'elemType (Var v)) s
 323   in
 324   (fst $ unFlatten (mk'replicateP (exprType e) len e ) s, s)
 325
 326 -- pick those variables of the given set that occur (if albeit in lifted form)
 327 -- in the current parallel context (EXPORTED)
 328 --
 329 --  * the variables returned are from the given set and *not* the corresponding
 330 --   context variables
 331 --
 332 intersectWithContext    :: VarSet -> Flatten [Var]
 333 intersectWithContext vs  = Flatten $ \s ->
 334   let
 335     vs' = filter (`elemVarEnv` ctxtEnv s) (varSetElems vs)
 336   in
 337   (vs', s)
 338
 339
 340 -- construct applications of prelude functions
 341 -- -------------------------------------------
 342
 343 -- NB: keep all the used names listed in `FlattenInfo.namesNeededForFlattening'
 344
 345 -- generate an application of `fst' (EXPORTED)
 346 --
 347 mk'fst           :: Type -> Type -> CoreExpr -> Flatten CoreExpr
 348 mk'fst ty1 ty2 a  = mkFunApp fstName [Type ty1, Type ty2, a]
 349
 350 -- generate an application of `&&' (EXPORTED)
 351 --
 352 mk'and       :: CoreExpr -> CoreExpr -> Flatten CoreExpr
 353 mk'and a1 a2  = mkFunApp andName [a1, a2]
 354
 355 -- generate an application of `||' (EXPORTED)
 356 --
 357 mk'or       :: CoreExpr -> CoreExpr -> Flatten CoreExpr
 358 mk'or a1 a2  = mkFunApp orName [a1, a2]
 359
 360 -- generate an application of `==' where the arguments may only be literals
 361 -- that may occur in a Core case expression (i.e., `Char', `Int', `Float', and
 362 -- `Double') (EXPORTED)
 363 --
 364 mk'eq          :: Type -> CoreExpr -> CoreExpr -> Flatten CoreExpr
 365 mk'eq ty a1 a2  = return (mkApps (Var eqName) [a1, a2])
 366                   where
 367                     tc = tyConAppTyCon ty
 368                     --
 369                     eqName | tc == charPrimTyCon   = primOpId CharEqOp
 370                            | tc == intPrimTyCon    = primOpId IntEqOp
 371                            | tc == floatPrimTyCon  = primOpId FloatEqOp
 372                            | tc == doublePrimTyCon = primOpId DoubleEqOp
 373                            | otherwise                   =
 374                              pprPanic "FlattenMonad.mk'eq: " (ppr ty)
 375
 376 -- generate an application of `==' where the arguments may only be literals
 377 -- that may occur in a Core case expression (i.e., `Char', `Int', `Float', and
 378 -- `Double') (EXPORTED)
 379 --
 380 mk'neq          :: Type -> CoreExpr -> CoreExpr -> Flatten CoreExpr
 381 mk'neq ty a1 a2  = return (mkApps (Var neqName) [a1, a2])
 382                    where
 383                      tc = tyConAppTyCon ty
 384                      --
 385                      neqName {-  | name == charPrimTyConName   = neqCharName -}
 386                              | tc == intPrimTyCon             = primOpId IntNeOp
 387                              {-  | name == floatPrimTyConName  = neqFloatName -}
 388                              {-  | name == doublePrimTyConName = neqDoubleName -}
 389                              | otherwise                   =
 390                                pprPanic "FlattenMonad.mk'neq: " (ppr ty)
 391
 392 -- generate an application of `lengthP' (EXPORTED)
 393 --
 394 mk'lengthP      :: Type -> CoreExpr -> Flatten CoreExpr
 395 mk'lengthP ty a  = mkFunApp lengthPName [Type ty, a]
 396
 397 -- generate an application of `replicateP' (EXPORTED)
 398 --
 399 mk'replicateP          :: Type -> CoreExpr -> CoreExpr -> Flatten CoreExpr
 400 mk'replicateP ty a1 a2  = mkFunApp replicatePName [Type ty, a1, a2]
 401
 402 -- generate an application of `replicateP' (EXPORTED)
 403 --
 404 mk'mapP :: Type -> Type -> CoreExpr -> CoreExpr -> Flatten CoreExpr
 405 mk'mapP ty1 ty2 a1 a2  = mkFunApp mapPName [Type ty1, Type ty2, a1, a2]
 406
 407 -- generate an application of `bpermuteP' (EXPORTED)
 408 --
 409 mk'bpermuteP          :: Type -> CoreExpr -> CoreExpr -> Flatten CoreExpr
 410 mk'bpermuteP ty a1 a2  = mkFunApp bpermutePName [Type ty, a1, a2]
 411
 412 -- generate an application of `bpermuteDftP' (EXPORTED)
 413 --
 414 mk'bpermuteDftP :: Type -> CoreExpr -> CoreExpr -> CoreExpr -> Flatten CoreExpr
 415 mk'bpermuteDftP ty a1 a2 a3 = mkFunApp bpermuteDftPName [Type ty, a1, a2, a3]
 416
 417 -- generate an application of `indexOfP' (EXPORTED)
 418 --
 419 mk'indexOfP          :: Type -> CoreExpr -> CoreExpr -> Flatten CoreExpr
 420 mk'indexOfP ty a1 a2  = mkFunApp indexOfPName [Type ty, a1, a2]
 421
 422
 423 -- auxilliary functions
 424 -- --------------------
 425
 426 -- obtain the context variable, aborting if it is not available (as this
 427 -- signals an internal error in the usage of the `Flatten' monad)
 428 --
 429 ctxtVarErr   :: FlattenState -> Var
 430 ctxtVarErr s  = case ctxtVar s of
 431                   Nothing -> panic "FlattenMonad.ctxtVarErr: No context variable available!"
 432                   Just v  -> v
 433
 434 -- given the name of a known function and a set of arguments (needs to include
 435 -- all needed type arguments), build a Core expression that applies the named
 436 -- function to those arguments
 437 --
 438 mkFunApp           :: Name -> [CoreExpr] -> Flatten CoreExpr
 439 mkFunApp name args  =
 440   do
 441     fun <- lookupName name
 442     return $ mkApps (Var fun) args
 443
 444 -- get the `Id' of a known `Name'
 445 --
 446 --  * this can be the `Name' of any function that's visible on the toplevel of
 447 --   the current compilation unit
 448 --
 449 lookupName      :: Name -> Flatten Id
 450 lookupName name  = Flatten $ \s ->
 451   (env s name, s)