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
101
102
103
104
105
|
{-# OPTIONS -fno-warn-tabs #-}
-- The above warning supression flag is a temporary kludge.
-- While working on this module you are encouraged to remove it and
-- detab the module (please do the detabbing in a separate patch). See
-- http://hackage.haskell.org/trac/ghc/wiki/Commentary/CodingStyle#TabsvsSpaces
-- for details
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 enviroment 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)
|
