compiler/vectorise/Vectorise/Utils/PADict.hs

   1
   2 module Vectorise.Utils.PADict (
   3         mkPADictType,
   4         paDictArgType,
   5         paDictOfType,
   6         paDFunType,
   7         paDFunApply,
   8         paMethod
   9 )
  10 where
  11 import Vectorise.Monad
  12 import Vectorise.Builtins
  13 import Vectorise.Utils.Base
  14
  15 import CoreSyn
  16 import Type
  17 import TypeRep
  18 import TyCon
  19 import Var
  20 import Outputable
  21 import FastString
  22 import Control.Monad
  23
  24
  25 mkPADictType :: Type -> VM Type
  26 mkPADictType ty = mkBuiltinTyConApp paTyCon [ty]
  27
  28
  29 paDictArgType :: TyVar -> VM (Maybe Type)
  30 paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
  31   where
  32     go ty k | Just k' <- kindView k = go ty k'
  33     go ty (FunTy k1 k2)
  34       = do
  35           tv   <- newTyVar (fsLit "a") k1
  36           mty1 <- go (TyVarTy tv) k1
  37           case mty1 of
  38             Just ty1 -> do
  39                           mty2 <- go (AppTy ty (TyVarTy tv)) k2
  40                           return $ fmap (ForAllTy tv . FunTy ty1) mty2
  41             Nothing  -> go ty k2
  42
  43     go ty k
  44       | isLiftedTypeKind k
  45       = liftM Just (mkPADictType ty)
  46
  47     go _ _ = return Nothing
  48
  49
  50 -- | Get the PA dictionary for some type, or `Nothing` if there isn't one.
  51 paDictOfType :: Type -> VM (Maybe CoreExpr)
  52 paDictOfType ty
  53   = paDictOfTyApp ty_fn ty_args
  54   where
  55     (ty_fn, ty_args) = splitAppTys ty
  56
  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
  61
  62     paDictOfTyApp (TyVarTy tv) ty_args
  63      = do dfun <- maybeV (lookupTyVarPA tv)
  64           liftM Just $ paDFunApply dfun ty_args
  65
  66     paDictOfTyApp (TyConApp tc _) ty_args
  67      = do mdfun <- lookupTyConPA tc
  68           case mdfun of
  69             Nothing
  70              -> pprTrace "VectUtils.paDictOfType"
  71                          (vcat [ text "No PA dictionary"
  72                                , text "for tycon: " <> ppr tc
  73                                , text "in type:   " <> ppr ty])
  74              $ return Nothing
  75
  76             Just dfun   -> liftM Just $ paDFunApply (Var dfun) ty_args
  77
  78     paDictOfTyApp ty _
  79      = cantVectorise "Can't construct PA dictionary for type" (ppr ty)
  80
  81
  82
  83 paDFunType :: TyCon -> VM Type
  84 paDFunType tc
  85   = do
  86       margs <- mapM paDictArgType tvs
  87       res   <- mkPADictType (mkTyConApp tc arg_tys)
  88       return . mkForAllTys tvs
  89              $ mkFunTys [arg | Just arg <- margs] res
  90   where
  91     tvs = tyConTyVars tc
  92     arg_tys = mkTyVarTys tvs
  93
  94 paDFunApply :: CoreExpr -> [Type] -> VM CoreExpr
  95 paDFunApply dfun tys
  96  = do Just dicts <- liftM sequence $ mapM paDictOfType tys
  97       return $ mkApps (mkTyApps dfun tys) dicts
  98
  99
 100 paMethod :: (Builtins -> Var) -> String -> Type -> VM CoreExpr
 101 paMethod _ name ty
 102   | Just tycon <- splitPrimTyCon ty
 103   = liftM Var
 104   . maybeCantVectoriseM "No PA method" (text name <+> text "for" <+> ppr tycon)
 105   $ lookupPrimMethod tycon name
 106
 107 paMethod method _ ty
 108   = do
 109       fn        <- builtin method
 110       Just dict <- paDictOfType ty
 111       return $ mkApps (Var fn) [Type ty, dict]
 112