module Vectorise.Utils.PADict (
mkPADictType,
paDictArgType,
paDictOfType,
paDFunType,
paDFunApply,
paMethod
)
where
import Vectorise.Monad
import Vectorise.Builtins
import Vectorise.Utils.Base
import CoreSyn
import Type
import TypeRep
import TyCon
import Var
import Outputable
import FastString
import Control.Monad
mkPADictType :: Type -> VM Type
mkPADictType ty = mkBuiltinTyConApp paTyCon [ty]
paDictArgType :: TyVar -> VM (Maybe Type)
paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
where
go ty k | Just k' <- kindView k = go ty k'
go ty (FunTy k1 k2)
= do
tv <- newTyVar (fsLit "a") k1
mty1 <- go (TyVarTy tv) k1
case mty1 of
Just ty1 -> do
mty2 <- go (AppTy ty (TyVarTy tv)) k2
return $ fmap (ForAllTy tv . FunTy ty1) mty2
Nothing -> go ty k2
go ty k
| isLiftedTypeKind k
= liftM Just (mkPADictType ty)
go _ _ = return Nothing
paDictOfType :: Type -> VM (Maybe CoreExpr)
paDictOfType ty
= paDictOfTyApp ty_fn ty_args
where
(ty_fn, ty_args) = splitAppTys ty
paDictOfTyApp :: Type -> [Type] -> VM (Maybe CoreExpr)
paDictOfTyApp ty_fn ty_args
| Just ty_fn' <- coreView ty_fn
= paDictOfTyApp ty_fn' ty_args
paDictOfTyApp (TyVarTy tv) ty_args
= do dfun <- maybeV (lookupTyVarPA tv)
liftM Just $ paDFunApply dfun ty_args
paDictOfTyApp (TyConApp tc _) ty_args
= do mdfun <- lookupTyConPA tc
case mdfun of
Nothing
-> pprTrace "VectUtils.paDictOfType"
(vcat [ text "No PA dictionary"
, text "for tycon: " <> ppr tc
, text "in type: " <> ppr ty])
$ return Nothing
Just dfun -> liftM Just $ paDFunApply (Var dfun) ty_args
paDictOfTyApp ty _
= cantVectorise "Can't construct PA dictionary for type" (ppr ty)
paDFunType :: TyCon -> VM Type
paDFunType tc
= do
margs <- mapM paDictArgType tvs
res <- mkPADictType (mkTyConApp tc arg_tys)
return . mkForAllTys tvs
$ mkFunTys [arg | Just arg <- margs] res
where
tvs = tyConTyVars tc
arg_tys = mkTyVarTys tvs
paDFunApply :: CoreExpr -> [Type] -> VM CoreExpr
paDFunApply dfun tys
= do Just dicts <- liftM sequence $ mapM paDictOfType tys
return $ mkApps (mkTyApps dfun tys) dicts
paMethod :: (Builtins -> Var) -> String -> Type -> VM CoreExpr
paMethod _ name ty
| Just tycon <- splitPrimTyCon ty
= liftM Var
. maybeCantVectoriseM "No PA method" (text name <+> text "for" <+> ppr tycon)
$ lookupPrimMethod tycon name
paMethod method _ ty
= do
fn <- builtin method
Just dict <- paDictOfType ty
return $ mkApps (Var fn) [Type ty, dict]