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