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