compiler/vectorise/Vectorise/Type/Repr.hs

   1
   2 -- | Representation of Algebraic Data Types.
   3 module Vectorise.Type.Repr
   4         ( CompRepr      (..)
   5         , ProdRepr      (..)
   6         , ConRepr       (..)
   7         , SumRepr       (..)
   8         , tyConRepr
   9         , sumReprType
  10         , conReprType
  11         , prodReprType
  12         , compReprType
  13         , compOrigType)
  14 where
  15 import Vectorise.Utils
  16 import Vectorise.Monad
  17 import Vectorise.Builtins
  18
  19 import CoreSyn
  20 import DataCon
  21 import TyCon
  22 import Type
  23 import Control.Monad
  24
  25
  26 data CompRepr = Keep Type
  27                      CoreExpr     -- PR dictionary for the type
  28               | Wrap Type
  29
  30 data ProdRepr = EmptyProd
  31               | UnaryProd CompRepr
  32               | Prod { repr_tup_tc   :: TyCon  -- representation tuple tycon
  33                      , repr_ptup_tc  :: TyCon  -- PData representation tycon
  34                      , repr_comp_tys :: [Type] -- representation types of
  35                      , repr_comps    :: [CompRepr]          -- components
  36                      }
  37 data ConRepr  = ConRepr DataCon ProdRepr
  38
  39 data SumRepr  = EmptySum
  40               | UnarySum ConRepr
  41               | Sum  { repr_sum_tc   :: TyCon  -- representation sum tycon
  42                      , repr_psum_tc  :: TyCon  -- PData representation tycon
  43                      , repr_sel_ty   :: Type   -- type of selector
  44                      , repr_con_tys :: [Type]  -- representation types of
  45                      , repr_cons     :: [ConRepr]           -- components
  46                      }
  47
  48 tyConRepr :: TyCon -> VM SumRepr
  49 tyConRepr tc = sum_repr (tyConDataCons tc)
  50   where
  51     sum_repr []    = return EmptySum
  52     sum_repr [con] = liftM UnarySum (con_repr con)
  53     sum_repr cons  = do
  54                        rs     <- mapM con_repr cons
  55                        sum_tc <- builtin (sumTyCon arity)
  56                        tys    <- mapM conReprType rs
  57                        (psum_tc, _) <- pdataReprTyCon (mkTyConApp sum_tc tys)
  58                        sel_ty <- builtin (selTy arity)
  59                        return $ Sum { repr_sum_tc  = sum_tc
  60                                     , repr_psum_tc = psum_tc
  61                                     , repr_sel_ty  = sel_ty
  62                                     , repr_con_tys = tys
  63                                     , repr_cons    = rs
  64                                     }
  65       where
  66         arity = length cons
  67
  68     con_repr con = liftM (ConRepr con) (prod_repr (dataConRepArgTys con))
  69
  70     prod_repr []   = return EmptyProd
  71     prod_repr [ty] = liftM UnaryProd (comp_repr ty)
  72     prod_repr tys  = do
  73                        rs <- mapM comp_repr tys
  74                        tup_tc <- builtin (prodTyCon arity)
  75                        tys'    <- mapM compReprType rs
  76                        (ptup_tc, _) <- pdataReprTyCon (mkTyConApp tup_tc tys')
  77                        return $ Prod { repr_tup_tc   = tup_tc
  78                                      , repr_ptup_tc  = ptup_tc
  79                                      , repr_comp_tys = tys'
  80                                      , repr_comps    = rs
  81                                      }
  82       where
  83         arity = length tys
  84
  85     comp_repr ty = liftM (Keep ty) (prDictOfType ty)
  86                    `orElseV` return (Wrap ty)
  87
  88 sumReprType :: SumRepr -> VM Type
  89 sumReprType EmptySum = voidType
  90 sumReprType (UnarySum r) = conReprType r
  91 sumReprType (Sum { repr_sum_tc  = sum_tc, repr_con_tys = tys })
  92   = return $ mkTyConApp sum_tc tys
  93
  94 conReprType :: ConRepr -> VM Type
  95 conReprType (ConRepr _ r) = prodReprType r
  96
  97 prodReprType :: ProdRepr -> VM Type
  98 prodReprType EmptyProd = voidType
  99 prodReprType (UnaryProd r) = compReprType r
 100 prodReprType (Prod { repr_tup_tc = tup_tc, repr_comp_tys = tys })
 101   = return $ mkTyConApp tup_tc tys
 102
 103 compReprType :: CompRepr -> VM Type
 104 compReprType (Keep ty _) = return ty
 105 compReprType (Wrap ty) = do
 106                              wrap_tc <- builtin wrapTyCon
 107                              return $ mkTyConApp wrap_tc [ty]
 108
 109 compOrigType :: CompRepr -> Type
 110 compOrigType (Keep ty _) = ty
 111 compOrigType (Wrap ty) = ty