Use lengthIs and friends in more places
[ghc.git] / compiler / vectorise / Vectorise / Utils / PADict.hs
1 module Vectorise.Utils.PADict (
2 paDictArgType,
3 paDictOfType,
4 paMethod,
5 prDictOfReprType,
6 prDictOfPReprInstTyCon
7 ) where
8
9 import Vectorise.Monad
10 import Vectorise.Builtins
11 import Vectorise.Utils.Base
12
13 import CoreSyn
14 import CoreUtils
15 import FamInstEnv
16 import Coercion
17 import Type
18 import TyCoRep
19 import TyCon
20 import CoAxiom
21 import Var
22 import Outputable
23 import DynFlags
24 import FastString
25 import Util
26 import Control.Monad
27
28
29 -- |Construct the PA argument type for the tyvar. For the tyvar (v :: *) it's
30 -- just PA v. For (v :: (* -> *) -> *) it's
31 --
32 -- > forall (a :: * -> *). (forall (b :: *). PA b -> PA (a b)) -> PA (v a)
33 --
34 paDictArgType :: TyVar -> VM (Maybe Type)
35 paDictArgType tv = go (mkTyVarTy tv) (tyVarKind tv)
36 where
37 go ty (FunTy k1 k2)
38 = do
39 tv <- if isCoercionType k1
40 then newCoVar (fsLit "c") k1
41 else newTyVar (fsLit "a") k1
42 mty1 <- go (mkTyVarTy tv) k1
43 case mty1 of
44 Just ty1 -> do
45 mty2 <- go (mkAppTy ty (mkTyVarTy tv)) k2
46 return $ fmap (mkInvForAllTy tv . mkFunTy ty1) mty2
47 Nothing -> go ty k2
48
49 go ty k
50 | isLiftedTypeKind k
51 = do
52 pa_cls <- builtin paClass
53 return $ Just $ mkClassPred pa_cls [ty]
54
55 go _ _ = return Nothing
56
57
58 -- |Get the PA dictionary for some type
59 --
60 paDictOfType :: Type -> VM CoreExpr
61 paDictOfType ty
62 = paDictOfTyApp ty_fn ty_args
63 where
64 (ty_fn, ty_args) = splitAppTys ty
65
66 paDictOfTyApp :: Type -> [Type] -> VM CoreExpr
67 paDictOfTyApp ty_fn ty_args
68 | Just ty_fn' <- coreView ty_fn
69 = paDictOfTyApp ty_fn' ty_args
70
71 -- for type variables, look up the dfun and apply to the PA dictionaries
72 -- of the type arguments
73 paDictOfTyApp (TyVarTy tv) ty_args
74 = do
75 { dfun <- maybeCantVectoriseM "No PA dictionary for type variable"
76 (ppr tv <+> text "in" <+> ppr ty)
77 $ lookupTyVarPA tv
78 ; dicts <- mapM paDictOfType ty_args
79 ; return $ dfun `mkTyApps` ty_args `mkApps` dicts
80 }
81
82 -- for tycons, we also need to apply the dfun to the PR dictionary of
83 -- the representation type if the tycon is polymorphic
84 paDictOfTyApp (TyConApp tc []) ty_args
85 = do
86 { dfun <- maybeCantVectoriseM noPADictErr (ppr tc <+> text "in" <+> ppr ty)
87 $ lookupTyConPA tc
88 ; super <- super_dict tc ty_args
89 ; dicts <- mapM paDictOfType ty_args
90 ; return $ Var dfun `mkTyApps` ty_args `mkApps` super `mkApps` dicts
91 }
92 where
93 noPADictErr = "No PA dictionary for type constructor (did you import 'Data.Array.Parallel'?)"
94
95 super_dict _ [] = return []
96 super_dict tycon ty_args
97 = do
98 { pr <- prDictOfPReprInst (TyConApp tycon ty_args)
99 ; return [pr]
100 }
101
102 paDictOfTyApp _ _ = getDynFlags >>= failure
103
104 failure dflags = cantVectorise dflags "Can't construct PA dictionary for type" (ppr ty)
105
106 -- |Produce code that refers to a method of the 'PA' class.
107 --
108 paMethod :: (Builtins -> Var) -> (TyCon -> Builtins -> Var) -> Type -> VM CoreExpr
109 paMethod _ query ty
110 | Just tycon <- splitPrimTyCon ty -- Is 'ty' from 'GHC.Prim' (e.g., 'Int#')?
111 = liftM Var $ builtin (query tycon)
112 paMethod method _ ty
113 = do
114 { fn <- builtin method
115 ; dict <- paDictOfType ty
116 ; return $ mkApps (Var fn) [Type ty, dict]
117 }
118
119 -- |Given a type @ty@, return the PR dictionary for @PRepr ty@.
120 --
121 prDictOfPReprInst :: Type -> VM CoreExpr
122 prDictOfPReprInst ty
123 = do
124 { (FamInstMatch { fim_instance = prepr_fam, fim_tys = prepr_args })
125 <- preprFamInst ty
126 ; prDictOfPReprInstTyCon ty (famInstAxiom prepr_fam) prepr_args
127 }
128
129 -- |Given a type @ty@, its PRepr synonym tycon and its type arguments,
130 -- return the PR @PRepr ty@. Suppose we have:
131 --
132 -- > type instance PRepr (T a1 ... an) = t
133 --
134 -- which is internally translated into
135 --
136 -- > type :R:PRepr a1 ... an = t
137 --
138 -- and the corresponding coercion. Then,
139 --
140 -- > prDictOfPReprInstTyCon (T a1 ... an) :R:PRepr u1 ... un = PR (T u1 ... un)
141 --
142 -- Note that @ty@ is only used for error messages
143 --
144 prDictOfPReprInstTyCon :: Type -> CoAxiom Unbranched -> [Type] -> VM CoreExpr
145 prDictOfPReprInstTyCon _ty prepr_ax prepr_args
146 = do
147 let rhs = mkUnbranchedAxInstRHS prepr_ax prepr_args []
148 dict <- prDictOfReprType' rhs
149 pr_co <- mkBuiltinCo prTyCon
150 let co = mkAppCo pr_co
151 $ mkSymCo
152 $ mkUnbranchedAxInstCo Nominal prepr_ax prepr_args []
153 return $ mkCast dict co
154
155 -- |Get the PR dictionary for a type. The argument must be a representation
156 -- type.
157 --
158 prDictOfReprType :: Type -> VM CoreExpr
159 prDictOfReprType ty
160 | Just (tycon, tyargs) <- splitTyConApp_maybe ty
161 = do
162 prepr <- builtin preprTyCon
163 if tycon == prepr
164 then do
165 let [ty'] = tyargs
166 pa <- paDictOfType ty'
167 sel <- builtin paPRSel
168 return $ Var sel `App` Type ty' `App` pa
169 else do
170 -- a representation tycon must have a PR instance
171 dfun <- maybeV (text "look up PR dictionary for" <+> ppr tycon) $
172 lookupTyConPR tycon
173 prDFunApply dfun tyargs
174
175 | otherwise
176 = do
177 -- it is a tyvar or an application of a tyvar
178 -- determine the PR dictionary from its PA dictionary
179 --
180 -- NOTE: This assumes that PRepr t ~ t is for all representation types
181 -- t
182 --
183 -- FIXME: This doesn't work for kinds other than * at the moment. We'd
184 -- have to simply abstract the term over the missing type arguments.
185 pa <- paDictOfType ty
186 prsel <- builtin paPRSel
187 return $ Var prsel `mkApps` [Type ty, pa]
188
189 prDictOfReprType' :: Type -> VM CoreExpr
190 prDictOfReprType' ty = prDictOfReprType ty `orElseV`
191 do dflags <- getDynFlags
192 cantVectorise dflags "No PR dictionary for representation type"
193 (ppr ty)
194
195 -- | Apply a tycon's PR dfun to dictionary arguments (PR or PA) corresponding
196 -- to the argument types.
197 prDFunApply :: Var -> [Type] -> VM CoreExpr
198 prDFunApply dfun tys
199 | Just [] <- ctxs -- PR (a :-> b) doesn't have a context
200 = return $ Var dfun `mkTyApps` tys
201
202 | Just tycons <- ctxs
203 , tycons `equalLength` tys
204 = do
205 pa <- builtin paTyCon
206 pr <- builtin prTyCon
207 dflags <- getDynFlags
208 args <- zipWithM (dictionary dflags pa pr) tys tycons
209 return $ Var dfun `mkTyApps` tys `mkApps` args
210
211 | otherwise = do dflags <- getDynFlags
212 invalid dflags
213 where
214 -- the dfun's contexts - if its type is (PA a, PR b) => PR (C a b) then
215 -- ctxs is Just [PA, PR]
216 ctxs = fmap (map fst)
217 $ sequence
218 $ map splitTyConApp_maybe
219 $ fst
220 $ splitFunTys
221 $ snd
222 $ splitForAllTys
223 $ varType dfun
224
225 dictionary dflags pa pr ty tycon
226 | tycon == pa = paDictOfType ty
227 | tycon == pr = prDictOfReprType ty
228 | otherwise = invalid dflags
229
230 invalid dflags = cantVectorise dflags "Invalid PR dfun type" (ppr (varType dfun) <+> ppr tys)