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 -- | Construct the PA argument type for the tyvar. For the tyvar (v :: *) it's
  25 -- just PA v. For (v :: (* -> *) -> *) it's
  26 --
  27 -- > forall (a :: * -> *). (forall (b :: *). PA b -> PA (a b)) -> PA (v a)
  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 (mkBuiltinTyConApp paTyCon [ty])
  46
  47     go _ _ = return Nothing
  48
  49
  50 -- | Get the PA dictionary for some type
  51 --
  52 paDictOfType :: Type -> VM 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 CoreExpr
  59     paDictOfTyApp ty_fn ty_args
  60         | Just ty_fn' <- coreView ty_fn
  61         = paDictOfTyApp ty_fn' ty_args
  62
  63     -- for type variables, look up the dfun and apply to the PA dictionaries
  64     -- of the type arguments
  65     paDictOfTyApp (TyVarTy tv) ty_args
  66      = do dfun <- maybeCantVectoriseM "No PA dictionary for type variable"
  67                                       (ppr tv <+> text "in" <+> ppr ty)
  68                 $ lookupTyVarPA tv
  69           dicts <- mapM paDictOfType ty_args
  70           return $ dfun `mkTyApps` ty_args `mkApps` dicts
  71
  72     -- for tycons, we also need to apply the dfun to the PR dictionary of
  73     -- the representation type if the tycon is polymorphic
  74     paDictOfTyApp (TyConApp tc []) ty_args
  75      = do
  76          dfun <- maybeCantVectoriseM "No PA dictionary for type constructor"
  77                                       (ppr tc <+> text "in" <+> ppr ty)
  78                 $ lookupTyConPA tc
  79          super <- super_dict tc ty_args
  80          dicts <- mapM paDictOfType ty_args
  81          return $ Var dfun `mkTyApps` ty_args `mkApps` super `mkApps` dicts
  82
  83     paDictOfTyApp _ _ = failure
  84
  85     super_dict _ [] = return []
  86     super_dict tycon ty_args
  87       = do
  88           pr <- prDictOfPReprInst (TyConApp tycon ty_args)
  89           return [pr]
  90
  91     failure = cantVectorise "Can't construct PA dictionary for type" (ppr ty)
  92
  93 paMethod :: (Builtins -> Var) -> String -> Type -> VM CoreExpr
  94 paMethod _ name ty
  95   | Just tycon <- splitPrimTyCon ty
  96   = liftM Var
  97   . maybeCantVectoriseM "No PA method" (text name <+> text "for" <+> ppr tycon)
  98   $ lookupPrimMethod tycon name
  99
 100 paMethod method _ ty
 101   = do
 102       fn   <- builtin method
 103       dict <- paDictOfType ty
 104       return $ mkApps (Var fn) [Type ty, dict]
 105
 106 -- | Given a type @ty@, return the PR dictionary for @PRepr ty@.
 107 prDictOfPReprInst :: Type -> VM CoreExpr
 108 prDictOfPReprInst ty
 109   = do
 110       (prepr_tc, prepr_args) <- preprSynTyCon ty
 111       case coreView (mkTyConApp prepr_tc prepr_args) of
 112         Just rhs -> do
 113                       dict <- prDictOfReprType rhs
 114                       pr_co <- mkBuiltinCo prTyCon
 115                       let Just arg_co = tyConFamilyCoercion_maybe prepr_tc
 116                       let co = mkAppCoercion pr_co
 117                              $ mkSymCoercion
 118                              $ mkTyConApp arg_co prepr_args
 119                       return $ mkCoerce co dict
 120         Nothing  -> cantVectorise "Invalid PRepr type instance"
 121                                   $ ppr ty
 122
 123 -- | Get the PR dictionary for a type. The argument must be a representation
 124 -- type.
 125 prDictOfReprType :: Type -> VM CoreExpr
 126 prDictOfReprType ty
 127   | Just (tycon, tyargs) <- splitTyConApp_maybe ty
 128     = do
 129         prepr <- builtin preprTyCon
 130         if tycon == prepr
 131           then do
 132                  [ty'] <- return tyargs
 133                  prDictOfPReprInst ty'
 134           else do
 135                  -- a representation tycon must have a PR instance
 136                  dfun <- maybeCantVectoriseM
 137                            "No PR dictionary for type constructor"
 138                            (ppr tycon <+> text "in" <+> ppr ty)
 139                        $ lookupTyConPR tycon
 140                  prDFunApply dfun tyargs
 141
 142   | otherwise
 143     = do
 144         -- it is a tyvar or an application of a tyvar
 145         -- determine the PR dictionary from its PA dictionary
 146         --
 147         -- NOTE: This assumes that PRepr t ~ t is for all representation types
 148         -- t
 149         --
 150         -- FIXME: This doesn't work for kinds other than * at the moment. We'd
 151         -- have to simply abstract the term over the missing type arguments.
 152         pa    <- paDictOfType ty
 153         prsel <- builtin paPRSel
 154         return $ Var prsel `mkApps` [Type ty, pa]
 155
 156 -- | Apply a tycon's PR dfun to dictionary arguments (PR or PA) corresponding
 157 -- to the argument types.
 158 prDFunApply :: Var -> [Type] -> VM CoreExpr
 159 prDFunApply dfun tys
 160   | Just [] <- ctxs    -- PR (a :-> b) doesn't have a context
 161   = return $ Var dfun `mkTyApps` tys
 162
 163   | Just tycons <- ctxs
 164   , length tycons == length tys
 165   = do
 166       pa <- builtin paTyCon
 167       pr <- builtin prTyCon
 168       args <- zipWithM (dictionary pa pr) tys tycons
 169       return $ Var dfun `mkTyApps` tys `mkApps` args
 170
 171   | otherwise = invalid
 172   where
 173     -- the dfun's contexts - if its type is (PA a, PR b) => PR (C a b) then
 174     -- ctxs is Just [PA, PR]
 175     ctxs = fmap (map fst)
 176          $ sequence
 177          $ map splitTyConApp_maybe
 178          $ fst
 179          $ splitFunTys
 180          $ snd
 181          $ splitForAllTys
 182          $ varType dfun
 183
 184     dictionary pa pr ty tycon
 185       | tycon == pa = paDictOfType ty
 186       | tycon == pr = prDictOfReprType ty
 187       | otherwise   = invalid
 188
 189     invalid = cantVectorise "Invalid PR dfun type" (ppr (varType dfun) <+> ppr tys)
 190