Major refactoring of CoAxioms
[ghc.git] / compiler / vectorise / Vectorise / Type / Env.hs
1 -- Vectorise a modules type and class declarations.
2 --
3 -- This produces new type constructors and family instances top be included in the module toplevel
4 -- as well as bindings for worker functions, dfuns, and the like.
5
6 module Vectorise.Type.Env (
7 vectTypeEnv,
8 ) where
9
10 #include "HsVersions.h"
11
12 import Vectorise.Env
13 import Vectorise.Vect
14 import Vectorise.Monad
15 import Vectorise.Builtins
16 import Vectorise.Type.TyConDecl
17 import Vectorise.Type.Classify
18 import Vectorise.Generic.PADict
19 import Vectorise.Generic.PAMethods
20 import Vectorise.Generic.PData
21 import Vectorise.Generic.Description
22 import Vectorise.Utils
23
24 import CoreSyn
25 import CoreUtils
26 import CoreUnfold
27 import DataCon
28 import TyCon
29 import Type
30 import FamInstEnv
31 import Id
32 import MkId
33 import NameEnv
34 import NameSet
35 import OccName
36
37 import Util
38 import Outputable
39 import FastString
40 import MonadUtils
41
42 import Control.Monad
43 import Data.Maybe
44 import Data.List
45
46
47 -- Note [Pragmas to vectorise tycons]
48 -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
49 --
50 -- VECTORISE pragmas for type constructors cover three different flavours of vectorising data type
51 -- constructors:
52 --
53 -- (1) Data type constructor 'T' that may be used in vectorised code, where 'T' represents itself,
54 -- but the representation of 'T' is opaque in vectorised code.
55 --
56 -- An example is the treatment of 'Int'. 'Int's can be used in vectorised code and remain
57 -- unchanged by vectorisation. However, the representation of 'Int' by the 'I#' data
58 -- constructor wrapping an 'Int#' is not exposed in vectorised code. Instead, computations
59 -- involving the representation need to be confined to scalar code.
60 --
61 -- 'PData' and 'PRepr' instances need to be explicitly supplied for 'T' (they are not generated
62 -- by the vectoriser).
63 --
64 -- Type constructors declared with {-# VECTORISE SCALAR type T #-} are treated in this manner.
65 -- (The vectoriser never treats a type constructor automatically in this manner.)
66 --
67 -- (2) Data type constructor 'T' that may be used in vectorised code, where 'T' is represented by an
68 -- explicitly given 'Tv', but the representation of 'T' is opaque in vectorised code.
69 --
70 -- An example is the treatment of '[::]'. '[::]'s can be used in vectorised code and is
71 -- vectorised to 'PArray'. However, the representation of '[::]' is not exposed in vectorised
72 -- code. Instead, computations involving the representation need to be confined to scalar code.
73 --
74 -- 'PData' and 'PRepr' instances need to be explicitly supplied for 'T' (they are not generated
75 -- by the vectoriser).
76 --
77 -- Type constructors declared with {-# VECTORISE SCALAR type T = T' #-} are treated in this
78 -- manner. (The vectoriser never treats a type constructor automatically in this manner.)
79 --
80 -- (3) Data type constructor 'T' that together with its constructors 'Cn' may be used in vectorised
81 -- code, where 'T' and the 'Cn' are automatically vectorised in the same manner as data types
82 -- declared in a vectorised module. This includes the case where the vectoriser determines that
83 -- the original representation of 'T' may be used in vectorised code (as it does not embed any
84 -- parallel arrays.) This case is for type constructors that are *imported* from a non-
85 -- vectorised module, but that we want to use with full vectorisation support.
86 --
87 -- An example is the treatment of 'Ordering' and '[]'. The former remains unchanged by
88 -- vectorisation, whereas the latter is fully vectorised.
89
90 -- 'PData' and 'PRepr' instances are automatically generated by the vectoriser.
91 --
92 -- Type constructors declared with {-# VECTORISE type T #-} are treated in this manner.
93 --
94 -- (4) Data type constructor 'T' that together with its constructors 'Cn' may be used in vectorised
95 -- code, where 'T' is represented by an explicitly given 'Tv' whose constructors 'Cvn' represent
96 -- the original constructors in vectorised code. As a special case, we can have 'Tv = T'
97 --
98 -- An example is the treatment of 'Bool', which is represented by itself in vectorised code
99 -- (as it cannot embed any parallel arrays). However, we do not want any automatic generation
100 -- of class and family instances, which is why Case (3) does not apply.
101 --
102 -- 'PData' and 'PRepr' instances need to be explicitly supplied for 'T' (they are not generated
103 -- by the vectoriser).
104 --
105 -- Type constructors declared with {-# VECTORISE type T = T' #-} are treated in this manner.
106 --
107 -- In addition, we have also got a single pragma form for type classes: {-# VECTORISE class C #-}.
108 -- It implies that the class type constructor may be used in vectorised code together with its data
109 -- constructor. We generally produce a vectorised version of the data type and data constructor.
110 -- We do not generate 'PData' and 'PRepr' instances for class type constructors. This pragma is the
111 -- default for all type classes declared in this module, but the pragma can also be used explitly on
112 -- imported classes.
113
114 -- Note [Vectorising classes]
115 -- ~~~~~~~~~~~~~~~~~~~~~~~~~~
116 --
117 -- We vectorise classes essentially by just vectorising their desugared Core representation, but we
118 -- do generate a 'Class' structure along the way (see 'Vectorise.Type.TyConDecl.vectTyConDecl').
119 --
120 -- Here is an example illustrating the mapping — assume
121 --
122 -- class Num a where
123 -- (+) :: a -> a -> a
124 --
125 -- It desugars to
126 --
127 -- data Num a = D:Num { (+) :: a -> a -> a }
128 --
129 -- which we vectorise to
130 --
131 -- data V:Num a = D:V:Num { ($v+) :: PArray a :-> PArray a :-> PArray a }
132 --
133 -- while adding the following entries to the vectorisation map:
134 --
135 -- tycon : Num --> V:Num
136 -- datacon: D:Num --> D:V:Num
137 -- var : (+) --> ($v+)
138
139 -- |Vectorise type constructor including class type constructors.
140 --
141 vectTypeEnv :: [TyCon] -- Type constructors defined in this module
142 -> [CoreVect] -- All 'VECTORISE [SCALAR] type' declarations in this module
143 -> [CoreVect] -- All 'VECTORISE class' declarations in this module
144 -> VM ( [TyCon] -- old TyCons ++ new TyCons
145 , [FamInst] -- New type family instances.
146 , [(Var, CoreExpr)]) -- New top level bindings.
147 vectTypeEnv tycons vectTypeDecls vectClassDecls
148 = do { traceVt "** vectTypeEnv" $ ppr tycons
149
150 -- Build a map containing all vectorised type constructor. If they are scalar, they are
151 -- mapped to 'False' (vectorised type constructor == original type constructor).
152 ; allScalarTyConNames <- globalScalarTyCons -- covers both current and imported modules
153 ; vectTyCons <- globalVectTyCons
154 ; let vectTyConBase = mapNameEnv (const True) vectTyCons -- by default fully vectorised
155 vectTyConFlavour = foldNameSet (\n env -> extendNameEnv env n False) vectTyConBase
156 allScalarTyConNames
157
158 ; let -- {-# VECTORISE SCALAR type T -#} (imported and local tycons)
159 localAbstractTyCons = [tycon | VectType True tycon Nothing <- vectTypeDecls]
160
161 -- {-# VECTORISE type T -#} (ONLY the imported tycons)
162 impVectTyCons = ( [tycon | VectType False tycon Nothing <- vectTypeDecls]
163 ++ [tycon | VectClass tycon <- vectClassDecls])
164 \\ tycons
165
166 -- {-# VECTORISE [SCALAR] type T = T' -#} (imported and local tycons)
167 vectTyConsWithRHS = [ (tycon, rhs, isAbstract)
168 | VectType isAbstract tycon (Just rhs) <- vectTypeDecls]
169
170 -- filter VECTORISE SCALAR tycons and VECTORISE tycons with explicit rhses
171 vectSpecialTyConNames = mkNameSet . map tyConName $
172 localAbstractTyCons ++ map fst3 vectTyConsWithRHS
173 notVectSpecialTyCon tc = not $ (tyConName tc) `elemNameSet` vectSpecialTyConNames
174
175 -- Split the list of 'TyCons' into the ones (1) that we must vectorise and those (2)
176 -- that we could, but don't need to vectorise. Type constructors that are not data
177 -- type constructors or use non-Haskell98 features are being dropped. They may not
178 -- appear in vectorised code. (We also drop the local type constructors appearing in a
179 -- VECTORISE SCALAR pragma or a VECTORISE pragma with an explicit right-hand side, as
180 -- these are being handled separately. NB: Some type constructors may be marked SCALAR
181 -- /and/ have an explicit right-hand side.)
182 --
183 -- Furthermore, 'drop_tcs' are those type constructors that we cannot vectorise.
184 ; let maybeVectoriseTyCons = filter notVectSpecialTyCon tycons ++ impVectTyCons
185 (conv_tcs, keep_tcs, drop_tcs) = classifyTyCons vectTyConFlavour maybeVectoriseTyCons
186
187 ; traceVt " VECT SCALAR : " $ ppr localAbstractTyCons
188 ; traceVt " VECT [class] : " $ ppr impVectTyCons
189 ; traceVt " VECT with rhs : " $ ppr (map fst3 vectTyConsWithRHS)
190 ; traceVt " -- after classification (local and VECT [class] tycons) --" empty
191 ; traceVt " reuse : " $ ppr keep_tcs
192 ; traceVt " convert : " $ ppr conv_tcs
193
194 -- warn the user about unvectorised type constructors
195 ; let explanation = ptext (sLit "(They use unsupported language extensions") $$
196 ptext (sLit "or depend on type constructors that are not vectorised)")
197 drop_tcs_nosyn = filter (not . isSynTyCon) drop_tcs
198 ; unless (null drop_tcs_nosyn) $
199 emitVt "Warning: cannot vectorise these type constructors:" $
200 pprQuotedList drop_tcs_nosyn $$ explanation
201
202 ; mapM_ addGlobalScalarTyCon keep_tcs
203
204 ; let mapping =
205 -- Type constructors that we don't need to vectorise, use the same
206 -- representation in both unvectorised and vectorised code; they are not
207 -- abstract.
208 [(tycon, tycon, False) | tycon <- keep_tcs]
209 -- We do the same for type constructors declared VECTORISE SCALAR /without/
210 -- an explicit right-hand side, but ignore their representation (data
211 -- constructors) as they are abstract.
212 ++ [(tycon, tycon, True) | tycon <- localAbstractTyCons]
213 -- Type constructors declared VECTORISE /with/ an explicit vectorised type,
214 -- we map from the original to the given type; whether they are abstract depends
215 -- on whether the vectorisation declaration was SCALAR.
216 ++ vectTyConsWithRHS
217 ; syn_tcs <- catMaybes <$> mapM defTyConDataCons mapping
218
219 -- Vectorise all the data type declarations that we can and must vectorise (enter the
220 -- type and data constructors into the vectorisation map on-the-fly.)
221 ; new_tcs <- vectTyConDecls conv_tcs
222
223 -- We don't need new representation types for dictionary constructors. The constructors
224 -- are always fully applied, and we don't need to lift them to arrays as a dictionary
225 -- of a particular type always has the same value.
226 ; let orig_tcs = filter (not . isClassTyCon) $ keep_tcs ++ conv_tcs
227 vect_tcs = filter (not . isClassTyCon) $ keep_tcs ++ new_tcs
228
229 -- Build 'PRepr' and 'PData' instance type constructors and family instances for all
230 -- type constructors with vectorised representations.
231 ; reprs <- mapM tyConRepr vect_tcs
232 ; repr_fis <- zipWith3M buildPReprTyCon orig_tcs vect_tcs reprs
233 ; pdata_fis <- zipWith3M buildPDataTyCon orig_tcs vect_tcs reprs
234 ; pdatas_fis <- zipWith3M buildPDatasTyCon orig_tcs vect_tcs reprs
235
236 ; let fam_insts = repr_fis ++ pdata_fis ++ pdatas_fis
237 repr_axs = map famInstAxiom repr_fis
238 pdata_tcs = famInstsRepTyCons pdata_fis
239 pdatas_tcs = famInstsRepTyCons pdatas_fis
240
241 ; updGEnv $ extendFamEnv fam_insts
242
243 -- Generate workers for the vectorised data constructors, dfuns for the 'PA' instances of
244 -- the vectorised type constructors, and associate the type constructors with their dfuns
245 -- in the global environment. We get back the dfun bindings (which we will subsequently
246 -- inject into the modules toplevel).
247 ; (_, binds) <- fixV $ \ ~(dfuns, _) ->
248 do { defTyConPAs (zipLazy vect_tcs dfuns)
249
250 -- Query the 'PData' instance type constructors for type constructors that have a
251 -- VECTORISE pragma with an explicit right-hand side (this is Item (4) of
252 -- "Note [Pragmas to vectorise tycons]" above).
253 ; let (withRHS_non_abstract, vwithRHS_non_abstract)
254 = unzip [(tycon, vtycon) | (tycon, vtycon, False) <- vectTyConsWithRHS]
255 ; pdata_withRHS_tcs <- mapM pdataReprTyConExact withRHS_non_abstract
256
257 -- Build workers for all vectorised data constructors (except abstract ones)
258 ; sequence_ $
259 zipWith3 vectDataConWorkers (orig_tcs ++ withRHS_non_abstract)
260 (vect_tcs ++ vwithRHS_non_abstract)
261 (pdata_tcs ++ pdata_withRHS_tcs)
262
263 -- Build a 'PA' dictionary for all type constructors (except abstract ones & those
264 -- defined with an explicit right-hand side where the dictionary is user-supplied)
265 ; dfuns <- sequence $
266 zipWith4 buildTyConPADict
267 vect_tcs
268 repr_axs
269 pdata_tcs
270 pdatas_tcs
271
272 ; binds <- takeHoisted
273 ; return (dfuns, binds)
274 }
275
276 -- Return the vectorised variants of type constructors as well as the generated instance
277 -- type constructors, family instances, and dfun bindings.
278 ; return ( new_tcs ++ pdata_tcs ++ pdatas_tcs ++ syn_tcs
279 , fam_insts, binds)
280 }
281 where
282 fst3 (a, _, _) = a
283
284 -- Add a mapping from the original to vectorised type constructor to the vectorisation map.
285 -- Unless the type constructor is abstract, also mappings from the orignal's data constructors
286 -- to the vectorised type's data constructors.
287 --
288 -- We have three cases: (1) original and vectorised type constructor are the same, (2) the
289 -- name of the vectorised type constructor is canonical (as prescribed by 'mkVectTyConOcc'), or
290 -- (3) the name is not canonical. In the third case, we additionally introduce a type synonym
291 -- with the canonical name that is set equal to the non-canonical name (so that we find the
292 -- right type constructor when reading vectorisation information from interface files).
293 --
294 defTyConDataCons (origTyCon, vectTyCon, isAbstract)
295 = do { canonName <- mkLocalisedName mkVectTyConOcc origName
296 ; if origName == vectName -- Case (1)
297 || vectName == canonName -- Case (2)
298 then do
299 { defTyCon origTyCon vectTyCon -- T --> vT
300 ; defDataCons -- Ci --> vCi
301 ; return Nothing
302 }
303 else do -- Case (3)
304 { let synTyCon = mkSyn canonName (mkTyConTy vectTyCon) -- type S = vT
305 ; defTyCon origTyCon synTyCon -- T --> S
306 ; defDataCons -- Ci --> vCi
307 ; return $ Just synTyCon
308 }
309 }
310 where
311 origName = tyConName origTyCon
312 vectName = tyConName vectTyCon
313
314 mkSyn canonName ty = mkSynTyCon canonName (typeKind ty) [] (SynonymTyCon ty) NoParentTyCon
315
316 defDataCons
317 | isAbstract = return ()
318 | otherwise
319 = do { MASSERT(length (tyConDataCons origTyCon) == length (tyConDataCons vectTyCon))
320 ; zipWithM_ defDataCon (tyConDataCons origTyCon) (tyConDataCons vectTyCon)
321 }
322
323
324 -- Helpers --------------------------------------------------------------------
325
326 buildTyConPADict :: TyCon -> CoAxiom -> TyCon -> TyCon -> VM Var
327 buildTyConPADict vect_tc prepr_ax pdata_tc pdatas_tc
328 = tyConRepr vect_tc >>= buildPADict vect_tc prepr_ax pdata_tc pdatas_tc
329
330 -- Produce a custom-made worker for the data constructors of a vectorised data type. This includes
331 -- all data constructors that may be used in vetcorised code — i.e., all data constructors of data
332 -- types other than scalar ones. Also adds a mapping from the original to vectorised worker into
333 -- the vectorisation map.
334 --
335 -- FIXME: It's not nice that we need create a special worker after the data constructors has
336 -- already been constructed. Also, I don't think the worker is properly added to the data
337 -- constructor. Seems messy.
338 vectDataConWorkers :: TyCon -> TyCon -> TyCon -> VM ()
339 vectDataConWorkers orig_tc vect_tc arr_tc
340 = do { traceVt "Building vectorised worker for datatype" (ppr orig_tc)
341
342 ; bs <- sequence
343 . zipWith3 def_worker (tyConDataCons orig_tc) rep_tys
344 $ zipWith4 mk_data_con (tyConDataCons vect_tc)
345 rep_tys
346 (inits rep_tys)
347 (tail $ tails rep_tys)
348 ; mapM_ (uncurry hoistBinding) bs
349 }
350 where
351 tyvars = tyConTyVars vect_tc
352 var_tys = mkTyVarTys tyvars
353 ty_args = map Type var_tys
354 res_ty = mkTyConApp vect_tc var_tys
355
356 cons = tyConDataCons vect_tc
357 arity = length cons
358 [arr_dc] = tyConDataCons arr_tc
359
360 rep_tys = map dataConRepArgTys $ tyConDataCons vect_tc
361
362 mk_data_con con tys pre post
363 = liftM2 (,) (vect_data_con con)
364 (lift_data_con tys pre post (mkDataConTag con))
365
366 sel_replicate len tag
367 | arity > 1 = do
368 rep <- builtin (selReplicate arity)
369 return [rep `mkApps` [len, tag]]
370
371 | otherwise = return []
372
373 vect_data_con con = return $ mkConApp con ty_args
374 lift_data_con tys pre_tys post_tys tag
375 = do
376 len <- builtin liftingContext
377 args <- mapM (newLocalVar (fsLit "xs"))
378 =<< mapM mkPDataType tys
379
380 sel <- sel_replicate (Var len) tag
381
382 pre <- mapM emptyPD (concat pre_tys)
383 post <- mapM emptyPD (concat post_tys)
384
385 return . mkLams (len : args)
386 . wrapFamInstBody arr_tc var_tys
387 . mkConApp arr_dc
388 $ ty_args ++ sel ++ pre ++ map Var args ++ post
389
390 def_worker data_con arg_tys mk_body
391 = do
392 arity <- polyArity tyvars
393 body <- closedV
394 . inBind orig_worker
395 . polyAbstract tyvars $ \args ->
396 liftM (mkLams (tyvars ++ args) . vectorised)
397 $ buildClosures tyvars [] [] arg_tys res_ty mk_body
398
399 raw_worker <- mkVectId orig_worker (exprType body)
400 let vect_worker = raw_worker `setIdUnfolding`
401 mkInlineUnfolding (Just arity) body
402 defGlobalVar orig_worker vect_worker
403 return (vect_worker, body)
404 where
405 orig_worker = dataConWorkId data_con