compiler/vectorise/VectMonad.hs

   1 module VectMonad (
   2   Scope(..),
   3   VM,
   4
   5   noV, tryV, maybeV, orElseV, fixV, localV, closedV, initV,
   6   cloneName, newLocalVar, newTyVar,
   7
   8   Builtins(..), paDictTyCon,
   9   builtin,
  10
  11   GlobalEnv(..),
  12   readGEnv, setGEnv, updGEnv,
  13
  14   LocalEnv(..),
  15   readLEnv, setLEnv, updLEnv,
  16
  17   defGlobalVar, lookupVar,
  18   lookupTyCon,
  19   lookupTyVarPA, defLocalTyVar, defLocalTyVarWithPA, localTyVars,
  20
  21   lookupInst, lookupFamInst
  22 ) where
  23
  24 #include "HsVersions.h"
  25
  26 import HscTypes
  27 import CoreSyn
  28 import Class
  29 import TyCon
  30 import DataCon
  31 import Type
  32 import Var
  33 import VarEnv
  34 import Id
  35 import OccName
  36 import Name
  37 import NameEnv
  38
  39 import DsMonad
  40 import PrelNames
  41
  42 import InstEnv
  43 import FamInstEnv
  44
  45 import Panic
  46 import Outputable
  47 import FastString
  48
  49 import Control.Monad ( liftM )
  50
  51 data Scope a b = Global a | Local b
  52
  53 -- ----------------------------------------------------------------------------
  54 -- Vectorisation monad
  55
  56 data Builtins = Builtins {
  57                   parrayTyCon      :: TyCon
  58                 , paClass          :: Class
  59                 , closureTyCon     :: TyCon
  60                 , mkClosureVar     :: Var
  61                 , applyClosureVar  :: Var
  62                 , mkClosurePVar    :: Var
  63                 , applyClosurePVar :: Var
  64                 , lengthPAVar      :: Var
  65                 , replicatePAVar   :: Var
  66                 }
  67
  68 paDictTyCon :: Builtins -> TyCon
  69 paDictTyCon = classTyCon . paClass
  70
  71 initBuiltins :: DsM Builtins
  72 initBuiltins
  73   = do
  74       parrayTyCon  <- dsLookupTyCon parrayTyConName
  75       paClass      <- dsLookupClass paClassName
  76       closureTyCon <- dsLookupTyCon closureTyConName
  77
  78       mkClosureVar     <- dsLookupGlobalId mkClosureName
  79       applyClosureVar  <- dsLookupGlobalId applyClosureName
  80       mkClosurePVar    <- dsLookupGlobalId mkClosurePName
  81       applyClosurePVar <- dsLookupGlobalId applyClosurePName
  82       lengthPAVar      <- dsLookupGlobalId lengthPAName
  83       replicatePAVar   <- dsLookupGlobalId replicatePAName
  84
  85       return $ Builtins {
  86                  parrayTyCon      = parrayTyCon
  87                , paClass          = paClass
  88                , closureTyCon     = closureTyCon
  89                , mkClosureVar     = mkClosureVar
  90                , applyClosureVar  = applyClosureVar
  91                , mkClosurePVar    = mkClosurePVar
  92                , applyClosurePVar = applyClosurePVar
  93                , lengthPAVar      = lengthPAVar
  94                , replicatePAVar   = replicatePAVar
  95                }
  96
  97 data GlobalEnv = GlobalEnv {
  98                   -- Mapping from global variables to their vectorised versions.
  99                   --
 100                   global_vars :: VarEnv CoreExpr
 101
 102                   -- Exported variables which have a vectorised version
 103                   --
 104                 , global_exported_vars :: VarEnv (Var, Var)
 105
 106                   -- Mapping from TyCons to their vectorised versions.
 107                   -- TyCons which do not have to be vectorised are mapped to
 108                   -- themselves.
 109                   --
 110                 , global_tycons :: NameEnv TyCon
 111
 112                   -- Mapping from TyCons to their PA dictionaries
 113                   --
 114                 , global_tycon_pa :: NameEnv CoreExpr
 115
 116                   -- Mapping from DataCons to their vectorised versions
 117                   --
 118                 , global_datacons :: NameEnv DataCon
 119
 120                 -- External package inst-env & home-package inst-env for class
 121                 -- instances
 122                 --
 123                 , global_inst_env :: (InstEnv, InstEnv)
 124
 125                 -- External package inst-env & home-package inst-env for family
 126                 -- instances
 127                 --
 128                 , global_fam_inst_env :: FamInstEnvs
 129
 130                 -- Hoisted bindings
 131                 , global_bindings :: [(Var, CoreExpr)]
 132                 }
 133
 134 data LocalEnv = LocalEnv {
 135                  -- Mapping from local variables to their vectorised and
 136                  -- lifted versions
 137                  --
 138                  local_vars :: VarEnv (CoreExpr, CoreExpr)
 139
 140                  -- In-scope type variables
 141                  --
 142                , local_tyvars :: [TyVar]
 143
 144                  -- Mapping from tyvars to their PA dictionaries
 145                , local_tyvar_pa :: VarEnv CoreExpr
 146                }
 147
 148
 149 initGlobalEnv :: VectInfo -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
 150 initGlobalEnv info instEnvs famInstEnvs
 151   = GlobalEnv {
 152       global_vars          = mapVarEnv  (Var . snd) $ vectInfoVar   info
 153     , global_exported_vars = emptyVarEnv
 154     , global_tycons        = mapNameEnv snd $ vectInfoTyCon info
 155     , global_tycon_pa      = emptyNameEnv
 156     , global_datacons      = mapNameEnv snd $ vectInfoDataCon info
 157     , global_inst_env      = instEnvs
 158     , global_fam_inst_env  = famInstEnvs
 159     , global_bindings      = []
 160     }
 161
 162 emptyLocalEnv = LocalEnv {
 163                    local_vars     = emptyVarEnv
 164                  , local_tyvars   = []
 165                  , local_tyvar_pa = emptyVarEnv
 166                  }
 167
 168 -- FIXME
 169 updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
 170 updVectInfo env tyenv info
 171   = info {
 172       vectInfoVar     = global_exported_vars env
 173     , vectInfoTyCon   = tc_env
 174     , vectInfoDataCon = dc_env
 175     }
 176   where
 177     tc_env = mkNameEnv [(tc_name, (tc,tc'))
 178                | tc <- typeEnvTyCons tyenv
 179                , let tc_name = tyConName tc
 180                , Just tc' <- [lookupNameEnv (global_tycons env) tc_name]]
 181
 182     dc_env = mkNameEnv [(dc_name, (dc,dc'))
 183                | dc <- typeEnvDataCons tyenv
 184                , let dc_name = dataConName dc
 185                , Just dc' <- [lookupNameEnv (global_datacons env) dc_name]]
 186
 187 data VResult a = Yes GlobalEnv LocalEnv a | No
 188
 189 newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
 190
 191 instance Monad VM where
 192   return x   = VM $ \bi genv lenv -> return (Yes genv lenv x)
 193   VM p >>= f = VM $ \bi genv lenv -> do
 194                                       r <- p bi genv lenv
 195                                       case r of
 196                                         Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
 197                                         No                -> return No
 198
 199 noV :: VM a
 200 noV = VM $ \_ _ _ -> return No
 201
 202 tryV :: VM a -> VM (Maybe a)
 203 tryV (VM p) = VM $ \bi genv lenv ->
 204   do
 205     r <- p bi genv lenv
 206     case r of
 207       Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
 208       No                -> return (Yes genv  lenv  Nothing)
 209
 210 maybeV :: VM (Maybe a) -> VM a
 211 maybeV p = maybe noV return =<< p
 212
 213 orElseV :: VM a -> VM a -> VM a
 214 orElseV p q = maybe q return =<< tryV p
 215
 216 fixV :: (a -> VM a) -> VM a
 217 fixV f = VM (\bi genv lenv -> fixDs $ \r -> runVM (f (unYes r)) bi genv lenv )
 218   where
 219     unYes (Yes _ _ x) = x
 220
 221 localV :: VM a -> VM a
 222 localV p = do
 223              env <- readLEnv id
 224              x <- p
 225              setLEnv env
 226              return x
 227
 228 closedV :: VM a -> VM a
 229 closedV p = do
 230               env <- readLEnv id
 231               setLEnv emptyLocalEnv
 232               x <- p
 233               setLEnv env
 234               return x
 235
 236 liftDs :: DsM a -> VM a
 237 liftDs p = VM $ \bi genv lenv -> do { x <- p; return (Yes genv lenv x) }
 238
 239 builtin :: (Builtins -> a) -> VM a
 240 builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
 241
 242 readGEnv :: (GlobalEnv -> a) -> VM a
 243 readGEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f genv))
 244
 245 setGEnv :: GlobalEnv -> VM ()
 246 setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
 247
 248 updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
 249 updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
 250
 251 readLEnv :: (LocalEnv -> a) -> VM a
 252 readLEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f lenv))
 253
 254 setLEnv :: LocalEnv -> VM ()
 255 setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
 256
 257 updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
 258 updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
 259
 260 getInstEnv :: VM (InstEnv, InstEnv)
 261 getInstEnv = readGEnv global_inst_env
 262
 263 getFamInstEnv :: VM FamInstEnvs
 264 getFamInstEnv = readGEnv global_fam_inst_env
 265
 266 cloneName :: (OccName -> OccName) -> Name -> VM Name
 267 cloneName mk_occ name = liftM make (liftDs newUnique)
 268   where
 269     occ_name = mk_occ (nameOccName name)
 270
 271     make u | isExternalName name = mkExternalName u (nameModule name)
 272                                                     occ_name
 273                                                     (nameSrcSpan name)
 274            | otherwise           = mkSystemName u occ_name
 275
 276 newLocalVar :: FastString -> Type -> VM Var
 277 newLocalVar fs ty
 278   = do
 279       u <- liftDs newUnique
 280       return $ mkSysLocal fs u ty
 281
 282 newTyVar :: FastString -> Kind -> VM Var
 283 newTyVar fs k
 284   = do
 285       u <- liftDs newUnique
 286       return $ mkTyVar (mkSysTvName u fs) k
 287
 288 defGlobalVar :: Var -> Var -> VM ()
 289 defGlobalVar v v' = updGEnv $ \env ->
 290   env { global_vars = extendVarEnv (global_vars env) v (Var v')
 291       , global_exported_vars = upd (global_exported_vars env)
 292       }
 293   where
 294     upd env | isExportedId v = extendVarEnv env v (v, v')
 295             | otherwise      = env
 296
 297 lookupVar :: Var -> VM (Scope CoreExpr (CoreExpr, CoreExpr))
 298 lookupVar v
 299   = do
 300       r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
 301       case r of
 302         Just e  -> return (Local e)
 303         Nothing -> liftM Global
 304                  $  maybeV (readGEnv $ \env -> lookupVarEnv (global_vars env) v)
 305
 306 lookupTyCon :: TyCon -> VM (Maybe TyCon)
 307 lookupTyCon tc = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
 308
 309 lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
 310 lookupTyVarPA tv = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv
 311
 312 defLocalTyVar :: TyVar -> VM ()
 313 defLocalTyVar tv = updLEnv $ \env ->
 314   env { local_tyvars   = tv : local_tyvars env
 315       , local_tyvar_pa = local_tyvar_pa env `delVarEnv` tv
 316       }
 317
 318 defLocalTyVarWithPA :: TyVar -> CoreExpr -> VM ()
 319 defLocalTyVarWithPA tv pa = updLEnv $ \env ->
 320   env { local_tyvars   = tv : local_tyvars env
 321       , local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa
 322       }
 323
 324 localTyVars :: VM [TyVar]
 325 localTyVars = readLEnv (reverse . local_tyvars)
 326
 327 -- Look up the dfun of a class instance.
 328 --
 329 -- The match must be unique - ie, match exactly one instance - but the
 330 -- type arguments used for matching may be more specific than those of
 331 -- the class instance declaration.  The found class instances must not have
 332 -- any type variables in the instance context that do not appear in the
 333 -- instances head (i.e., no flexi vars); for details for what this means,
 334 -- see the docs at InstEnv.lookupInstEnv.
 335 --
 336 lookupInst :: Class -> [Type] -> VM (DFunId, [Type])
 337 lookupInst cls tys
 338   = do { instEnv <- getInstEnv
 339        ; case lookupInstEnv instEnv cls tys of
 340            ([(inst, inst_tys)], _)
 341              | noFlexiVar -> return (instanceDFunId inst, inst_tys')
 342              | otherwise  -> pprPanic "VectMonad.lookupInst: flexi var: "
 343                                       (ppr $ mkTyConApp (classTyCon cls) tys)
 344              where
 345                inst_tys'  = [ty | Right ty <- inst_tys]
 346                noFlexiVar = all isRight inst_tys
 347            _other         -> noV
 348        }
 349   where
 350     isRight (Left  _) = False
 351     isRight (Right _) = True
 352
 353 -- Look up the representation tycon of a family instance.
 354 --
 355 -- The match must be unique - ie, match exactly one instance - but the
 356 -- type arguments used for matching may be more specific than those of
 357 -- the family instance declaration.
 358 --
 359 -- Return the instance tycon and its type instance.  For example, if we have
 360 --
 361 --  lookupFamInst 'T' '[Int]' yields (':R42T', 'Int')
 362 --
 363 -- then we have a coercion (ie, type instance of family instance coercion)
 364 --
 365 --  :Co:R42T Int :: T [Int] ~ :R42T Int
 366 --
 367 -- which implies that :R42T was declared as 'data instance T [a]'.
 368 --
 369 lookupFamInst :: TyCon -> [Type] -> VM (TyCon, [Type])
 370 lookupFamInst tycon tys
 371   = ASSERT( isOpenTyCon tycon )
 372     do { instEnv <- getFamInstEnv
 373        ; case lookupFamInstEnv instEnv tycon tys of
 374            [(fam_inst, rep_tys)] -> return (famInstTyCon fam_inst, rep_tys)
 375            _other                ->
 376              pprPanic "VectMonad.lookupFamInst: not found: "
 377                       (ppr $ mkTyConApp tycon tys)
 378        }
 379
 380 initV :: HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
 381 initV hsc_env guts info p
 382   = do
 383       eps <- hscEPS hsc_env
 384       let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
 385       let instEnvs    = (eps_inst_env     eps, mg_inst_env     guts)
 386
 387       Just r <- initDs hsc_env (mg_module guts)
 388                                (mg_rdr_env guts)
 389                                (mg_types guts)
 390                                (go instEnvs famInstEnvs)
 391       return r
 392   where
 393
 394     go instEnvs famInstEnvs =
 395       do
 396         builtins <- initBuiltins
 397         r <- runVM p builtins (initGlobalEnv info instEnvs famInstEnvs)
 398                    emptyLocalEnv
 399         case r of
 400           Yes genv _ x -> return $ Just (new_info genv, x)
 401           No           -> return Nothing
 402
 403     new_info genv = updVectInfo genv (mg_types guts) info
 404