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)
34 go ty k | Just k' <- kindView k = go ty k'
37 tv <- newTyVar (fsLit "a") k1
38 mty1 <- go (TyVarTy tv) k1
41 mty2 <- go (AppTy ty (TyVarTy tv)) k2
42 return $ fmap (ForAllTy tv . FunTy ty1) mty2
48 pa_cls <- builtin paClass
49 return $ Just $ PredTy $ ClassP pa_cls [ty]
51 go _ _ = return Nothing
54 -- | Get the PA dictionary for some type
56 paDictOfType :: Type -> VM CoreExpr
58 = paDictOfTyApp ty_fn ty_args
60 (ty_fn, ty_args) = splitAppTys ty
62 paDictOfTyApp :: Type -> [Type] -> VM CoreExpr
63 paDictOfTyApp ty_fn ty_args
64 | Just ty_fn' <- coreView ty_fn
65 = paDictOfTyApp ty_fn' ty_args
67 -- for type variables, look up the dfun and apply to the PA dictionaries
68 -- of the type arguments
69 paDictOfTyApp (TyVarTy tv) ty_args
70 = do dfun <- maybeCantVectoriseM "No PA dictionary for type variable"
71 (ppr tv <+> text "in" <+> ppr ty)
73 dicts <- mapM paDictOfType ty_args
74 return $ dfun `mkTyApps` ty_args `mkApps` dicts
76 -- for tycons, we also need to apply the dfun to the PR dictionary of
77 -- the representation type if the tycon is polymorphic
78 paDictOfTyApp (TyConApp tc []) ty_args
80 dfun <- maybeCantVectoriseM "No PA dictionary for type constructor"
81 (ppr tc <+> text "in" <+> ppr ty)
83 super <- super_dict tc ty_args
84 dicts <- mapM paDictOfType ty_args
85 return $ Var dfun `mkTyApps` ty_args `mkApps` super `mkApps` dicts
87 paDictOfTyApp _ _ = failure
89 super_dict _ [] = return []
90 super_dict tycon ty_args
92 pr <- prDictOfPReprInst (TyConApp tycon ty_args)
95 failure = cantVectorise "Can't construct PA dictionary for type" (ppr ty)
97 paMethod :: (Builtins -> Var) -> String -> Type -> VM CoreExpr
99 | Just tycon <- splitPrimTyCon ty
101 . maybeCantVectoriseM "No PA method" (text name <+> text "for" <+> ppr tycon)
102 $ lookupPrimMethod tycon name
107 dict <- paDictOfType ty
108 return $ mkApps (Var fn) [Type ty, dict]
110 -- | Given a type @ty@, return the PR dictionary for @PRepr ty@.
111 prDictOfPReprInst :: Type -> VM CoreExpr
114 (prepr_tc, prepr_args) <- preprSynTyCon ty
115 prDictOfPReprInstTyCon ty prepr_tc prepr_args
117 -- | Given a type @ty@, its PRepr synonym tycon and its type arguments,
118 -- return the PR @PRepr ty@. Suppose we have:
120 -- > type instance PRepr (T a1 ... an) = t
122 -- which is internally translated into
124 -- > type :R:PRepr a1 ... an = t
126 -- and the corresponding coercion. Then,
128 -- > prDictOfPReprInstTyCon (T a1 ... an) :R:PRepr u1 ... un = PR (T u1 ... un)
130 -- Note that @ty@ is only used for error messages
132 prDictOfPReprInstTyCon :: Type -> TyCon -> [Type] -> VM CoreExpr
133 prDictOfPReprInstTyCon ty prepr_tc prepr_args
134 | Just rhs <- coreView (mkTyConApp prepr_tc prepr_args)
136 dict <- prDictOfReprType' rhs
137 pr_co <- mkBuiltinCo prTyCon
138 let Just arg_co = tyConFamilyCoercion_maybe prepr_tc
139 let co = mkAppCoercion pr_co
141 $ mkTyConApp arg_co prepr_args
142 return $ mkCoerce co dict
144 | otherwise = cantVectorise "Invalid PRepr type instance" (ppr ty)
146 -- | Get the PR dictionary for a type. The argument must be a representation
148 prDictOfReprType :: Type -> VM CoreExpr
150 | Just (tycon, tyargs) <- splitTyConApp_maybe ty
152 prepr <- builtin preprTyCon
156 pa <- paDictOfType ty'
157 sel <- builtin paPRSel
158 return $ Var sel `App` Type ty' `App` pa
160 -- a representation tycon must have a PR instance
161 dfun <- maybeV $ lookupTyConPR tycon
162 prDFunApply dfun tyargs
166 -- it is a tyvar or an application of a tyvar
167 -- determine the PR dictionary from its PA dictionary
169 -- NOTE: This assumes that PRepr t ~ t is for all representation types
172 -- FIXME: This doesn't work for kinds other than * at the moment. We'd
173 -- have to simply abstract the term over the missing type arguments.
174 pa <- paDictOfType ty
175 prsel <- builtin paPRSel
176 return $ Var prsel `mkApps` [Type ty, pa]
178 prDictOfReprType' :: Type -> VM CoreExpr
179 prDictOfReprType' ty = prDictOfReprType ty `orElseV`
180 cantVectorise "No PR dictionary for representation type"
183 -- | Apply a tycon's PR dfun to dictionary arguments (PR or PA) corresponding
184 -- to the argument types.
185 prDFunApply :: Var -> [Type] -> VM CoreExpr
187 | Just [] <- ctxs -- PR (a :-> b) doesn't have a context
188 = return $ Var dfun `mkTyApps` tys
190 | Just tycons <- ctxs
191 , length tycons == length tys
193 pa <- builtin paTyCon
194 pr <- builtin prTyCon
195 args <- zipWithM (dictionary pa pr) tys tycons
196 return $ Var dfun `mkTyApps` tys `mkApps` args
198 | otherwise = invalid
200 -- the dfun's contexts - if its type is (PA a, PR b) => PR (C a b) then
201 -- ctxs is Just [PA, PR]
202 ctxs = fmap (map fst)
204 $ map splitTyConApp_maybe
211 dictionary pa pr ty tycon
212 | tycon == pa = paDictOfType ty
213 | tycon == pr = prDictOfReprType ty
214 | otherwise = invalid
216 invalid = cantVectorise "Invalid PR dfun type" (ppr (varType dfun) <+> ppr tys)