2 module Vectorise.Utils.PADict (
11 import Vectorise.Monad
12 import Vectorise.Builtins
13 import Vectorise.Utils.Base
26 mkPADictType :: Type -> VM Type
27 mkPADictType ty = mkBuiltinTyConApp paTyCon [ty]
30 paDictArgType :: TyVar -> VM (Maybe Type)
31 paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
33 go ty k | Just k' <- kindView k = go ty k'
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
46 = liftM Just (mkPADictType ty)
48 go _ _ = return Nothing
51 -- | Get the PA dictionary for some type, or `Nothing` if there isn't one.
52 paDictOfType :: Type -> VM (Maybe CoreExpr)
54 = paDictOfTyApp ty_fn ty_args
56 (ty_fn, ty_args) = splitAppTys ty
58 paDictOfTyApp :: Type -> [Type] -> VM (Maybe CoreExpr)
59 paDictOfTyApp ty_fn ty_args
60 | Just ty_fn' <- coreView ty_fn
61 = paDictOfTyApp ty_fn' ty_args
63 paDictOfTyApp (TyVarTy tv) ty_args
64 = do dfun <- maybeV (lookupTyVarPA tv)
65 liftM Just $ paDFunApply dfun ty_args
67 paDictOfTyApp (TyConApp tc _) ty_args
68 = do mdfun <- lookupTyConPA tc
71 -> pprTrace "VectUtils.paDictOfType"
72 (vcat [ text "No PA dictionary"
73 , text "for tycon: " <> ppr tc
74 , text "in type: " <> ppr ty])
77 Just dfun -> liftM Just $ paDFunApply (Var dfun) ty_args
80 = cantVectorise "Can't construct PA dictionary for type" (ppr ty)
84 paDFunType :: TyCon -> VM Type
87 margs <- mapM paDictArgType tvs
88 res <- mkPADictType (mkTyConApp tc arg_tys)
89 return . mkForAllTys tvs
90 $ mkFunTys [arg | Just arg <- margs] res
93 arg_tys = mkTyVarTys tvs
95 paDFunApply :: CoreExpr -> [Type] -> VM CoreExpr
97 = do Just dicts <- liftM sequence $ mapM paDictOfType tys
98 return $ mkApps (mkTyApps dfun tys) dicts
101 paMethod :: (Builtins -> Var) -> String -> Type -> VM CoreExpr
103 | Just tycon <- splitPrimTyCon ty
105 . maybeCantVectoriseM "No PA method" (text name <+> text "for" <+> ppr tycon)
106 $ lookupPrimMethod tycon name
111 Just dict <- paDictOfType ty
112 return $ mkApps (Var fn) [Type ty, dict]