compiler/vectorise/Vectorise/Type/Classify.hs

   1
   2 module Vectorise.Type.Classify
   3         ( TyConGroup
   4         , classifyTyCons
   5         , tyConGroups)
   6 where
   7 import UniqSet
   8 import UniqFM
   9 import DataCon
  10 import TyCon
  11 import TypeRep
  12 import Type
  13 import Digraph
  14 import Outputable
  15
  16 type TyConGroup = ([TyCon], UniqSet TyCon)
  17
  18 -- | Split the given tycons into two sets depending on whether they have to be
  19 --   converted (first list) or not (second list). The first argument contains
  20 --   information about the conversion status of external tycons:
  21 --
  22 --   * tycons which have converted versions are mapped to True
  23 --   * tycons which are not changed by vectorisation are mapped to False
  24 --   * tycons which can't be converted are not elements of the map
  25 --
  26 classifyTyCons
  27         :: UniqFM Bool
  28         -> [TyConGroup]
  29         -> ([TyCon], [TyCon])
  30
  31 classifyTyCons = classify [] []
  32   where
  33     classify conv keep _  [] = (conv, keep)
  34     classify conv keep cs ((tcs, ds) : rs)
  35       | can_convert && must_convert
  36         = classify (tcs ++ conv) keep (cs `addListToUFM` [(tc,True) | tc <- tcs]) rs
  37       | can_convert
  38         = classify conv (tcs ++ keep) (cs `addListToUFM` [(tc,False) | tc <- tcs]) rs
  39       | otherwise
  40         = classify conv keep cs rs
  41       where
  42         refs = ds `delListFromUniqSet` tcs
  43
  44         can_convert  = isNullUFM (refs `minusUFM` cs) && all convertable tcs
  45         must_convert = foldUFM (||) False (intersectUFM_C const cs refs)
  46
  47         convertable tc = isDataTyCon tc && all isVanillaDataCon (tyConDataCons tc)
  48
  49
  50 -- | Compute mutually recursive groups of tycons in topological order
  51 tyConGroups :: [TyCon] -> [TyConGroup]
  52 tyConGroups tcs = map mk_grp (stronglyConnCompFromEdgedVertices edges)
  53   where
  54     edges = [((tc, ds), tc, uniqSetToList ds) | tc <- tcs
  55                                 , let ds = tyConsOfTyCon tc]
  56
  57     mk_grp (AcyclicSCC (tc, ds)) = ([tc], ds)
  58     mk_grp (CyclicSCC els)       = (tcs, unionManyUniqSets dss)
  59       where
  60         (tcs, dss) = unzip els
  61
  62
  63 -- | Collect the set of TyCons used by the representation of some data type.
  64 tyConsOfTyCon :: TyCon -> UniqSet TyCon
  65 tyConsOfTyCon
  66   = tyConsOfTypes . concatMap dataConRepArgTys . tyConDataCons
  67
  68
  69 -- | Collect the set of TyCons that occur in these types.
  70 tyConsOfTypes :: [Type] -> UniqSet TyCon
  71 tyConsOfTypes = unionManyUniqSets . map tyConsOfType
  72
  73
  74 -- | Collect the set of TyCons that occur in this type.
  75 tyConsOfType :: Type -> UniqSet TyCon
  76 tyConsOfType ty
  77   | Just ty' <- coreView ty    = tyConsOfType ty'
  78 tyConsOfType (TyVarTy _)       = emptyUniqSet
  79 tyConsOfType (TyConApp tc tys) = extend (tyConsOfTypes tys)
  80   where
  81     extend |  isUnLiftedTyCon tc
  82            || isTupleTyCon   tc = id
  83
  84            | otherwise          = (`addOneToUniqSet` tc)
  85
  86 tyConsOfType (AppTy a b)       = tyConsOfType a `unionUniqSets` tyConsOfType b
  87 tyConsOfType (FunTy a b)       = (tyConsOfType a `unionUniqSets` tyConsOfType b)
  88                                  `addOneToUniqSet` funTyCon
  89 tyConsOfType (ForAllTy _ ty)   = tyConsOfType ty
  90 tyConsOfType other             = pprPanic "ClosureConv.tyConsOfType" $ ppr other
  91