VectMonad.hs 16.8 KB
Newer Older
1
module VectMonad (
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
2
  Scope(..),
3 4
  VM,

5
  noV, tryV, maybeV, traceMaybeV, orElseV, fixV, localV, closedV, initV,
6
  liftDs,
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
7
  cloneName, cloneId, cloneVar,
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
8
  newExportedVar, newLocalVar, newDummyVar, newTyVar,
9
  
10
  Builtins(..), sumTyCon, prodTyCon,
11
  combinePAVar,
12
  builtin, builtins,
13 14

  GlobalEnv(..),
15
  setFamInstEnv,
16 17 18 19 20
  readGEnv, setGEnv, updGEnv,

  LocalEnv(..),
  readLEnv, setLEnv, updLEnv,

21 22
  getBindName, inBind,

23 24 25
  lookupVar, defGlobalVar,
  lookupTyCon, defTyCon,
  lookupDataCon, defDataCon,
26
  lookupTyConPA, defTyConPA, defTyConPAs,
27
  lookupTyConPR,
28
  lookupBoxedTyCon,
29
  lookupPrimMethod, lookupPrimPArray,
30
  lookupTyVarPA, defLocalTyVar, defLocalTyVarWithPA, localTyVars,
31

32
  {-lookupInst,-} lookupFamInst
33 34 35 36
) where

#include "HsVersions.h"

37 38
import VectBuiltIn

39
import HscTypes
40
import Module           ( PackageId )
41 42
import CoreSyn
import TyCon
43
import DataCon
44 45 46 47 48 49
import Type
import Var
import VarEnv
import Id
import Name
import NameEnv
50
import IOEnv         ( liftIO )
51 52 53

import DsMonad

54 55 56 57
import InstEnv
import FamInstEnv

import Outputable
58
import FastString
59
import SrcLoc        ( noSrcSpan )
60

twanvl's avatar
twanvl committed
61
import Control.Monad
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
62 63 64

data Scope a b = Global a | Local b

65 66 67 68 69 70
-- ----------------------------------------------------------------------------
-- Vectorisation monad

data GlobalEnv = GlobalEnv {
                  -- Mapping from global variables to their vectorised versions.
                  -- 
71
                  global_vars :: VarEnv Var
72 73 74 75 76 77 78 79 80 81 82

                  -- Exported variables which have a vectorised version
                  --
                , global_exported_vars :: VarEnv (Var, Var)

                  -- Mapping from TyCons to their vectorised versions.
                  -- TyCons which do not have to be vectorised are mapped to
                  -- themselves.
                  --
                , global_tycons :: NameEnv TyCon

83 84 85 86
                  -- Mapping from DataCons to their vectorised versions
                  --
                , global_datacons :: NameEnv DataCon

87 88 89 90
                  -- Mapping from TyCons to their PA dfuns
                  --
                , global_pa_funs :: NameEnv Var

91 92 93
                  -- Mapping from TyCons to their PR dfuns
                , global_pr_funs :: NameEnv Var

94 95 96
                  -- Mapping from unboxed TyCons to their boxed versions
                , global_boxed_tycons :: NameEnv TyCon

97 98 99 100 101 102 103 104 105
                -- External package inst-env & home-package inst-env for class
                -- instances
                --
                , global_inst_env :: (InstEnv, InstEnv)

                -- External package inst-env & home-package inst-env for family
                -- instances
                --
                , global_fam_inst_env :: FamInstEnvs
106 107 108

                -- Hoisted bindings
                , global_bindings :: [(Var, CoreExpr)]
109 110 111 112 113 114
                }

data LocalEnv = LocalEnv {
                 -- Mapping from local variables to their vectorised and
                 -- lifted versions
                 --
115
                 local_vars :: VarEnv (Var, Var)
116

117 118 119 120
                 -- In-scope type variables
                 --
               , local_tyvars :: [TyVar]

121 122
                 -- Mapping from tyvars to their PA dictionaries
               , local_tyvar_pa :: VarEnv CoreExpr
123 124 125

                 -- Local binding name
               , local_bind_name :: FastString
126 127
               }

