2 module Vectorise.Utils.PADict (
10 import Vectorise.Monad
11 import Vectorise.Builtins
12 import Vectorise.Utils.Base
26 -- | Construct the PA argument type for the tyvar. For the tyvar (v :: *) it's
27 -- just PA v. For (v :: (* -> *) -> *) it's
29 -- > forall (a :: * -> *). (forall (b :: *). PA b -> PA (a b)) -> PA (v a)
31 paDictArgType :: TyVar -> VM (Maybe Type)
32 paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
36 tv <- newTyVar (fsLit "a") k1
37 mty1 <- go (TyVarTy tv) k1
40 mty2 <- go (AppTy ty (TyVarTy tv)) k2
41 return $ fmap (ForAllTy tv . FunTy ty1) mty2
47 pa_cls <- builtin paClass
48 return $ Just $ PredTy $ ClassP pa_cls [ty]
50 go _ _ = return Nothing
53 -- | Get the PA dictionary for some type
55 paDictOfType :: Type -> VM CoreExpr
57 = paDictOfTyApp ty_fn ty_args
59 (ty_fn, ty_args) = splitAppTys ty
61 paDictOfTyApp :: Type -> [Type] -> VM CoreExpr
62 paDictOfTyApp ty_fn ty_args
63 | Just ty_fn' <- coreView ty_fn
64 = paDictOfTyApp ty_fn' ty_args
66 -- for type variables, look up the dfun and apply to the PA dictionaries
67 -- of the type arguments
68 paDictOfTyApp (TyVarTy tv) ty_args
69 = do dfun <- maybeCantVectoriseM "No PA dictionary for type variable"
70 (ppr tv <+> text "in" <+> ppr ty)
72 dicts <- mapM paDictOfType ty_args
73 return $ dfun `mkTyApps` ty_args `mkApps` dicts
75 -- for tycons, we also need to apply the dfun to the PR dictionary of
76 -- the representation type if the tycon is polymorphic
77 paDictOfTyApp (TyConApp tc []) ty_args
79 dfun <- maybeCantVectoriseM "No PA dictionary for type constructor"
80 (ppr tc <+> text "in" <+> ppr ty)
82 super <- super_dict tc ty_args
83 dicts <- mapM paDictOfType ty_args
84 return $ Var dfun `mkTyApps` ty_args `mkApps` super `mkApps` dicts
86 paDictOfTyApp _ _ = failure
88 super_dict _ [] = return []
89 super_dict tycon ty_args
91 pr <- prDictOfPReprInst (TyConApp tycon ty_args)
94 failure = cantVectorise "Can't construct PA dictionary for type" (ppr ty)
96 paMethod :: (Builtins -> Var) -> String -> Type -> VM CoreExpr
98 | Just tycon <- splitPrimTyCon ty
100 . maybeCantVectoriseM "No PA method" (text name <+> text "for" <+> ppr tycon)
101 $ lookupPrimMethod tycon name
106 dict <- paDictOfType ty
107 return $ mkApps (Var fn) [Type ty, dict]
109 -- | Given a type @ty@, return the PR dictionary for @PRepr ty@.
110 prDictOfPReprInst :: Type -> VM CoreExpr
113 (prepr_tc, prepr_args) <- preprSynTyCon ty
114 prDictOfPReprInstTyCon ty prepr_tc prepr_args
116 -- | Given a type @ty@, its PRepr synonym tycon and its type arguments,
117 -- return the PR @PRepr ty@. Suppose we have:
119 -- > type instance PRepr (T a1 ... an) = t
121 -- which is internally translated into
123 -- > type :R:PRepr a1 ... an = t
125 -- and the corresponding coercion. Then,
127 -- > prDictOfPReprInstTyCon (T a1 ... an) :R:PRepr u1 ... un = PR (T u1 ... un)
129 -- Note that @ty@ is only used for error messages
131 prDictOfPReprInstTyCon :: Type -> TyCon -> [Type] -> VM CoreExpr
132 prDictOfPReprInstTyCon ty prepr_tc prepr_args
133 | Just rhs <- coreView (mkTyConApp prepr_tc prepr_args)
135 dict <- prDictOfReprType' rhs
136 pr_co <- mkBuiltinCo prTyCon
137 let Just arg_co = tyConFamilyCoercion_maybe prepr_tc
138 let co = mkAppCo pr_co
140 $ mkAxInstCo arg_co prepr_args
141 return $ mkCoerce co dict
143 | otherwise = cantVectorise "Invalid PRepr type instance" (ppr ty)
145 -- | Get the PR dictionary for a type. The argument must be a representation
147 prDictOfReprType :: Type -> VM CoreExpr
149 | Just (tycon, tyargs) <- splitTyConApp_maybe ty
151 prepr <- builtin preprTyCon
155 pa <- paDictOfType ty'
156 sel <- builtin paPRSel
157 return $ Var sel `App` Type ty' `App` pa
159 -- a representation tycon must have a PR instance
160 dfun <- maybeV $ lookupTyConPR tycon
161 prDFunApply dfun tyargs
165 -- it is a tyvar or an application of a tyvar
166 -- determine the PR dictionary from its PA dictionary
168 -- NOTE: This assumes that PRepr t ~ t is for all representation types
171 -- FIXME: This doesn't work for kinds other than * at the moment. We'd
172 -- have to simply abstract the term over the missing type arguments.
173 pa <- paDictOfType ty
174 prsel <- builtin paPRSel
175 return $ Var prsel `mkApps` [Type ty, pa]
177 prDictOfReprType' :: Type -> VM CoreExpr
178 prDictOfReprType' ty = prDictOfReprType ty `orElseV`
179 cantVectorise "No PR dictionary for representation type"
182 -- | Apply a tycon's PR dfun to dictionary arguments (PR or PA) corresponding
183 -- to the argument types.
184 prDFunApply :: Var -> [Type] -> VM CoreExpr
186 | Just [] <- ctxs -- PR (a :-> b) doesn't have a context
187 = return $ Var dfun `mkTyApps` tys
189 | Just tycons <- ctxs
190 , length tycons == length tys
192 pa <- builtin paTyCon
193 pr <- builtin prTyCon
194 args <- zipWithM (dictionary pa pr) tys tycons
195 return $ Var dfun `mkTyApps` tys `mkApps` args
197 | otherwise = invalid
199 -- the dfun's contexts - if its type is (PA a, PR b) => PR (C a b) then
200 -- ctxs is Just [PA, PR]
201 ctxs = fmap (map fst)
203 $ map splitTyConApp_maybe
210 dictionary pa pr ty tycon
211 | tycon == pa = paDictOfType ty
212 | tycon == pr = prDictOfReprType ty
213 | otherwise = invalid
215 invalid = cantVectorise "Invalid PR dfun type" (ppr (varType dfun) <+> ppr tys)