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