compiler/vectorise/Vectorise/Type/Repr.hs

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