compiler/vectorise/VectType.hs

   1 module VectType ( vectTyCon, vectType, vectTypeEnv )
   2 where
   3
   4 #include "HsVersions.h"
   5
   6 import VectMonad
   7 import VectUtils
   8
   9 import HscTypes          ( TypeEnv, extendTypeEnvList, typeEnvTyCons )
  10 import DataCon
  11 import TyCon
  12 import Type
  13 import TypeRep
  14 import Coercion
  15 import OccName
  16 import MkId
  17 import BasicTypes        ( StrictnessMark(..), boolToRecFlag )
  18 import NameEnv
  19 import TysWiredIn
  20
  21 import Unique
  22 import UniqFM
  23 import UniqSet
  24 import Digraph           ( SCC(..), stronglyConnComp )
  25
  26 import Outputable
  27
  28 import Control.Monad  ( liftM2, zipWithM_ )
  29
  30 -- ----------------------------------------------------------------------------
  31 -- Types
  32
  33 vectTyCon :: TyCon -> VM TyCon
  34 vectTyCon tc
  35   | isFunTyCon tc        = builtin closureTyCon
  36   | isBoxedTupleTyCon tc = return tc
  37   | isUnLiftedTyCon tc   = return tc
  38   | otherwise = do
  39                   r <- lookupTyCon tc
  40                   case r of
  41                     Just tc' -> return tc'
  42
  43                     -- FIXME: just for now
  44                     Nothing  -> pprTrace "ccTyCon:" (ppr tc) $ return tc
  45
  46 vectType :: Type -> VM Type
  47 vectType ty | Just ty' <- coreView ty = vectType ty'
  48 vectType (TyVarTy tv) = return $ TyVarTy tv
  49 vectType (AppTy ty1 ty2) = liftM2 AppTy (vectType ty1) (vectType ty2)
  50 vectType (TyConApp tc tys) = liftM2 TyConApp (vectTyCon tc) (mapM vectType tys)
  51 vectType (FunTy ty1 ty2)   = liftM2 TyConApp (builtin closureTyCon)
  52                                              (mapM vectType [ty1,ty2])
  53 vectType ty@(ForAllTy _ _)
  54   = do
  55       mdicts   <- mapM paDictArgType tyvars
  56       mono_ty' <- vectType mono_ty
  57       return $ tyvars `mkForAllTys` ([dict | Just dict <- mdicts] `mkFunTys` mono_ty')
  58   where
  59     (tyvars, mono_ty) = splitForAllTys ty
  60
  61 vectType ty = pprPanic "vectType:" (ppr ty)
  62
  63 -- ----------------------------------------------------------------------------
  64 -- Type definitions
  65
  66 type TyConGroup = ([TyCon], UniqSet TyCon)
  67
  68 vectTypeEnv :: TypeEnv -> VM TypeEnv
  69 vectTypeEnv env
  70   = do
  71       cs <- readGEnv $ mk_map . global_tycons
  72       let (conv_tcs, keep_tcs) = classifyTyCons cs groups
  73           keep_dcs             = concatMap tyConDataCons keep_tcs
  74       zipWithM_ defTyCon   keep_tcs keep_tcs
  75       zipWithM_ defDataCon keep_dcs keep_dcs
  76       vect_tcs <- vectTyConDecls conv_tcs
  77       parr_tcs <- mapM buildPArrayTyCon (keep_tcs ++ vect_tcs)
  78       let new_tcs = vect_tcs ++ parr_tcs
  79       return $ extendTypeEnvList env
  80                  (map ATyCon new_tcs ++ [ADataCon dc | tc <- new_tcs
  81                                                      , dc <- tyConDataCons tc])
  82   where
  83     tycons = typeEnvTyCons env
  84     groups = tyConGroups tycons
  85
  86     mk_map env = listToUFM_Directly [(u, getUnique n /= u) | (u,n) <- nameEnvUniqueElts env]
  87
  88     keep_tc tc = let dcs = tyConDataCons tc
  89                  in
  90                  defTyCon tc tc >> zipWithM_ defDataCon dcs dcs
  91
  92
  93 vectTyConDecls :: [TyCon] -> VM [TyCon]
  94 vectTyConDecls tcs = fixV $ \tcs' ->
  95   do
  96     mapM_ (uncurry defTyCon) (lazy_zip tcs tcs')
  97     mapM vectTyConDecl tcs
  98   where
  99     lazy_zip [] _ = []
 100     lazy_zip (x:xs) ~(y:ys) = (x,y) : lazy_zip xs ys
 101
 102 vectTyConDecl :: TyCon -> VM TyCon
 103 vectTyConDecl tc
 104   = do
 105       name' <- cloneName mkVectTyConOcc name
 106       rhs'  <- vectAlgTyConRhs (algTyConRhs tc)
 107
 108       return $ mkAlgTyCon name'
 109                           kind
 110                           tyvars
 111                           []              -- no stupid theta
 112                           rhs'
 113                           []              -- no selector ids
 114                           NoParentTyCon   -- FIXME
 115                           rec_flag        -- FIXME: is this ok?
 116                           False           -- FIXME: no generics
 117                           False           -- not GADT syntax
 118   where
 119     name   = tyConName tc
 120     kind   = tyConKind tc
 121     tyvars = tyConTyVars tc
 122     rec_flag = boolToRecFlag (isRecursiveTyCon tc)
 123
 124 vectAlgTyConRhs :: AlgTyConRhs -> VM AlgTyConRhs
 125 vectAlgTyConRhs (DataTyCon { data_cons = data_cons
 126                            , is_enum   = is_enum
 127                            })
 128   = do
 129       data_cons' <- mapM vectDataCon data_cons
 130       zipWithM_ defDataCon data_cons data_cons'
 131       return $ DataTyCon { data_cons = data_cons'
 132                          , is_enum   = is_enum
 133                          }
 134
 135 vectDataCon :: DataCon -> VM DataCon
 136 vectDataCon dc
 137   | not . null $ dataConExTyVars dc = pprPanic "vectDataCon: existentials" (ppr dc)
 138   | not . null $ dataConEqSpec   dc = pprPanic "vectDataCon: eq spec" (ppr dc)
 139   | otherwise
 140   = do
 141       name'    <- cloneName mkVectDataConOcc name
 142       tycon'   <- vectTyCon tycon
 143       arg_tys  <- mapM vectType rep_arg_tys
 144       wrk_name <- cloneName mkDataConWorkerOcc name'
 145
 146       let ids      = mkDataConIds (panic "vectDataCon: wrapper id")
 147                                   wrk_name
 148                                   data_con
 149           data_con = mkDataCon name'
 150                                False           -- not infix
 151                                (map (const NotMarkedStrict) arg_tys)
 152                                []              -- no labelled fields
 153                                univ_tvs
 154                                []              -- no existential tvs for now
 155                                []              -- no eq spec for now
 156                                []              -- no theta
 157                                arg_tys
 158                                tycon'
 159                                []              -- no stupid theta
 160                                ids
 161       return data_con
 162   where
 163     name        = dataConName dc
 164     univ_tvs    = dataConUnivTyVars dc
 165     rep_arg_tys = dataConRepArgTys dc
 166     tycon       = dataConTyCon dc
 167
 168 buildPArrayTyCon :: TyCon -> VM TyCon
 169 buildPArrayTyCon orig_tc = fixV $ \repr_tc ->
 170   do
 171     name'  <- cloneName mkPArrayTyConOcc name
 172     parent <- buildPArrayParentInfo orig_tc repr_tc
 173     rhs    <- buildPArrayTyConRhs orig_tc repr_tc
 174
 175     return $ mkAlgTyCon name'
 176                         kind
 177                         tyvars
 178                         []              -- no stupid theta
 179                         rhs
 180                         []              -- no selector ids
 181                         parent
 182                         rec_flag        -- FIXME: is this ok?
 183                         False           -- FIXME: no generics
 184                         False           -- not GADT syntax
 185   where
 186     name   = tyConName orig_tc
 187     kind   = tyConKind orig_tc
 188     tyvars = tyConTyVars orig_tc
 189     rec_flag = boolToRecFlag (isRecursiveTyCon orig_tc)
 190
 191
 192 buildPArrayParentInfo :: TyCon -> TyCon -> VM TyConParent
 193 buildPArrayParentInfo orig_tc repr_tc
 194   = do
 195       parray_tc <- builtin parrayTyCon
 196       co_name <- cloneName mkInstTyCoOcc (tyConName repr_tc)
 197
 198       let inst_tys = [mkTyConApp orig_tc (map mkTyVarTy tyvars)]
 199
 200       return . FamilyTyCon parray_tc inst_tys
 201              $ mkFamInstCoercion co_name
 202                                  tyvars
 203                                  parray_tc
 204                                  inst_tys
 205                                  repr_tc
 206   where
 207     tyvars = tyConTyVars orig_tc
 208
 209 buildPArrayTyConRhs :: TyCon -> TyCon -> VM AlgTyConRhs
 210 buildPArrayTyConRhs orig_tc repr_tc
 211   = do
 212       data_con <- buildPArrayDataCon orig_tc repr_tc
 213       return $ DataTyCon { data_cons = [data_con], is_enum = False }
 214
 215 buildPArrayDataCon :: TyCon -> TyCon -> VM DataCon
 216 buildPArrayDataCon orig_tc repr_tc
 217   = do
 218       name     <- cloneName mkPArrayDataConOcc (tyConName orig_tc)
 219       shape_ty <- mkPArrayType intTy   -- FIXME: we want to unbox this!
 220       repr_tys <- mapM mkPArrayType types
 221       wrk_name <- cloneName mkDataConWorkerOcc name
 222       wrp_name <- cloneName mkDataConWrapperOcc name
 223
 224       let ids      = mkDataConIds wrp_name wrk_name data_con
 225           data_con = mkDataCon name
 226                                False
 227                                (MarkedStrict : map (const NotMarkedStrict) repr_tys)
 228                                []
 229                                (tyConTyVars orig_tc)
 230                                []
 231                                []
 232                                []
 233                                (shape_ty : repr_tys)
 234                                repr_tc
 235                                []
 236                                ids
 237
 238       return data_con
 239   where
 240     types = [ty | dc <- tyConDataCons orig_tc
 241                 , ty <- dataConRepArgTys dc]
 242
 243 -- | Split the given tycons into two sets depending on whether they have to be
 244 -- converted (first list) or not (second list). The first argument contains
 245 -- information about the conversion status of external tycons:
 246 --
 247 --   * tycons which have converted versions are mapped to True
 248 --   * tycons which are not changed by vectorisation are mapped to False
 249 --   * tycons which can't be converted are not elements of the map
 250 --
 251 classifyTyCons :: UniqFM Bool -> [TyConGroup] -> ([TyCon], [TyCon])
 252 classifyTyCons = classify [] []
 253   where
 254     classify conv keep cs [] = (conv, keep)
 255     classify conv keep cs ((tcs, ds) : rs)
 256       | can_convert && must_convert
 257         = classify (tcs ++ conv) keep (cs `addListToUFM` [(tc,True) | tc <- tcs]) rs
 258       | can_convert
 259         = classify conv (tcs ++ keep) (cs `addListToUFM` [(tc,False) | tc <- tcs]) rs
 260       | otherwise
 261         = classify conv keep cs rs
 262       where
 263         refs = ds `delListFromUniqSet` tcs
 264
 265         can_convert  = isNullUFM (refs `minusUFM` cs) && all convertable tcs
 266         must_convert = foldUFM (||) False (intersectUFM_C const cs refs)
 267
 268         convertable tc = isDataTyCon tc && all isVanillaDataCon (tyConDataCons tc)
 269
 270 -- | Compute mutually recursive groups of tycons in topological order
 271 --
 272 tyConGroups :: [TyCon] -> [TyConGroup]
 273 tyConGroups tcs = map mk_grp (stronglyConnComp edges)
 274   where
 275     edges = [((tc, ds), tc, uniqSetToList ds) | tc <- tcs
 276                                 , let ds = tyConsOfTyCon tc]
 277
 278     mk_grp (AcyclicSCC (tc, ds)) = ([tc], ds)
 279     mk_grp (CyclicSCC els)       = (tcs, unionManyUniqSets dss)
 280       where
 281         (tcs, dss) = unzip els
 282
 283 tyConsOfTyCon :: TyCon -> UniqSet TyCon
 284 tyConsOfTyCon
 285   = tyConsOfTypes . concatMap dataConRepArgTys . tyConDataCons
 286
 287 tyConsOfType :: Type -> UniqSet TyCon
 288 tyConsOfType ty
 289   | Just ty' <- coreView ty    = tyConsOfType ty'
 290 tyConsOfType (TyVarTy v)       = emptyUniqSet
 291 tyConsOfType (TyConApp tc tys) = extend (tyConsOfTypes tys)
 292   where
 293     extend | isUnLiftedTyCon tc
 294            || isTupleTyCon   tc = id
 295
 296            | otherwise          = (`addOneToUniqSet` tc)
 297
 298 tyConsOfType (AppTy a b)       = tyConsOfType a `unionUniqSets` tyConsOfType b
 299 tyConsOfType (FunTy a b)       = (tyConsOfType a `unionUniqSets` tyConsOfType b)
 300                                  `addOneToUniqSet` funTyCon
 301 tyConsOfType (ForAllTy _ ty)   = tyConsOfType ty
 302 tyConsOfType other             = pprPanic "ClosureConv.tyConsOfType" $ ppr other
 303
 304 tyConsOfTypes :: [Type] -> UniqSet TyCon
 305 tyConsOfTypes = unionManyUniqSets . map tyConsOfType
 306