module Vectorise.Type.Type
( vectTyCon
, vectAndLiftType
, vectType)
where
import Vectorise.Utils
import Vectorise.Monad
import Vectorise.Builtins
import TypeRep
import Type
import TyCon
import Outputable
import Control.Monad
import Data.List
import Data.Maybe
vectTyCon :: TyCon -> VM TyCon
vectTyCon tc
| isFunTyCon tc = builtin closureTyCon
| isBoxedTupleTyCon tc = return tc
| isUnLiftedTyCon tc = return tc
| otherwise
= maybeCantVectoriseM "Tycon not vectorised: " (ppr tc)
$ lookupTyCon tc
vectAndLiftType :: Type -> VM (Type, Type)
vectAndLiftType ty | Just ty' <- coreView ty = vectAndLiftType ty'
vectAndLiftType ty
= do
mdicts <- mapM paDictArgType (reverse tyvars)
let dicts = [dict | Just dict <- mdicts]
vmono_ty <- vectType mono_ty
lmono_ty <- mkPDataType vmono_ty
return (abstractType tyvars dicts vmono_ty,
abstractType tyvars dicts lmono_ty)
where
(tyvars, mono_ty) = splitForAllTys ty
vectType :: Type -> VM Type
vectType ty
| Just ty' <- coreView ty
= vectType ty'
vectType (TyVarTy tv) = return $ TyVarTy tv
vectType (AppTy ty1 ty2) = liftM2 AppTy (vectType ty1) (vectType ty2)
vectType (TyConApp tc tys) = liftM2 TyConApp (vectTyCon tc) (mapM vectType tys)
vectType (FunTy ty1 ty2) = liftM2 TyConApp (builtin closureTyCon)
(mapM vectAndBoxType [ty1,ty2])
vectType ty@(ForAllTy _ _)
= do
let (tyvars, tyBody) = splitForAllTys ty
let (tyArgs, tyResult) = splitFunTys tyBody
let (tyArgs_dict, tyArgs_regular)
= partition isDictType tyArgs
let tyBody' = mkFunTys tyArgs_regular tyResult
tyBody'' <- vectType tyBody'
dictsVect <- mapM vectType tyArgs_dict
dictsPA <- liftM catMaybes $ mapM paDictArgType tyvars
traceVt "vect ForAllTy: " $ ppr (abstractType tyvars (dictsPA ++ dictsVect) tyBody'')
return $ abstractType tyvars (dictsPA ++ dictsVect) tyBody''
vectType ty = cantVectorise "Can't vectorise type" (ppr ty)
abstractType :: [TyVar] -> [Type] -> Type -> Type
abstractType tyvars dicts = mkForAllTys tyvars . mkFunTys dicts
isDictType :: Type -> Bool
isDictType ty
= case splitTyConApp_maybe ty of
Just (tyCon, _) -> isClassTyCon tyCon
_ -> False
vectAndBoxType :: Type -> VM Type
vectAndBoxType ty = vectType ty >>= boxType
boxType :: Type -> VM Type
boxType ty
| Just (tycon, []) <- splitTyConApp_maybe ty
, isUnLiftedTyCon tycon
= do
r <- lookupBoxedTyCon tycon
case r of
Just tycon' -> return $ mkTyConApp tycon' []
Nothing -> return ty
| otherwise = return ty