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