2 module Vectorise.Utils.PADict (
11 import Vectorise.Monad
12 import Vectorise.Builtins
13 import Vectorise.Utils.Base
25 mkPADictType :: Type -> VM Type
26 mkPADictType ty = mkBuiltinTyConApp paTyCon [ty]
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 (mkPADictType ty)
47 go _ _ = return Nothing
50 -- | Get the PA dictionary for some type, or `Nothing` if there isn't one.
51 paDictOfType :: Type -> VM (Maybe CoreExpr)
53 = paDictOfTyApp ty_fn ty_args
55 (ty_fn, ty_args) = splitAppTys ty
57 paDictOfTyApp :: Type -> [Type] -> VM (Maybe CoreExpr)
58 paDictOfTyApp ty_fn ty_args
59 | Just ty_fn' <- coreView ty_fn
60 = paDictOfTyApp ty_fn' ty_args
62 paDictOfTyApp (TyVarTy tv) ty_args
63 = do dfun <- maybeV (lookupTyVarPA tv)
64 liftM Just $ paDFunApply dfun ty_args
66 paDictOfTyApp (TyConApp tc _) ty_args
67 = do mdfun <- lookupTyConPA tc
70 -> pprTrace "VectUtils.paDictOfType"
71 (vcat [ text "No PA dictionary"
72 , text "for tycon: " <> ppr tc
73 , text "in type: " <> ppr ty])
76 Just dfun -> liftM Just $ paDFunApply (Var dfun) ty_args
79 = cantVectorise "Can't construct PA dictionary for type" (ppr ty)
83 paDFunType :: TyCon -> VM Type
86 margs <- mapM paDictArgType tvs
87 res <- mkPADictType (mkTyConApp tc arg_tys)
88 return . mkForAllTys tvs
89 $ mkFunTys [arg | Just arg <- margs] res
92 arg_tys = mkTyVarTys tvs
94 paDFunApply :: CoreExpr -> [Type] -> VM CoreExpr
96 = do Just dicts <- liftM sequence $ mapM paDictOfType tys
97 return $ mkApps (mkTyApps dfun tys) dicts
100 paMethod :: (Builtins -> Var) -> String -> Type -> VM CoreExpr
102 | Just tycon <- splitPrimTyCon ty
104 . maybeCantVectoriseM "No PA method" (text name <+> text "for" <+> ppr tycon)
105 $ lookupPrimMethod tycon name
110 Just dict <- paDictOfType ty
111 return $ mkApps (Var fn) [Type ty, dict]