compiler/vectorise/Vectorise/Utils/PADict.hs

   1
   2 module Vectorise.Utils.PADict (
   3         paDictArgType,
   4         paDictOfType,
   5         paMethod
   6 )
   7 where
   8 import Vectorise.Monad
   9 import Vectorise.Builtins
  10 import Vectorise.Utils.Base
  11
  12 import CoreSyn
  13 import CoreUtils
  14 import Coercion
  15 import Type
  16 import TypeRep
  17 import TyCon
  18 import Var
  19 import Outputable
  20 import FastString
  21 import Control.Monad
  22
  23
  24 paDictArgType :: TyVar -> VM (Maybe Type)
  25 paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
  26   where
  27     go ty k | Just k' <- kindView k = go ty k'
  28     go ty (FunTy k1 k2)
  29       = do
  30           tv   <- newTyVar (fsLit "a") k1
  31           mty1 <- go (TyVarTy tv) k1
  32           case mty1 of
  33             Just ty1 -> do
  34                           mty2 <- go (AppTy ty (TyVarTy tv)) k2
  35                           return $ fmap (ForAllTy tv . FunTy ty1) mty2
  36             Nothing  -> go ty k2
  37
  38     go ty k
  39       | isLiftedTypeKind k
  40       = liftM Just (mkBuiltinTyConApp paTyCon [ty])
  41
  42     go _ _ = return Nothing
  43
  44
  45 -- | Get the PA dictionary for some type
  46 paDictOfType :: Type -> VM CoreExpr
  47 paDictOfType ty
  48   = paDictOfTyApp ty_fn ty_args
  49   where
  50     (ty_fn, ty_args) = splitAppTys ty
  51
  52     paDictOfTyApp :: Type -> [Type] -> VM CoreExpr
  53     paDictOfTyApp ty_fn ty_args
  54         | Just ty_fn' <- coreView ty_fn
  55         = paDictOfTyApp ty_fn' ty_args
  56
  57     -- for type variables, look up the dfun and apply to the PA dictionaries
  58     -- of the type arguments
  59     paDictOfTyApp (TyVarTy tv) ty_args
  60      = do dfun <- maybeV (lookupTyVarPA tv)
  61           dicts <- mapM paDictOfType ty_args
  62           return $ dfun `mkTyApps` ty_args `mkApps` dicts
  63
  64     -- for tycons, we also need to apply the dfun to the PR dictionary of
  65     -- the representation type
  66     paDictOfTyApp (TyConApp tc []) ty_args
  67      = do
  68          dfun <- maybeV $ lookupTyConPA tc
  69          pr <- prDictOfPRepr tc ty_args
  70          dicts <- mapM paDictOfType ty_args
  71          return $ Var dfun `mkTyApps` ty_args `mkApps` (pr:dicts)
  72
  73     paDictOfTyApp _ _ = failure
  74
  75     failure = cantVectorise "Can't construct PA dictionary for type" (ppr ty)
  76
  77 paMethod :: (Builtins -> Var) -> String -> Type -> VM CoreExpr
  78 paMethod _ name ty
  79   | Just tycon <- splitPrimTyCon ty
  80   = liftM Var
  81   . maybeCantVectoriseM "No PA method" (text name <+> text "for" <+> ppr tycon)
  82   $ lookupPrimMethod tycon name
  83
  84 paMethod method _ ty
  85   = do
  86       fn   <- builtin method
  87       dict <- paDictOfType ty
  88       return $ mkApps (Var fn) [Type ty, dict]
  89
  90 -- | Get the PR (PRepr t) dictionary, where t is the tycon applied to the type
  91 -- arguments
  92 prDictOfPRepr :: TyCon -> [Type] -> VM CoreExpr
  93 prDictOfPRepr tycon tys
  94   = do
  95       (prepr_tc, prepr_args) <- preprSynTyCon (mkTyConApp tycon tys)
  96       case coreView (mkTyConApp prepr_tc prepr_args) of
  97         Just rhs -> do
  98                       dict <- prDictOfReprType rhs
  99                       pr_co <- mkBuiltinCo prTyCon
 100                       let Just arg_co = tyConFamilyCoercion_maybe prepr_tc
 101                       let co = mkAppCoercion pr_co
 102                              $ mkSymCoercion
 103                              $ mkTyConApp arg_co prepr_args
 104                       return $ mkCoerce co dict
 105         Nothing  -> cantVectorise "Invalid PRepr type instance"
 106                                   $ ppr $ mkTyConApp prepr_tc prepr_args
 107
 108 -- | Get the PR dictionary for a type. The argument must be a representation
 109 -- type.
 110 prDictOfReprType :: Type -> VM CoreExpr
 111 prDictOfReprType ty
 112   | Just (tycon, tyargs) <- splitTyConApp_maybe ty
 113     = do
 114         -- a representation tycon must have a PR instance
 115         dfun <- maybeV $ lookupTyConPR tycon
 116         prDFunApply dfun tyargs
 117
 118   | otherwise
 119     = do
 120         -- it is a tyvar or an application of a tyvar
 121         -- determine the PR dictionary from its PA dictionary
 122         --
 123         -- NOTE: This assumes that PRepr t ~ t is for all representation types
 124         -- t
 125         --
 126         -- FIXME: This doesn't work for kinds other than * at the moment. We'd
 127         -- have to simply abstract the term over the missing type arguments.
 128         pa    <- paDictOfType ty
 129         prsel <- builtin paPRSel
 130         return $ Var prsel `mkApps` [Type ty, pa]
 131
 132 -- | Apply a tycon's PR dfun to dictionary arguments (PR or PA) corresponding
 133 -- to the argument types.
 134 prDFunApply :: Var -> [Type] -> VM CoreExpr
 135 prDFunApply dfun tys
 136   | Just [] <- ctxs    -- PR (a :-> b) doesn't have a context
 137   = return $ Var dfun `mkTyApps` tys
 138
 139   | Just tycons <- ctxs
 140   , length tycons == length tys
 141   = do
 142       pa <- builtin paTyCon
 143       pr <- builtin prTyCon
 144       args <- zipWithM (dictionary pa pr) tys tycons
 145       return $ Var dfun `mkTyApps` tys `mkApps` args
 146
 147   | otherwise = invalid
 148   where
 149     -- the dfun's contexts - if its type is (PA a, PR b) => PR (C a b) then
 150     -- ctxs is Just [PA, PR]
 151     ctxs = fmap (map fst)
 152          $ sequence
 153          $ map splitTyConApp_maybe
 154          $ fst
 155          $ splitFunTys
 156          $ snd
 157          $ splitForAllTys
 158          $ varType dfun
 159
 160     dictionary pa pr ty tycon
 161       | tycon == pa = paDictOfType ty
 162       | tycon == pr = prDictOfReprType ty
 163       | otherwise   = invalid
 164
 165     invalid = cantVectorise "Invalid PR dfun type" (ppr (varType dfun) <+> ppr tys)
 166