compiler/vectorise/VectMonad.hs

   1 module VectMonad (
   2   Scope(..),
   3   VM,
   4
   5   noV, tryV, maybeV, orElseV, fixV, localV, closedV, initV,
   6   cloneName, cloneId,
   7   newExportedVar, newLocalVar, newDummyVar, newTyVar,
   8
   9   Builtins(..), paDictTyCon, paDictDataCon,
  10   builtin,
  11
  12   GlobalEnv(..),
  13   setInstEnvs,
  14   readGEnv, setGEnv, updGEnv,
  15
  16   LocalEnv(..),
  17   readLEnv, setLEnv, updLEnv,
  18
  19   getBindName, inBind,
  20
  21   lookupVar, defGlobalVar,
  22   lookupTyCon, defTyCon,
  23   lookupDataCon, defDataCon,
  24   lookupTyVarPA, defLocalTyVar, defLocalTyVarWithPA, localTyVars,
  25
  26   lookupInst, lookupFamInst
  27 ) where
  28
  29 #include "HsVersions.h"
  30
  31 import HscTypes
  32 import CoreSyn
  33 import Class
  34 import TyCon
  35 import DataCon
  36 import Type
  37 import Var
  38 import VarEnv
  39 import Id
  40 import OccName
  41 import Name
  42 import NameEnv
  43 import TysPrim       ( intPrimTy )
  44
  45 import DsMonad
  46 import PrelNames
  47
  48 import InstEnv
  49 import FamInstEnv
  50
  51 import Panic
  52 import Outputable
  53 import FastString
  54 import SrcLoc        ( noSrcSpan )
  55
  56 import Control.Monad ( liftM )
  57
  58 data Scope a b = Global a | Local b
  59
  60 -- ----------------------------------------------------------------------------
  61 -- Vectorisation monad
  62
  63 data Builtins = Builtins {
  64                   parrayTyCon      :: TyCon
  65                 , paClass          :: Class
  66                 , closureTyCon     :: TyCon
  67                 , mkClosureVar     :: Var
  68                 , applyClosureVar  :: Var
  69                 , mkClosurePVar    :: Var
  70                 , applyClosurePVar :: Var
  71                 , lengthPAVar      :: Var
  72                 , replicatePAVar   :: Var
  73                 , emptyPAVar       :: Var
  74                 , liftingContext   :: Var
  75                 }
  76
  77 paDictTyCon :: Builtins -> TyCon
  78 paDictTyCon = classTyCon . paClass
  79
  80 paDictDataCon :: Builtins -> DataCon
  81 paDictDataCon = classDataCon . paClass
  82
  83 initBuiltins :: DsM Builtins
  84 initBuiltins
  85   = do
  86       parrayTyCon  <- dsLookupTyCon parrayTyConName
  87       paClass      <- dsLookupClass paClassName
  88       closureTyCon <- dsLookupTyCon closureTyConName
  89
  90       mkClosureVar     <- dsLookupGlobalId mkClosureName
  91       applyClosureVar  <- dsLookupGlobalId applyClosureName
  92       mkClosurePVar    <- dsLookupGlobalId mkClosurePName
  93       applyClosurePVar <- dsLookupGlobalId applyClosurePName
  94       lengthPAVar      <- dsLookupGlobalId lengthPAName
  95       replicatePAVar   <- dsLookupGlobalId replicatePAName
  96       emptyPAVar       <- dsLookupGlobalId emptyPAName
  97
  98       liftingContext <- liftM (\u -> mkSysLocal FSLIT("lc") u intPrimTy)
  99                               newUnique
 100
 101       return $ Builtins {
 102                  parrayTyCon      = parrayTyCon
 103                , paClass          = paClass
 104                , closureTyCon     = closureTyCon
 105                , mkClosureVar     = mkClosureVar
 106                , applyClosureVar  = applyClosureVar
 107                , mkClosurePVar    = mkClosurePVar
 108                , applyClosurePVar = applyClosurePVar
 109                , lengthPAVar      = lengthPAVar
 110                , replicatePAVar   = replicatePAVar
 111                , emptyPAVar       = emptyPAVar
 112                , liftingContext   = liftingContext
 113                }
 114
 115 data GlobalEnv = GlobalEnv {
 116                   -- Mapping from global variables to their vectorised versions.
 117                   --
 118                   global_vars :: VarEnv Var
 119
 120                   -- Exported variables which have a vectorised version
 121                   --
 122                 , global_exported_vars :: VarEnv (Var, Var)
 123
 124                   -- Mapping from TyCons to their vectorised versions.
 125                   -- TyCons which do not have to be vectorised are mapped to
 126                   -- themselves.
 127                   --
 128                 , global_tycons :: NameEnv TyCon
 129
 130                   -- Mapping from DataCons to their vectorised versions
 131                   --
 132                 , global_datacons :: NameEnv DataCon
 133
 134                 -- External package inst-env & home-package inst-env for class
 135                 -- instances
 136                 --
 137                 , global_inst_env :: (InstEnv, InstEnv)
 138
 139                 -- External package inst-env & home-package inst-env for family
 140                 -- instances
 141                 --
 142                 , global_fam_inst_env :: FamInstEnvs
 143
 144                 -- Hoisted bindings
 145                 , global_bindings :: [(Var, CoreExpr)]
 146                 }
 147
 148 data LocalEnv = LocalEnv {
 149                  -- Mapping from local variables to their vectorised and
 150                  -- lifted versions
 151                  --
 152                  local_vars :: VarEnv (Var, Var)
 153
 154                  -- In-scope type variables
 155                  --
 156                , local_tyvars :: [TyVar]
 157
 158                  -- Mapping from tyvars to their PA dictionaries
 159                , local_tyvar_pa :: VarEnv CoreExpr
 160
 161                  -- Local binding name
 162                , local_bind_name :: FastString
 163                }
 164
 165
 166 initGlobalEnv :: VectInfo -> (InstEnv, InstEnv) -> FamInstEnvs -> Builtins -> GlobalEnv
 167 initGlobalEnv info instEnvs famInstEnvs bi
 168   = GlobalEnv {
 169       global_vars          = mapVarEnv snd $ vectInfoVar info
 170     , global_exported_vars = emptyVarEnv
 171     , global_tycons        = extendNameEnv (mapNameEnv snd (vectInfoTyCon info))
 172                                            (tyConName funTyCon) (closureTyCon bi)
 173
 174     , global_datacons      = mapNameEnv snd $ vectInfoDataCon info
 175     , global_inst_env      = instEnvs
 176     , global_fam_inst_env  = famInstEnvs
 177     , global_bindings      = []
 178     }
 179
 180 setInstEnvs :: InstEnv -> FamInstEnv -> GlobalEnv -> GlobalEnv
 181 setInstEnvs l_inst l_fam_inst genv
 182   | (g_inst,     _) <- global_inst_env genv
 183   , (g_fam_inst, _) <- global_fam_inst_env genv
 184   = genv { global_inst_env     = (g_inst, l_inst)
 185          , global_fam_inst_env = (g_fam_inst, l_fam_inst) }
 186
 187 emptyLocalEnv = LocalEnv {
 188                    local_vars     = emptyVarEnv
 189                  , local_tyvars   = []
 190                  , local_tyvar_pa = emptyVarEnv
 191                  , local_bind_name  = FSLIT("fn")
 192                  }
 193
 194 -- FIXME
 195 updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
 196 updVectInfo env tyenv info
 197   = info {
 198       vectInfoVar     = global_exported_vars env
 199     , vectInfoTyCon   = mk_env typeEnvTyCons global_tycons
 200     , vectInfoDataCon = mk_env typeEnvDataCons global_datacons
 201     }
 202   where
 203     mk_env from_tyenv from_env = mkNameEnv [(name, (from,to))
 204                                    | from <- from_tyenv tyenv
 205                                    , let name = getName from
 206                                    , Just to <- [lookupNameEnv (from_env env) name]]
 207
 208 data VResult a = Yes GlobalEnv LocalEnv a | No
 209
 210 newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
 211
 212 instance Monad VM where
 213   return x   = VM $ \bi genv lenv -> return (Yes genv lenv x)
 214   VM p >>= f = VM $ \bi genv lenv -> do
 215                                       r <- p bi genv lenv
 216                                       case r of
 217                                         Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
 218                                         No                -> return No
 219
 220 noV :: VM a
 221 noV = VM $ \_ _ _ -> return No
 222
 223 traceNoV :: String -> SDoc -> VM a
 224 traceNoV s d = pprTrace s d noV
 225
 226 tryV :: VM a -> VM (Maybe a)
 227 tryV (VM p) = VM $ \bi genv lenv ->
 228   do
 229     r <- p bi genv lenv
 230     case r of
 231       Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
 232       No                -> return (Yes genv  lenv  Nothing)
 233
 234 maybeV :: VM (Maybe a) -> VM a
 235 maybeV p = maybe noV return =<< p
 236
 237 traceMaybeV :: String -> SDoc -> VM (Maybe a) -> VM a
 238 traceMaybeV s d p = maybe (traceNoV s d) return =<< p
 239
 240 orElseV :: VM a -> VM a -> VM a
 241 orElseV p q = maybe q return =<< tryV p
 242
 243 fixV :: (a -> VM a) -> VM a
 244 fixV f = VM (\bi genv lenv -> fixDs $ \r -> runVM (f (unYes r)) bi genv lenv )
 245   where
 246     unYes (Yes _ _ x) = x
 247
 248 localV :: VM a -> VM a
 249 localV p = do
 250              env <- readLEnv id
 251              x <- p
 252              setLEnv env
 253              return x
 254
 255 closedV :: VM a -> VM a
 256 closedV p = do
 257               env <- readLEnv id
 258               setLEnv (emptyLocalEnv { local_bind_name = local_bind_name env })
 259               x <- p
 260               setLEnv env
 261               return x
 262
 263 liftDs :: DsM a -> VM a
 264 liftDs p = VM $ \bi genv lenv -> do { x <- p; return (Yes genv lenv x) }
 265
 266 builtin :: (Builtins -> a) -> VM a
 267 builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
 268
 269 readGEnv :: (GlobalEnv -> a) -> VM a
 270 readGEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f genv))
 271
 272 setGEnv :: GlobalEnv -> VM ()
 273 setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
 274
 275 updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
 276 updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
 277
 278 readLEnv :: (LocalEnv -> a) -> VM a
 279 readLEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f lenv))
 280
 281 setLEnv :: LocalEnv -> VM ()
 282 setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
 283
 284 updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
 285 updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
 286
 287 getInstEnv :: VM (InstEnv, InstEnv)
 288 getInstEnv = readGEnv global_inst_env
 289
 290 getFamInstEnv :: VM FamInstEnvs
 291 getFamInstEnv = readGEnv global_fam_inst_env
 292
 293 getBindName :: VM FastString
 294 getBindName = readLEnv local_bind_name
 295
 296 inBind :: Id -> VM a -> VM a
 297 inBind id p
 298   = do updLEnv $ \env -> env { local_bind_name = occNameFS (getOccName id) }
 299        p
 300
 301 cloneName :: (OccName -> OccName) -> Name -> VM Name
 302 cloneName mk_occ name = liftM make (liftDs newUnique)
 303   where
 304     occ_name = mk_occ (nameOccName name)
 305
 306     make u | isExternalName name = mkExternalName u (nameModule name)
 307                                                     occ_name
 308                                                     (nameSrcSpan name)
 309            | otherwise           = mkSystemName u occ_name
 310
 311 cloneId :: (OccName -> OccName) -> Id -> Type -> VM Id
 312 cloneId mk_occ id ty
 313   = do
 314       name <- cloneName mk_occ (getName id)
 315       let id' | isExportedId id = Id.mkExportedLocalId name ty
 316               | otherwise       = Id.mkLocalId         name ty
 317       return id'
 318
 319 newExportedVar :: OccName -> Type -> VM Var
 320 newExportedVar occ_name ty
 321   = do
 322       mod <- liftDs getModuleDs
 323       u   <- liftDs newUnique
 324
 325       let name = mkExternalName u mod occ_name noSrcSpan
 326
 327       return $ Id.mkExportedLocalId name ty
 328
 329 newLocalVar :: FastString -> Type -> VM Var
 330 newLocalVar fs ty
 331   = do
 332       u <- liftDs newUnique
 333       return $ mkSysLocal fs u ty
 334
 335 newDummyVar :: Type -> VM Var
 336 newDummyVar = newLocalVar FSLIT("ds")
 337
 338 newTyVar :: FastString -> Kind -> VM Var
 339 newTyVar fs k
 340   = do
 341       u <- liftDs newUnique
 342       return $ mkTyVar (mkSysTvName u fs) k
 343
 344 defGlobalVar :: Var -> Var -> VM ()
 345 defGlobalVar v v' = updGEnv $ \env ->
 346   env { global_vars = extendVarEnv (global_vars env) v v'
 347       , global_exported_vars = upd (global_exported_vars env)
 348       }
 349   where
 350     upd env | isExportedId v = extendVarEnv env v (v, v')
 351             | otherwise      = env
 352
 353 lookupVar :: Var -> VM (Scope Var (Var, Var))
 354 lookupVar v
 355   = do
 356       r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
 357       case r of
 358         Just e  -> return (Local e)
 359         Nothing -> liftM Global
 360                  $  traceMaybeV "lookupVar" (ppr v)
 361                                 (readGEnv $ \env -> lookupVarEnv (global_vars env) v)
 362
 363 lookupTyCon :: TyCon -> VM (Maybe TyCon)
 364 lookupTyCon tc
 365   | isUnLiftedTyCon tc || isTupleTyCon tc = return (Just tc)
 366
 367   | otherwise = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
 368
 369 defTyCon :: TyCon -> TyCon -> VM ()
 370 defTyCon tc tc' = updGEnv $ \env ->
 371   env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }
 372
 373 lookupDataCon :: DataCon -> VM (Maybe DataCon)
 374 lookupDataCon dc = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
 375
 376 defDataCon :: DataCon -> DataCon -> VM ()
 377 defDataCon dc dc' = updGEnv $ \env ->
 378   env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }
 379
 380 lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
 381 lookupTyVarPA tv = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv
 382
 383 defLocalTyVar :: TyVar -> VM ()
 384 defLocalTyVar tv = updLEnv $ \env ->
 385   env { local_tyvars   = tv : local_tyvars env
 386       , local_tyvar_pa = local_tyvar_pa env `delVarEnv` tv
 387       }
 388
 389 defLocalTyVarWithPA :: TyVar -> CoreExpr -> VM ()
 390 defLocalTyVarWithPA tv pa = updLEnv $ \env ->
 391   env { local_tyvars   = tv : local_tyvars env
 392       , local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa
 393       }
 394
 395 localTyVars :: VM [TyVar]
 396 localTyVars = readLEnv (reverse . local_tyvars)
 397
 398 -- Look up the dfun of a class instance.
 399 --
 400 -- The match must be unique - ie, match exactly one instance - but the
 401 -- type arguments used for matching may be more specific than those of
 402 -- the class instance declaration.  The found class instances must not have
 403 -- any type variables in the instance context that do not appear in the
 404 -- instances head (i.e., no flexi vars); for details for what this means,
 405 -- see the docs at InstEnv.lookupInstEnv.
 406 --
 407 lookupInst :: Class -> [Type] -> VM (DFunId, [Type])
 408 lookupInst cls tys
 409   = do { instEnv <- getInstEnv
 410        ; case lookupInstEnv instEnv cls tys of
 411            ([(inst, inst_tys)], _)
 412              | noFlexiVar -> return (instanceDFunId inst, inst_tys')
 413              | otherwise  -> pprPanic "VectMonad.lookupInst: flexi var: "
 414                                       (ppr $ mkTyConApp (classTyCon cls) tys)
 415              where
 416                inst_tys'  = [ty | Right ty <- inst_tys]
 417                noFlexiVar = all isRight inst_tys
 418            _other         -> traceNoV "lookupInst" (ppr cls <+> ppr tys)
 419        }
 420   where
 421     isRight (Left  _) = False
 422     isRight (Right _) = True
 423
 424 -- Look up the representation tycon of a family instance.
 425 --
 426 -- The match must be unique - ie, match exactly one instance - but the
 427 -- type arguments used for matching may be more specific than those of
 428 -- the family instance declaration.
 429 --
 430 -- Return the instance tycon and its type instance.  For example, if we have
 431 --
 432 --  lookupFamInst 'T' '[Int]' yields (':R42T', 'Int')
 433 --
 434 -- then we have a coercion (ie, type instance of family instance coercion)
 435 --
 436 --  :Co:R42T Int :: T [Int] ~ :R42T Int
 437 --
 438 -- which implies that :R42T was declared as 'data instance T [a]'.
 439 --
 440 lookupFamInst :: TyCon -> [Type] -> VM (TyCon, [Type])
 441 lookupFamInst tycon tys
 442   = ASSERT( isOpenTyCon tycon )
 443     do { instEnv <- getFamInstEnv
 444        ; case lookupFamInstEnv instEnv tycon tys of
 445            [(fam_inst, rep_tys)] -> return (famInstTyCon fam_inst, rep_tys)
 446            _other                ->
 447              pprPanic "VectMonad.lookupFamInst: not found: "
 448                       (ppr $ mkTyConApp tycon tys)
 449        }
 450
 451 initV :: HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
 452 initV hsc_env guts info p
 453   = do
 454       eps <- hscEPS hsc_env
 455       let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
 456       let instEnvs    = (eps_inst_env     eps, mg_inst_env     guts)
 457
 458       Just r <- initDs hsc_env (mg_module guts)
 459                                (mg_rdr_env guts)
 460                                (mg_types guts)
 461                                (go instEnvs famInstEnvs)
 462       return r
 463   where
 464
 465     go instEnvs famInstEnvs =
 466       do
 467         builtins <- initBuiltins
 468         r <- runVM p builtins (initGlobalEnv info instEnvs famInstEnvs builtins)
 469                    emptyLocalEnv
 470         case r of
 471           Yes genv _ x -> return $ Just (new_info genv, x)
 472           No           -> return Nothing
 473
 474     new_info genv = updVectInfo genv (mg_types guts) info
 475