compiler/vectorise/VectMonad.hs

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