128 129
initGlobalEnv :: VectInfo -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
initGlobalEnv info instEnvs famInstEnvs
130
  = GlobalEnv {
131
      global_vars          = mapVarEnv snd $ vectInfoVar info
132
    , global_exported_vars = emptyVarEnv
133
    , global_tycons        = mapNameEnv snd $ vectInfoTyCon info
134
    , global_datacons      = mapNameEnv snd $ vectInfoDataCon info
135
    , global_pa_funs       = mapNameEnv snd $ vectInfoPADFun info
136
    , global_pr_funs       = emptyNameEnv
137
    , global_boxed_tycons  = emptyNameEnv
138 139
    , global_inst_env      = instEnvs
    , global_fam_inst_env  = famInstEnvs
140
    , global_bindings      = []
141 142
    }

143 144 145 146
extendImportedVarsEnv :: [(Var, Var)] -> GlobalEnv -> GlobalEnv
extendImportedVarsEnv ps genv
  = genv { global_vars = extendVarEnvList (global_vars genv) ps }

147 148 149 150 151
setFamInstEnv :: FamInstEnv -> GlobalEnv -> GlobalEnv
setFamInstEnv l_fam_inst genv
  = genv { global_fam_inst_env = (g_fam_inst, l_fam_inst) }
  where
    (g_fam_inst, _) = global_fam_inst_env genv
152

153 154 155 156
extendTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
extendTyConsEnv ps genv
  = genv { global_tycons = extendNameEnvList (global_tycons genv) ps }

157 158 159 160
extendDataConsEnv :: [(Name, DataCon)] -> GlobalEnv -> GlobalEnv
extendDataConsEnv ps genv
  = genv { global_datacons = extendNameEnvList (global_datacons genv) ps }

161 162 163 164
extendPAFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
extendPAFunsEnv ps genv
  = genv { global_pa_funs = extendNameEnvList (global_pa_funs genv) ps }

165 166 167 168
setPRFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
setPRFunsEnv ps genv
  = genv { global_pr_funs = mkNameEnv ps }

169 170 171 172
setBoxedTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
setBoxedTyConsEnv ps genv
  = genv { global_boxed_tycons = mkNameEnv ps }

twanvl's avatar
twanvl committed
173
emptyLocalEnv :: LocalEnv
174 175
emptyLocalEnv = LocalEnv {
                   local_vars     = emptyVarEnv
176
                 , local_tyvars   = []
177
                 , local_tyvar_pa = emptyVarEnv
Ian Lynagh's avatar
Ian Lynagh committed
178
                 , local_bind_name  = fsLit "fn"
179 180 181 182 183 184
                 }

-- FIXME
updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
updVectInfo env tyenv info
  = info {
185
      vectInfoVar     = global_exported_vars env
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
186 187
    , vectInfoTyCon   = mk_env typeEnvTyCons global_tycons
    , vectInfoDataCon = mk_env typeEnvDataCons global_datacons
188
    , vectInfoPADFun  = mk_env typeEnvTyCons global_pa_funs
189 190
    }
  where
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
191 192 193 194
    mk_env from_tyenv from_env = mkNameEnv [(name, (from,to))
                                   | from <- from_tyenv tyenv
                                   , let name = getName from
                                   , Just to <- [lookupNameEnv (from_env env) name]]
195

196 197 198 199 200
data VResult a = Yes GlobalEnv LocalEnv a | No

newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }

instance Monad VM where
twanvl's avatar
twanvl committed
201
  return x   = VM $ \_  genv lenv -> return (Yes genv lenv x)
202 203 204 205 206 207 208 209 210
  VM p >>= f = VM $ \bi genv lenv -> do
                                      r <- p bi genv lenv
                                      case r of
                                        Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
                                        No                -> return No

noV :: VM a
noV = VM $ \_ _ _ -> return No

211 212 213
traceNoV :: String -> SDoc -> VM a
traceNoV s d = pprTrace s d noV

214 215 216 217 218 219 220 221 222 223 224
tryV :: VM a -> VM (Maybe a)
tryV (VM p) = VM $ \bi genv lenv ->
  do
    r <- p bi genv lenv
    case r of
      Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
      No                -> return (Yes genv  lenv  Nothing)

