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,
  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       = emptyVarEnv
 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 newExportedVar :: OccName -> Type -> VM Var
 305 newExportedVar occ_name ty
 306   = do
 307       mod <- liftDs getModuleDs
 308       u   <- liftDs newUnique
 309
 310       let name = mkExternalName u mod occ_name noSrcSpan
 311
 312       return $ Id.mkExportedLocalId name ty
 313
 314 newLocalVar :: FastString -> Type -> VM Var
 315 newLocalVar fs ty
 316   = do
 317       u <- liftDs newUnique
 318       return $ mkSysLocal fs u ty
 319
 320 newDummyVar :: Type -> VM Var
 321 newDummyVar = newLocalVar FSLIT("ds")
 322
 323 newTyVar :: FastString -> Kind -> VM Var
 324 newTyVar fs k
 325   = do
 326       u <- liftDs newUnique
 327       return $ mkTyVar (mkSysTvName u fs) k
 328
 329 defGlobalVar :: Var -> Var -> VM ()
 330 defGlobalVar v v' = updGEnv $ \env ->
 331   env { global_vars = extendVarEnv (global_vars env) v v'
 332       , global_exported_vars = upd (global_exported_vars env)
 333       }
 334   where
 335     upd env | isExportedId v = extendVarEnv env v (v, v')
 336             | otherwise      = env
 337
 338 lookupVar :: Var -> VM (Scope Var (Var, Var))
 339 lookupVar v
 340   = do
 341       r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
 342       case r of
 343         Just e  -> return (Local e)
 344         Nothing -> liftM Global
 345                  $  traceMaybeV "lookupVar" (ppr v)
 346                                 (readGEnv $ \env -> lookupVarEnv (global_vars env) v)
 347
 348 lookupTyCon :: TyCon -> VM (Maybe TyCon)
 349 lookupTyCon tc
 350   | isUnLiftedTyCon tc || isTupleTyCon tc = return (Just tc)
 351
 352   | otherwise = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
 353
 354 defTyCon :: TyCon -> TyCon -> VM ()
 355 defTyCon tc tc' = updGEnv $ \env ->
 356   env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }
 357
 358 lookupDataCon :: DataCon -> VM (Maybe DataCon)
 359 lookupDataCon dc = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
 360
 361 defDataCon :: DataCon -> DataCon -> VM ()
 362 defDataCon dc dc' = updGEnv $ \env ->
 363   env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }
 364
 365 lookupPrimPArray :: TyCon -> VM (Maybe TyCon)
 366 lookupPrimPArray = liftDs . primPArray
 367
 368 lookupPrimMethod :: TyCon -> String -> VM (Maybe Var)
 369 lookupPrimMethod tycon = liftDs . primMethod tycon
 370
 371 lookupTyConPA :: TyCon -> VM (Maybe Var)
 372 lookupTyConPA tc = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)
 373
 374 defTyConPA :: TyCon -> Var -> VM ()
 375 defTyConPA tc pa = updGEnv $ \env ->
 376   env { global_pa_funs = extendNameEnv (global_pa_funs env) (tyConName tc) pa }
 377
 378 defTyConPAs :: [(TyCon, Var)] -> VM ()
 379 defTyConPAs ps = updGEnv $ \env ->
 380   env { global_pa_funs = extendNameEnvList (global_pa_funs env)
 381                                            [(tyConName tc, pa) | (tc, pa) <- ps] }
 382
 383 lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
 384 lookupTyVarPA tv = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv
 385
 386 lookupTyConPR :: TyCon -> VM (Maybe Var)
 387 lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc)
 388
 389 defLocalTyVar :: TyVar -> VM ()
 390 defLocalTyVar tv = updLEnv $ \env ->
 391   env { local_tyvars   = tv : local_tyvars env
 392       , local_tyvar_pa = local_tyvar_pa env `delVarEnv` tv
 393       }
 394
 395 defLocalTyVarWithPA :: TyVar -> CoreExpr -> VM ()
 396 defLocalTyVarWithPA tv pa = updLEnv $ \env ->
 397   env { local_tyvars   = tv : local_tyvars env
 398       , local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa
 399       }
 400
 401 localTyVars :: VM [TyVar]
 402 localTyVars = readLEnv (reverse . local_tyvars)
 403
 404 -- Look up the dfun of a class instance.
 405 --
 406 -- The match must be unique - ie, match exactly one instance - but the
 407 -- type arguments used for matching may be more specific than those of
 408 -- the class instance declaration.  The found class instances must not have
 409 -- any type variables in the instance context that do not appear in the
 410 -- instances head (i.e., no flexi vars); for details for what this means,
 411 -- see the docs at InstEnv.lookupInstEnv.
 412 --
 413 {-
 414 lookupInst :: Class -> [Type] -> VM (DFunId, [Type])
 415 lookupInst cls tys
 416   = do { instEnv <- getInstEnv
 417        ; case lookupInstEnv instEnv cls tys of
 418            ([(inst, inst_tys)], _)
 419              | noFlexiVar -> return (instanceDFunId inst, inst_tys')
 420              | otherwise  -> pprPanic "VectMonad.lookupInst: flexi var: "
 421                                       (ppr $ mkTyConApp (classTyCon cls) tys)
 422              where
 423                inst_tys'  = [ty | Right ty <- inst_tys]
 424                noFlexiVar = all isRight inst_tys
 425            _other         -> traceNoV "lookupInst" (ppr cls <+> ppr tys)
 426        }
 427   where
 428     isRight (Left  _) = False
 429     isRight (Right _) = True
 430 -}
 431
 432 -- Look up the representation tycon of a family instance.
 433 --
 434 -- The match must be unique - ie, match exactly one instance - but the
 435 -- type arguments used for matching may be more specific than those of
 436 -- the family instance declaration.
 437 --
 438 -- Return the instance tycon and its type instance.  For example, if we have
 439 --
 440 --  lookupFamInst 'T' '[Int]' yields (':R42T', 'Int')
 441 --
 442 -- then we have a coercion (ie, type instance of family instance coercion)
 443 --
 444 --  :Co:R42T Int :: T [Int] ~ :R42T Int
 445 --
 446 -- which implies that :R42T was declared as 'data instance T [a]'.
 447 --
 448 lookupFamInst :: TyCon -> [Type] -> VM (TyCon, [Type])
 449 lookupFamInst tycon tys
 450   = ASSERT( isOpenTyCon tycon )
 451     do { instEnv <- getFamInstEnv
 452        ; case lookupFamInstEnv instEnv tycon tys of
 453            [(fam_inst, rep_tys)] -> return (famInstTyCon fam_inst, rep_tys)
 454            _other                ->
 455              pprPanic "VectMonad.lookupFamInst: not found: "
 456                       (ppr $ mkTyConApp tycon tys)
 457        }
 458
 459 initV :: HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
 460 initV hsc_env guts info p
 461   = do
 462       Just r <- initDs hsc_env (mg_module guts)
 463                                (mg_rdr_env guts)
 464                                (mg_types guts)
 465                                go
 466       return r
 467   where
 468
 469     go =
 470       do
 471         builtins       <- initBuiltins
 472         let builtin_tycons = initBuiltinTyCons builtins
 473         builtin_pas    <- initBuiltinPAs builtins
 474         builtin_prs    <- initBuiltinPRs builtins
 475
 476         eps <- ioToIOEnv $ hscEPS hsc_env
 477         let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
 478             instEnvs    = (eps_inst_env     eps, mg_inst_env     guts)
 479
 480         let genv = extendTyConsEnv builtin_tycons
 481                  . extendPAFunsEnv builtin_pas
 482                  . setPRFunsEnv    builtin_prs
 483                  $ initGlobalEnv info instEnvs famInstEnvs
 484
 485         r <- runVM p builtins genv emptyLocalEnv
 486         case r of
 487           Yes genv _ x -> return $ Just (new_info genv, x)
 488           No           -> return Nothing
 489
 490     new_info genv = updVectInfo genv (mg_types guts) info
 491