2 module Vectorise.Utils.PADict (
9 import Vectorise.Builtins
10 import Vectorise.Utils.Base
24 -- | Construct the PA argument type for the tyvar. For the tyvar (v :: *) it's
25 -- just PA v. For (v :: (* -> *) -> *) it's
27 -- > forall (a :: * -> *). (forall (b :: *). PA b -> PA (a b)) -> PA (v a)
29 paDictArgType :: TyVar -> VM (Maybe Type)
30 paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
32 go ty k | Just k' <- kindView k = go ty k'
35 tv <- newTyVar (fsLit "a") k1
36 mty1 <- go (TyVarTy tv) k1
39 mty2 <- go (AppTy ty (TyVarTy tv)) k2
40 return $ fmap (ForAllTy tv . FunTy ty1) mty2
45 = liftM Just (mkBuiltinTyConApp paTyCon [ty])
47 go _ _ = return Nothing
50 -- | Get the PA dictionary for some type
52 paDictOfType :: Type -> VM CoreExpr
54 = paDictOfTyApp ty_fn ty_args
56 (ty_fn, ty_args) = splitAppTys ty
58 paDictOfTyApp :: Type -> [Type] -> VM CoreExpr
59 paDictOfTyApp ty_fn ty_args
60 | Just ty_fn' <- coreView ty_fn
61 = paDictOfTyApp ty_fn' ty_args
63 -- for type variables, look up the dfun and apply to the PA dictionaries
64 -- of the type arguments
65 paDictOfTyApp (TyVarTy tv) ty_args
66 = do dfun <- maybeCantVectoriseM "No PA dictionary for type variable"
67 (ppr tv <+> text "in" <+> ppr ty)
69 dicts <- mapM paDictOfType ty_args
70 return $ dfun `mkTyApps` ty_args `mkApps` dicts
72 -- for tycons, we also need to apply the dfun to the PR dictionary of
73 -- the representation type if the tycon is polymorphic
74 paDictOfTyApp (TyConApp tc []) ty_args
76 dfun <- maybeCantVectoriseM "No PA dictionary for type constructor"
77 (ppr tc <+> text "in" <+> ppr ty)
79 super <- super_dict tc ty_args
80 dicts <- mapM paDictOfType ty_args
81 return $ Var dfun `mkTyApps` ty_args `mkApps` super `mkApps` dicts
83 paDictOfTyApp _ _ = failure
85 super_dict _ [] = return []
86 super_dict tycon ty_args
88 pr <- prDictOfPReprInst (TyConApp tycon ty_args)
91 failure = cantVectorise "Can't construct PA dictionary for type" (ppr ty)
93 paMethod :: (Builtins -> Var) -> String -> Type -> VM CoreExpr
95 | Just tycon <- splitPrimTyCon ty
97 . maybeCantVectoriseM "No PA method" (text name <+> text "for" <+> ppr tycon)
98 $ lookupPrimMethod tycon name
103 dict <- paDictOfType ty
104 return $ mkApps (Var fn) [Type ty, dict]
106 -- | Given a type @ty@, return the PR dictionary for @PRepr ty@.
107 prDictOfPReprInst :: Type -> VM CoreExpr
110 (prepr_tc, prepr_args) <- preprSynTyCon ty
111 case coreView (mkTyConApp prepr_tc prepr_args) of
113 dict <- prDictOfReprType rhs
114 pr_co <- mkBuiltinCo prTyCon
115 let Just arg_co = tyConFamilyCoercion_maybe prepr_tc
116 let co = mkAppCoercion pr_co
118 $ mkTyConApp arg_co prepr_args
119 return $ mkCoerce co dict
120 Nothing -> cantVectorise "Invalid PRepr type instance"
123 -- | Get the PR dictionary for a type. The argument must be a representation
125 prDictOfReprType :: Type -> VM CoreExpr
127 | Just (tycon, tyargs) <- splitTyConApp_maybe ty
129 prepr <- builtin preprTyCon
132 [ty'] <- return tyargs
133 prDictOfPReprInst ty'
135 -- a representation tycon must have a PR instance
136 dfun <- maybeCantVectoriseM
137 "No PR dictionary for type constructor"
138 (ppr tycon <+> text "in" <+> ppr ty)
139 $ lookupTyConPR tycon
140 prDFunApply dfun tyargs
144 -- it is a tyvar or an application of a tyvar
145 -- determine the PR dictionary from its PA dictionary
147 -- NOTE: This assumes that PRepr t ~ t is for all representation types
150 -- FIXME: This doesn't work for kinds other than * at the moment. We'd
151 -- have to simply abstract the term over the missing type arguments.
152 pa <- paDictOfType ty
153 prsel <- builtin paPRSel
154 return $ Var prsel `mkApps` [Type ty, pa]
156 -- | Apply a tycon's PR dfun to dictionary arguments (PR or PA) corresponding
157 -- to the argument types.
158 prDFunApply :: Var -> [Type] -> VM CoreExpr
160 | Just [] <- ctxs -- PR (a :-> b) doesn't have a context
161 = return $ Var dfun `mkTyApps` tys
163 | Just tycons <- ctxs
164 , length tycons == length tys
166 pa <- builtin paTyCon
167 pr <- builtin prTyCon
168 args <- zipWithM (dictionary pa pr) tys tycons
169 return $ Var dfun `mkTyApps` tys `mkApps` args
171 | otherwise = invalid
173 -- the dfun's contexts - if its type is (PA a, PR b) => PR (C a b) then
174 -- ctxs is Just [PA, PR]
175 ctxs = fmap (map fst)
177 $ map splitTyConApp_maybe
184 dictionary pa pr ty tycon
185 | tycon == pa = paDictOfType ty
186 | tycon == pr = prDictOfReprType ty
187 | otherwise = invalid
189 invalid = cantVectorise "Invalid PR dfun type" (ppr (varType dfun) <+> ppr tys)