More refactoring
[ghc.git] / compiler / vectorise / VectUtils.hs
1 module VectUtils (
2 collectAnnTypeBinders, collectAnnTypeArgs, isAnnTypeArg,
3 splitClosureTy,
4 mkPADictType, mkPArrayType,
5 paDictArgType, paDictOfType,
6 paMethod, lengthPA, replicatePA,
7 lookupPArrayFamInst,
8 hoistExpr, takeHoisted
9 ) where
10
11 #include "HsVersions.h"
12
13 import VectMonad
14
15 import CoreSyn
16 import CoreUtils
17 import Type
18 import TypeRep
19 import TyCon
20 import Var
21 import PrelNames
22
23 import Outputable
24 import FastString
25
26 import Control.Monad ( liftM )
27
28 collectAnnTypeArgs :: AnnExpr b ann -> (AnnExpr b ann, [Type])
29 collectAnnTypeArgs expr = go expr []
30 where
31 go (_, AnnApp f (_, AnnType ty)) tys = go f (ty : tys)
32 go e tys = (e, tys)
33
34 collectAnnTypeBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann)
35 collectAnnTypeBinders expr = go [] expr
36 where
37 go bs (_, AnnLam b e) | isTyVar b = go (b:bs) e
38 go bs e = (reverse bs, e)
39
40 isAnnTypeArg :: AnnExpr b ann -> Bool
41 isAnnTypeArg (_, AnnType t) = True
42 isAnnTypeArg _ = False
43
44 isClosureTyCon :: TyCon -> Bool
45 isClosureTyCon tc = tyConUnique tc == closureTyConKey
46
47 splitClosureTy :: Type -> (Type, Type)
48 splitClosureTy ty
49 | Just (tc, [arg_ty, res_ty]) <- splitTyConApp_maybe ty
50 , isClosureTyCon tc
51 = (arg_ty, res_ty)
52
53 | otherwise = pprPanic "splitClosureTy" (ppr ty)
54
55 mkPADictType :: Type -> VM Type
56 mkPADictType ty
57 = do
58 tc <- builtin paDictTyCon
59 return $ TyConApp tc [ty]
60
61 mkPArrayType :: Type -> VM Type
62 mkPArrayType ty
63 = do
64 tc <- builtin parrayTyCon
65 return $ TyConApp tc [ty]
66
67 paDictArgType :: TyVar -> VM (Maybe Type)
68 paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
69 where
70 go ty k | Just k' <- kindView k = go ty k'
71 go ty (FunTy k1 k2)
72 = do
73 tv <- newTyVar FSLIT("a") k1
74 mty1 <- go (TyVarTy tv) k1
75 case mty1 of
76 Just ty1 -> do
77 mty2 <- go (AppTy ty (TyVarTy tv)) k2
78 return $ fmap (ForAllTy tv . FunTy ty1) mty2
79 Nothing -> go ty k2
80
81 go ty k
82 | isLiftedTypeKind k
83 = liftM Just (mkPADictType ty)
84
85 go ty k = return Nothing
86
87 paDictOfType :: Type -> VM CoreExpr
88 paDictOfType ty = paDictOfTyApp ty_fn ty_args
89 where
90 (ty_fn, ty_args) = splitAppTys ty
91
92 paDictOfTyApp :: Type -> [Type] -> VM CoreExpr
93 paDictOfTyApp ty_fn ty_args
94 | Just ty_fn' <- coreView ty_fn = paDictOfTyApp ty_fn' ty_args
95 paDictOfTyApp (TyVarTy tv) ty_args
96 = do
97 dfun <- maybeV (lookupTyVarPA tv)
98 paDFunApply dfun ty_args
99 paDictOfTyApp (TyConApp tc _) ty_args
100 = do
101 pa_class <- builtin paClass
102 (dfun, ty_args') <- lookupInst pa_class [TyConApp tc ty_args]
103 paDFunApply (Var dfun) ty_args'
104 paDictOfTyApp ty ty_args = pprPanic "paDictOfTyApp" (ppr ty)
105
106 paDFunApply :: CoreExpr -> [Type] -> VM CoreExpr
107 paDFunApply dfun tys
108 = do
109 dicts <- mapM paDictOfType tys
110 return $ mkApps (mkTyApps dfun tys) dicts
111
112 paMethod :: (Builtins -> Var) -> Type -> VM CoreExpr
113 paMethod method ty
114 = do
115 fn <- builtin method
116 dict <- paDictOfType ty
117 return $ mkApps (Var fn) [Type ty, dict]
118
119 lengthPA :: CoreExpr -> VM CoreExpr
120 lengthPA x = liftM (`App` x) (paMethod lengthPAVar (exprType x))
121
122 replicatePA :: CoreExpr -> CoreExpr -> VM CoreExpr
123 replicatePA len x = liftM (`mkApps` [len,x])
124 (paMethod replicatePAVar (exprType x))
125
126 lookupPArrayFamInst :: Type -> VM (TyCon, [Type])
127 lookupPArrayFamInst ty = builtin parrayTyCon >>= (`lookupFamInst` [ty])
128
129 hoistExpr :: FastString -> CoreExpr -> VM Var
130 hoistExpr fs expr
131 = do
132 var <- newLocalVar fs (exprType expr)
133 updGEnv $ \env ->
134 env { global_bindings = (var, expr) : global_bindings env }
135 return var
136
137 takeHoisted :: VM [(Var, CoreExpr)]
138 takeHoisted
139 = do
140 env <- readGEnv id
141 setGEnv $ env { global_bindings = [] }
142 return $ global_bindings env
143