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, combinePAVar,
  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   lookupBoxedTyCon,
  35   lookupPrimMethod, lookupPrimPArray,
  36   lookupTyVarPA, defLocalTyVar, defLocalTyVarWithPA, localTyVars,
  37
  38   {-lookupInst,-} lookupFamInst
  39 ) where
  40
  41 #include "HsVersions.h"
  42
  43 import VectBuiltIn
  44
  45 import HscTypes
  46 import CoreSyn
  47 import TyCon
  48 import DataCon
  49 import Type
  50 import Class
  51 import Var
  52 import VarEnv
  53 import Id
  54 import OccName
  55 import Name
  56 import NameEnv
  57 import TysPrim       ( intPrimTy )
  58 import Module
  59 import IfaceEnv
  60 import IOEnv         ( ioToIOEnv )
  61
  62 import DsMonad
  63 import PrelNames
  64
  65 import InstEnv
  66 import FamInstEnv
  67
  68 import Panic
  69 import Outputable
  70 import FastString
  71 import SrcLoc        ( noSrcSpan )
  72
  73 import Control.Monad ( liftM, zipWithM )
  74
  75 data Scope a b = Global a | Local b
  76
  77 -- ----------------------------------------------------------------------------
  78 -- Vectorisation monad
  79
  80 data GlobalEnv = GlobalEnv {
  81                   -- Mapping from global variables to their vectorised versions.
  82                   --
  83                   global_vars :: VarEnv Var
  84
  85                   -- Exported variables which have a vectorised version
  86                   --
  87                 , global_exported_vars :: VarEnv (Var, Var)
  88
  89                   -- Mapping from TyCons to their vectorised versions.
  90                   -- TyCons which do not have to be vectorised are mapped to
  91                   -- themselves.
  92                   --
  93                 , global_tycons :: NameEnv TyCon
  94
  95                   -- Mapping from DataCons to their vectorised versions
  96                   --
  97                 , global_datacons :: NameEnv DataCon
  98
  99                   -- Mapping from TyCons to their PA dfuns
 100                   --
 101                 , global_pa_funs :: NameEnv Var
 102
 103                   -- Mapping from TyCons to their PR dfuns
 104                 , global_pr_funs :: NameEnv Var
 105
 106                   -- Mapping from unboxed TyCons to their boxed versions
 107                 , global_boxed_tycons :: NameEnv TyCon
 108
 109                 -- External package inst-env & home-package inst-env for class
 110                 -- instances
 111                 --
 112                 , global_inst_env :: (InstEnv, InstEnv)
 113
 114                 -- External package inst-env & home-package inst-env for family
 115                 -- instances
 116                 --
 117                 , global_fam_inst_env :: FamInstEnvs
 118
 119                 -- Hoisted bindings
 120                 , global_bindings :: [(Var, CoreExpr)]
 121                 }
 122
 123 data LocalEnv = LocalEnv {
 124                  -- Mapping from local variables to their vectorised and
 125                  -- lifted versions
 126                  --
 127                  local_vars :: VarEnv (Var, Var)
 128
 129                  -- In-scope type variables
 130                  --
 131                , local_tyvars :: [TyVar]
 132
 133                  -- Mapping from tyvars to their PA dictionaries
 134                , local_tyvar_pa :: VarEnv CoreExpr
 135
 136                  -- Local binding name
 137                , local_bind_name :: FastString
 138                }
 139
 140 initGlobalEnv :: VectInfo -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
 141 initGlobalEnv info instEnvs famInstEnvs
 142   = GlobalEnv {
 143       global_vars          = mapVarEnv snd $ vectInfoVar info
 144     , global_exported_vars = emptyVarEnv
 145     , global_tycons        = mapNameEnv snd $ vectInfoTyCon info
 146     , global_datacons      = mapNameEnv snd $ vectInfoDataCon info
 147     , global_pa_funs       = mapNameEnv snd $ vectInfoPADFun info
 148     , global_pr_funs       = emptyNameEnv
 149     , global_boxed_tycons  = emptyNameEnv
 150     , global_inst_env      = instEnvs
 151     , global_fam_inst_env  = famInstEnvs
 152     , global_bindings      = []
 153     }
 154
 155 extendImportedVarsEnv :: [(Var, Var)] -> GlobalEnv -> GlobalEnv
 156 extendImportedVarsEnv ps genv
 157   = genv { global_vars = extendVarEnvList (global_vars genv) ps }
 158
 159 setFamInstEnv :: FamInstEnv -> GlobalEnv -> GlobalEnv
 160 setFamInstEnv l_fam_inst genv
 161   = genv { global_fam_inst_env = (g_fam_inst, l_fam_inst) }
 162   where
 163     (g_fam_inst, _) = global_fam_inst_env genv
 164
 165 extendTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
 166 extendTyConsEnv ps genv
 167   = genv { global_tycons = extendNameEnvList (global_tycons genv) ps }
 168
 169 extendDataConsEnv :: [(Name, DataCon)] -> GlobalEnv -> GlobalEnv
 170 extendDataConsEnv ps genv
 171   = genv { global_datacons = extendNameEnvList (global_datacons genv) ps }
 172
 173 extendPAFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
 174 extendPAFunsEnv ps genv
 175   = genv { global_pa_funs = extendNameEnvList (global_pa_funs genv) ps }
 176
 177 setPRFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
 178 setPRFunsEnv ps genv
 179   = genv { global_pr_funs = mkNameEnv ps }
 180
 181 setBoxedTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
 182 setBoxedTyConsEnv ps genv
 183   = genv { global_boxed_tycons = mkNameEnv ps }
 184
 185 emptyLocalEnv = LocalEnv {
 186                    local_vars     = emptyVarEnv
 187                  , local_tyvars   = []
 188                  , local_tyvar_pa = emptyVarEnv
 189                  , local_bind_name  = FSLIT("fn")
 190                  }
 191
 192 -- FIXME
 193 updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
 194 updVectInfo env tyenv info
 195   = info {
 196       vectInfoVar     = global_exported_vars env
 197     , vectInfoTyCon   = mk_env typeEnvTyCons global_tycons
 198     , vectInfoDataCon = mk_env typeEnvDataCons global_datacons
 199     , vectInfoPADFun  = mk_env typeEnvTyCons global_pa_funs
 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 builtins :: (a -> Builtins -> b) -> VM (a -> b)
 269 builtins f = VM $ \bi genv lenv -> return (Yes genv lenv (`f` bi))
 270
 271 readGEnv :: (GlobalEnv -> a) -> VM a
 272 readGEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f genv))
 273
 274 setGEnv :: GlobalEnv -> VM ()
 275 setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
 276
 277 updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
 278 updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
 279
 280 readLEnv :: (LocalEnv -> a) -> VM a
 281 readLEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f lenv))
 282
 283 setLEnv :: LocalEnv -> VM ()
 284 setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
 285
 286 updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
 287 updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
 288
 289 getInstEnv :: VM (InstEnv, InstEnv)
 290 getInstEnv = readGEnv global_inst_env
 291
 292 getFamInstEnv :: VM FamInstEnvs
 293 getFamInstEnv = readGEnv global_fam_inst_env
 294
 295 getBindName :: VM FastString
 296 getBindName = readLEnv local_bind_name
 297
 298 inBind :: Id -> VM a -> VM a
 299 inBind id p
 300   = do updLEnv $ \env -> env { local_bind_name = occNameFS (getOccName id) }
 301        p
 302
 303 cloneName :: (OccName -> OccName) -> Name -> VM Name
 304 cloneName mk_occ name = liftM make (liftDs newUnique)
 305   where
 306     occ_name = mk_occ (nameOccName name)
 307
 308     make u | isExternalName name = mkExternalName u (nameModule name)
 309                                                     occ_name
 310                                                     (nameSrcSpan name)
 311            | otherwise           = mkSystemName u occ_name
 312
 313 cloneId :: (OccName -> OccName) -> Id -> Type -> VM Id
 314 cloneId mk_occ id ty
 315   = do
 316       name <- cloneName mk_occ (getName id)
 317       let id' | isExportedId id = Id.mkExportedLocalId name ty
 318               | otherwise       = Id.mkLocalId         name ty
 319       return id'
 320
 321 cloneVar :: Var -> VM Var
 322 cloneVar var = liftM (setIdUnique var) (liftDs newUnique)
 323
 324 newExportedVar :: OccName -> Type -> VM Var
 325 newExportedVar occ_name ty
 326   = do
 327       mod <- liftDs getModuleDs
 328       u   <- liftDs newUnique
 329
 330       let name = mkExternalName u mod occ_name noSrcSpan
 331
 332       return $ Id.mkExportedLocalId name ty
 333
 334 newLocalVar :: FastString -> Type -> VM Var
 335 newLocalVar fs ty
 336   = do
 337       u <- liftDs newUnique
 338       return $ mkSysLocal fs u ty
 339
 340 newDummyVar :: Type -> VM Var
 341 newDummyVar = newLocalVar FSLIT("ds")
 342
 343 newTyVar :: FastString -> Kind -> VM Var
 344 newTyVar fs k
 345   = do
 346       u <- liftDs newUnique
 347       return $ mkTyVar (mkSysTvName u fs) k
 348
 349 defGlobalVar :: Var -> Var -> VM ()
 350 defGlobalVar v v' = updGEnv $ \env ->
 351   env { global_vars = extendVarEnv (global_vars env) v v'
 352       , global_exported_vars = upd (global_exported_vars env)
 353       }
 354   where
 355     upd env | isExportedId v = extendVarEnv env v (v, v')
 356             | otherwise      = env
 357
 358 lookupVar :: Var -> VM (Scope Var (Var, Var))
 359 lookupVar v
 360   = do
 361       r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
 362       case r of
 363         Just e  -> return (Local e)
 364         Nothing -> liftM Global
 365                  $  traceMaybeV "lookupVar" (ppr v)
 366                                 (readGEnv $ \env -> lookupVarEnv (global_vars env) v)
 367
 368 lookupTyCon :: TyCon -> VM (Maybe TyCon)
 369 lookupTyCon tc
 370   | isUnLiftedTyCon tc || isTupleTyCon tc = return (Just tc)
 371
 372   | otherwise = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
 373
 374 defTyCon :: TyCon -> TyCon -> VM ()
 375 defTyCon tc tc' = updGEnv $ \env ->
 376   env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }
 377
 378 lookupDataCon :: DataCon -> VM (Maybe DataCon)
 379 lookupDataCon dc
 380   | isTupleTyCon (dataConTyCon dc) = return (Just dc)
 381   | otherwise = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
 382
 383 defDataCon :: DataCon -> DataCon -> VM ()
 384 defDataCon dc dc' = updGEnv $ \env ->
 385   env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }
 386
 387 lookupPrimPArray :: TyCon -> VM (Maybe TyCon)
 388 lookupPrimPArray = liftDs . primPArray
 389
 390 lookupPrimMethod :: TyCon -> String -> VM (Maybe Var)
 391 lookupPrimMethod tycon = liftDs . primMethod tycon
 392
 393 lookupTyConPA :: TyCon -> VM (Maybe Var)
 394 lookupTyConPA tc = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)
 395
 396 defTyConPA :: TyCon -> Var -> VM ()
 397 defTyConPA tc pa = updGEnv $ \env ->
 398   env { global_pa_funs = extendNameEnv (global_pa_funs env) (tyConName tc) pa }
 399
 400 defTyConPAs :: [(TyCon, Var)] -> VM ()
 401 defTyConPAs ps = updGEnv $ \env ->
 402   env { global_pa_funs = extendNameEnvList (global_pa_funs env)
 403                                            [(tyConName tc, pa) | (tc, pa) <- ps] }
 404
 405 lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
 406 lookupTyVarPA tv = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv
 407
 408 lookupTyConPR :: TyCon -> VM (Maybe Var)
 409 lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc)
 410
 411 lookupBoxedTyCon :: TyCon -> VM (Maybe TyCon)
 412 lookupBoxedTyCon tc = readGEnv $ \env -> lookupNameEnv (global_boxed_tycons env)
 413                                                        (tyConName tc)
 414
 415 defLocalTyVar :: TyVar -> VM ()
 416 defLocalTyVar tv = updLEnv $ \env ->
 417   env { local_tyvars   = tv : local_tyvars env
 418       , local_tyvar_pa = local_tyvar_pa env `delVarEnv` tv
 419       }
 420
 421 defLocalTyVarWithPA :: TyVar -> CoreExpr -> VM ()
 422 defLocalTyVarWithPA tv pa = updLEnv $ \env ->
 423   env { local_tyvars   = tv : local_tyvars env
 424       , local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa
 425       }
 426
 427 localTyVars :: VM [TyVar]
 428 localTyVars = readLEnv (reverse . local_tyvars)
 429
 430 -- Look up the dfun of a class instance.
 431 --
 432 -- The match must be unique - ie, match exactly one instance - but the
 433 -- type arguments used for matching may be more specific than those of
 434 -- the class instance declaration.  The found class instances must not have
 435 -- any type variables in the instance context that do not appear in the
 436 -- instances head (i.e., no flexi vars); for details for what this means,
 437 -- see the docs at InstEnv.lookupInstEnv.
 438 --
 439 {-
 440 lookupInst :: Class -> [Type] -> VM (DFunId, [Type])
 441 lookupInst cls tys
 442   = do { instEnv <- getInstEnv
 443        ; case lookupInstEnv instEnv cls tys of
 444            ([(inst, inst_tys)], _)
 445              | noFlexiVar -> return (instanceDFunId inst, inst_tys')
 446              | otherwise  -> pprPanic "VectMonad.lookupInst: flexi var: "
 447                                       (ppr $ mkTyConApp (classTyCon cls) tys)
 448              where
 449                inst_tys'  = [ty | Right ty <- inst_tys]
 450                noFlexiVar = all isRight inst_tys
 451            _other         -> traceNoV "lookupInst" (ppr cls <+> ppr tys)
 452        }
 453   where
 454     isRight (Left  _) = False
 455     isRight (Right _) = True
 456 -}
 457
 458 -- Look up the representation tycon of a family instance.
 459 --
 460 -- The match must be unique - ie, match exactly one instance - but the
 461 -- type arguments used for matching may be more specific than those of
 462 -- the family instance declaration.
 463 --
 464 -- Return the instance tycon and its type instance.  For example, if we have
 465 --
 466 --  lookupFamInst 'T' '[Int]' yields (':R42T', 'Int')
 467 --
 468 -- then we have a coercion (ie, type instance of family instance coercion)
 469 --
 470 --  :Co:R42T Int :: T [Int] ~ :R42T Int
 471 --
 472 -- which implies that :R42T was declared as 'data instance T [a]'.
 473 --
 474 lookupFamInst :: TyCon -> [Type] -> VM (TyCon, [Type])
 475 lookupFamInst tycon tys
 476   = ASSERT( isOpenTyCon tycon )
 477     do { instEnv <- getFamInstEnv
 478        ; case lookupFamInstEnv instEnv tycon tys of
 479            [(fam_inst, rep_tys)] -> return (famInstTyCon fam_inst, rep_tys)
 480            _other                ->
 481              pprPanic "VectMonad.lookupFamInst: not found: "
 482                       (ppr $ mkTyConApp tycon tys)
 483        }
 484
 485 initV :: HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
 486 initV hsc_env guts info p
 487   = do
 488       Just r <- initDs hsc_env (mg_module guts)
 489                                (mg_rdr_env guts)
 490                                (mg_types guts)
 491                                go
 492       return r
 493   where
 494
 495     go =
 496       do
 497         builtins       <- initBuiltins
 498         builtin_vars   <- initBuiltinVars builtins
 499         builtin_tycons <- initBuiltinTyCons builtins
 500         let builtin_datacons = initBuiltinDataCons builtins
 501         builtin_pas    <- initBuiltinPAs builtins
 502         builtin_prs    <- initBuiltinPRs builtins
 503         builtin_boxed  <- initBuiltinBoxedTyCons builtins
 504
 505         eps <- ioToIOEnv $ hscEPS hsc_env
 506         let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
 507             instEnvs    = (eps_inst_env     eps, mg_inst_env     guts)
 508
 509         let genv = extendImportedVarsEnv builtin_vars
 510                  . extendTyConsEnv builtin_tycons
 511                  . extendDataConsEnv builtin_datacons
 512                  . extendPAFunsEnv builtin_pas
 513                  . setPRFunsEnv    builtin_prs
 514                  . setBoxedTyConsEnv builtin_boxed
 515                  $ initGlobalEnv info instEnvs famInstEnvs
 516
 517         r <- runVM p builtins genv emptyLocalEnv
 518         case r of
 519           Yes genv _ x -> return $ Just (new_info genv, x)
 520           No           -> return Nothing
 521
 522     new_info genv = updVectInfo genv (mg_types guts) info
 523