VectMonad.hs 21.3 KB
Newer Older
1
{-# LANGUAGE NamedFieldPuns #-}
benl@ouroborus.net's avatar
benl@ouroborus.net committed
2
3

-- | The Vectorisation monad.
4
module VectMonad (
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
5
  Scope(..),
6
7
  VM,

8
9
  noV, traceNoV, ensureV, traceEnsureV, tryV, maybeV, traceMaybeV, orElseV,
  onlyIfV, fixV, localV, closedV,
10
  initV, cantVectorise, maybeCantVectorise, maybeCantVectoriseM,
11
  liftDs,
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
12
  cloneName, cloneId, cloneVar,
13
  newExportedVar, newLocalVar, newLocalVars, newDummyVar, newTyVar,
14
  
15
  Builtins(..), sumTyCon, prodTyCon, prodDataCon,
16
  selTy, selReplicate, selPick, selTags, selElements,
17
  combinePDVar, scalarZip, closureCtrFun,
18
  builtin, builtins,
19
20

  GlobalEnv(..),
21
  setFamInstEnv,
22
23
24
25
26
  readGEnv, setGEnv, updGEnv,

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

27
28
  getBindName, inBind,

29
  lookupVar, defGlobalVar, globalScalars,
30
31
  lookupTyCon, defTyCon,
  lookupDataCon, defDataCon,
32
  lookupTyConPA, defTyConPA, defTyConPAs,
33
  lookupTyConPR,
34
  lookupBoxedTyCon,
35
  lookupPrimMethod, lookupPrimPArray,
36
  lookupTyVarPA, defLocalTyVar, defLocalTyVarWithPA, localTyVars,
37

38
  lookupInst, lookupFamInst
39
40
41
42
) where

#include "HsVersions.h"

43
44
import VectBuiltIn

45
import HscTypes hiding  ( MonadThings(..) )
46
import Module           ( PackageId )
47
import CoreSyn
48
import Class
49
import TyCon
50
import DataCon
51
52
import Type
import Var
53
import VarSet
54
55
56
57
58
59
60
import VarEnv
import Id
import Name
import NameEnv

import DsMonad

61
62
63
64
import InstEnv
import FamInstEnv

import Outputable
65
import FastString
66
import SrcLoc        ( noSrcSpan )
67

twanvl's avatar
twanvl committed
68
import Control.Monad
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
69

benl@ouroborus.net's avatar
benl@ouroborus.net committed
70
-- | Indicates what scope something (a variable) is in.
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
71
72
data Scope a b = Global a | Local b

73

benl@ouroborus.net's avatar
benl@ouroborus.net committed
74
-- | The global environment.
75
data GlobalEnv = GlobalEnv {
benl@ouroborus.net's avatar
benl@ouroborus.net committed
76
                  -- | Mapping from global variables to their vectorised versions.
77
                  -- 
78
                  global_vars :: VarEnv Var
79

benl@ouroborus.net's avatar
benl@ouroborus.net committed
80
81
                  -- | Purely scalar variables. Code which mentions only these
                  --   variables doesn't have to be lifted.
82
83
                , global_scalars :: VarSet

benl@ouroborus.net's avatar
benl@ouroborus.net committed
84
                  -- | Exported variables which have a vectorised version
85
86
87
                  --
                , global_exported_vars :: VarEnv (Var, Var)

benl@ouroborus.net's avatar
benl@ouroborus.net committed
88
89
90
                  -- | Mapping from TyCons to their vectorised versions.
                  --   TyCons which do not have to be vectorised are mapped to
                  --   themselves.
91
92
93
                  --
                , global_tycons :: NameEnv TyCon

benl@ouroborus.net's avatar
benl@ouroborus.net committed
94
                  -- | Mapping from DataCons to their vectorised versions
95
96
97
                  --
                , global_datacons :: NameEnv DataCon

benl@ouroborus.net's avatar
benl@ouroborus.net committed
98
                  -- | Mapping from TyCons to their PA dfuns
99
100
101
                  --
                , global_pa_funs :: NameEnv Var

benl@ouroborus.net's avatar
benl@ouroborus.net committed
102
                  -- | Mapping from TyCons to their PR dfuns
103
104
                , global_pr_funs :: NameEnv Var

benl@ouroborus.net's avatar
benl@ouroborus.net committed
105
                  -- | Mapping from unboxed TyCons to their boxed versions
106
107
                , global_boxed_tycons :: NameEnv TyCon

benl@ouroborus.net's avatar
benl@ouroborus.net committed
108
109
                -- | External package inst-env & home-package inst-env for class
                --   instances
110
111
112
                --
                , global_inst_env :: (InstEnv, InstEnv)

benl@ouroborus.net's avatar
benl@ouroborus.net committed
113
114
                -- | External package inst-env & home-package inst-env for family
                --   instances
115
116
                --
                , global_fam_inst_env :: FamInstEnvs
117

benl@ouroborus.net's avatar
benl@ouroborus.net committed
118
                -- | Hoisted bindings
119
                , global_bindings :: [(Var, CoreExpr)]
120
121
                }

benl@ouroborus.net's avatar
benl@ouroborus.net committed
122
-- | The local environment.
123
124
125
126
data LocalEnv = LocalEnv {
                 -- Mapping from local variables to their vectorised and
                 -- lifted versions
                 --
127
                 local_vars :: VarEnv (Var, Var)
128

129
130
131
132
                 -- In-scope type variables
                 --
               , local_tyvars :: [TyVar]

133
134
                 -- Mapping from tyvars to their PA dictionaries
               , local_tyvar_pa :: VarEnv CoreExpr
135
136
137

                 -- Local binding name
               , local_bind_name :: FastString
138
139
               }

benl@ouroborus.net's avatar
benl@ouroborus.net committed
140
141

-- | Create an initial global environment
142
143
initGlobalEnv :: VectInfo -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
initGlobalEnv info instEnvs famInstEnvs
144
  = GlobalEnv {
145
      global_vars          = mapVarEnv snd $ vectInfoVar info
146
    , global_scalars   = emptyVarSet
147
    , global_exported_vars = emptyVarEnv
148
    , global_tycons        = mapNameEnv snd $ vectInfoTyCon info
149
    , global_datacons      = mapNameEnv snd $ vectInfoDataCon info
150
    , global_pa_funs       = mapNameEnv snd $ vectInfoPADFun info
151
    , global_pr_funs       = emptyNameEnv
152
    , global_boxed_tycons  = emptyNameEnv
153
154
    , global_inst_env      = instEnvs
    , global_fam_inst_env  = famInstEnvs
155
    , global_bindings      = []
156
157
    }

benl@ouroborus.net's avatar
benl@ouroborus.net committed
158
159

-- Operators on Global Environments -------------------------------------------
160
161
162
163
extendImportedVarsEnv :: [(Var, Var)] -> GlobalEnv -> GlobalEnv
extendImportedVarsEnv ps genv
  = genv { global_vars = extendVarEnvList (global_vars genv) ps }

164
165
166
167
extendScalars :: [Var] -> GlobalEnv -> GlobalEnv
extendScalars vs genv
  = genv { global_scalars = extendVarSetList (global_scalars genv) vs }

168
169
170
171
172
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
173

174
175
176
177
extendTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
extendTyConsEnv ps genv
  = genv { global_tycons = extendNameEnvList (global_tycons genv) ps }

178
179
180
181
extendDataConsEnv :: [(Name, DataCon)] -> GlobalEnv -> GlobalEnv
extendDataConsEnv ps genv
  = genv { global_datacons = extendNameEnvList (global_datacons genv) ps }

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

186
187
188
189
setPRFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
setPRFunsEnv ps genv
  = genv { global_pr_funs = mkNameEnv ps }

190
191
192
193
setBoxedTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
setBoxedTyConsEnv ps genv
  = genv { global_boxed_tycons = mkNameEnv ps }

benl@ouroborus.net's avatar
benl@ouroborus.net committed
194
195

-- | Create an empty local environment.
twanvl's avatar
twanvl committed
196
emptyLocalEnv :: LocalEnv
197
198
emptyLocalEnv = LocalEnv {
                   local_vars     = emptyVarEnv
199
                 , local_tyvars   = []
200
                 , local_tyvar_pa = emptyVarEnv
Ian Lynagh's avatar
Ian Lynagh committed
201
                 , local_bind_name  = fsLit "fn"
202
203
204
205
206
207
                 }

-- FIXME
updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
updVectInfo env tyenv info
  = info {
208
      vectInfoVar     = global_exported_vars env
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
209
210
    , vectInfoTyCon   = mk_env typeEnvTyCons global_tycons
    , vectInfoDataCon = mk_env typeEnvDataCons global_datacons
211
    , vectInfoPADFun  = mk_env typeEnvTyCons global_pa_funs
212
213
    }
  where
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
214
215
216
217
    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]]
