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 setFamInstEnv :: FamInstEnv -> GlobalEnv -> GlobalEnv
 156 setFamInstEnv l_fam_inst genv
 157   = genv { global_fam_inst_env = (g_fam_inst, l_fam_inst) }
 158   where
 159     (g_fam_inst, _) = global_fam_inst_env genv
 160
 161 extendTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
 162 extendTyConsEnv ps genv
 163   = genv { global_tycons = extendNameEnvList (global_tycons genv) ps }
 164
 165 extendPAFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
 166 extendPAFunsEnv ps genv
 167   = genv { global_pa_funs = extendNameEnvList (global_pa_funs genv) ps }
 168
 169 setPRFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
 170 setPRFunsEnv ps genv
 171   = genv { global_pr_funs = mkNameEnv ps }
 172
 173 setBoxedTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
 174 setBoxedTyConsEnv ps genv
 175   = genv { global_boxed_tycons = mkNameEnv ps }
 176
 177 emptyLocalEnv = LocalEnv {
 178                    local_vars     = emptyVarEnv
 179                  , local_tyvars   = []
 180                  , local_tyvar_pa = emptyVarEnv
 181                  , local_bind_name  = FSLIT("fn")
 182                  }
 183
 184 -- FIXME
 185 updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
 186 updVectInfo env tyenv info
 187   = info {
 188       vectInfoVar     = global_exported_vars env
 189     , vectInfoTyCon   = mk_env typeEnvTyCons global_tycons
 190     , vectInfoDataCon = mk_env typeEnvDataCons global_datacons
 191     , vectInfoPADFun  = mk_env typeEnvTyCons global_pa_funs
 192     }
 193   where
 194     mk_env from_tyenv from_env = mkNameEnv [(name, (from,to))
 195                                    | from <- from_tyenv tyenv
 196                                    , let name = getName from
 197                                    , Just to <- [lookupNameEnv (from_env env) name]]
 198
 199 data VResult a = Yes GlobalEnv LocalEnv a | No
 200
 201 newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
 202
 203 instance Monad VM where
 204   return x   = VM $ \bi genv lenv -> return (Yes genv lenv x)
 205   VM p >>= f = VM $ \bi genv lenv -> do
 206                                       r <- p bi genv lenv
 207                                       case r of
 208                                         Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
 209                                         No                -> return No
 210
 211 noV :: VM a
 212 noV = VM $ \_ _ _ -> return No
 213
 214 traceNoV :: String -> SDoc -> VM a
 215 traceNoV s d = pprTrace s d noV
 216
 217 tryV :: VM a -> VM (Maybe a)
 218 tryV (VM p) = VM $ \bi genv lenv ->
 219   do
 220     r <- p bi genv lenv
 221     case r of
 222       Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
 223       No                -> return (Yes genv  lenv  Nothing)
 224
 225 maybeV :: VM (Maybe a) -> VM a
 226 maybeV p = maybe noV return =<< p
 227
 228 traceMaybeV :: String -> SDoc -> VM (Maybe a) -> VM a
 229 traceMaybeV s d p = maybe (traceNoV s d) return =<< p
 230
 231 orElseV :: VM a -> VM a -> VM a
 232 orElseV p q = maybe q return =<< tryV p
 233
 234 fixV :: (a -> VM a) -> VM a
 235 fixV f = VM (\bi genv lenv -> fixDs $ \r -> runVM (f (unYes r)) bi genv lenv )
 236   where
 237     unYes (Yes _ _ x) = x
 238
 239 localV :: VM a -> VM a
 240 localV p = do
 241              env <- readLEnv id
 242              x <- p
 243              setLEnv env
 244              return x
 245
 246 closedV :: VM a -> VM a
 247 closedV p = do
 248               env <- readLEnv id
 249               setLEnv (emptyLocalEnv { local_bind_name = local_bind_name env })
 250               x <- p
 251               setLEnv env
 252               return x
 253
 254 liftDs :: DsM a -> VM a
 255 liftDs p = VM $ \bi genv lenv -> do { x <- p; return (Yes genv lenv x) }
 256
 257 builtin :: (Builtins -> a) -> VM a
 258 builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
 259
 260 builtins :: (a -> Builtins -> b) -> VM (a -> b)
 261 builtins f = VM $ \bi genv lenv -> return (Yes genv lenv (`f` bi))
 262
 263 readGEnv :: (GlobalEnv -> a) -> VM a
 264 readGEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f genv))
 265
 266 setGEnv :: GlobalEnv -> VM ()
 267 setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
 268
 269 updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
 270 updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
 271
 272 readLEnv :: (LocalEnv -> a) -> VM a
 273 readLEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f lenv))
 274
 275 setLEnv :: LocalEnv -> VM ()
 276 setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
 277
 278 updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
 279 updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
 280
 281 getInstEnv :: VM (InstEnv, InstEnv)
 282 getInstEnv = readGEnv global_inst_env
 283
 284 getFamInstEnv :: VM FamInstEnvs
 285 getFamInstEnv = readGEnv global_fam_inst_env
 286
 287 getBindName :: VM FastString
 288 getBindName = readLEnv local_bind_name
 289
 290 inBind :: Id -> VM a -> VM a
 291 inBind id p
 292   = do updLEnv $ \env -> env { local_bind_name = occNameFS (getOccName id) }
 293        p
 294
 295 cloneName :: (OccName -> OccName) -> Name -> VM Name
 296 cloneName mk_occ name = liftM make (liftDs newUnique)
 297   where
 298     occ_name = mk_occ (nameOccName name)
 299
 300     make u | isExternalName name = mkExternalName u (nameModule name)
 301                                                     occ_name
 302                                                     (nameSrcSpan name)
 303            | otherwise           = mkSystemName u occ_name
 304
 305 cloneId :: (OccName -> OccName) -> Id -> Type -> VM Id
 306 cloneId mk_occ id ty
 307   = do
 308       name <- cloneName mk_occ (getName id)
 309       let id' | isExportedId id = Id.mkExportedLocalId name ty
 310               | otherwise       = Id.mkLocalId         name ty
 311       return id'
 312
 313 cloneVar :: Var -> VM Var
 314 cloneVar var = liftM (setIdUnique var) (liftDs newUnique)
 315
 316 newExportedVar :: OccName -> Type -> VM Var
 317 newExportedVar occ_name ty
 318   = do
 319       mod <- liftDs getModuleDs
 320       u   <- liftDs newUnique
 321
 322       let name = mkExternalName u mod occ_name noSrcSpan
 323
 324       return $ Id.mkExportedLocalId name ty
 325
 326 newLocalVar :: FastString -> Type -> VM Var
 327 newLocalVar fs ty
 328   = do
 329       u <- liftDs newUnique
 330       return $ mkSysLocal fs u ty
 331
 332 newDummyVar :: Type -> VM Var
 333 newDummyVar = newLocalVar FSLIT("ds")
 334
 335 newTyVar :: FastString -> Kind -> VM Var
 336 newTyVar fs k
 337   = do
 338       u <- liftDs newUnique
 339       return $ mkTyVar (mkSysTvName u fs) k
 340
 341 defGlobalVar :: Var -> Var -> VM ()
 342 defGlobalVar v v' = updGEnv $ \env ->
 343   env { global_vars = extendVarEnv (global_vars env) v v'
 344       , global_exported_vars = upd (global_exported_vars env)
 345       }
 346   where
 347     upd env | isExportedId v = extendVarEnv env v (v, v')
 348             | otherwise      = env
 349
 350 lookupVar :: Var -> VM (Scope Var (Var, Var))
 351 lookupVar v
 352   = do
 353       r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
 354       case r of
 355         Just e  -> return (Local e)
 356         Nothing -> liftM Global
 357                  $  traceMaybeV "lookupVar" (ppr v)
 358                                 (readGEnv $ \env -> lookupVarEnv (global_vars env) v)
 359
 360 lookupTyCon :: TyCon -> VM (Maybe TyCon)
 361 lookupTyCon tc
 362   | isUnLiftedTyCon tc || isTupleTyCon tc = return (Just tc)
 363
 364   | otherwise = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
 365
 366 defTyCon :: TyCon -> TyCon -> VM ()
 367 defTyCon tc tc' = updGEnv $ \env ->
 368   env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }
 369
 370 lookupDataCon :: DataCon -> VM (Maybe DataCon)
 371 lookupDataCon dc = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
 372
 373 defDataCon :: DataCon -> DataCon -> VM ()
 374 defDataCon dc dc' = updGEnv $ \env ->
 375   env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }
 376
 377 lookupPrimPArray :: TyCon -> VM (Maybe TyCon)
 378 lookupPrimPArray = liftDs . primPArray
 379
 380 lookupPrimMethod :: TyCon -> String -> VM (Maybe Var)
 381 lookupPrimMethod tycon = liftDs . primMethod tycon
 382
 383 lookupTyConPA :: TyCon -> VM (Maybe Var)
 384 lookupTyConPA tc = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)
 385
 386 defTyConPA :: TyCon -> Var -> VM ()
 387 defTyConPA tc pa = updGEnv $ \env ->
 388   env { global_pa_funs = extendNameEnv (global_pa_funs env) (tyConName tc) pa }
 389
 390 defTyConPAs :: [(TyCon, Var)] -> VM ()
 391 defTyConPAs ps = updGEnv $ \env ->
 392   env { global_pa_funs = extendNameEnvList (global_pa_funs env)
 393                                            [(tyConName tc, pa) | (tc, pa) <- ps] }
 394
 395 lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
 396 lookupTyVarPA tv = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv
 397
 398 lookupTyConPR :: TyCon -> VM (Maybe Var)
 399 lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc)
 400
 401 lookupBoxedTyCon :: TyCon -> VM (Maybe TyCon)
 402 lookupBoxedTyCon tc = readGEnv $ \env -> lookupNameEnv (global_boxed_tycons env)
 403                                                        (tyConName tc)
 404
 405 defLocalTyVar :: TyVar -> VM ()
 406 defLocalTyVar tv = updLEnv $ \env ->
 407   env { local_tyvars   = tv : local_tyvars env
 408       , local_tyvar_pa = local_tyvar_pa env `delVarEnv` tv
 409       }
 410
 411 defLocalTyVarWithPA :: TyVar -> CoreExpr -> VM ()
 412 defLocalTyVarWithPA tv pa = updLEnv $ \env ->
 413   env { local_tyvars   = tv : local_tyvars env
 414       , local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa
 415       }
 416
 417 localTyVars :: VM [TyVar]
 418 localTyVars = readLEnv (reverse . local_tyvars)
 419
 420 -- Look up the dfun of a class instance.
 421 --
 422 -- The match must be unique - ie, match exactly one instance - but the
 423 -- type arguments used for matching may be more specific than those of
 424 -- the class instance declaration.  The found class instances must not have
 425 -- any type variables in the instance context that do not appear in the
 426 -- instances head (i.e., no flexi vars); for details for what this means,
 427 -- see the docs at InstEnv.lookupInstEnv.
 428 --
 429 {-
 430 lookupInst :: Class -> [Type] -> VM (DFunId, [Type])
 431 lookupInst cls tys
 432   = do { instEnv <- getInstEnv
 433        ; case lookupInstEnv instEnv cls tys of
 434            ([(inst, inst_tys)], _)
 435              | noFlexiVar -> return (instanceDFunId inst, inst_tys')
 436              | otherwise  -> pprPanic "VectMonad.lookupInst: flexi var: "
 437                                       (ppr $ mkTyConApp (classTyCon cls) tys)
 438              where
 439                inst_tys'  = [ty | Right ty <- inst_tys]
 440                noFlexiVar = all isRight inst_tys
 441            _other         -> traceNoV "lookupInst" (ppr cls <+> ppr tys)
 442        }
 443   where
 444     isRight (Left  _) = False
 445     isRight (Right _) = True
 446 -}
 447
 448 -- Look up the representation tycon of a family instance.
 449 --
 450 -- The match must be unique - ie, match exactly one instance - but the
 451 -- type arguments used for matching may be more specific than those of
 452 -- the family instance declaration.
 453 --
 454 -- Return the instance tycon and its type instance.  For example, if we have
 455 --
 456 --  lookupFamInst 'T' '[Int]' yields (':R42T', 'Int')
 457 --
 458 -- then we have a coercion (ie, type instance of family instance coercion)
 459 --
 460 --  :Co:R42T Int :: T [Int] ~ :R42T Int
 461 --
 462 -- which implies that :R42T was declared as 'data instance T [a]'.
 463 --
 464 lookupFamInst :: TyCon -> [Type] -> VM (TyCon, [Type])
 465 lookupFamInst tycon tys
 466   = ASSERT( isOpenTyCon tycon )
 467     do { instEnv <- getFamInstEnv
 468        ; case lookupFamInstEnv instEnv tycon tys of
 469            [(fam_inst, rep_tys)] -> return (famInstTyCon fam_inst, rep_tys)
 470            _other                ->
 471              pprPanic "VectMonad.lookupFamInst: not found: "
 472                       (ppr $ mkTyConApp tycon tys)
 473        }
 474
 475 initV :: HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
 476 initV hsc_env guts info p
 477   = do
 478       Just r <- initDs hsc_env (mg_module guts)
 479                                (mg_rdr_env guts)
 480                                (mg_types guts)
 481                                go
 482       return r
 483   where
 484
 485     go =
 486       do
 487         builtins       <- initBuiltins
 488         let builtin_tycons = initBuiltinTyCons builtins
 489         builtin_pas    <- initBuiltinPAs builtins
 490         builtin_prs    <- initBuiltinPRs builtins
 491         builtin_boxed  <- initBuiltinBoxedTyCons builtins
 492
 493         eps <- ioToIOEnv $ hscEPS hsc_env
 494         let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
 495             instEnvs    = (eps_inst_env     eps, mg_inst_env     guts)
 496
 497         let genv = extendTyConsEnv builtin_tycons
 498                  . extendPAFunsEnv builtin_pas
 499                  . setPRFunsEnv    builtin_prs
 500                  . setBoxedTyConsEnv builtin_boxed
 501                  $ initGlobalEnv info instEnvs famInstEnvs
 502
 503         r <- runVM p builtins genv emptyLocalEnv
 504         case r of
 505           Yes genv _ x -> return $ Just (new_info genv, x)
 506           No           -> return Nothing
 507
 508     new_info genv = updVectInfo genv (mg_types guts) info
 509