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