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                   -- Mapping from TyCons to their PA dfuns
 135                   --
 136                 , global_pa_funs :: NameEnv Var
 137
 138                 -- External package inst-env & home-package inst-env for class
 139                 -- instances
 140                 --
 141                 , global_inst_env :: (InstEnv, InstEnv)
 142
 143                 -- External package inst-env & home-package inst-env for family
 144                 -- instances
 145                 --
 146                 , global_fam_inst_env :: FamInstEnvs
 147
 148                 -- Hoisted bindings
 149                 , global_bindings :: [(Var, CoreExpr)]
 150                 }
 151
 152 data LocalEnv = LocalEnv {
 153                  -- Mapping from local variables to their vectorised and
 154                  -- lifted versions
 155                  --
 156                  local_vars :: VarEnv (Var, Var)
 157
 158                  -- In-scope type variables
 159                  --
 160                , local_tyvars :: [TyVar]
 161
 162                  -- Mapping from tyvars to their PA dictionaries
 163                , local_tyvar_pa :: VarEnv CoreExpr
 164
 165                  -- Local binding name
 166                , local_bind_name :: FastString
 167                }
 168
 169
 170 initGlobalEnv :: VectInfo -> (InstEnv, InstEnv) -> FamInstEnvs -> Builtins -> GlobalEnv
 171 initGlobalEnv info instEnvs famInstEnvs bi
 172   = GlobalEnv {
 173       global_vars          = mapVarEnv snd $ vectInfoVar info
 174     , global_exported_vars = emptyVarEnv
 175     , global_tycons        = extendNameEnv (mapNameEnv snd (vectInfoTyCon info))
 176                                            (tyConName funTyCon) (closureTyCon bi)
 177
 178     , global_datacons      = mapNameEnv snd $ vectInfoDataCon info
 179     , global_pa_funs       = mapNameEnv snd $ vectInfoPADFun info
 180     , global_inst_env      = instEnvs
 181     , global_fam_inst_env  = famInstEnvs
 182     , global_bindings      = []
 183     }
 184
 185 setInstEnvs :: InstEnv -> FamInstEnv -> GlobalEnv -> GlobalEnv
 186 setInstEnvs l_inst l_fam_inst genv
 187   | (g_inst,     _) <- global_inst_env genv
 188   , (g_fam_inst, _) <- global_fam_inst_env genv
 189   = genv { global_inst_env     = (g_inst, l_inst)
 190          , global_fam_inst_env = (g_fam_inst, l_fam_inst) }
 191
 192 emptyLocalEnv = LocalEnv {
 193                    local_vars     = emptyVarEnv
 194                  , local_tyvars   = []
 195                  , local_tyvar_pa = emptyVarEnv
 196                  , local_bind_name  = FSLIT("fn")
 197                  }
 198
 199 -- FIXME
 200 updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
 201 updVectInfo env tyenv info
 202   = info {
 203       vectInfoVar     = global_exported_vars env
 204     , vectInfoTyCon   = mk_env typeEnvTyCons global_tycons
 205     , vectInfoDataCon = mk_env typeEnvDataCons global_datacons
 206     , vectInfoPADFun  = mk_env typeEnvTyCons global_pa_funs
 207     }
 208   where
 209     mk_env from_tyenv from_env = mkNameEnv [(name, (from,to))
 210                                    | from <- from_tyenv tyenv
 211                                    , let name = getName from
 212                                    , Just to <- [lookupNameEnv (from_env env) name]]
 213
 214 data VResult a = Yes GlobalEnv LocalEnv a | No
 215
 216 newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
 217
 218 instance Monad VM where
 219   return x   = VM $ \bi genv lenv -> return (Yes genv lenv x)
 220   VM p >>= f = VM $ \bi genv lenv -> do
 221                                       r <- p bi genv lenv
 222                                       case r of
 223                                         Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
 224                                         No                -> return No
 225
 226 noV :: VM a
 227 noV = VM $ \_ _ _ -> return No
 228
 229 traceNoV :: String -> SDoc -> VM a
 230 traceNoV s d = pprTrace s d noV
 231
 232 tryV :: VM a -> VM (Maybe a)
 233 tryV (VM p) = VM $ \bi genv lenv ->
 234   do
 235     r <- p bi genv lenv
 236     case r of
 237       Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
 238       No                -> return (Yes genv  lenv  Nothing)
 239
 240 maybeV :: VM (Maybe a) -> VM a
 241 maybeV p = maybe noV return =<< p
 242
 243 traceMaybeV :: String -> SDoc -> VM (Maybe a) -> VM a
 244 traceMaybeV s d p = maybe (traceNoV s d) return =<< p
 245
 246 orElseV :: VM a -> VM a -> VM a
 247 orElseV p q = maybe q return =<< tryV p
 248
 249 fixV :: (a -> VM a) -> VM a
 250 fixV f = VM (\bi genv lenv -> fixDs $ \r -> runVM (f (unYes r)) bi genv lenv )
 251   where
 252     unYes (Yes _ _ x) = x
 253
 254 localV :: VM a -> VM a
 255 localV p = do
 256              env <- readLEnv id
 257              x <- p
 258              setLEnv env
 259              return x
 260
 261 closedV :: VM a -> VM a
 262 closedV p = do
 263               env <- readLEnv id
 264               setLEnv (emptyLocalEnv { local_bind_name = local_bind_name env })
 265               x <- p
 266               setLEnv env
 267               return x
 268
 269 liftDs :: DsM a -> VM a
 270 liftDs p = VM $ \bi genv lenv -> do { x <- p; return (Yes genv lenv x) }
 271
 272 builtin :: (Builtins -> a) -> VM a
 273 builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
 274
 275 readGEnv :: (GlobalEnv -> a) -> VM a
 276 readGEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f genv))
 277
 278 setGEnv :: GlobalEnv -> VM ()
 279 setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
 280
 281 updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
 282 updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
 283
 284 readLEnv :: (LocalEnv -> a) -> VM a
 285 readLEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f lenv))
 286
 287 setLEnv :: LocalEnv -> VM ()
 288 setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
 289
 290 updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
 291 updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
 292
 293 getInstEnv :: VM (InstEnv, InstEnv)
 294 getInstEnv = readGEnv global_inst_env
 295
 296 getFamInstEnv :: VM FamInstEnvs
 297 getFamInstEnv = readGEnv global_fam_inst_env
 298
 299 getBindName :: VM FastString
 300 getBindName = readLEnv local_bind_name
 301
 302 inBind :: Id -> VM a -> VM a
 303 inBind id p
 304   = do updLEnv $ \env -> env { local_bind_name = occNameFS (getOccName id) }
 305        p
 306
 307 cloneName :: (OccName -> OccName) -> Name -> VM Name
 308 cloneName mk_occ name = liftM make (liftDs newUnique)
 309   where
 310     occ_name = mk_occ (nameOccName name)
 311
 312     make u | isExternalName name = mkExternalName u (nameModule name)
 313                                                     occ_name
 314                                                     (nameSrcSpan name)
 315            | otherwise           = mkSystemName u occ_name
 316
 317 cloneId :: (OccName -> OccName) -> Id -> Type -> VM Id
 318 cloneId mk_occ id ty
 319   = do
 320       name <- cloneName mk_occ (getName id)
 321       let id' | isExportedId id = Id.mkExportedLocalId name ty
 322               | otherwise       = Id.mkLocalId         name ty
 323       return id'
 324
 325 newExportedVar :: OccName -> Type -> VM Var
 326 newExportedVar occ_name ty
 327   = do
 328       mod <- liftDs getModuleDs
 329       u   <- liftDs newUnique
 330
 331       let name = mkExternalName u mod occ_name noSrcSpan
 332
 333       return $ Id.mkExportedLocalId name ty
 334
 335 newLocalVar :: FastString -> Type -> VM Var
 336 newLocalVar fs ty
 337   = do
 338       u <- liftDs newUnique
 339       return $ mkSysLocal fs u ty
 340
 341 newDummyVar :: Type -> VM Var
 342 newDummyVar = newLocalVar FSLIT("ds")
 343
 344 newTyVar :: FastString -> Kind -> VM Var
 345 newTyVar fs k
 346   = do
 347       u <- liftDs newUnique
 348       return $ mkTyVar (mkSysTvName u fs) k
 349
 350 defGlobalVar :: Var -> Var -> VM ()
 351 defGlobalVar v v' = updGEnv $ \env ->
 352   env { global_vars = extendVarEnv (global_vars env) v v'
 353       , global_exported_vars = upd (global_exported_vars env)
 354       }
 355   where
 356     upd env | isExportedId v = extendVarEnv env v (v, v')
 357             | otherwise      = env
 358
 359 lookupVar :: Var -> VM (Scope Var (Var, Var))
 360 lookupVar v
 361   = do
 362       r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
 363       case r of
 364         Just e  -> return (Local e)
 365         Nothing -> liftM Global
 366                  $  traceMaybeV "lookupVar" (ppr v)
 367                                 (readGEnv $ \env -> lookupVarEnv (global_vars env) v)
 368
 369 lookupTyCon :: TyCon -> VM (Maybe TyCon)
 370 lookupTyCon tc
 371   | isUnLiftedTyCon tc || isTupleTyCon tc = return (Just tc)
 372
 373   | otherwise = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
 374
 375 defTyCon :: TyCon -> TyCon -> VM ()
 376 defTyCon tc tc' = updGEnv $ \env ->
 377   env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }
 378
 379 lookupDataCon :: DataCon -> VM (Maybe DataCon)
 380 lookupDataCon dc = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
 381
 382 defDataCon :: DataCon -> DataCon -> VM ()
 383 defDataCon dc dc' = updGEnv $ \env ->
 384   env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }
 385
 386 lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
 387 lookupTyVarPA tv = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv
 388
 389 defLocalTyVar :: TyVar -> VM ()
 390 defLocalTyVar tv = updLEnv $ \env ->
 391   env { local_tyvars   = tv : local_tyvars env
 392       , local_tyvar_pa = local_tyvar_pa env `delVarEnv` tv
 393       }
 394
 395 defLocalTyVarWithPA :: TyVar -> CoreExpr -> VM ()
 396 defLocalTyVarWithPA tv pa = updLEnv $ \env ->
 397   env { local_tyvars   = tv : local_tyvars env
 398       , local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa
 399       }
 400
 401 localTyVars :: VM [TyVar]
 402 localTyVars = readLEnv (reverse . local_tyvars)
 403
 404 -- Look up the dfun of a class instance.
 405 --
 406 -- The match must be unique - ie, match exactly one instance - but the
 407 -- type arguments used for matching may be more specific than those of
 408 -- the class instance declaration.  The found class instances must not have
 409 -- any type variables in the instance context that do not appear in the
 410 -- instances head (i.e., no flexi vars); for details for what this means,
 411 -- see the docs at InstEnv.lookupInstEnv.
 412 --
 413 lookupInst :: Class -> [Type] -> VM (DFunId, [Type])
 414 lookupInst cls tys
 415   = do { instEnv <- getInstEnv
 416        ; case lookupInstEnv instEnv cls tys of
 417            ([(inst, inst_tys)], _)
 418              | noFlexiVar -> return (instanceDFunId inst, inst_tys')
 419              | otherwise  -> pprPanic "VectMonad.lookupInst: flexi var: "
 420                                       (ppr $ mkTyConApp (classTyCon cls) tys)
 421              where
 422                inst_tys'  = [ty | Right ty <- inst_tys]
 423                noFlexiVar = all isRight inst_tys
 424            _other         -> traceNoV "lookupInst" (ppr cls <+> ppr tys)
 425        }
 426   where
 427     isRight (Left  _) = False
 428     isRight (Right _) = True
 429
 430 -- Look up the representation tycon of a family instance.
 431 --
 432 -- The match must be unique - ie, match exactly one instance - but the
 433 -- type arguments used for matching may be more specific than those of
 434 -- the family instance declaration.
 435 --
 436 -- Return the instance tycon and its type instance.  For example, if we have
 437 --
 438 --  lookupFamInst 'T' '[Int]' yields (':R42T', 'Int')
 439 --
 440 -- then we have a coercion (ie, type instance of family instance coercion)
 441 --
 442 --  :Co:R42T Int :: T [Int] ~ :R42T Int
 443 --
 444 -- which implies that :R42T was declared as 'data instance T [a]'.
 445 --
 446 lookupFamInst :: TyCon -> [Type] -> VM (TyCon, [Type])
 447 lookupFamInst tycon tys
 448   = ASSERT( isOpenTyCon tycon )
 449     do { instEnv <- getFamInstEnv
 450        ; case lookupFamInstEnv instEnv tycon tys of
 451            [(fam_inst, rep_tys)] -> return (famInstTyCon fam_inst, rep_tys)
 452            _other                ->
 453              pprPanic "VectMonad.lookupFamInst: not found: "
 454                       (ppr $ mkTyConApp tycon tys)
 455        }
 456
 457 initV :: HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
 458 initV hsc_env guts info p
 459   = do
 460       eps <- hscEPS hsc_env
 461       let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
 462       let instEnvs    = (eps_inst_env     eps, mg_inst_env     guts)
 463
 464       Just r <- initDs hsc_env (mg_module guts)
 465                                (mg_rdr_env guts)
 466                                (mg_types guts)
 467                                (go instEnvs famInstEnvs)
 468       return r
 469   where
 470
 471     go instEnvs famInstEnvs =
 472       do
 473         builtins <- initBuiltins
 474         r <- runVM p builtins (initGlobalEnv info instEnvs famInstEnvs builtins)
 475                    emptyLocalEnv
 476         case r of
 477           Yes genv _ x -> return $ Just (new_info genv, x)
 478           No           -> return Nothing
 479
 480     new_info genv = updVectInfo genv (mg_types guts) info
 481