maybeV :: VM (Maybe a) -> VM a
maybeV p = maybe noV return =<< p

225 226 227
traceMaybeV :: String -> SDoc -> VM (Maybe a) -> VM a
traceMaybeV s d p = maybe (traceNoV s d) return =<< p

228 229 230
orElseV :: VM a -> VM a -> VM a
orElseV p q = maybe q return =<< tryV p

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
231
fixV :: (a -> VM a) -> VM a
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
232 233 234 235 236 237
fixV f = VM (\bi genv lenv -> fixDs $ \r -> runVM (f (unYes r)) bi genv lenv )
  where
    -- NOTE: It is essential that we are lazy in r above so do not replace
    --       calls to this function by an explicit case.
    unYes (Yes _ _ x) = x
    unYes No          = panic "VectMonad.fixV: no result"
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
238

239 240 241 242 243 244 245
localV :: VM a -> VM a
localV p = do
             env <- readLEnv id
             x <- p
             setLEnv env
             return x

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
246 247 248
closedV :: VM a -> VM a
closedV p = do
              env <- readLEnv id
249
              setLEnv (emptyLocalEnv { local_bind_name = local_bind_name env })
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
250 251 252 253
              x <- p
              setLEnv env
              return x

254
liftDs :: DsM a -> VM a
twanvl's avatar
twanvl committed
255
liftDs p = VM $ \_ genv lenv -> do { x <- p; return (Yes genv lenv x) }
256

257 258 259
liftBuiltinDs :: (Builtins -> DsM a) -> VM a
liftBuiltinDs p = VM $ \bi genv lenv -> do { x <- p bi; return (Yes genv lenv x)}

260 261 262
builtin :: (Builtins -> a) -> VM a
builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))

263 264 265
builtins :: (a -> Builtins -> b) -> VM (a -> b)
builtins f = VM $ \bi genv lenv -> return (Yes genv lenv (`f` bi))

266
readGEnv :: (GlobalEnv -> a) -> VM a
twanvl's avatar
twanvl committed
267
readGEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f genv))
268 269 270 271 272 273 274 275

setGEnv :: GlobalEnv -> VM ()
setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())

updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())

readLEnv :: (LocalEnv -> a) -> VM a
twanvl's avatar
twanvl committed
276
readLEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f lenv))
277 278 279 280 281 282 283

setLEnv :: LocalEnv -> VM ()
setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())

updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())

twanvl's avatar
twanvl committed
284
{-
285 286
getInstEnv :: VM (InstEnv, InstEnv)
getInstEnv = readGEnv global_inst_env
twanvl's avatar
twanvl committed
287
-}
288 289 290 291

getFamInstEnv :: VM FamInstEnvs
getFamInstEnv = readGEnv global_fam_inst_env

292 293 294 295 296 297 298 299
getBindName :: VM FastString
getBindName = readLEnv local_bind_name

inBind :: Id -> VM a -> VM a
inBind id p
  = do updLEnv $ \env -> env { local_bind_name = occNameFS (getOccName id) }
       p

300 301 302 303 304 305 306 307 308 309
cloneName :: (OccName -> OccName) -> Name -> VM Name
cloneName mk_occ name = liftM make (liftDs newUnique)
  where
    occ_name = mk_occ (nameOccName name)

    make u | isExternalName name = mkExternalName u (nameModule name)
                                                    occ_name
                                                    (nameSrcSpan name)
           | otherwise           = mkSystemName u occ_name

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
310 311 312 313 314 315 316 317
cloneId :: (OccName -> OccName) -> Id -> Type -> VM Id
cloneId mk_occ id ty
  = do
      name <- cloneName mk_occ (getName id)
      let id' | isExportedId id = Id.mkExportedLocalId name ty
              | otherwise       = Id.mkLocalId         name ty
      return id'

rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
318 319 320
cloneVar :: Var -> VM Var
cloneVar var = liftM (setIdUnique var) (liftDs newUnique)

321 322 323 324 325 326 327 328 329 330
newExportedVar :: OccName -> Type -> VM Var
newExportedVar occ_name ty 
  = do
      mod <- liftDs getModuleDs
      u   <- liftDs newUnique

      let name = mkExternalName u mod occ_name noSrcSpan
      
      return $ Id.mkExportedLocalId name ty