218

benl@ouroborus.net's avatar
benl@ouroborus.net committed
219
220
221
222
223
224

-- The Vectorisation Monad ----------------------------------------------------

-- Vectorisation can either succeed with new envionment and a value,
-- or return with failure.
--
225
226
227
228
229
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
230
  return x   = VM $ \_  genv lenv -> return (Yes genv lenv x)
231
232
233
234
235
236
  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

237

benl@ouroborus.net's avatar
benl@ouroborus.net committed
238
-- | Throw an error saying we can't vectorise something
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
cantVectorise :: String -> SDoc -> a
cantVectorise s d = pgmError
                  . showSDocDump
                  $ vcat [text "*** Vectorisation error ***",
                          nest 4 $ sep [text s, nest 4 d]]

maybeCantVectorise :: String -> SDoc -> Maybe a -> a
maybeCantVectorise s d Nothing  = cantVectorise s d
maybeCantVectorise _ _ (Just x) = x

maybeCantVectoriseM :: Monad m => String -> SDoc -> m (Maybe a) -> m a
maybeCantVectoriseM s d p
  = do
      r <- p
      case r of
        Just x  -> return x
        Nothing -> cantVectorise s d

benl@ouroborus.net's avatar
benl@ouroborus.net committed
257
258
259

-- Control --------------------------------------------------------------------
-- | Return some result saying we've failed.
260
261
262
noV :: VM a
noV = VM $ \_ _ _ -> return No

