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