331 332 333 334 335 336
newLocalVar :: FastString -> Type -> VM Var
newLocalVar fs ty
  = do
      u <- liftDs newUnique
      return $ mkSysLocal fs u ty

337
newDummyVar :: Type -> VM Var
Ian Lynagh's avatar
Ian Lynagh committed
338
newDummyVar = newLocalVar (fsLit "ds")
339

340 341 342 343 344 345
newTyVar :: FastString -> Kind -> VM Var
newTyVar fs k
  = do
      u <- liftDs newUnique
      return $ mkTyVar (mkSysTvName u fs) k

346 347
defGlobalVar :: Var -> Var -> VM ()
defGlobalVar v v' = updGEnv $ \env ->
348
  env { global_vars = extendVarEnv (global_vars env) v v'
349 350 351 352 353
      , global_exported_vars = upd (global_exported_vars env)
      }
  where
    upd env | isExportedId v = extendVarEnv env v (v, v')
            | otherwise      = env
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
354

355
lookupVar :: Var -> VM (Scope Var (Var, Var))
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
356 357 358 359 360 361
lookupVar v
  = do
      r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
      case r of
        Just e  -> return (Local e)
        Nothing -> liftM Global
362 363
                 $  traceMaybeV "lookupVar" (ppr v)
                                (readGEnv $ \env -> lookupVarEnv (global_vars env) v)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
364

365
lookupTyCon :: TyCon -> VM (Maybe TyCon)
366 367 368 369
lookupTyCon tc
  | isUnLiftedTyCon tc || isTupleTyCon tc = return (Just tc)

  | otherwise = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
370

371 372 373 374 375
defTyCon :: TyCon -> TyCon -> VM ()
defTyCon tc tc' = updGEnv $ \env ->
  env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }

lookupDataCon :: DataCon -> VM (Maybe DataCon)
376 377 378
lookupDataCon dc
  | isTupleTyCon (dataConTyCon dc) = return (Just dc)
  | otherwise = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
379 380 381 382 383

defDataCon :: DataCon -> DataCon -> VM ()
defDataCon dc dc' = updGEnv $ \env ->
  env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }

384
lookupPrimPArray :: TyCon -> VM (Maybe TyCon)
385
lookupPrimPArray = liftBuiltinDs . primPArray
386

387
lookupPrimMethod :: TyCon -> String -> VM (Maybe Var)
388
lookupPrimMethod tycon = liftBuiltinDs . primMethod tycon
389

390 391 392 393 394 395 396
lookupTyConPA :: TyCon -> VM (Maybe Var)
lookupTyConPA tc = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)

defTyConPA :: TyCon -> Var -> VM ()
defTyConPA tc pa = updGEnv $ \env ->
  env { global_pa_funs = extendNameEnv (global_pa_funs env) (tyConName tc) pa }

397 398 399 400 401
defTyConPAs :: [(TyCon, Var)] -> VM ()
defTyConPAs ps = updGEnv $ \env ->
  env { global_pa_funs = extendNameEnvList (global_pa_funs env)
                                           [(tyConName tc, pa) | (tc, pa) <- ps] }

402
lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
403 404 405 406
lookupTyVarPA tv = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv

lookupTyConPR :: TyCon -> VM (Maybe Var)
lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc)
407

408 409 410 411
lookupBoxedTyCon :: TyCon -> VM (Maybe TyCon)
lookupBoxedTyCon tc = readGEnv $ \env -> lookupNameEnv (global_boxed_tycons env)
                                                       (tyConName tc)

412 413 414 415 416 417 418 419 420 421 422
defLocalTyVar :: TyVar -> VM ()
defLocalTyVar tv = updLEnv $ \env ->
  env { local_tyvars   = tv : local_tyvars env
      , local_tyvar_pa = local_tyvar_pa env `delVarEnv` tv
      }

defLocalTyVarWithPA :: TyVar -> CoreExpr -> VM ()
defLocalTyVarWithPA tv pa = updLEnv $ \env ->
  env { local_tyvars   = tv : local_tyvars env
      , local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa
      }
