64a4a22dab921cd5f5e67aa9fb6691aa04767955
[ghc.git] / compiler / vectorise / Vectorise / Type / Type.hs
1 -- Apply the vectorisation transformation to types. This is the \mathcal{L}_t scheme in HtM.
2
3 module Vectorise.Type.Type (
4 vectTyCon,
5 vectAndLiftType,
6 vectType
7 ) where
8
9 import Vectorise.Utils
10 import Vectorise.Monad
11 import Vectorise.Builtins
12 import TypeRep
13 import Type
14 import TyCon
15 import Outputable
16 import Control.Monad
17 import Data.List
18 import Data.Maybe
19
20 -- | Vectorise a type constructor.
21 --
22 vectTyCon :: TyCon -> VM TyCon
23 vectTyCon tc
24 | isFunTyCon tc = builtin closureTyCon
25 | isBoxedTupleTyCon tc = return tc
26 | isUnLiftedTyCon tc = return tc
27 | otherwise
28 = maybeCantVectoriseM "Tycon not vectorised: " (ppr tc)
29 $ lookupTyCon tc
30
31 -- |Produce the vectorised and lifted versions of a type.
32 --
33 vectAndLiftType :: Type -> VM (Type, Type)
34 vectAndLiftType ty | Just ty' <- coreView ty = vectAndLiftType ty'
35 vectAndLiftType ty
36 = do
37 mdicts <- mapM paDictArgType (reverse tyvars)
38 let dicts = [dict | Just dict <- mdicts]
39 vmono_ty <- vectType mono_ty
40 lmono_ty <- mkPDataType vmono_ty
41 return (abstractType tyvars dicts vmono_ty,
42 abstractType tyvars dicts lmono_ty)
43 where
44 (tyvars, mono_ty) = splitForAllTys ty
45
46 -- |Vectorise a type.
47 --
48 vectType :: Type -> VM Type
49 vectType ty
50 | Just ty' <- coreView ty
51 = vectType ty'
52
53 vectType (TyVarTy tv) = return $ TyVarTy tv
54 vectType (AppTy ty1 ty2) = liftM2 AppTy (vectType ty1) (vectType ty2)
55 vectType (TyConApp tc tys) = liftM2 TyConApp (vectTyCon tc) (mapM vectType tys)
56 vectType (FunTy ty1 ty2) = liftM2 TyConApp (builtin closureTyCon) (mapM vectType [ty1,ty2])
57
58 -- For each quantified var we need to add a PA dictionary out the front of the type.
59 -- So forall a. C a => a -> a
60 -- turns into forall a. Cv a => PA a => a :-> a
61 vectType ty@(ForAllTy _ _)
62 = do
63 -- split the type into the quantified vars, its dictionaries and the body.
64 let (tyvars, tyBody) = splitForAllTys ty
65 let (tyArgs, tyResult) = splitFunTys tyBody
66
67 let (tyArgs_dict, tyArgs_regular)
68 = partition isDictTy tyArgs
69
70 -- vectorise the body.
71 let tyBody' = mkFunTys tyArgs_regular tyResult
72 tyBody'' <- vectType tyBody'
73
74 -- vectorise the dictionary parameters.
75 dictsVect <- mapM vectType tyArgs_dict
76
77 -- make a PA dictionary for each of the type variables.
78 dictsPA <- liftM catMaybes $ mapM paDictArgType tyvars
79
80 -- pack it all back together.
81 traceVt "vect ForAllTy: " $ ppr (abstractType tyvars (dictsPA ++ dictsVect) tyBody'')
82 return $ abstractType tyvars (dictsPA ++ dictsVect) tyBody''
83
84 -- |Add quantified vars and dictionary parameters to the front of a type.
85 --
86 abstractType :: [TyVar] -> [Type] -> Type -> Type
87 abstractType tyvars dicts = mkForAllTys tyvars . mkFunTys dicts