compiler/vectorise/VectUtils.hs

   1 module VectUtils (
   2   splitClosureTy,
   3   mkPADictType, mkPArrayType,
   4   paDictArgType, paDictOfType
   5 ) where
   6
   7 #include "HsVersions.h"
   8
   9 import VectMonad
  10
  11 import CoreSyn
  12 import Type
  13 import TypeRep
  14 import TyCon
  15 import Var
  16 import PrelNames
  17
  18 import Outputable
  19
  20 import Control.Monad         ( liftM )
  21
  22 isClosureTyCon :: TyCon -> Bool
  23 isClosureTyCon tc = tyConUnique tc == closureTyConKey
  24
  25 splitClosureTy :: Type -> (Type, Type)
  26 splitClosureTy ty
  27   | Just (tc, [arg_ty, res_ty]) <- splitTyConApp_maybe ty
  28   , isClosureTyCon tc
  29   = (arg_ty, res_ty)
  30
  31   | otherwise = pprPanic "splitClosureTy" (ppr ty)
  32
  33 mkPADictType :: Type -> VM Type
  34 mkPADictType ty
  35   = do
  36       tc <- builtin paDictTyCon
  37       return $ TyConApp tc [ty]
  38
  39 mkPArrayType :: Type -> VM Type
  40 mkPArrayType ty
  41   = do
  42       tc <- builtin parrayTyCon
  43       return $ TyConApp tc [ty]
  44
  45 paDictArgType :: TyVar -> VM (Maybe Type)
  46 paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
  47   where
  48     go ty k | Just k' <- kindView k = go ty k'
  49     go ty (FunTy k1 k2)
  50       = do
  51           tv   <- newTyVar FSLIT("a") k1
  52           mty1 <- go (TyVarTy tv) k1
  53           case mty1 of
  54             Just ty1 -> do
  55                           mty2 <- go (AppTy ty (TyVarTy tv)) k2
  56                           return $ fmap (ForAllTy tv . FunTy ty1) mty2
  57             Nothing  -> go ty k2
  58
  59     go ty k
  60       | isLiftedTypeKind k
  61       = liftM Just (mkPADictType ty)
  62
  63     go ty k = return Nothing
  64
  65 paDictOfType :: Type -> VM CoreExpr
  66 paDictOfType ty = paDictOfTyApp ty_fn ty_args
  67   where
  68     (ty_fn, ty_args) = splitAppTys ty
  69
  70 paDictOfTyApp :: Type -> [Type] -> VM CoreExpr
  71 paDictOfTyApp ty_fn ty_args
  72   | Just ty_fn' <- coreView ty_fn = paDictOfTyApp ty_fn' ty_args
  73 paDictOfTyApp (TyVarTy tv) ty_args
  74   = do
  75       dfun <- maybeV (lookupTyVarPA tv)
  76       paDFunApply dfun ty_args
  77 paDictOfTyApp (TyConApp tc _) ty_args
  78   = do
  79       pa_class <- builtin paClass
  80       (dfun, ty_args') <- lookupInst pa_class [TyConApp tc ty_args]
  81       paDFunApply (Var dfun) ty_args'
  82 paDictOfTyApp ty ty_args = pprPanic "paDictOfTyApp" (ppr ty)
  83
  84 paDFunApply :: CoreExpr -> [Type] -> VM CoreExpr
  85 paDFunApply dfun tys
  86   = do
  87       dicts <- mapM paDictOfType tys
  88       return $ mkApps (mkTyApps dfun tys) dicts
  89