Move code
[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, unitDataConId )
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 (TyConRepr {
216                 repr_tys         = repr_tys
217               , repr_prod_tycons = prod_tycons
218               , repr_prod_tys    = prod_tys
219               , repr_sum_tycon   = repr_sum_tycon
220               , repr_type        = repr_type
221               })
222               vect_tc prepr_tc _
223   = do
224       arg  <- newLocalVar FSLIT("x") arg_ty
225       vars <- mapM (mapM (newLocalVar FSLIT("x"))) repr_tys
226
227       return . Lam arg
228              . wrapFamInstBody prepr_tc var_tys
229              . Case (Var arg) (mkWildId arg_ty) repr_type
230              . mk_alts data_cons vars
231              . zipWith3 mk_prod prod_tycons repr_tys $ map (map Var) vars
232   where
233     var_tys   = mkTyVarTys $ tyConTyVars vect_tc
234     arg_ty    = mkTyConApp vect_tc var_tys
235     data_cons = tyConDataCons vect_tc
236
237     Just sum_tycon = repr_sum_tycon
238     sum_datacons   = tyConDataCons sum_tycon
239
240     mk_alts _    _      []     = [(DEFAULT, [], Var unitDataConId)]
241     mk_alts [dc] [vars] [expr] = [(DataAlt dc, vars, expr)]
242     mk_alts dcs  vars   exprs  = zipWith4 mk_alt dcs vars sum_datacons exprs 
243
244     mk_alt dc vars sum_dc expr = (DataAlt dc, vars,
245                                   mkConApp sum_dc (map Type prod_tys ++ [expr]))
246
247     mk_prod _         _   []     = Var unitDataConId
248     mk_prod _         _   [expr] = expr
249     mk_prod (Just tc) tys exprs  = mkConApp dc (map Type tys ++ exprs)
250       where
251         [dc] = tyConDataCons tc
252
253 buildToArrPRepr :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
254 buildToArrPRepr _ vect_tc prepr_tc arr_tc
255   = do
256       arg_ty  <- mkPArrayType el_ty
257       rep_tys <- mapM (mapM mkPArrayType) rep_el_tys
258
259       arg     <- newLocalVar FSLIT("xs") arg_ty
260       bndrss  <- mapM (mapM (newLocalVar FSLIT("ys"))) rep_tys
261       len     <- newLocalVar FSLIT("len") intPrimTy
262       sel     <- newLocalVar FSLIT("sel") =<< mkPArrayType intTy
263
264       let add_sel xs | has_selector = sel : xs
265                      | otherwise    = xs
266
267           all_bndrs = len : add_sel (concat bndrss)
268
269       res      <- parrayCoerce prepr_tc var_tys
270                 =<< mkToArrPRepr (Var len) (Var sel) (map (map Var) bndrss)
271       res_ty   <- mkPArrayType =<< mkPReprType el_ty
272
273       return . Lam arg
274              $ Case (unwrapFamInstScrut arr_tc var_tys (Var arg))
275                     (mkWildId (mkTyConApp arr_tc var_tys))
276                     res_ty
277                     [(DataAlt arr_dc, all_bndrs, res)]
278   where
279     var_tys    = mkTyVarTys $ tyConTyVars vect_tc
280     el_ty      = mkTyConApp vect_tc var_tys
281     data_cons  = tyConDataCons vect_tc
282     rep_el_tys = map dataConRepArgTys data_cons
283
284     [arr_dc]   = tyConDataCons arr_tc
285
286     has_selector | [_] <- data_cons = False
287                  | otherwise        = True
288
289
290 buildFromPRepr :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
291 buildFromPRepr _ vect_tc prepr_tc _
292   = do
293       arg_ty <- mkPReprType res_ty
294       arg <- newLocalVar FSLIT("x") arg_ty
295       alts <- mapM mk_alt data_cons
296       body <- mkFromPRepr (unwrapFamInstScrut prepr_tc var_tys (Var arg))
297                           res_ty alts
298       return $ Lam arg body
299   where
300     var_tys   = mkTyVarTys $ tyConTyVars vect_tc
301     res_ty    = mkTyConApp vect_tc var_tys
302     data_cons = tyConDataCons vect_tc
303
304     mk_alt dc = do
305                   bndrs <- mapM (newLocalVar FSLIT("x")) $ dataConRepArgTys dc
306                   return (bndrs, mkConApp dc (map Type var_tys ++ map Var bndrs))
307
308 buildFromArrPRepr :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
309 buildFromArrPRepr _ vect_tc prepr_tc arr_tc
310   = mkFromArrPRepr undefined undefined undefined undefined undefined undefined
311
312 buildPRDict :: TyConRepr -> TyCon -> TyCon -> TyCon -> VM CoreExpr
313 buildPRDict _ vect_tc prepr_tc _
314   = prCoerce prepr_tc var_tys
315   =<< prDictOfType (mkTyConApp prepr_tc var_tys)
316   where
317     var_tys = mkTyVarTys $ tyConTyVars vect_tc
318
319 buildPArrayTyCon :: TyCon -> TyCon -> VM TyCon
320 buildPArrayTyCon orig_tc vect_tc = fixV $ \repr_tc ->
321   do
322     name'  <- cloneName mkPArrayTyConOcc orig_name
323     rhs    <- buildPArrayTyConRhs orig_name vect_tc repr_tc
324     parray <- builtin parrayTyCon
325
326     liftDs $ buildAlgTyCon name'
327                            tyvars
328                            []          -- no stupid theta
329                            rhs
330                            rec_flag    -- FIXME: is this ok?
331                            False       -- FIXME: no generics
332                            False       -- not GADT syntax
333                            (Just $ mk_fam_inst parray vect_tc)
334   where
335     orig_name = tyConName orig_tc
336     tyvars = tyConTyVars vect_tc
337     rec_flag = boolToRecFlag (isRecursiveTyCon vect_tc)
338     
339
340 buildPArrayTyConRhs :: Name -> TyCon -> TyCon -> VM AlgTyConRhs
341 buildPArrayTyConRhs orig_name vect_tc repr_tc
342   = do
343       data_con <- buildPArrayDataCon orig_name vect_tc repr_tc
344       return $ DataTyCon { data_cons = [data_con], is_enum = False }
345
346 buildPArrayDataCon :: Name -> TyCon -> TyCon -> VM DataCon
347 buildPArrayDataCon orig_name vect_tc repr_tc
348   = do
349       dc_name  <- cloneName mkPArrayDataConOcc orig_name
350       shape    <- tyConShape vect_tc
351       repr_tys <- mapM mkPArrayType types
352
353       liftDs $ buildDataCon dc_name
354                             False                  -- not infix
355                             (shapeStrictness shape ++ map (const NotMarkedStrict) repr_tys)
356                             []                     -- no field labels
357                             (tyConTyVars vect_tc)
358                             []                     -- no existentials
359                             []                     -- no eq spec
360                             []                     -- no context
361                             (shapeReprTys shape ++ repr_tys)
362                             repr_tc
363   where
364     types = [ty | dc <- tyConDataCons vect_tc
365                 , ty <- dataConRepArgTys dc]
366
367 mkPADFun :: TyCon -> VM Var
368 mkPADFun vect_tc
369   = newExportedVar (mkPADFunOcc $ getOccName vect_tc) =<< paDFunType vect_tc
370
371 data Shape = Shape {
372                shapeReprTys    :: [Type]
373              , shapeStrictness :: [StrictnessMark]
374              , shapeLength     :: [CoreExpr] -> VM CoreExpr
375              , shapeReplicate  :: CoreExpr -> CoreExpr -> VM [CoreExpr]
376              }
377
378 tyConShape :: TyCon -> VM Shape
379 tyConShape vect_tc
380   | isProductTyCon vect_tc
381   = return $ Shape {
382                 shapeReprTys    = [intPrimTy]
383               , shapeStrictness = [NotMarkedStrict]
384               , shapeLength     = \[len] -> return len
385               , shapeReplicate  = \len _ -> return [len]
386               }
387
388   | otherwise
389   = do
390       repr_ty <- mkPArrayType intTy   -- FIXME: we want to unbox this
391       return $ Shape {
392                  shapeReprTys    = [repr_ty]
393                , shapeStrictness = [MarkedStrict]
394                , shapeLength     = \[sel] -> lengthPA sel
395                , shapeReplicate  = \len n -> do
396                                                e <- replicatePA len n
397                                                return [e]
398                }
399
400 buildTyConBindings :: TyCon -> TyCon -> TyCon -> TyCon -> Var
401                    -> VM [(Var, CoreExpr)]
402 buildTyConBindings orig_tc vect_tc prepr_tc arr_tc dfun
403   = do
404       shape <- tyConShape vect_tc
405       repr  <- mkTyConRepr vect_tc
406       sequence_ (zipWith4 (vectDataConWorker shape vect_tc arr_tc arr_dc)
407                           orig_dcs
408                           vect_dcs
409                           (inits repr_tys)
410                           (tails repr_tys))
411       dict <- buildPADict repr vect_tc prepr_tc arr_tc dfun
412       binds <- takeHoisted
413       return $ (dfun, dict) : binds
414   where
415     orig_dcs = tyConDataCons orig_tc
416     vect_dcs = tyConDataCons vect_tc
417     [arr_dc] = tyConDataCons arr_tc
418
419     repr_tys = map dataConRepArgTys vect_dcs
420
421 vectDataConWorker :: Shape -> TyCon -> TyCon -> DataCon
422                   -> DataCon -> DataCon -> [[Type]] -> [[Type]]
423                   -> VM ()
424 vectDataConWorker shape vect_tc arr_tc arr_dc orig_dc vect_dc pre (dc_tys : post)
425   = do
426       clo <- closedV
427            . inBind orig_worker
428            . polyAbstract tvs $ \abstract ->
429              liftM (abstract . vectorised)
430            $ buildClosures tvs [] dc_tys res_ty (liftM2 (,) mk_vect mk_lift)
431
432       worker <- cloneId mkVectOcc orig_worker (exprType clo)
433       hoistBinding worker clo
434       defGlobalVar orig_worker worker
435       return ()
436   where
437     tvs     = tyConTyVars vect_tc
438     arg_tys = mkTyVarTys tvs
439     res_ty  = mkTyConApp vect_tc arg_tys
440
441     orig_worker = dataConWorkId orig_dc
442
443     mk_vect = return . mkConApp vect_dc $ map Type arg_tys
444     mk_lift = do
445                 len     <- newLocalVar FSLIT("n") intPrimTy
446                 arr_tys <- mapM mkPArrayType dc_tys
447                 args    <- mapM (newLocalVar FSLIT("xs")) arr_tys
448                 shapes  <- shapeReplicate shape
449                                           (Var len)
450                                           (mkDataConTag vect_dc)
451                 
452                 empty_pre  <- mapM emptyPA (concat pre)
453                 empty_post <- mapM emptyPA (concat post)
454
455                 return . mkLams (len : args)
456                        . wrapFamInstBody arr_tc arg_tys
457                        . mkConApp arr_dc
458                        $ map Type arg_tys ++ shapes
459                                           ++ empty_pre
460                                           ++ map Var args
461                                           ++ empty_post
462
463 buildPADict :: TyConRepr -> TyCon -> TyCon -> TyCon -> Var -> VM CoreExpr
464 buildPADict repr vect_tc prepr_tc arr_tc dfun
465   = polyAbstract tvs $ \abstract ->
466     do
467       meth_binds <- mapM (mk_method repr) paMethods
468       let meth_exprs = map (Var . fst) meth_binds
469
470       pa_dc <- builtin paDataCon
471       let dict = mkConApp pa_dc (Type (mkTyConApp vect_tc arg_tys) : meth_exprs)
472           body = Let (Rec meth_binds) dict
473       return . mkInlineMe $ abstract body
474   where
475     tvs = tyConTyVars arr_tc
476     arg_tys = mkTyVarTys tvs
477
478     mk_method repr (name, build)
479       = localV
480       $ do
481           body <- build repr vect_tc prepr_tc arr_tc
482           var  <- newLocalVar name (exprType body)
483           return (var, mkInlineMe body)
484           
485 paMethods = [(FSLIT("toPRepr"),      buildToPRepr),
486              (FSLIT("fromPRepr"),    buildFromPRepr),
487              (FSLIT("toArrPRepr"),   buildToArrPRepr),
488              (FSLIT("fromArrPRepr"), buildFromArrPRepr),
489              (FSLIT("dictPRepr"),    buildPRDict)]
490
491 -- | Split the given tycons into two sets depending on whether they have to be
492 -- converted (first list) or not (second list). The first argument contains
493 -- information about the conversion status of external tycons:
494 -- 
495 --   * tycons which have converted versions are mapped to True
496 --   * tycons which are not changed by vectorisation are mapped to False
497 --   * tycons which can't be converted are not elements of the map
498 --
499 classifyTyCons :: UniqFM Bool -> [TyConGroup] -> ([TyCon], [TyCon])
500 classifyTyCons = classify [] []
501   where
502     classify conv keep cs [] = (conv, keep)
503     classify conv keep cs ((tcs, ds) : rs)
504       | can_convert && must_convert
505         = classify (tcs ++ conv) keep (cs `addListToUFM` [(tc,True) | tc <- tcs]) rs
506       | can_convert
507         = classify conv (tcs ++ keep) (cs `addListToUFM` [(tc,False) | tc <- tcs]) rs
508       | otherwise
509         = classify conv keep cs rs
510       where
511         refs = ds `delListFromUniqSet` tcs
512
513         can_convert  = isNullUFM (refs `minusUFM` cs) && all convertable tcs
514         must_convert = foldUFM (||) False (intersectUFM_C const cs refs)
515
516         convertable tc = isDataTyCon tc && all isVanillaDataCon (tyConDataCons tc)
517     
518 -- | Compute mutually recursive groups of tycons in topological order
519 --
520 tyConGroups :: [TyCon] -> [TyConGroup]
521 tyConGroups tcs = map mk_grp (stronglyConnComp edges)
522   where
523     edges = [((tc, ds), tc, uniqSetToList ds) | tc <- tcs
524                                 , let ds = tyConsOfTyCon tc]
525
526     mk_grp (AcyclicSCC (tc, ds)) = ([tc], ds)
527     mk_grp (CyclicSCC els)       = (tcs, unionManyUniqSets dss)
528       where
529         (tcs, dss) = unzip els
530
531 tyConsOfTyCon :: TyCon -> UniqSet TyCon
532 tyConsOfTyCon 
533   = tyConsOfTypes . concatMap dataConRepArgTys . tyConDataCons
534
535 tyConsOfType :: Type -> UniqSet TyCon
536 tyConsOfType ty
537   | Just ty' <- coreView ty    = tyConsOfType ty'
538 tyConsOfType (TyVarTy v)       = emptyUniqSet
539 tyConsOfType (TyConApp tc tys) = extend (tyConsOfTypes tys)
540   where
541     extend | isUnLiftedTyCon tc
542            || isTupleTyCon   tc = id
543
544            | otherwise          = (`addOneToUniqSet` tc)
545
546 tyConsOfType (AppTy a b)       = tyConsOfType a `unionUniqSets` tyConsOfType b
547 tyConsOfType (FunTy a b)       = (tyConsOfType a `unionUniqSets` tyConsOfType b)
548                                  `addOneToUniqSet` funTyCon
549 tyConsOfType (ForAllTy _ ty)   = tyConsOfType ty
550 tyConsOfType other             = pprPanic "ClosureConv.tyConsOfType" $ ppr other
551
552 tyConsOfTypes :: [Type] -> UniqSet TyCon
553 tyConsOfTypes = unionManyUniqSets . map tyConsOfType
554