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