compiler/vectorise/VectUtils.hs

   1 module VectUtils (
   2   collectAnnTypeBinders, collectAnnTypeArgs, isAnnTypeArg,
   3   splitClosureTy,
   4   mkPADictType, mkPArrayType,
   5   paDictArgType, paDictOfType,
   6   paMethod, lengthPA, replicatePA, emptyPA,
   7   lookupPArrayFamInst,
   8   hoistExpr, takeHoisted
   9 ) where
  10
  11 #include "HsVersions.h"
  12
  13 import VectMonad
  14
  15 import CoreSyn
  16 import CoreUtils
  17 import Type
  18 import TypeRep
  19 import TyCon
  20 import Var
  21 import PrelNames
  22
  23 import Outputable
  24 import FastString
  25
  26 import Control.Monad         ( liftM )
  27
  28 collectAnnTypeArgs :: AnnExpr b ann -> (AnnExpr b ann, [Type])
  29 collectAnnTypeArgs expr = go expr []
  30   where
  31     go (_, AnnApp f (_, AnnType ty)) tys = go f (ty : tys)
  32     go e                             tys = (e, tys)
  33
  34 collectAnnTypeBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann)
  35 collectAnnTypeBinders expr = go [] expr
  36   where
  37     go bs (_, AnnLam b e) | isTyVar b = go (b:bs) e
  38     go bs e                           = (reverse bs, e)
  39
  40 isAnnTypeArg :: AnnExpr b ann -> Bool
  41 isAnnTypeArg (_, AnnType t) = True
  42 isAnnTypeArg _              = False
  43
  44 isClosureTyCon :: TyCon -> Bool
  45 isClosureTyCon tc = tyConUnique tc == closureTyConKey
  46
  47 splitClosureTy :: Type -> (Type, Type)
  48 splitClosureTy ty
  49   | Just (tc, [arg_ty, res_ty]) <- splitTyConApp_maybe ty
  50   , isClosureTyCon tc
  51   = (arg_ty, res_ty)
  52
  53   | otherwise = pprPanic "splitClosureTy" (ppr ty)
  54
  55 mkPADictType :: Type -> VM Type
  56 mkPADictType ty
  57   = do
  58       tc <- builtin paDictTyCon
  59       return $ TyConApp tc [ty]
  60
  61 mkPArrayType :: Type -> VM Type
  62 mkPArrayType ty
  63   = do
  64       tc <- builtin parrayTyCon
  65       return $ TyConApp tc [ty]
  66
  67 paDictArgType :: TyVar -> VM (Maybe Type)
  68 paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
  69   where
  70     go ty k | Just k' <- kindView k = go ty k'
  71     go ty (FunTy k1 k2)
  72       = do
  73           tv   <- newTyVar FSLIT("a") k1
  74           mty1 <- go (TyVarTy tv) k1
  75           case mty1 of
  76             Just ty1 -> do
  77                           mty2 <- go (AppTy ty (TyVarTy tv)) k2
  78                           return $ fmap (ForAllTy tv . FunTy ty1) mty2
  79             Nothing  -> go ty k2
  80
  81     go ty k
  82       | isLiftedTypeKind k
  83       = liftM Just (mkPADictType ty)
  84
  85     go ty k = return Nothing
  86
  87 paDictOfType :: Type -> VM CoreExpr
  88 paDictOfType ty = paDictOfTyApp ty_fn ty_args
  89   where
  90     (ty_fn, ty_args) = splitAppTys ty
  91
  92 paDictOfTyApp :: Type -> [Type] -> VM CoreExpr
  93 paDictOfTyApp ty_fn ty_args
  94   | Just ty_fn' <- coreView ty_fn = paDictOfTyApp ty_fn' ty_args
  95 paDictOfTyApp (TyVarTy tv) ty_args
  96   = do
  97       dfun <- maybeV (lookupTyVarPA tv)
  98       paDFunApply dfun ty_args
  99 paDictOfTyApp (TyConApp tc _) ty_args
 100   = do
 101       pa_class <- builtin paClass
 102       (dfun, ty_args') <- lookupInst pa_class [TyConApp tc ty_args]
 103       paDFunApply (Var dfun) ty_args'
 104 paDictOfTyApp ty ty_args = pprPanic "paDictOfTyApp" (ppr ty)
 105
 106 paDFunApply :: CoreExpr -> [Type] -> VM CoreExpr
 107 paDFunApply dfun tys
 108   = do
 109       dicts <- mapM paDictOfType tys
 110       return $ mkApps (mkTyApps dfun tys) dicts
 111
 112 paMethod :: (Builtins -> Var) -> Type -> VM CoreExpr
 113 paMethod method ty
 114   = do
 115       fn   <- builtin method
 116       dict <- paDictOfType ty
 117       return $ mkApps (Var fn) [Type ty, dict]
 118
 119 lengthPA :: CoreExpr -> VM CoreExpr
 120 lengthPA x = liftM (`App` x) (paMethod lengthPAVar (exprType x))
 121
 122 replicatePA :: CoreExpr -> CoreExpr -> VM CoreExpr
 123 replicatePA len x = liftM (`mkApps` [len,x])
 124                           (paMethod replicatePAVar (exprType x))
 125
 126 emptyPA :: Type -> VM CoreExpr
 127 emptyPA = paMethod emptyPAVar
 128
 129 lookupPArrayFamInst :: Type -> VM (TyCon, [Type])
 130 lookupPArrayFamInst ty = builtin parrayTyCon >>= (`lookupFamInst` [ty])
 131
 132 hoistExpr :: FastString -> CoreExpr -> VM Var
 133 hoistExpr fs expr
 134   = do
 135       var <- newLocalVar fs (exprType expr)
 136       updGEnv $ \env ->
 137         env { global_bindings = (var, expr) : global_bindings env }
 138       return var
 139
 140 takeHoisted :: VM [(Var, CoreExpr)]
 141 takeHoisted
 142   = do
 143       env <- readGEnv id
 144       setGEnv $ env { global_bindings = [] }
 145       return $ global_bindings env
 146