compiler/vectorise/VectUtils.hs

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