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