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 Outputable
  14 import Control.Monad
  15 import Data.List
  16 import Data.Maybe
  17
  18
  19 -- | Vectorise a type constructor.
  20 vectTyCon :: TyCon -> VM TyCon
  21 vectTyCon tc
  22   | isFunTyCon tc        = builtin closureTyCon
  23   | isBoxedTupleTyCon tc = return tc
  24   | isUnLiftedTyCon tc   = return tc
  25   | otherwise
  26   = maybeCantVectoriseM "Tycon not vectorised: " (ppr tc)
  27         $ lookupTyCon tc
  28
  29
  30 -- | Produce the vectorised and lifted versions of a type.
  31 vectAndLiftType :: Type -> VM (Type, Type)
  32 vectAndLiftType ty | Just ty' <- coreView ty = vectAndLiftType ty'
  33 vectAndLiftType ty
  34   = do
  35       mdicts   <- mapM paDictArgType (reverse tyvars)
  36       let dicts = [dict | Just dict <- mdicts]
  37       vmono_ty <- vectType mono_ty
  38       lmono_ty <- mkPDataType vmono_ty
  39       return (abstractType tyvars dicts vmono_ty,
  40               abstractType tyvars dicts lmono_ty)
  41   where
  42     (tyvars, mono_ty) = splitForAllTys ty
  43
  44
  45 -- | Vectorise a type.
  46 vectType :: Type -> VM Type
  47 vectType ty
  48         | Just ty'      <- coreView ty
  49         = vectType ty'
  50
  51 vectType (TyVarTy tv)           = return $ TyVarTy tv
  52 vectType (AppTy ty1 ty2)        = liftM2 AppTy    (vectType ty1) (vectType ty2)
  53 vectType (TyConApp tc tys)      = liftM2 TyConApp (vectTyCon tc) (mapM vectType tys)
  54 vectType (FunTy ty1 ty2)        = liftM2 TyConApp (builtin closureTyCon)
  55                                                   (mapM vectAndBoxType [ty1,ty2])
  56
  57 -- For each quantified var we need to add a PA dictionary out the front of the type.
  58 -- So          forall a. C  a => a -> a
  59 -- turns into  forall a. Cv a => PA a => a :-> a
  60 vectType ty@(ForAllTy _ _)
  61  = do
  62       -- split the type into the quantified vars, its dictionaries and the body.
  63       let (tyvars, tyBody)   = splitForAllTys ty
  64       let (tyArgs, tyResult) = splitFunTys    tyBody
  65
  66       let (tyArgs_dict, tyArgs_regular)
  67                   = partition isDictType tyArgs
  68
  69       -- vectorise the body.
  70       let tyBody' = mkFunTys tyArgs_regular tyResult
  71       tyBody''    <- vectType tyBody'
  72
  73       -- vectorise the dictionary parameters.
  74       dictsVect   <- mapM vectType tyArgs_dict
  75
  76       -- make a PA dictionary for each of the type variables.
  77       dictsPA     <- liftM catMaybes $ mapM paDictArgType tyvars
  78
  79       -- pack it all back together.
  80       traceVt "vect ForAllTy: " $ ppr (abstractType tyvars (dictsPA ++ dictsVect) tyBody'')
  81       return $ abstractType tyvars (dictsPA ++ dictsVect) 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