423

424 425
localTyVars :: VM [TyVar]
localTyVars = readLEnv (reverse . local_tyvars)
426

427 428 429 430 431 432 433 434 435
-- Look up the dfun of a class instance.
--
-- The match must be unique - ie, match exactly one instance - but the 
-- type arguments used for matching may be more specific than those of 
-- the class instance declaration.  The found class instances must not have
-- any type variables in the instance context that do not appear in the
-- instances head (i.e., no flexi vars); for details for what this means,
-- see the docs at InstEnv.lookupInstEnv.
--
436
{-
437 438 439 440 441 442 443 444 445 446 447
lookupInst :: Class -> [Type] -> VM (DFunId, [Type])
lookupInst cls tys
  = do { instEnv <- getInstEnv
       ; case lookupInstEnv instEnv cls tys of
	   ([(inst, inst_tys)], _) 
             | noFlexiVar -> return (instanceDFunId inst, inst_tys')
             | otherwise  -> pprPanic "VectMonad.lookupInst: flexi var: " 
                                      (ppr $ mkTyConApp (classTyCon cls) tys)
             where
               inst_tys'  = [ty | Right ty <- inst_tys]
               noFlexiVar = all isRight inst_tys
448
	   _other         -> traceNoV "lookupInst" (ppr cls <+> ppr tys)
449 450 451 452
       }
  where
    isRight (Left  _) = False
    isRight (Right _) = True
453
-}
454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481

-- Look up the representation tycon of a family instance.
--
-- The match must be unique - ie, match exactly one instance - but the 
-- type arguments used for matching may be more specific than those of 
-- the family instance declaration.
--
-- Return the instance tycon and its type instance.  For example, if we have
--
--  lookupFamInst 'T' '[Int]' yields (':R42T', 'Int')
--
-- then we have a coercion (ie, type instance of family instance coercion)
--
--  :Co:R42T Int :: T [Int] ~ :R42T Int
--
-- which implies that :R42T was declared as 'data instance T [a]'.
--
lookupFamInst :: TyCon -> [Type] -> VM (TyCon, [Type])
lookupFamInst tycon tys
  = ASSERT( isOpenTyCon tycon )
    do { instEnv <- getFamInstEnv
       ; case lookupFamInstEnv instEnv tycon tys of
	   [(fam_inst, rep_tys)] -> return (famInstTyCon fam_inst, rep_tys)
	   _other                -> 
             pprPanic "VectMonad.lookupFamInst: not found: " 
                      (ppr $ mkTyConApp tycon tys)
       }

482 483
initV :: PackageId -> HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
initV pkg hsc_env guts info p
484 485 486 487
  = do
      Just r <- initDs hsc_env (mg_module guts)
                               (mg_rdr_env guts)
                               (mg_types guts)
488
                               go
489 490
      return r
  where
491

492
    go =
493
      do
494
        builtins       <- initBuiltins pkg
495 496 497
        builtin_vars   <- initBuiltinVars builtins
        builtin_tycons <- initBuiltinTyCons builtins
        let builtin_datacons = initBuiltinDataCons builtins
498
        builtin_pas    <- initBuiltinPAs builtins
499
        builtin_prs    <- initBuiltinPRs builtins
500
        builtin_boxed  <- initBuiltinBoxedTyCons builtins
501

502
        eps <- liftIO $ hscEPS hsc_env
503 504 505
        let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
            instEnvs    = (eps_inst_env     eps, mg_inst_env     guts)

506 507
        let genv = extendImportedVarsEnv builtin_vars
                 . extendTyConsEnv builtin_tycons
508
                 . extendDataConsEnv builtin_datacons
509
                 . extendPAFunsEnv builtin_pas
510
                 . setPRFunsEnv    builtin_prs
511
                 . setBoxedTyConsEnv builtin_boxed
512 513 514
                 $ initGlobalEnv info instEnvs famInstEnvs

        r <- runVM p builtins genv emptyLocalEnv
515 516 517
        case r of
          Yes genv _ x -> return $ Just (new_info genv, x)
          No           -> return Nothing
518 519 520

    new_info genv = updVectInfo genv (mg_types guts) info