263
264
265
traceNoV :: String -> SDoc -> VM a
traceNoV s d = pprTrace s d noV

benl@ouroborus.net's avatar
benl@ouroborus.net committed
266
267

-- | If True then carry on, otherwise fail.
268
269
270
271
ensureV :: Bool -> VM ()
ensureV False = noV
ensureV True  = return ()

benl@ouroborus.net's avatar
benl@ouroborus.net committed
272
273

-- | If True then return the first argument, otherwise fail.
274
275
276
277
278
onlyIfV :: Bool -> VM a -> VM a
onlyIfV b p = ensureV b >> p

traceEnsureV :: String -> SDoc -> Bool -> VM ()
traceEnsureV s d False = traceNoV s d
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
279
traceEnsureV _ _ True  = return ()
280

benl@ouroborus.net's avatar
benl@ouroborus.net committed
281
282
283
284

-- | Try some vectorisation computaton.
--	If it succeeds then return Just the result,
--	otherwise return Nothing.
285
286
287
288
289
290
291
292
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)

benl@ouroborus.net's avatar
benl@ouroborus.net committed
293

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

297
298
299
traceMaybeV :: String -> SDoc -> VM (Maybe a) -> VM a
traceMaybeV s d p = maybe (traceNoV s d) return =<< p

300
301
302
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
303
fixV :: (a -> VM a) -> VM a
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
304
305
306
307
308
309
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
310

benl@ouroborus.net's avatar
benl@ouroborus.net committed
311
312
313
314

-- Local Environments ---------------------------------------------------------
-- | Perform a computation in its own local environment.
--	This does not alter the environment of the current state.
315
316
317
318
319
320
321
localV :: VM a -> VM a
localV p = do
             env <- readLEnv id
             x <- p
             setLEnv env
             return x

benl@ouroborus.net's avatar
benl@ouroborus.net committed
322
-- | Perform a computation in an empty local environment.
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
323
324
325
closedV :: VM a -> VM a
closedV p = do
              env <- readLEnv id
326
              setLEnv (emptyLocalEnv { local_bind_name = local_bind_name env })
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
327
328
329
330
              x <- p
              setLEnv env
              return x

benl@ouroborus.net's avatar
benl@ouroborus.net committed
331
332
-- Lifting --------------------------------------------------------------------
-- | Lift a desugaring computation into the vectorisation monad.
333
liftDs :: DsM a -> VM a
twanvl's avatar
twanvl committed
334
liftDs p = VM $ \_ genv lenv -> do { x <- p; return (Yes genv lenv x) }
335

benl@ouroborus.net's avatar
benl@ouroborus.net committed
336
337
338
339


