Type.hs 2.77 KB
Newer Older
1 2
-- Apply the vectorisation transformation to types. This is the \mathcal{L}_t scheme in HtM.

3 4 5 6 7 8
module Vectorise.Type.Type
  ( vectTyCon
  , vectAndLiftType
  , vectType
  ) 
where
9

10
import Vectorise.Utils
11 12
import Vectorise.Monad
import Vectorise.Builtins
13
import TcType
14
import Type
15
import TypeRep
16 17
import TyCon
import Control.Monad
18
import Control.Applicative
19
import Data.Maybe
20
import Prelude -- avoid redundant import warning due to AMP
21 22 23

-- |Vectorise a type constructor. Unless there is a vectorised version (stripped of embedded
-- parallel arrays), the vectorised version is the same as the original.
24
--
25
vectTyCon :: TyCon -> VM TyCon
26
vectTyCon tc = maybe tc id <$> lookupTyCon tc
27

28 29
-- |Produce the vectorised and lifted versions of a type.
--
30 31 32
-- NB: Here we are limited to properly handle predicates at the toplevel only.  Anything embedded
--     in what is called the 'body_ty' below will end up as an argument to the type family 'PData'.
--
33 34 35
vectAndLiftType :: Type -> VM (Type, Type)
vectAndLiftType ty | Just ty' <- coreView ty = vectAndLiftType ty'
vectAndLiftType ty
36 37 38 39 40 41
  = do { padicts  <- liftM catMaybes $ mapM paDictArgType tyvars
       ; vmono_ty <- vectType mono_ty
       ; lmono_ty <- mkPDataType vmono_ty
       ; return (abstractType tyvars (padicts ++ theta) vmono_ty,
                 abstractType tyvars (padicts ++ theta) lmono_ty)
       }
42
  where
43 44
    (tyvars, phiTy)  = splitForAllTys ty
    (theta, mono_ty) = tcSplitPhiTy phiTy 
45

46 47
-- |Vectorise a type.
--
48 49 50 51
-- For each quantified var we need to add a PA dictionary out the front of the type.
-- So          forall a.         C  a => a -> a   
-- turns into  forall a. PA a => Cv a => a :-> a
--
52 53
vectType :: Type -> VM Type
vectType ty
54 55 56
  | Just ty'  <- coreView ty
  = vectType ty'
vectType (TyVarTy tv)      = return $ TyVarTy tv
57
vectType (LitTy l)         = return $ LitTy l
58 59 60 61 62 63 64
vectType (AppTy ty1 ty2)   = AppTy <$> vectType ty1 <*> vectType ty2
vectType (TyConApp tc tys) = TyConApp <$> vectTyCon tc <*> mapM vectType tys
vectType (FunTy ty1 ty2)   
  | isPredTy ty1
  = FunTy <$> vectType ty1 <*> vectType ty2   -- don't build a closure for dictionary abstraction
  | otherwise
  = TyConApp <$> builtin closureTyCon <*> mapM vectType [ty1, ty2]
65
vectType ty@(ForAllTy _ _)
66 67
 = do {   -- strip off consecutive foralls
      ; let (tyvars, tyBody) = splitForAllTys ty
68

69 70
          -- vectorise the body
      ; vtyBody <- vectType tyBody
71

72 73
          -- make a PA dictionary for each of the type variables
      ; dictsPA <- liftM catMaybes $ mapM paDictArgType tyvars
74

75 76 77
          -- add the PA dictionaries after the foralls
      ; return $ abstractType tyvars dictsPA vtyBody
      }
78

79 80
-- |Add quantified vars and dictionary parameters to the front of a type.
--
81 82
abstractType :: [TyVar] -> [Type] -> Type -> Type
abstractType tyvars dicts = mkForAllTys tyvars . mkFunTys dicts