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 traceNoV :: String -> SDoc -> VM a
 217 traceNoV s d = pprTrace s d noV
 218
 219 tryV :: VM a -> VM (Maybe a)
 220 tryV (VM p) = VM $ \bi genv lenv ->
 221   do
 222     r <- p bi genv lenv
 223     case r of
 224       Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
 225       No                -> return (Yes genv  lenv  Nothing)
 226
 227 maybeV :: VM (Maybe a) -> VM a
 228 maybeV p = maybe noV return =<< p
 229
 230 traceMaybeV :: String -> SDoc -> VM (Maybe a) -> VM a
 231 traceMaybeV s d p = maybe (traceNoV s d) return =<< p
 232
 233 orElseV :: VM a -> VM a -> VM a
 234 orElseV p q = maybe q return =<< tryV p
 235
 236 fixV :: (a -> VM a) -> VM a
 237 fixV f = VM (\bi genv lenv -> fixDs $ \r -> runVM (f (unYes r)) bi genv lenv )
 238   where
 239     unYes (Yes _ _ x) = x
 240
 241 localV :: VM a -> VM a
 242 localV p = do
 243              env <- readLEnv id
 244              x <- p
 245              setLEnv env
 246              return x
 247
 248 closedV :: VM a -> VM a
 249 closedV p = do
 250               env <- readLEnv id
 251               setLEnv (emptyLocalEnv { local_bind_name = local_bind_name env })
 252               x <- p
 253               setLEnv env
 254               return x
 255
 256 liftDs :: DsM a -> VM a
 257 liftDs p = VM $ \bi genv lenv -> do { x <- p; return (Yes genv lenv x) }
 258
 259 builtin :: (Builtins -> a) -> VM a
 260 builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
 261
 262 readGEnv :: (GlobalEnv -> a) -> VM a
 263 readGEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f genv))
 264
 265 setGEnv :: GlobalEnv -> VM ()
 266 setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
 267
 268 updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
 269 updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
 270
 271 readLEnv :: (LocalEnv -> a) -> VM a
 272 readLEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f lenv))
 273
 274 setLEnv :: LocalEnv -> VM ()
 275 setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
 276
 277 updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
 278 updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
 279
 280 getInstEnv :: VM (InstEnv, InstEnv)
 281 getInstEnv = readGEnv global_inst_env
 282
 283 getFamInstEnv :: VM FamInstEnvs
 284 getFamInstEnv = readGEnv global_fam_inst_env
 285
 286 getBindName :: VM FastString
 287 getBindName = readLEnv local_bind_name
 288
 289 inBind :: Id -> VM a -> VM a
 290 inBind id p
 291   = do updLEnv $ \env -> env { local_bind_name = occNameFS (getOccName id) }
 292        p
 293
 294 cloneName :: (OccName -> OccName) -> Name -> VM Name
 295 cloneName mk_occ name = liftM make (liftDs newUnique)
 296   where
 297     occ_name = mk_occ (nameOccName name)
 298
 299     make u | isExternalName name = mkExternalName u (nameModule name)
 300                                                     occ_name
 301                                                     (nameSrcSpan name)
 302            | otherwise           = mkSystemName u occ_name
 303
 304 newExportedVar :: OccName -> Type -> VM Var
 305 newExportedVar occ_name ty
 306   = do
 307       mod <- liftDs getModuleDs
 308       u   <- liftDs newUnique
 309
 310       let name = mkExternalName u mod occ_name noSrcSpan
 311
 312       return $ Id.mkExportedLocalId name ty
 313
 314 newLocalVar :: FastString -> Type -> VM Var
 315 newLocalVar fs ty
 316   = do
 317       u <- liftDs newUnique
 318       return $ mkSysLocal fs u ty
 319
 320 newDummyVar :: Type -> VM Var
 321 newDummyVar = newLocalVar FSLIT("ds")
 322
 323 newTyVar :: FastString -> Kind -> VM Var
 324 newTyVar fs k
 325   = do
 326       u <- liftDs newUnique
 327       return $ mkTyVar (mkSysTvName u fs) k
 328
 329 defGlobalVar :: Var -> Var -> VM ()
 330 defGlobalVar v v' = updGEnv $ \env ->
 331   env { global_vars = extendVarEnv (global_vars env) v v'
 332       , global_exported_vars = upd (global_exported_vars env)
 333       }
 334   where
 335     upd env | isExportedId v = extendVarEnv env v (v, v')
 336             | otherwise      = env
 337
 338 lookupVar :: Var -> VM (Scope Var (Var, Var))
 339 lookupVar v
 340   = do
 341       r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
 342       case r of
 343         Just e  -> return (Local e)
 344         Nothing -> liftM Global
 345                  $  traceMaybeV "lookupVar" (ppr v)
 346                                 (readGEnv $ \env -> lookupVarEnv (global_vars env) v)
 347
 348 lookupTyCon :: TyCon -> VM (Maybe TyCon)
 349 lookupTyCon tc
 350   | isUnLiftedTyCon tc || isTupleTyCon tc = return (Just tc)
 351
 352   | otherwise = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
 353
 354 defTyCon :: TyCon -> TyCon -> VM ()
 355 defTyCon tc tc' = updGEnv $ \env ->
 356   env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }
 357
 358 lookupDataCon :: DataCon -> VM (Maybe DataCon)
 359 lookupDataCon dc = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
 360
 361 defDataCon :: DataCon -> DataCon -> VM ()
 362 defDataCon dc dc' = updGEnv $ \env ->
 363   env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }
 364
 365 lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
 366 lookupTyVarPA tv = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv
 367
 368 defLocalTyVar :: TyVar -> VM ()
 369 defLocalTyVar tv = updLEnv $ \env ->
 370   env { local_tyvars   = tv : local_tyvars env
 371       , local_tyvar_pa = local_tyvar_pa env `delVarEnv` tv
 372       }
 373
 374 defLocalTyVarWithPA :: TyVar -> CoreExpr -> VM ()
 375 defLocalTyVarWithPA tv pa = updLEnv $ \env ->
 376   env { local_tyvars   = tv : local_tyvars env
 377       , local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa
 378       }
 379
 380 localTyVars :: VM [TyVar]
 381 localTyVars = readLEnv (reverse . local_tyvars)
 382
 383 -- Look up the dfun of a class instance.
 384 --
 385 -- The match must be unique - ie, match exactly one instance - but the
 386 -- type arguments used for matching may be more specific than those of
 387 -- the class instance declaration.  The found class instances must not have
 388 -- any type variables in the instance context that do not appear in the
 389 -- instances head (i.e., no flexi vars); for details for what this means,
 390 -- see the docs at InstEnv.lookupInstEnv.
 391 --
 392 lookupInst :: Class -> [Type] -> VM (DFunId, [Type])
 393 lookupInst cls tys
 394   = do { instEnv <- getInstEnv
 395        ; case lookupInstEnv instEnv cls tys of
 396            ([(inst, inst_tys)], _)
 397              | noFlexiVar -> return (instanceDFunId inst, inst_tys')
 398              | otherwise  -> pprPanic "VectMonad.lookupInst: flexi var: "
 399                                       (ppr $ mkTyConApp (classTyCon cls) tys)
 400              where
 401                inst_tys'  = [ty | Right ty <- inst_tys]
 402                noFlexiVar = all isRight inst_tys
 403            _other         -> traceNoV "lookupInst" (ppr cls <+> ppr tys)
 404        }
 405   where
 406     isRight (Left  _) = False
 407     isRight (Right _) = True
 408
 409 -- Look up the representation tycon of a family instance.
 410 --
 411 -- The match must be unique - ie, match exactly one instance - but the
 412 -- type arguments used for matching may be more specific than those of
 413 -- the family instance declaration.
 414 --
 415 -- Return the instance tycon and its type instance.  For example, if we have
 416 --
 417 --  lookupFamInst 'T' '[Int]' yields (':R42T', 'Int')
 418 --
 419 -- then we have a coercion (ie, type instance of family instance coercion)
 420 --
 421 --  :Co:R42T Int :: T [Int] ~ :R42T Int
 422 --
 423 -- which implies that :R42T was declared as 'data instance T [a]'.
 424 --
 425 lookupFamInst :: TyCon -> [Type] -> VM (TyCon, [Type])
 426 lookupFamInst tycon tys
 427   = ASSERT( isOpenTyCon tycon )
 428     do { instEnv <- getFamInstEnv
 429        ; case lookupFamInstEnv instEnv tycon tys of
 430            [(fam_inst, rep_tys)] -> return (famInstTyCon fam_inst, rep_tys)
 431            _other                ->
 432              pprPanic "VectMonad.lookupFamInst: not found: "
 433                       (ppr $ mkTyConApp tycon tys)
 434        }
 435
 436 initV :: HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
 437 initV hsc_env guts info p
 438   = do
 439       eps <- hscEPS hsc_env
 440       let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
 441       let instEnvs    = (eps_inst_env     eps, mg_inst_env     guts)
 442
 443       Just r <- initDs hsc_env (mg_module guts)
 444                                (mg_rdr_env guts)
 445                                (mg_types guts)
 446                                (go instEnvs famInstEnvs)
 447       return r
 448   where
 449
 450     go instEnvs famInstEnvs =
 451       do
 452         builtins <- initBuiltins
 453         r <- runVM p builtins (initGlobalEnv info instEnvs famInstEnvs builtins)
 454                    emptyLocalEnv
 455         case r of
 456           Yes genv _ x -> return $ Just (new_info genv, x)
 457           No           -> return Nothing
 458
 459     new_info genv = updVectInfo genv (mg_types guts) info
 460