VectMonad.hs 15.2 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 8
  cloneName, cloneId,
  newExportedVar, newLocalVar, newDummyVar, newTyVar,
9
  
10 11
  Builtins(..), sumTyCon, prodTyCon,
  builtin, builtins,
12 13

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

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

20 21
  getBindName, inBind,

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

30
  {-lookupInst,-} lookupFamInst
31 32 33 34
) where

#include "HsVersions.h"

35 36
import VectBuiltIn

37 38 39
import HscTypes
import CoreSyn
import TyCon
40
import DataCon
41
import Type
42
import Class
43 44 45
import Var
import VarEnv
import Id
46
import OccName
47 48
import Name
import NameEnv
49
import TysPrim       ( intPrimTy )
50 51
import Module
import IfaceEnv
52
import IOEnv         ( ioToIOEnv )
53 54 55 56

import DsMonad
import PrelNames

57 58 59 60 61
import InstEnv
import FamInstEnv

import Panic
import Outputable
62
import FastString
63
import SrcLoc        ( noSrcSpan )
64

65
import Control.Monad ( liftM, zipWithM )
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
66 67 68

data Scope a b = Global a | Local b

69 70 71 72 73 74
-- ----------------------------------------------------------------------------
-- Vectorisation monad

data GlobalEnv = GlobalEnv {
                  -- Mapping from global variables to their vectorised versions.
                  -- 
75
                  global_vars :: VarEnv Var
76 77 78 79 80 81 82 83 84 85 86

                  -- 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

87 88 89 90
                  -- Mapping from DataCons to their vectorised versions
                  --
                , global_datacons :: NameEnv DataCon

91 92 93 94
                  -- Mapping from TyCons to their PA dfuns
                  --
                , global_pa_funs :: NameEnv Var

95 96 97
                  -- Mapping from TyCons to their PR dfuns
                , global_pr_funs :: NameEnv Var

98 99 100 101 102 103 104 105 106
                -- 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
107 108 109

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

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

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

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

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

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

143 144 145 146 147
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
148

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

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

157 158 159 160
setPRFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
setPRFunsEnv ps genv
  = genv { global_pr_funs = mkNameEnv ps }

161 162
emptyLocalEnv = LocalEnv {
                   local_vars     = emptyVarEnv
163
                 , local_tyvars   = []
164
                 , local_tyvar_pa = emptyVarEnv
165
                 , local_bind_name  = FSLIT("fn")
166 167 168 169 170 171
                 }

-- FIXME
updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
updVectInfo env tyenv info
  = info {
172
      vectInfoVar     = global_exported_vars env
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
173 174
    , vectInfoTyCon   = mk_env typeEnvTyCons global_tycons
    , vectInfoDataCon = mk_env typeEnvDataCons global_datacons
175
    , vectInfoPADFun  = mk_env typeEnvTyCons global_pa_funs
176 177
    }
  where
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
178 179 180 181
    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]]
182

183 184 185 186 187 188 189 190 191 192 193 194 195 196 197
data VResult a = Yes GlobalEnv LocalEnv a | No

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

instance Monad VM where
  return x   = VM $ \bi genv lenv -> return (Yes genv lenv x)
  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

198 199 200
traceNoV :: String -> SDoc -> VM a
traceNoV s d = pprTrace s d noV

201 202 203 204 205 206 207 208 209 210 211
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

212 213 214
traceMaybeV :: String -> SDoc -> VM (Maybe a) -> VM a
traceMaybeV s d p = maybe (traceNoV s d) return =<< p

215 216 217
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
218 219 220 221 222
fixV :: (a -> VM a) -> VM a
fixV f = VM (\bi genv lenv -> fixDs $ \r -> runVM (f (unYes r)) bi genv lenv )
  where
    unYes (Yes _ _ x) = x

223 224 225 226 227 228 229
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
230 231 232
closedV :: VM a -> VM a
closedV p = do
              env <- readLEnv id
233
              setLEnv (emptyLocalEnv { local_bind_name = local_bind_name env })
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
234 235 236 237
              x <- p
              setLEnv env
              return x

238 239 240 241 242 243
liftDs :: DsM a -> VM a
liftDs p = VM $ \bi genv lenv -> do { x <- p; return (Yes genv lenv x) }

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

244 245 246
builtins :: (a -> Builtins -> b) -> VM (a -> b)
builtins f = VM $ \bi genv lenv -> return (Yes genv lenv (`f` bi))

247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264
readGEnv :: (GlobalEnv -> a) -> VM a
readGEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f genv))

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
readLEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f lenv))

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) ())

