compiler/vectorise/Vectorise/Utils/PADict.hs

   1
   2 module Vectorise.Utils.PADict (
   3         paDictArgType,
   4         paDictOfType,
   5         paMethod,
   6         prDictOfReprType,
   7         prDictOfPReprInstTyCon
   8 )
   9 where
  10 import Vectorise.Monad
  11 import Vectorise.Builtins
  12 import Vectorise.Utils.Base
  13
  14 import CoreSyn
  15 import CoreUtils
  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 -- | Construct the PA argument type for the tyvar. For the tyvar (v :: *) it's
  27 -- just PA v. For (v :: (* -> *) -> *) it's
  28 --
  29 -- > forall (a :: * -> *). (forall (b :: *). PA b -> PA (a b)) -> PA (v a)
  30 --
  31 paDictArgType :: TyVar -> VM (Maybe Type)
  32 paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
  33   where
  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       = do
  47           pa_cls <- builtin paClass
  48           return $ Just $ PredTy $ ClassP pa_cls [ty]
  49
  50     go _ _ = return Nothing
  51
  52
  53 -- | Get the PA dictionary for some type
  54 --
  55 paDictOfType :: Type -> VM CoreExpr
  56 paDictOfType ty
  57   = paDictOfTyApp ty_fn ty_args
  58   where
  59     (ty_fn, ty_args) = splitAppTys ty
  60
  61     paDictOfTyApp :: Type -> [Type] -> VM CoreExpr
  62     paDictOfTyApp ty_fn ty_args
  63         | Just ty_fn' <- coreView ty_fn
  64         = paDictOfTyApp ty_fn' ty_args
  65
  66     -- for type variables, look up the dfun and apply to the PA dictionaries
  67     -- of the type arguments
  68     paDictOfTyApp (TyVarTy tv) ty_args
  69      = do dfun <- maybeCantVectoriseM "No PA dictionary for type variable"
  70                                       (ppr tv <+> text "in" <+> ppr ty)
  71                 $ lookupTyVarPA tv
  72           dicts <- mapM paDictOfType ty_args
  73           return $ dfun `mkTyApps` ty_args `mkApps` dicts
  74
  75     -- for tycons, we also need to apply the dfun to the PR dictionary of
  76     -- the representation type if the tycon is polymorphic
  77     paDictOfTyApp (TyConApp tc []) ty_args
  78      = do
  79          dfun <- maybeCantVectoriseM "No PA dictionary for type constructor"
  80                                       (ppr tc <+> text "in" <+> ppr ty)
  81                 $ lookupTyConPA tc
  82          super <- super_dict tc ty_args
  83          dicts <- mapM paDictOfType ty_args
  84          return $ Var dfun `mkTyApps` ty_args `mkApps` super `mkApps` dicts
  85
  86     paDictOfTyApp _ _ = failure
  87
  88     super_dict _ [] = return []
  89     super_dict tycon ty_args
  90       = do
  91           pr <- prDictOfPReprInst (TyConApp tycon ty_args)
  92           return [pr]
  93
  94     failure = cantVectorise "Can't construct PA dictionary for type" (ppr ty)
  95
  96 paMethod :: (Builtins -> Var) -> String -> Type -> VM CoreExpr
  97 paMethod _ name ty
  98   | Just tycon <- splitPrimTyCon ty
  99   = liftM Var
 100   . maybeCantVectoriseM "No PA method" (text name <+> text "for" <+> ppr tycon)
 101   $ lookupPrimMethod tycon name
 102
 103 paMethod method _ ty
 104   = do
 105       fn   <- builtin method
 106       dict <- paDictOfType ty
 107       return $ mkApps (Var fn) [Type ty, dict]
 108
 109 -- | Given a type @ty@, return the PR dictionary for @PRepr ty@.
 110 prDictOfPReprInst :: Type -> VM CoreExpr
 111 prDictOfPReprInst ty
 112   = do
 113       (prepr_tc, prepr_args) <- preprSynTyCon ty
 114       prDictOfPReprInstTyCon ty prepr_tc prepr_args
 115
 116 -- | Given a type @ty@, its PRepr synonym tycon and its type arguments,
 117 -- return the PR @PRepr ty@. Suppose we have:
 118 --
 119 -- > type instance PRepr (T a1 ... an) = t
 120 --
 121 -- which is internally translated into
 122 --
 123 -- > type :R:PRepr a1 ... an = t
 124 --
 125 -- and the corresponding coercion. Then,
 126 --
 127 -- > prDictOfPReprInstTyCon (T a1 ... an) :R:PRepr u1 ... un = PR (T u1 ... un)
 128 --
 129 -- Note that @ty@ is only used for error messages
 130 --
 131 prDictOfPReprInstTyCon :: Type -> TyCon -> [Type] -> VM CoreExpr
 132 prDictOfPReprInstTyCon ty prepr_tc prepr_args
 133   | Just rhs <- coreView (mkTyConApp prepr_tc prepr_args)
 134   = do
 135       dict <- prDictOfReprType' rhs
 136       pr_co <- mkBuiltinCo prTyCon
 137       let Just arg_co = tyConFamilyCoercion_maybe prepr_tc
 138       let co = mkAppCo pr_co
 139              $ mkSymCo
 140              $ mkAxInstCo arg_co prepr_args
 141       return $ mkCoerce co dict
 142
 143   | otherwise = cantVectorise "Invalid PRepr type instance" (ppr ty)
 144
 145 -- | Get the PR dictionary for a type. The argument must be a representation
 146 -- type.
 147 prDictOfReprType :: Type -> VM CoreExpr
 148 prDictOfReprType ty
 149   | Just (tycon, tyargs) <- splitTyConApp_maybe ty
 150     = do
 151         prepr <- builtin preprTyCon
 152         if tycon == prepr
 153           then do
 154                  let [ty'] = tyargs
 155                  pa <- paDictOfType ty'
 156                  sel <- builtin paPRSel
 157                  return $ Var sel `App` Type ty' `App` pa
 158           else do
 159                  -- a representation tycon must have a PR instance
 160                  dfun <- maybeV $ lookupTyConPR tycon
 161                  prDFunApply dfun tyargs
 162
 163   | otherwise
 164     = do
 165         -- it is a tyvar or an application of a tyvar
 166         -- determine the PR dictionary from its PA dictionary
 167         --
 168         -- NOTE: This assumes that PRepr t ~ t is for all representation types
 169         -- t
 170         --
 171         -- FIXME: This doesn't work for kinds other than * at the moment. We'd
 172         -- have to simply abstract the term over the missing type arguments.
 173         pa    <- paDictOfType ty
 174         prsel <- builtin paPRSel
 175         return $ Var prsel `mkApps` [Type ty, pa]
 176
 177 prDictOfReprType' :: Type -> VM CoreExpr
 178 prDictOfReprType' ty = prDictOfReprType ty `orElseV`
 179                        cantVectorise "No PR dictionary for representation type"
 180                                      (ppr ty)
 181
 182 -- | Apply a tycon's PR dfun to dictionary arguments (PR or PA) corresponding
 183 -- to the argument types.
 184 prDFunApply :: Var -> [Type] -> VM CoreExpr
 185 prDFunApply dfun tys
 186   | Just [] <- ctxs    -- PR (a :-> b) doesn't have a context
 187   = return $ Var dfun `mkTyApps` tys
 188
 189   | Just tycons <- ctxs
 190   , length tycons == length tys
 191   = do
 192       pa <- builtin paTyCon
 193       pr <- builtin prTyCon
 194       args <- zipWithM (dictionary pa pr) tys tycons
 195       return $ Var dfun `mkTyApps` tys `mkApps` args
 196
 197   | otherwise = invalid
 198   where
 199     -- the dfun's contexts - if its type is (PA a, PR b) => PR (C a b) then
 200     -- ctxs is Just [PA, PR]
 201     ctxs = fmap (map fst)
 202          $ sequence
 203          $ map splitTyConApp_maybe
 204          $ fst
 205          $ splitFunTys
 206          $ snd
 207          $ splitForAllTys
 208          $ varType dfun
 209
 210     dictionary pa pr ty tycon
 211       | tycon == pa = paDictOfType ty
 212       | tycon == pr = prDictOfReprType ty
 213       | otherwise   = invalid
 214
 215     invalid = cantVectorise "Invalid PR dfun type" (ppr (varType dfun) <+> ppr tys)
 216