-- Builtins -------------------------------------------------------------------
-- Operations on Builtins
340
341
342
liftBuiltinDs :: (Builtins -> DsM a) -> VM a
liftBuiltinDs p = VM $ \bi genv lenv -> do { x <- p bi; return (Yes genv lenv x)}

benl@ouroborus.net's avatar
benl@ouroborus.net committed
343
344

-- | Project something from the set of builtins.
345
346
347
builtin :: (Builtins -> a) -> VM a
builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))

348
349
350
builtins :: (a -> Builtins -> b) -> VM (a -> b)
builtins f = VM $ \bi genv lenv -> return (Yes genv lenv (`f` bi))

benl@ouroborus.net's avatar
benl@ouroborus.net committed
351
352
353

-- Environments ---------------------------------------------------------------
-- | Project something from the global environment.
354
readGEnv :: (GlobalEnv -> a) -> VM a
twanvl's avatar
twanvl committed
355
readGEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f genv))
356
357
358
359
360
361
362

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

benl@ouroborus.net's avatar
benl@ouroborus.net committed
363
364

-- | Project something from the local environment.
365
readLEnv :: (LocalEnv -> a) -> VM a
twanvl's avatar
twanvl committed
366
readLEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f lenv))
367
368
369
370
371
372
373

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

benl@ouroborus.net's avatar
benl@ouroborus.net committed
374
375

-- InstEnv --------------------------------------------------------------------
376
377
378
379
380
381
getInstEnv :: VM (InstEnv, InstEnv)
getInstEnv = readGEnv global_inst_env

getFamInstEnv :: VM FamInstEnvs
getFamInstEnv = readGEnv global_fam_inst_env

benl@ouroborus.net's avatar
benl@ouroborus.net committed
382
383
384

-- Names ----------------------------------------------------------------------
-- | Get the name of the local binding currently being vectorised.
385
386
387
388
389
390
391
392
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

393
394
395
396
397
398
399
400
401
402
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
403
404
405
406
407
408
409
410
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'

benl@ouroborus.net's avatar
benl@ouroborus.net committed
411
-- Make a fresh instance of this var, with a new unique.
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
412
413
414
cloneVar :: Var -> VM Var
cloneVar var = liftM (setIdUnique var) (liftDs newUnique)

415
416
417
418
419
420
421
422
423
424
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

425
426
427
428
429
430
newLocalVar :: FastString -> Type -> VM Var
newLocalVar fs ty
  = do
      u <- liftDs newUnique
      return $ mkSysLocal fs u ty

431
432
433
newLocalVars :: FastString -> [Type] -> VM [Var]
newLocalVars fs = mapM (newLocalVar fs)

434
newDummyVar :: Type -> VM Var
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
435
newDummyVar = newLocalVar (fsLit "vv")
436

437
438
439
440
441
442
newTyVar :: FastString -> Kind -> VM Var
newTyVar fs k
  = do
      u <- liftDs newUnique
      return $ mkTyVar (mkSysTvName u fs) k

443
444

-- | Add a mapping between a global var and its vectorised version to the state.
445
446
defGlobalVar :: Var -> Var -> VM ()
defGlobalVar v v' = updGEnv $ \env ->
447
  env { global_vars = extendVarEnv (global_vars env) v v'
448
449
450
451
452
      , 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
453

454
-- Var ------------------------------------------------------------------------
benl@ouroborus.net's avatar
benl@ouroborus.net committed
455
456
457
458
-- | Lookup the vectorised and\/or lifted versions of this variable.
--	If it's in the global environment we get the vectorised version.
--      If it's in the local environment we get both the vectorised and lifted version.
--	
459
lookupVar :: Var -> VM (Scope Var (Var, Var))
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
460
lookupVar v
461
 = do r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
462
463
464
      case r of
        Just e  -> return (Local e)
        Nothing -> liftM Global
465
                . maybeCantVectoriseVarM v
466
                . readGEnv $ \env -> lookupVarEnv (global_vars env) v
rl@cse.unsw.edu.au's avatar
rl@cse.unsw.edu.au committed
467

468
469
470
471
472
473
474
475
476
maybeCantVectoriseVarM :: Monad m => Var -> m (Maybe Var) -> m Var
maybeCantVectoriseVarM v p
 = do r <- p
      case r of
        Just x  -> return x
        Nothing -> dumpVar v

dumpVar :: Var -> a
dumpVar var
benl@ouroborus.net's avatar
benl@ouroborus.net committed
477
	| Just _		<- isClassOpId_maybe var
478
479
480
481
482
483
	= cantVectorise "ClassOpId not vectorised:" (ppr var)

	| otherwise
	= cantVectorise "Variable not vectorised:" (ppr var)

-------------------------------------------------------------------------------
484
485
486
globalScalars :: VM VarSet
globalScalars = readGEnv global_scalars

487
lookupTyCon :: TyCon -> VM (Maybe TyCon)
488
489
490
491
lookupTyCon tc
  | isUnLiftedTyCon tc || isTupleTyCon tc = return (Just tc)

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

493
494
495
496
497
defTyCon :: TyCon -> TyCon -> VM ()
defTyCon tc tc' = updGEnv $ \env ->
  env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }

