Change buildToPRepr to work with the new representation scheme
[ghc-hetmet.git] / compiler / vectorise / VectType.hs
1 module VectType ( vectTyCon, vectType, vectTypeEnv,
2                    PAInstance, buildPADict )
3 where
4
5 #include "HsVersions.h"
6
7 import VectMonad
8 import VectUtils
9 import VectCore
10
11 import HscTypes          ( TypeEnv, extendTypeEnvList, typeEnvTyCons )
12 import CoreSyn
13 import CoreUtils
14 import BuildTyCl
15 import DataCon
16 import TyCon
17 import Type
18 import TypeRep
19 import Coercion
20 import FamInstEnv        ( FamInst, mkLocalFamInst )
21 import InstEnv           ( Instance, mkLocalInstance, instanceDFunId )
22 import OccName
23 import MkId
24 import BasicTypes        ( StrictnessMark(..), OverlapFlag(..), boolToRecFlag )
25 import Var               ( Var )
26 import Id                ( mkWildId )
27 import Name              ( Name, getOccName )
28 import NameEnv
29 import TysWiredIn        ( unitTy, intTy, intDataCon )
30 import TysPrim           ( intPrimTy )
31
32 import Unique
33 import UniqFM
34 import UniqSet
35 import Digraph           ( SCC(..), stronglyConnComp )
36
37 import Outputable
38
39 import Control.Monad  ( liftM, liftM2, zipWithM, zipWithM_ )
40 import Data.List      ( inits, tails, zipWith4, zipWith5 )
41
42 -- ----------------------------------------------------------------------------
43 -- Types
44
45 vectTyCon :: TyCon -> VM TyCon
46 vectTyCon tc
47   | isFunTyCon tc        = builtin closureTyCon
48   | isBoxedTupleTyCon tc = return tc
49   | isUnLiftedTyCon tc   = return tc
50   | otherwise = do
51                   r <- lookupTyCon tc
52                   case r of
53                     Just tc' -> return tc'
54
55                     -- FIXME: just for now
56                     Nothing  -> pprTrace "ccTyCon:" (ppr tc) $ return tc
57
58 vectType :: Type -> VM Type
59 vectType ty | Just ty' <- coreView ty = vectType ty'
60 vectType (TyVarTy tv) = return $ TyVarTy tv
61 vectType (AppTy ty1 ty2) = liftM2 AppTy (vectType ty1) (vectType ty2)
62 vectType (TyConApp tc tys) = liftM2 TyConApp (vectTyCon tc) (mapM vectType tys)
63 vectType (FunTy ty1 ty2)   = liftM2 TyConApp (builtin closureTyCon)
64                                              (mapM vectType [ty1,ty2])
65 vectType ty@(ForAllTy _ _)
66   = do
67       mdicts   <- mapM paDictArgType tyvars
68       mono_ty' <- vectType mono_ty
69       return $ tyvars `mkForAllTys` ([dict | Just dict <- mdicts] `mkFunTys` mono_ty')
70   where
71     (tyvars, mono_ty) = splitForAllTys ty
72
73 vectType ty = pprPanic "vectType:" (ppr ty)
74
75 -- ----------------------------------------------------------------------------
76 -- Type definitions
77
78 type TyConGroup = ([TyCon], UniqSet TyCon)
79
80 data PAInstance = PAInstance {
81                     painstDFun      :: Var
82                   , painstOrigTyCon :: TyCon
83                   , painstVectTyCon :: TyCon
84                   , painstArrTyCon  :: TyCon
85                   }
86
87 vectTypeEnv :: TypeEnv -> VM (TypeEnv, [FamInst], [(Var, CoreExpr)])
88 vectTypeEnv env
89   = do
90       cs <- readGEnv $ mk_map . global_tycons
91       let (conv_tcs, keep_tcs) = classifyTyCons cs groups
92           keep_dcs             = concatMap tyConDataCons keep_tcs
93       zipWithM_ defTyCon   keep_tcs keep_tcs
94       zipWithM_ defDataCon keep_dcs keep_dcs
95       new_tcs <- vectTyConDecls conv_tcs
96
97       let orig_tcs = keep_tcs ++ conv_tcs
98           vect_tcs  = keep_tcs ++ new_tcs
99
100       repr_tcs <- zipWithM buildPReprTyCon   orig_tcs vect_tcs
101       parr_tcs <- zipWithM buildPArrayTyCon orig_tcs vect_tcs
102       dfuns    <- mapM mkPADFun vect_tcs
103       defTyConPAs (zip vect_tcs dfuns)
104       binds    <- sequence (zipWith5 buildTyConBindings orig_tcs
105                                                         vect_tcs
106                                                         repr_tcs
107                                                         parr_tcs
108                                                         dfuns)
109
110       let all_new_tcs = new_tcs ++ repr_tcs ++ parr_tcs
111
112       let new_env = extendTypeEnvList env
113                        (map ATyCon all_new_tcs
114                         ++ [ADataCon dc | tc <- all_new_tcs
115                                         , dc <- tyConDataCons tc])
116
117       return (new_env, map mkLocalFamInst (repr_tcs ++ parr_tcs), concat binds)
118   where
119     tycons = typeEnvTyCons env
120     groups = tyConGroups tycons
121
122     mk_map env = listToUFM_Directly [(u, getUnique n /= u) | (u,n) <- nameEnvUniqueElts env]
123
124     keep_tc tc = let dcs = tyConDataCons tc
125                  in
126                  defTyCon tc tc >> zipWithM_ defDataCon dcs dcs
127
128
129 vectTyConDecls :: [TyCon] -> VM [TyCon]
130 vectTyConDecls tcs = fixV $ \tcs' ->
131   do
132     mapM_ (uncurry defTyCon) (lazy_zip tcs tcs')
133     mapM vectTyConDecl tcs
134   where
135     lazy_zip [] _ = []
136     lazy_zip (x:xs) ~(y:ys) = (x,y) : lazy_zip xs ys
137
138 vectTyConDecl :: TyCon -> VM TyCon
139 vectTyConDecl tc
140   = do
141       name' <- cloneName mkVectTyConOcc name
142       rhs'  <- vectAlgTyConRhs (algTyConRhs tc)
143
144       liftDs $ buildAlgTyCon name'
145                              tyvars
146                              []           -- no stupid theta
147                              rhs'
148                              rec_flag     -- FIXME: is this ok?
149                              False        -- FIXME: no generics
150                              False        -- not GADT syntax
151                              Nothing      -- not a family instance
152   where
153     name   = tyConName tc
154     tyvars = tyConTyVars tc
155     rec_flag = boolToRecFlag (isRecursiveTyCon tc)
156
157 vectAlgTyConRhs :: AlgTyConRhs -> VM AlgTyConRhs
158 vectAlgTyConRhs (DataTyCon { data_cons = data_cons
159                            , is_enum   = is_enum
160                            })
161   = do
162       data_cons' <- mapM vectDataCon data_cons
163       zipWithM_ defDataCon data_cons data_cons'
164       return $ DataTyCon { data_cons = data_cons'
165                          , is_enum   = is_enum
166                          }
167
168 vectDataCon :: DataCon -> VM DataCon
169 vectDataCon dc
170   | not . null $ dataConExTyVars dc = pprPanic "vectDataCon: existentials" (ppr dc)
171   | not . null $ dataConEqSpec   dc = pprPanic "vectDataCon: eq spec" (ppr dc)
172   | otherwise
173   = do
174       name'    <- cloneName mkVectDataConOcc name
175       tycon'   <- vectTyCon tycon
176       arg_tys  <- mapM vectType rep_arg_tys
177
178       liftDs $ buildDataCon name'
179                             False           -- not infix
180                             (map (const NotMarkedStrict) arg_tys)
181                             []              -- no labelled fields
182                             univ_tvs
183                             []              -- no existential tvs for now
184                             []              -- no eq spec for now
185                             []              -- no context
186                             arg_tys
187                             tycon'
188   where
189     name        = dataConName dc
190     univ_tvs    = dataConUnivTyVars dc
191     rep_arg_tys = dataConRepArgTys dc
192     tycon       = dataConTyCon dc
193
194 mk_fam_inst :: TyCon -> TyCon -> (TyCon, [Type])
195 mk_fam_inst fam_tc arg_tc
196   = (fam_tc, [mkTyConApp arg_tc . mkTyVarTys $ tyConTyVars arg_tc])
197
198 buildPReprTyCon :: TyCon -> TyCon -> VM TyCon
199 buildPReprTyCon orig_tc vect_tc
200   = do
201       name     <- cloneName mkPReprTyConOcc (tyConName orig_tc)
202       rhs_ty   <- buildPReprType vect_tc
203       prepr_tc <- builtin preprTyCon
204       liftDs $ buildSynTyCon name
205                              tyvars
206                              (SynonymTyCon rhs_ty)
207                              (Just $ mk_fam_inst prepr_tc vect_tc)
208   where
209     tyvars = tyConTyVars vect_tc
210
211 buildPReprType :: TyCon -> VM Type
212 buildPReprType = liftM repr_type . mkTyConRepr
213
214 buildToPRepr :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
215 buildToPRepr repr vect_tc prepr_tc _
216   = do
217       arg <- newLocalVar FSLIT("x") arg_ty
218       bndrss <- mapM (mapM (newLocalVar FSLIT("x")))
219                      (repr_tys repr)
220
221       return . Lam arg
222              . wrapFamInstBody prepr_tc var_tys
223              . Case (Var arg) (mkWildId arg_ty) (repr_type repr)
224              . zipWith3 mk_alt data_cons bndrss
225              . mkToPRepr repr $ map (map Var) bndrss
226   where
227     var_tys   = mkTyVarTys $ tyConTyVars vect_tc
228     arg_ty    = mkTyConApp vect_tc var_tys
229     data_cons = tyConDataCons vect_tc
230
231     mk_alt data_con bndrs body = (DataAlt data_con, bndrs, body)
232
233 buildToArrPRepr :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
234 buildToArrPRepr _ vect_tc prepr_tc arr_tc
235   = do
236       arg_ty  <- mkPArrayType el_ty
237       rep_tys <- mapM (mapM mkPArrayType) rep_el_tys
238
239       arg     <- newLocalVar FSLIT("xs") arg_ty
240       bndrss  <- mapM (mapM (newLocalVar FSLIT("ys"))) rep_tys
241       len     <- newLocalVar FSLIT("len") intPrimTy
242       sel     <- newLocalVar FSLIT("sel") =<< mkPArrayType intTy
243
244       let add_sel xs | has_selector = sel : xs
245                      | otherwise    = xs
246
247           all_bndrs = len : add_sel (concat bndrss)
248
249       res      <- parrayCoerce prepr_tc var_tys
250                 =<< mkToArrPRepr (Var len) (Var sel) (map (map Var) bndrss)
251       res_ty   <- mkPArrayType =<< mkPReprType el_ty
252
253       return . Lam arg
254              $ Case (unwrapFamInstScrut arr_tc var_tys (Var arg))
255                     (mkWildId (mkTyConApp arr_tc var_tys))
256                     res_ty
257                     [(DataAlt arr_dc, all_bndrs, res)]
258   where
259     var_tys    = mkTyVarTys $ tyConTyVars vect_tc
260     el_ty      = mkTyConApp vect_tc var_tys
261     data_cons  = tyConDataCons vect_tc
262     rep_el_tys = map dataConRepArgTys data_cons
263
264     [arr_dc]   = tyConDataCons arr_tc
265
266     has_selector | [_] <- data_cons = False
267                  | otherwise        = True
268
269
270 buildFromPRepr :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
271 buildFromPRepr _ vect_tc prepr_tc _
272   = do
273       arg_ty <- mkPReprType res_ty
274       arg <- newLocalVar FSLIT("x") arg_ty
275       alts <- mapM mk_alt data_cons
276       body <- mkFromPRepr (unwrapFamInstScrut prepr_tc var_tys (Var arg))
277                           res_ty alts
278       return $ Lam arg body
279   where
280     var_tys   = mkTyVarTys $ tyConTyVars vect_tc
281     res_ty    = mkTyConApp vect_tc var_tys
282     data_cons = tyConDataCons vect_tc
283
284     mk_alt dc = do
285                   bndrs <- mapM (newLocalVar FSLIT("x")) $ dataConRepArgTys dc
286                   return (bndrs, mkConApp dc (map Type var_tys ++ map Var bndrs))
287
288 buildFromArrPRepr :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
289 buildFromArrPRepr _ vect_tc prepr_tc arr_tc
290   = mkFromArrPRepr undefined undefined undefined undefined undefined undefined
291
292 buildPRDict :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
293 buildPRDict _ vect_tc prepr_tc _
294   = prCoerce prepr_tc var_tys
295   =<< prDictOfType (mkTyConApp prepr_tc var_tys)
296   where
297     var_tys = mkTyVarTys $ tyConTyVars vect_tc
298
299 buildPArrayTyCon :: TyCon -> TyCon -> VM TyCon
300 buildPArrayTyCon orig_tc vect_tc = fixV $ \repr_tc ->
301   do
302     name'  <- cloneName mkPArrayTyConOcc orig_name
303     rhs    <- buildPArrayTyConRhs orig_name vect_tc repr_tc
304     parray <- builtin parrayTyCon
305
306     liftDs $ buildAlgTyCon name'
307                            tyvars
308                            []          -- no stupid theta
309                            rhs
310                            rec_flag    -- FIXME: is this ok?
311                            False       -- FIXME: no generics
312                            False       -- not GADT syntax
313                            (Just $ mk_fam_inst parray vect_tc)
314   where
315     orig_name = tyConName orig_tc
316     tyvars = tyConTyVars vect_tc
317     rec_flag = boolToRecFlag (isRecursiveTyCon vect_tc)
318     
319
320 buildPArrayTyConRhs :: Name -> TyCon -> TyCon -> VM AlgTyConRhs
321 buildPArrayTyConRhs orig_name vect_tc repr_tc
322   = do
323       data_con <- buildPArrayDataCon orig_name vect_tc repr_tc
324       return $ DataTyCon { data_cons = [data_con], is_enum = False }
325
326 buildPArrayDataCon :: Name -> TyCon -> TyCon -> VM DataCon
327 buildPArrayDataCon orig_name vect_tc repr_tc
328   = do
329       dc_name  <- cloneName mkPArrayDataConOcc orig_name
330       shape    <- tyConShape vect_tc
331       repr_tys <- mapM mkPArrayType types
332
333       liftDs $ buildDataCon dc_name
334                             False                  -- not infix
335                             (shapeStrictness shape ++ map (const NotMarkedStrict) repr_tys)
336                             []                     -- no field labels
337                             (tyConTyVars vect_tc)
338                             []                     -- no existentials
339                             []                     -- no eq spec
340                             []                     -- no context
341                             (shapeReprTys shape ++ repr_tys)
342                             repr_tc
343   where
344     types = [ty | dc <- tyConDataCons vect_tc
345                 , ty <- dataConRepArgTys dc]
346
347 mkPADFun :: TyCon -> VM Var
348 mkPADFun vect_tc
349   = newExportedVar (mkPADFunOcc $ getOccName vect_tc) =<< paDFunType vect_tc
350
351 data Shape = Shape {
352                shapeReprTys    :: [Type]
353              , shapeStrictness :: [StrictnessMark]
354              , shapeLength     :: [CoreExpr] -> VM CoreExpr
355              , shapeReplicate  :: CoreExpr -> CoreExpr -> VM [CoreExpr]
356              }
357
358 tyConShape :: TyCon -> VM Shape
359 tyConShape vect_tc
360   | isProductTyCon vect_tc
361   = return $ Shape {
362                 shapeReprTys    = [intPrimTy]
363               , shapeStrictness = [NotMarkedStrict]
364               , shapeLength     = \[len] -> return len
365               , shapeReplicate  = \len _ -> return [len]
366               }
367
368   | otherwise
369   = do
370       repr_ty <- mkPArrayType intTy   -- FIXME: we want to unbox this
371       return $ Shape {
372                  shapeReprTys    = [repr_ty]
373                , shapeStrictness = [MarkedStrict]
374                , shapeLength     = \[sel] -> lengthPA sel
375                , shapeReplicate  = \len n -> do
376                                                e <- replicatePA len n
377                                                return [e]
378                }
379
380 buildTyConBindings :: TyCon -> TyCon -> TyCon -> TyCon -> Var
381                    -> VM [(Var, CoreExpr)]
382 buildTyConBindings orig_tc vect_tc prepr_tc arr_tc dfun
383   = do
384       shape <- tyConShape vect_tc
385       repr  <- mkTyConRepr vect_tc
386       sequence_ (zipWith4 (vectDataConWorker shape vect_tc arr_tc arr_dc)
387                           orig_dcs
388                           vect_dcs
389                           (inits repr_tys)
390                           (tails repr_tys))
391       dict <- buildPADict repr vect_tc prepr_tc arr_tc dfun
392       binds <- takeHoisted
393       return $ (dfun, dict) : binds
394   where
395     orig_dcs = tyConDataCons orig_tc
396     vect_dcs = tyConDataCons vect_tc
397     [arr_dc] = tyConDataCons arr_tc
398
399     repr_tys = map dataConRepArgTys vect_dcs
400
401 vectDataConWorker :: Shape -> TyCon -> TyCon -> DataCon
402                   -> DataCon -> DataCon -> [[Type]] -> [[Type]]
403                   -> VM ()
404 vectDataConWorker shape vect_tc arr_tc arr_dc orig_dc vect_dc pre (dc_tys : post)
405   = do
406       clo <- closedV
407            . inBind orig_worker
408            . polyAbstract tvs $ \abstract ->
409              liftM (abstract . vectorised)
410            $ buildClosures tvs [] dc_tys res_ty (liftM2 (,) mk_vect mk_lift)
411
412       worker <- cloneId mkVectOcc orig_worker (exprType clo)
413       hoistBinding worker clo
414       defGlobalVar orig_worker worker
415       return ()
416   where
417     tvs     = tyConTyVars vect_tc
418     arg_tys = mkTyVarTys tvs
419     res_ty  = mkTyConApp vect_tc arg_tys
420
421     orig_worker = dataConWorkId orig_dc
422
423     mk_vect = return . mkConApp vect_dc $ map Type arg_tys
424     mk_lift = do
425                 len     <- newLocalVar FSLIT("n") intPrimTy
426                 arr_tys <- mapM mkPArrayType dc_tys
427                 args    <- mapM (newLocalVar FSLIT("xs")) arr_tys
428                 shapes  <- shapeReplicate shape
429                                           (Var len)
430                                           (mkDataConTag vect_dc)
431                 
432                 empty_pre  <- mapM emptyPA (concat pre)
433                 empty_post <- mapM emptyPA (concat post)
434
435                 return . mkLams (len : args)
436                        . wrapFamInstBody arr_tc arg_tys
437                        . mkConApp arr_dc
438                        $ map Type arg_tys ++ shapes
439                                           ++ empty_pre
440                                           ++ map Var args
441                                           ++ empty_post
442
443 buildPADict :: TyConRepr -> TyCon -> TyCon -> TyCon -> Var -> VM CoreExpr
444 buildPADict repr vect_tc prepr_tc arr_tc dfun
445   = polyAbstract tvs $ \abstract ->
446     do
447       meth_binds <- mapM (mk_method repr) paMethods
448       let meth_exprs = map (Var . fst) meth_binds
449
450       pa_dc <- builtin paDataCon
451       let dict = mkConApp pa_dc (Type (mkTyConApp vect_tc arg_tys) : meth_exprs)
452           body = Let (Rec meth_binds) dict
453       return . mkInlineMe $ abstract body
454   where
455     tvs = tyConTyVars arr_tc
456     arg_tys = mkTyVarTys tvs
457
458     mk_method repr (name, build)
459       = localV
460       $ do
461           body <- build repr vect_tc prepr_tc arr_tc
462           var  <- newLocalVar name (exprType body)
463           return (var, mkInlineMe body)
464           
465 paMethods = [(FSLIT("toPRepr"),      buildToPRepr),
466              (FSLIT("fromPRepr"),    buildFromPRepr),
467              (FSLIT("toArrPRepr"),   buildToArrPRepr),
468              (FSLIT("fromArrPRepr"), buildFromArrPRepr),
469              (FSLIT("dictPRepr"),    buildPRDict)]
470
471 -- | Split the given tycons into two sets depending on whether they have to be
472 -- converted (first list) or not (second list). The first argument contains
473 -- information about the conversion status of external tycons:
474 -- 
475 --   * tycons which have converted versions are mapped to True
476 --   * tycons which are not changed by vectorisation are mapped to False
477 --   * tycons which can't be converted are not elements of the map
478 --
479 classifyTyCons :: UniqFM Bool -> [TyConGroup] -> ([TyCon], [TyCon])
480 classifyTyCons = classify [] []
481   where
482     classify conv keep cs [] = (conv, keep)
483     classify conv keep cs ((tcs, ds) : rs)
484       | can_convert && must_convert
485         = classify (tcs ++ conv) keep (cs `addListToUFM` [(tc,True) | tc <- tcs]) rs
486       | can_convert
487         = classify conv (tcs ++ keep) (cs `addListToUFM` [(tc,False) | tc <- tcs]) rs
488       | otherwise
489         = classify conv keep cs rs
490       where
491         refs = ds `delListFromUniqSet` tcs
492
493         can_convert  = isNullUFM (refs `minusUFM` cs) && all convertable tcs
494         must_convert = foldUFM (||) False (intersectUFM_C const cs refs)
495
496         convertable tc = isDataTyCon tc && all isVanillaDataCon (tyConDataCons tc)
497     
498 -- | Compute mutually recursive groups of tycons in topological order
499 --
500 tyConGroups :: [TyCon] -> [TyConGroup]
501 tyConGroups tcs = map mk_grp (stronglyConnComp edges)
502   where
503     edges = [((tc, ds), tc, uniqSetToList ds) | tc <- tcs
504                                 , let ds = tyConsOfTyCon tc]
505
506     mk_grp (AcyclicSCC (tc, ds)) = ([tc], ds)
507     mk_grp (CyclicSCC els)       = (tcs, unionManyUniqSets dss)
508       where
509         (tcs, dss) = unzip els
510
511 tyConsOfTyCon :: TyCon -> UniqSet TyCon
512 tyConsOfTyCon 
513   = tyConsOfTypes . concatMap dataConRepArgTys . tyConDataCons
514
515 tyConsOfType :: Type -> UniqSet TyCon
516 tyConsOfType ty
517   | Just ty' <- coreView ty    = tyConsOfType ty'
518 tyConsOfType (TyVarTy v)       = emptyUniqSet
519 tyConsOfType (TyConApp tc tys) = extend (tyConsOfTypes tys)
520   where
521     extend | isUnLiftedTyCon tc
522            || isTupleTyCon   tc = id
523
524            | otherwise          = (`addOneToUniqSet` tc)
525
526 tyConsOfType (AppTy a b)       = tyConsOfType a `unionUniqSets` tyConsOfType b
527 tyConsOfType (FunTy a b)       = (tyConsOfType a `unionUniqSets` tyConsOfType b)
528                                  `addOneToUniqSet` funTyCon
529 tyConsOfType (ForAllTy _ ty)   = tyConsOfType ty
530 tyConsOfType other             = pprPanic "ClosureConv.tyConsOfType" $ ppr other
531
532 tyConsOfTypes :: [Type] -> UniqSet TyCon
533 tyConsOfTypes = unionManyUniqSets . map tyConsOfType
534