compiler/vectorise/Vectorise/Type/Type.hs

   1
   2 module Vectorise.Type.Type
   3         ( vectTyCon
   4         , vectAndLiftType
   5         , vectType)
   6 where
   7 import Vectorise.Utils
   8 import Vectorise.Monad
   9 import Vectorise.Builtins
  10 import TypeRep
  11 import Type
  12 import TyCon
  13 import Var
  14 import Outputable
  15 import Control.Monad
  16 import Data.List
  17 import Data.Maybe
  18
  19
  20 -- | Vectorise a type constructor.
  21 vectTyCon :: TyCon -> VM TyCon
  22 vectTyCon tc
  23   | isFunTyCon tc        = builtin closureTyCon
  24   | isBoxedTupleTyCon tc = return tc
  25   | isUnLiftedTyCon tc   = return tc
  26   | otherwise
  27   = maybeCantVectoriseM "Tycon not vectorised: " (ppr tc)
  28         $ lookupTyCon tc
  29
  30
  31 -- | Produce the vectorised and lifted versions of a type.
  32 vectAndLiftType :: Type -> VM (Type, Type)
  33 vectAndLiftType ty | Just ty' <- coreView ty = vectAndLiftType ty'
  34 vectAndLiftType ty
  35   = do
  36       mdicts   <- mapM paDictArgType tyvars
  37       let dicts = [dict | Just dict <- mdicts]
  38       vmono_ty <- vectType mono_ty
  39       lmono_ty <- mkPDataType vmono_ty
  40       return (abstractType tyvars dicts vmono_ty,
  41               abstractType tyvars dicts lmono_ty)
  42   where
  43     (tyvars, mono_ty) = splitForAllTys ty
  44
  45
  46 -- | Vectorise a type.
  47 vectType :: Type -> VM Type
  48 vectType ty
  49         | Just ty'      <- coreView ty
  50         = vectType ty'
  51
  52 vectType (TyVarTy tv)           = return $ TyVarTy tv
  53 vectType (AppTy ty1 ty2)        = liftM2 AppTy    (vectType ty1) (vectType ty2)
  54 vectType (TyConApp tc tys)      = liftM2 TyConApp (vectTyCon tc) (mapM vectType tys)
  55 vectType (FunTy ty1 ty2)        = liftM2 TyConApp (builtin closureTyCon)
  56                                                   (mapM vectAndBoxType [ty1,ty2])
  57
  58 -- For each quantified var we need to add a PA dictionary out the front of the type.
  59 -- So          forall a. C  a => a -> a
  60 -- turns into  forall a. Cv a => PA a => a :-> a
  61 vectType ty@(ForAllTy _ _)
  62  = do
  63       -- split the type into the quantified vars, its dictionaries and the body.
  64       let (tyvars, tyBody)   = splitForAllTys ty
  65       let (tyArgs, tyResult) = splitFunTys    tyBody
  66
  67       let (tyArgs_dict, tyArgs_regular)
  68                   = partition isDictType tyArgs
  69
  70       -- vectorise the body.
  71       let tyBody' = mkFunTys tyArgs_regular tyResult
  72       tyBody''    <- vectType tyBody'
  73
  74       -- vectorise the dictionary parameters.
  75       dictsVect   <- mapM vectType tyArgs_dict
  76
  77       -- make a PA dictionary for each of the type variables.
  78       dictsPA     <- liftM catMaybes $ mapM paDictArgType tyvars
  79
  80       -- pack it all back together.
  81       return $ abstractType tyvars (dictsVect ++ dictsPA) tyBody''
  82
  83 vectType ty = cantVectorise "Can't vectorise type" (ppr ty)
  84
  85
  86 -- | Add quantified vars and dictionary parameters to the front of a type.
  87 abstractType :: [TyVar] -> [Type] -> Type -> Type
  88 abstractType tyvars dicts = mkForAllTys tyvars . mkFunTys dicts
  89
  90
  91 -- | Check if some type is a type class dictionary.
  92 isDictType :: Type -> Bool
  93 isDictType ty
  94  = case splitTyConApp_maybe ty of
  95         Just (tyCon, _)         -> isClassTyCon tyCon
  96         _                       -> False
  97
  98
  99 -- | Create the boxed version of a vectorised type.
 100 vectAndBoxType :: Type -> VM Type
 101 vectAndBoxType ty = vectType ty >>= boxType
 102
 103
 104 -- | Create the boxed version of a type.
 105 boxType :: Type -> VM Type
 106 boxType ty
 107   | Just (tycon, []) <- splitTyConApp_maybe ty
 108   , isUnLiftedTyCon tycon
 109   = do
 110       r <- lookupBoxedTyCon tycon
 111       case r of
 112         Just tycon' -> return $ mkTyConApp tycon' []
 113         Nothing     -> return ty
 114
 115   | otherwise   = return ty
 116
 117