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