Classify.hs 4.59 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18
-- Extract from a list of type constructors those (1) which need to be vectorised and (2) those
-- that could be, but need not be vectorised (as a scalar representation is sufficient and more
-- efficient).  The type constructors that cannot be vectorised will be dropped.
--
-- A type constructor will only be vectorised if it is
--
-- (1) a data type constructor, with vanilla data constructors (i.e., data constructors admitted by
--     Haskell 98) and
-- (2) at least one of the type constructors that appears in its definition is also vectorised.
--
-- If (1) is met, but not (2), the type constructor may appear in vectorised code, but there is no
-- need to vectorise that type constructor itself.  This holds, for example, for all enumeration
-- types.  As '([::])' is being vectorised, any type constructor whose definition involves
-- '([::])', either directly or indirectly, will be vectorised.

module Vectorise.Type.Classify (
  classifyTyCons
) where
19 20 21 22 23 24 25 26 27 28

import UniqSet
import UniqFM
import DataCon
import TyCon
import TypeRep
import Type
import Digraph


29
-- |From a list of type constructors, extract those that can be vectorised, returning them in two
30
-- sets, where the first result list /must be/ vectorised and the second result list /need not be/
31 32 33
-- vectorised.  The third result list are those type constructors that we cannot convert (either
-- because they use language extensions or because they dependent on type constructors for which
-- no vectorised version is available).
34 35

-- The first argument determines the /conversion status/ of external type constructors as follows:
36
--
37 38 39
-- * tycons which have converted versions are mapped to 'True'
-- * tycons which are not changed by vectorisation are mapped to 'False'
-- * tycons which can't be converted are not elements of the map
40
--
41 42 43 44
classifyTyCons :: UniqFM Bool                     -- ^type constructor conversion status
               -> [TyCon]                         -- ^type constructors that need to be classified
               -> ([TyCon], [TyCon], [TyCon])     -- ^tycons to be converted & not to be converted
classifyTyCons convStatus tcs = classify [] [] [] convStatus (tyConGroups tcs)
45
  where
46 47
    classify conv keep ignored _  [] = (conv, keep, ignored)
    classify conv keep ignored cs ((tcs, ds) : rs)
48
      | can_convert && must_convert
49
      = classify (tcs ++ conv) keep ignored (cs `addListToUFM` [(tc, True) | tc <- tcs]) rs
50
      | can_convert
51
      = classify conv (tcs ++ keep) ignored (cs `addListToUFM` [(tc, False) | tc <- tcs]) rs
52
      | otherwise
53
      = classify conv keep (tcs ++ ignored) cs rs
54 55 56 57 58 59
      where
        refs = ds `delListFromUniqSet` tcs

        can_convert  = isNullUFM (refs `minusUFM` cs) && all convertable tcs
        must_convert = foldUFM (||) False (intersectUFM_C const cs refs)

60 61 62 63 64
        -- We currently admit Haskell 2011-style data and newtype declarations as well as type
        -- constructors representing classes.
        convertable tc 
          = (isDataTyCon tc || isNewTyCon tc) && all isVanillaDataCon (tyConDataCons tc)
            || isClassTyCon tc
65

66 67 68
-- Used to group type constructors into mutually dependent groups.
--
type TyConGroup = ([TyCon], UniqSet TyCon)
69

70 71
-- Compute mutually recursive groups of tycons in topological order.
--
72 73 74 75 76 77 78 79 80 81 82
tyConGroups :: [TyCon] -> [TyConGroup]
tyConGroups tcs = map mk_grp (stronglyConnCompFromEdgedVertices edges)
  where
    edges = [((tc, ds), tc, uniqSetToList ds) | tc <- tcs
                                , let ds = tyConsOfTyCon tc]

    mk_grp (AcyclicSCC (tc, ds)) = ([tc], ds)
    mk_grp (CyclicSCC els)       = (tcs, unionManyUniqSets dss)
      where
        (tcs, dss) = unzip els

83 84
-- |Collect the set of TyCons used by the representation of some data type.
--
85
tyConsOfTyCon :: TyCon -> UniqSet TyCon
86
tyConsOfTyCon = tyConsOfTypes . concatMap dataConRepArgTys . tyConDataCons
87

88 89
-- |Collect the set of TyCons that occur in these types.
--
90 91 92
tyConsOfTypes :: [Type] -> UniqSet TyCon
tyConsOfTypes = unionManyUniqSets . map tyConsOfType

93 94
-- |Collect the set of TyCons that occur in this type.
--
95 96 97 98 99 100 101 102 103 104 105 106 107 108 109
tyConsOfType :: Type -> UniqSet TyCon
tyConsOfType ty
  | Just ty' <- coreView ty    = tyConsOfType ty'
tyConsOfType (TyVarTy _)       = emptyUniqSet
tyConsOfType (TyConApp tc tys) = extend (tyConsOfTypes tys)
  where
    extend |  isUnLiftedTyCon tc
           || isTupleTyCon   tc = id

           | otherwise          = (`addOneToUniqSet` tc)

tyConsOfType (AppTy a b)       = tyConsOfType a `unionUniqSets` tyConsOfType b
tyConsOfType (FunTy a b)       = (tyConsOfType a `unionUniqSets` tyConsOfType b)
                                 `addOneToUniqSet` funTyCon
tyConsOfType (ForAllTy _ ty)   = tyConsOfType ty