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