630c425f44f49bd7e07c3bbf83bfee522815ad59
[ghc.git] / compiler / vectorise / VectUtils.hs
1 module VectUtils (
2 collectAnnTypeBinders, collectAnnTypeArgs, isAnnTypeArg,
3 splitClosureTy,
4 mkPADictType, mkPArrayType,
5 paDictArgType, paDictOfType,
6 lookupPArrayFamInst
7 ) where
8
9 #include "HsVersions.h"
10
11 import VectMonad
12
13 import CoreSyn
14 import Type
15 import TypeRep
16 import TyCon
17 import Var
18 import PrelNames
19
20 import Outputable
21
22 import Control.Monad ( liftM )
23
24 collectAnnTypeArgs :: AnnExpr b ann -> (AnnExpr b ann, [Type])
25 collectAnnTypeArgs expr = go expr []
26 where
27 go (_, AnnApp f (_, AnnType ty)) tys = go f (ty : tys)
28 go e tys = (e, tys)
29
30 collectAnnTypeBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann)
31 collectAnnTypeBinders expr = go [] expr
32 where
33 go bs (_, AnnLam b e) | isTyVar b = go (b:bs) e
34 go bs e = (reverse bs, e)
35
36 isAnnTypeArg :: AnnExpr b ann -> Bool
37 isAnnTypeArg (_, AnnType t) = True
38 isAnnTypeArg _ = False
39
40 isClosureTyCon :: TyCon -> Bool
41 isClosureTyCon tc = tyConUnique tc == closureTyConKey
42
43 splitClosureTy :: Type -> (Type, Type)
44 splitClosureTy ty
45 | Just (tc, [arg_ty, res_ty]) <- splitTyConApp_maybe ty
46 , isClosureTyCon tc
47 = (arg_ty, res_ty)
48
49 | otherwise = pprPanic "splitClosureTy" (ppr ty)
50
51 mkPADictType :: Type -> VM Type
52 mkPADictType ty
53 = do
54 tc <- builtin paDictTyCon
55 return $ TyConApp tc [ty]
56
57 mkPArrayType :: Type -> VM Type
58 mkPArrayType ty
59 = do
60 tc <- builtin parrayTyCon
61 return $ TyConApp tc [ty]
62
63 paDictArgType :: TyVar -> VM (Maybe Type)
64 paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
65 where
66 go ty k | Just k' <- kindView k = go ty k'
67 go ty (FunTy k1 k2)
68 = do
69 tv <- newTyVar FSLIT("a") k1
70 mty1 <- go (TyVarTy tv) k1
71 case mty1 of
72 Just ty1 -> do
73 mty2 <- go (AppTy ty (TyVarTy tv)) k2
74 return $ fmap (ForAllTy tv . FunTy ty1) mty2
75 Nothing -> go ty k2
76
77 go ty k
78 | isLiftedTypeKind k
79 = liftM Just (mkPADictType ty)
80
81 go ty k = return Nothing
82
83 paDictOfType :: Type -> VM CoreExpr
84 paDictOfType ty = paDictOfTyApp ty_fn ty_args
85 where
86 (ty_fn, ty_args) = splitAppTys ty
87
88 paDictOfTyApp :: Type -> [Type] -> VM CoreExpr
89 paDictOfTyApp ty_fn ty_args
90 | Just ty_fn' <- coreView ty_fn = paDictOfTyApp ty_fn' ty_args
91 paDictOfTyApp (TyVarTy tv) ty_args
92 = do
93 dfun <- maybeV (lookupTyVarPA tv)
94 paDFunApply dfun ty_args
95 paDictOfTyApp (TyConApp tc _) ty_args
96 = do
97 pa_class <- builtin paClass
98 (dfun, ty_args') <- lookupInst pa_class [TyConApp tc ty_args]
99 paDFunApply (Var dfun) ty_args'
100 paDictOfTyApp ty ty_args = pprPanic "paDictOfTyApp" (ppr ty)
101
102 paDFunApply :: CoreExpr -> [Type] -> VM CoreExpr
103 paDFunApply dfun tys
104 = do
105 dicts <- mapM paDictOfType tys
106 return $ mkApps (mkTyApps dfun tys) dicts
107
108 lookupPArrayFamInst :: Type -> VM (TyCon, [Type])
109 lookupPArrayFamInst ty = builtin parrayTyCon >>= (`lookupFamInst` [ty])
110