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