1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
|
module Vectorise.Monad.Local
( readLEnv
, setLEnv
, updLEnv
, localV
, closedV
, getBindName
, inBind
, lookupTyVarPA
, defLocalTyVar
, defLocalTyVarWithPA
, localTyVars
)
where
import Vectorise.Monad.Base
import Vectorise.Env
import CoreSyn
import Name
import VarEnv
import Var
import FastString
-- Local Environment ----------------------------------------------------------
-- |Project something from the local environment.
--
readLEnv :: (LocalEnv -> a) -> VM a
readLEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f lenv))
-- |Set the local environment.
--
setLEnv :: LocalEnv -> VM ()
setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
-- |Update the environment using the provided function.
--
updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
-- |Perform a computation in its own local environment.
-- This does not alter the environment of the current state.
--
localV :: VM a -> VM a
localV p
= do
{ env <- readLEnv id
; x <- p
; setLEnv env
; return x
}
-- |Perform a computation in an empty local environment.
--
closedV :: VM a -> VM a
closedV p
= do
{ env <- readLEnv id
; setLEnv (emptyLocalEnv { local_bind_name = local_bind_name env })
; x <- p
; setLEnv env
; return x
}
-- |Get the name of the local binding currently being vectorised.
--
getBindName :: VM FastString
getBindName = readLEnv local_bind_name
-- |Run a vectorisation computation in a local environment,
-- with this id set as the current binding.
--
inBind :: Id -> VM a -> VM a
inBind id p
= do updLEnv $ \env -> env { local_bind_name = occNameFS (getOccName id) }
p
-- |Lookup a PA tyvars from the local environment.
lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
lookupTyVarPA tv
= readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv
-- |Add a tyvar to the local environment.
defLocalTyVar :: TyVar -> VM ()
defLocalTyVar tv = updLEnv $ \env ->
env { local_tyvars = tv : local_tyvars env
, local_tyvar_pa = local_tyvar_pa env `delVarEnv` tv
}
-- |Add mapping between a tyvar and pa dictionary to the local environment.
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
}
-- |Get the set of tyvars from the local environment.
localTyVars :: VM [TyVar]
localTyVars = readLEnv (reverse . local_tyvars)
|
