compiler/vectorise/VectUtils.hs

   1 module VectUtils (
   2   paDictArgType, paDictOfType
   3 ) where
   4
   5 #include "HsVersions.h"
   6
   7 import VectMonad
   8
   9 import CoreSyn
  10 import Type
  11 import TypeRep
  12 import Var
  13
  14 import Outputable
  15
  16 paDictArgType :: TyVar -> VM (Maybe Type)
  17 paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
  18   where
  19     go ty k | Just k' <- kindView k = go ty k'
  20     go ty (FunTy k1 k2)
  21       = do
  22           tv   <- newTyVar FSLIT("a") k1
  23           mty1 <- go (TyVarTy tv) k1
  24           case mty1 of
  25             Just ty1 -> do
  26                           mty2 <- go (AppTy ty (TyVarTy tv)) k2
  27                           return $ fmap (ForAllTy tv . FunTy ty1) mty2
  28             Nothing  -> go ty k2
  29
  30     go ty k
  31       | isLiftedTypeKind k
  32       = do
  33           tc <- builtin paDictTyCon
  34           return . Just $ TyConApp tc [ty]
  35
  36     go ty k = return Nothing
  37
  38 paDictOfType :: Type -> VM CoreExpr
  39 paDictOfType ty = paDictOfTyApp ty_fn ty_args
  40   where
  41     (ty_fn, ty_args) = splitAppTys ty
  42
  43 paDictOfTyApp :: Type -> [Type] -> VM CoreExpr
  44 paDictOfTyApp ty_fn ty_args
  45   | Just ty_fn' <- coreView ty_fn = paDictOfTyApp ty_fn' ty_args
  46 paDictOfTyApp (TyVarTy tv) ty_args
  47   = do
  48       dfun <- maybeV (lookupTyVarPA tv)
  49       paDFunApply dfun ty_args
  50 paDictOfTyApp (TyConApp tc _) ty_args
  51   = do
  52       pa_class <- builtin paClass
  53       (dfun, ty_args') <- lookupInst pa_class [TyConApp tc ty_args]
  54       paDFunApply (Var dfun) ty_args'
  55 paDictOfTyApp ty ty_args = pprPanic "paDictOfTyApp" (ppr ty)
  56
  57 paDFunApply :: CoreExpr -> [Type] -> VM CoreExpr
  58 paDFunApply dfun tys
  59   = do
  60       dicts <- mapM paDictOfType tys
  61       return $ mkApps (mkTyApps dfun tys) dicts
  62