compiler/vectorise/VectMonad.hs

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