compiler/vectorise/VectMonad.hs

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