def1ffa58c67a091d138ca4bfdf1dd22210d4b73
[ghc.git] / compiler / vectorise / Vectorise / Monad / Naming.hs
1 -- |Computations in the vectorisation monad concerned with naming and fresh variable generation.
2
3 module Vectorise.Monad.Naming
4 ( mkLocalisedName
5 , mkDerivedName
6 , mkVectId
7 , cloneVar
8 , newExportedVar
9 , newLocalVar
10 , newLocalVars
11 , newDummyVar
12 , newTyVar
13 )
14 where
15
16 import Vectorise.Monad.Base
17
18 import DsMonad
19 import TcType
20 import Type
21 import Var
22 import Module
23 import Name
24 import SrcLoc
25 import MkId
26 import Id
27 import FastString
28
29 import Control.Monad
30
31
32 -- Naming ---------------------------------------------------------------------
33
34 -- |Create a localised variant of a name, using the provided function to transform its `OccName`.
35 --
36 -- If the name external, encode the orignal name's module into the new 'OccName'. The result is
37 -- always an internal system name.
38 --
39 mkLocalisedName :: (Maybe String -> OccName -> OccName) -> Name -> VM Name
40 mkLocalisedName mk_occ name
41 = do { mod <- liftDs getModule
42 ; u <- liftDs newUnique
43 ; let occ_name = mkLocalisedOccName mod mk_occ name
44
45 new_name | isExternalName name = mkExternalName u mod occ_name (nameSrcSpan name)
46 | otherwise = mkSystemName u occ_name
47
48 ; return new_name }
49
50 mkDerivedName :: (OccName -> OccName) -> Name -> VM Name
51 -- Similar to mkLocalisedName, but assumes the
52 -- incoming name is from this module.
53 -- Works on External names only
54 mkDerivedName mk_occ name
55 = do { u <- liftDs newUnique
56 ; return (mkExternalName u (nameModule name)
57 (mk_occ (nameOccName name))
58 (nameSrcSpan name)) }
59
60 -- |Produce the vectorised variant of an `Id` with the given vectorised type, while taking care that
61 -- vectorised dfun ids must be dfuns again.
62 --
63 -- Force the new name to be a system name and, if the original was an external name, disambiguate
64 -- the new name with the module name of the original.
65 --
66 mkVectId :: Id -> Type -> VM Id
67 mkVectId id ty
68 = do { name <- mkLocalisedName mkVectOcc (getName id)
69 ; let id' | isDFunId id = MkId.mkDictFunId name tvs theta cls tys
70 | isExportedId id = Id.mkExportedLocalId name ty
71 | otherwise = Id.mkLocalId name ty
72 ; return id'
73 }
74 where
75 -- Decompose a dictionary function signature: \forall tvs. theta -> cls tys
76 -- NB: We do *not* use closures '(:->)' for vectorised predicate abstraction as dictionary
77 -- functions are always fully applied.
78 (tvs, theta, pty) = tcSplitSigmaTy ty
79 (cls, tys) = tcSplitDFunHead pty
80
81 -- |Make a fresh instance of this var, with a new unique.
82 --
83 cloneVar :: Var -> VM Var
84 cloneVar var = liftM (setIdUnique var) (liftDs newUnique)
85
86 -- |Make a fresh exported variable with the given type.
87 --
88 newExportedVar :: OccName -> Type -> VM Var
89 newExportedVar occ_name ty
90 = do mod <- liftDs getModule
91 u <- liftDs newUnique
92
93 let name = mkExternalName u mod occ_name noSrcSpan
94
95 return $ Id.mkExportedLocalId name ty
96
97 -- |Make a fresh local variable with the given type.
98 -- The variable's name is formed using the given string as the prefix.
99 --
100 newLocalVar :: FastString -> Type -> VM Var
101 newLocalVar fs ty
102 = do u <- liftDs newUnique
103 return $ mkSysLocal fs u ty
104
105 -- |Make several fresh local variables with the given types.
106 -- The variable's names are formed using the given string as the prefix.
107 --
108 newLocalVars :: FastString -> [Type] -> VM [Var]
109 newLocalVars fs = mapM (newLocalVar fs)
110
111 -- |Make a new local dummy variable.
112 --
113 newDummyVar :: Type -> VM Var
114 newDummyVar = newLocalVar (fsLit "vv")
115
116 -- |Make a fresh type variable with the given kind.
117 -- The variable's name is formed using the given string as the prefix.
118 --
119 newTyVar :: FastString -> Kind -> VM Var
120 newTyVar fs k
121 = do u <- liftDs newUnique
122 return $ mkTyVar (mkSysTvName u fs) k