265 266 267 268 269 270
getInstEnv :: VM (InstEnv, InstEnv)
getInstEnv = readGEnv global_inst_env

getFamInstEnv :: VM FamInstEnvs
getFamInstEnv = readGEnv global_fam_inst_env

271 272 273 274 275 276 277 278
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

279 280 281 282 283 284 285 286 287 288
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
289 290 291 292 293 294 295 296
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'

297 298 299 300 301 302 303 304 305 306
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

307 308 309 310 311 312
newLocalVar :: FastString -> Type -> VM Var
newLocalVar fs ty
  = do
      u <- liftDs newUnique
      return $ mkSysLocal fs u ty

313 314 315
newDummyVar :: Type -> VM Var
newDummyVar = newLocalVar FSLIT("ds")

316 317 318 319 320 321
newTyVar :: FastString -> Kind -> VM Var
newTyVar fs k
  = do
      u <- liftDs newUnique
      return $ mkTyVar (mkSysTvName u fs) k

322 323
defGlobalVar :: Var -> Var -> VM ()
defGlobalVar v v' = updGEnv $ \env ->
324
  env { global_vars = extendVarEnv (global_vars env) v v'
325 326 327 328 329
      , 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
330

331
lookupVar :: Var -> VM (Scope Var (Var, Var))
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
332 333 334 335 336 337
lookupVar v
  = do
      r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
      case r of
        Just e  -> return (Local e)
        Nothing -> liftM Global
338 339
                 $  traceMaybeV "lookupVar" (ppr v)
                                (readGEnv $ \env -> lookupVarEnv (global_vars env) v)
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
340

341
lookupTyCon :: TyCon -> VM (Maybe TyCon)
342 343 344 345
lookupTyCon tc
  | isUnLiftedTyCon tc || isTupleTyCon tc = return (Just tc)

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

347 348 349 350 351 352 353 354 355 356 357
defTyCon :: TyCon -> TyCon -> VM ()
defTyCon tc tc' = updGEnv $ \env ->
  env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }

lookupDataCon :: DataCon -> VM (Maybe DataCon)
lookupDataCon dc = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)

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

358 359 360
lookupPrimPArray :: TyCon -> VM (Maybe TyCon)
lookupPrimPArray = liftDs . primPArray

361
lookupPrimMethod :: TyCon -> String -> VM (Maybe Var)
362
lookupPrimMethod tycon = liftDs . primMethod tycon
363

364 365 366 367 368 369 370
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 }

371 372 373 374 375
defTyConPAs :: [(TyCon, Var)] -> VM ()
defTyConPAs ps = updGEnv $ \env ->
  env { global_pa_funs = extendNameEnvList (global_pa_funs env)
                                           [(tyConName tc, pa) | (tc, pa) <- ps] }

376
lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
377 378 379 380
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)
381

382 383 384 385 386 387 388 389 390 391 392
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
      }
393

394 395
localTyVars :: VM [TyVar]
localTyVars = readLEnv (reverse . local_tyvars)
396

397 398 399 400 401 402 403 404 405
-- 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.
--
406
{-
407 408 409 410 411 412 413 414 415 416 417
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
418
	   _other         -> traceNoV "lookupInst" (ppr cls <+> ppr tys)
419 420 421 422
       }
  where
    isRight (Left  _) = False
    isRight (Right _) = True
423
-}
424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451

-- 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)
       }

452 453 454 455 456 457
initV :: HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
initV hsc_env guts info p
  = do
      Just r <- initDs hsc_env (mg_module guts)
                               (mg_rdr_env guts)
                               (mg_types guts)
458
                               go
459 460
      return r
  where
461

462
    go =
463
      do
464
        builtins       <- initBuiltins
465
        builtin_tycons <- initBuiltinTyCons
466
        builtin_pas    <- initBuiltinPAs builtins
467
        builtin_prs    <- initBuiltinPRs builtins
468

469 470 471 472
        eps <- ioToIOEnv $ hscEPS hsc_env
        let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
            instEnvs    = (eps_inst_env     eps, mg_inst_env     guts)

473 474
        let genv = extendTyConsEnv builtin_tycons
                 . extendPAFunsEnv builtin_pas
475
                 . setPRFunsEnv    builtin_prs
476 477 478
                 $ initGlobalEnv info instEnvs famInstEnvs

        r <- runVM p builtins genv emptyLocalEnv
479 480 481
        case r of
          Yes genv _ x -> return $ Just (new_info genv, x)
          No           -> return Nothing
482 483 484

    new_info genv = updVectInfo genv (mg_types guts) info