Add kind equalities to GHC.
[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 )
8 where
9
10 import Vectorise.Utils
11 import Vectorise.Monad
12 import Vectorise.Builtins
13 import TcType
14 import Type
15 import TyCoRep
16 import TyCon
17 import Control.Monad
18 import Control.Applicative
19 import Data.Maybe
20 import Outputable
21 import Prelude -- avoid redundant import warning due to AMP
22
23 -- |Vectorise a type constructor. Unless there is a vectorised version (stripped of embedded
24 -- parallel arrays), the vectorised version is the same as the original.
25 --
26 vectTyCon :: TyCon -> VM TyCon
27 vectTyCon tc = maybe tc id <$> lookupTyCon tc
28
29 -- |Produce the vectorised and lifted versions of a type.
30 --
31 -- NB: Here we are limited to properly handle predicates at the toplevel only. Anything embedded
32 -- in what is called the 'body_ty' below will end up as an argument to the type family 'PData'.
33 --
34 vectAndLiftType :: Type -> VM (Type, Type)
35 vectAndLiftType ty | Just ty' <- coreView ty = vectAndLiftType ty'
36 vectAndLiftType ty
37 = do { padicts <- liftM catMaybes $ mapM paDictArgType tyvars
38 ; vmono_ty <- vectType mono_ty
39 ; lmono_ty <- mkPDataType vmono_ty
40 ; return (abstractType tyvars (padicts ++ theta) vmono_ty,
41 abstractType tyvars (padicts ++ theta) lmono_ty)
42 }
43 where
44 (tyvars, phiTy) = splitForAllTys ty
45 (theta, mono_ty) = tcSplitPhiTy phiTy
46
47 -- |Vectorise a type.
48 --
49 -- For each quantified var we need to add a PA dictionary out the front of the type.
50 -- So forall a. C a => a -> a
51 -- turns into forall a. PA a => Cv a => a :-> a
52 --
53 vectType :: Type -> VM Type
54 vectType ty
55 | Just ty' <- coreView ty
56 = vectType ty'
57 vectType (TyVarTy tv) = return $ TyVarTy tv
58 vectType (LitTy l) = return $ LitTy l
59 vectType (AppTy ty1 ty2) = AppTy <$> vectType ty1 <*> vectType ty2
60 vectType (TyConApp tc tys) = TyConApp <$> vectTyCon tc <*> mapM vectType tys
61 vectType (ForAllTy (Anon ty1) ty2)
62 | isPredTy ty1
63 = mkFunTy <$> vectType ty1 <*> vectType ty2 -- don't build a closure for dictionary abstraction
64 | otherwise
65 = TyConApp <$> builtin closureTyCon <*> mapM vectType [ty1, ty2]
66 vectType ty@(ForAllTy {})
67 = do { -- strip off consecutive foralls
68 ; let (tyvars, tyBody) = splitForAllTys ty
69
70 -- vectorise the body
71 ; vtyBody <- vectType tyBody
72
73 -- make a PA dictionary for each of the type variables
74 ; dictsPA <- liftM catMaybes $ mapM paDictArgType tyvars
75
76 -- add the PA dictionaries after the foralls
77 ; return $ abstractType tyvars dictsPA vtyBody
78 }
79 vectType ty@(CastTy {})
80 = pprSorry "Vectorise.Type.Type.vectType: CastTy" (ppr ty)
81 vectType ty@(CoercionTy {})
82 = pprSorry "Vectorise.Type.Type.vectType: CoercionTy" (ppr ty)
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 = mkInvForAllTys tyvars . mkFunTys dicts