lookupDataCon :: DataCon -> VM (Maybe DataCon)
498
499
500
lookupDataCon dc
  | isTupleTyCon (dataConTyCon dc) = return (Just dc)
  | otherwise = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
501
502
503
504
505

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

506
lookupPrimPArray :: TyCon -> VM (Maybe TyCon)
507
lookupPrimPArray = liftBuiltinDs . primPArray
508

509
lookupPrimMethod :: TyCon -> String -> VM (Maybe Var)
510
lookupPrimMethod tycon = liftBuiltinDs . primMethod tycon
511

512
513
514
515
516
517
518
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 }

519
520
521
522
523
defTyConPAs :: [(TyCon, Var)] -> VM ()
defTyConPAs ps = updGEnv $ \env ->
  env { global_pa_funs = extendNameEnvList (global_pa_funs env)
                                           [(tyConName tc, pa) | (tc, pa) <- ps] }

524
lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
525
526
527
528
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)
529

530
531
532
533
lookupBoxedTyCon :: TyCon -> VM (Maybe TyCon)
lookupBoxedTyCon tc = readGEnv $ \env -> lookupNameEnv (global_boxed_tycons env)
                                                       (tyConName tc)

534
535
536
537
538
539
540
541
542
543
544
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
      }
545

546
547
localTyVars :: VM [TyVar]
localTyVars = readLEnv (reverse . local_tyvars)
548

549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
-- 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.
--
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
569
570
	   _other         ->
             pprPanic "VectMonad.lookupInst: not found " (ppr cls <+> ppr tys)
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
       }
  where
    isRight (Left  _) = False
    isRight (Right _) = True

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

603
604

-- | Run a vectorisation computation.
605
606
initV :: PackageId -> HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
initV pkg hsc_env guts info p
607
  = do
608
609
610
         -- XXX: ignores error messages and warnings, check that this is
         -- indeed ok (the use of "Just r" suggests so)
      (_,Just r) <- initDs hsc_env (mg_module guts)
611
612
                               (mg_rdr_env guts)
                               (mg_types guts)
613
                               go
614
615
      return r
  where
616

617
    go =
618
      do
619
        builtins       <- initBuiltins pkg
620
621
622
        builtin_vars   <- initBuiltinVars builtins
        builtin_tycons <- initBuiltinTyCons builtins
        let builtin_datacons = initBuiltinDataCons builtins
623
        builtin_boxed  <- initBuiltinBoxedTyCons builtins
624
        builtin_scalars <- initBuiltinScalars builtins
625

626
        eps <- liftIO $ hscEPS hsc_env
627
628
629
        let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
            instEnvs    = (eps_inst_env     eps, mg_inst_env     guts)

630
631
632
        builtin_prs    <- initBuiltinPRs builtins instEnvs
        builtin_pas    <- initBuiltinPAs builtins instEnvs

633
        let genv = extendImportedVarsEnv builtin_vars
634
                 . extendScalars builtin_scalars
635
                 . extendTyConsEnv builtin_tycons
636
                 . extendDataConsEnv builtin_datacons
637
                 . extendPAFunsEnv builtin_pas
638
                 . setPRFunsEnv    builtin_prs
639
                 . setBoxedTyConsEnv builtin_boxed
640
641
642
                 $ initGlobalEnv info instEnvs famInstEnvs

        r <- runVM p builtins genv emptyLocalEnv
643
644
645
        case r of
          Yes genv _ x -> return $ Just (new_info genv, x)
          No           -> return Nothing
646
647
648

    new_info genv = updVectInfo genv (mg_types guts) info