Add VECTORISE [SCALAR] type pragma
[ghc.git] / compiler / vectorise / Vectorise / Type / Classify.hs
1 -- Extract from a list of type constructors those (1) which need to be vectorised and (2) those
2 -- that could be, but need not be vectorised (as a scalar representation is sufficient and more
3 -- efficient). The type constructors that cannot be vectorised will be dropped.
4 --
5 -- A type constructor will only be vectorised if it is
6 --
7 -- (1) a data type constructor, with vanilla data constructors (i.e., data constructors admitted by
8 -- Haskell 98) and
9 -- (2) at least one of the type constructors that appears in its definition is also vectorised.
10 --
11 -- If (1) is met, but not (2), the type constructor may appear in vectorised code, but there is no
12 -- need to vectorise that type constructor itself. This holds, for example, for all enumeration
13 -- types. As '([::])' is being vectorised, any type constructor whose definition involves
14 -- '([::])', either directly or indirectly, will be vectorised.
15
16 module Vectorise.Type.Classify (
17 classifyTyCons
18 ) where
19
20 import UniqSet
21 import UniqFM
22 import DataCon
23 import TyCon
24 import TypeRep
25 import Type
26 import Digraph
27 import Outputable
28
29
30 -- |From a list of type constructors, extract those thatcan be vectorised, returning them in two
31 -- sets, where the first result list /must be/ vectorised and the second result list /need not be/
32 -- vectroised.
33
34 -- The first argument determines the /conversion status/ of external type constructors as follows:
35 --
36 -- * tycons which have converted versions are mapped to 'True'
37 -- * tycons which are not changed by vectorisation are mapped to 'False'
38 -- * tycons which can't be converted are not elements of the map
39 --
40 classifyTyCons :: UniqFM Bool -- ^type constructor conversion status
41 -> [TyCon] -- ^type constructors that need to be classified
42 -> ([TyCon], [TyCon]) -- ^tycons to be converted & not to be converted
43 classifyTyCons convStatus tcs = classify [] [] convStatus (tyConGroups tcs)
44 where
45 classify conv keep _ [] = (conv, keep)
46 classify conv keep cs ((tcs, ds) : rs)
47 | can_convert && must_convert
48 = classify (tcs ++ conv) keep (cs `addListToUFM` [(tc, True) | tc <- tcs]) rs
49 | can_convert
50 = classify conv (tcs ++ keep) (cs `addListToUFM` [(tc, False) | tc <- tcs]) rs
51 | otherwise
52 = classify conv keep cs rs
53 where
54 refs = ds `delListFromUniqSet` tcs
55
56 can_convert = isNullUFM (refs `minusUFM` cs) && all convertable tcs
57 must_convert = foldUFM (||) False (intersectUFM_C const cs refs)
58
59 convertable tc = isDataTyCon tc && all isVanillaDataCon (tyConDataCons tc)
60
61 -- Used to group type constructors into mutually dependent groups.
62 --
63 type TyConGroup = ([TyCon], UniqSet TyCon)
64
65 -- Compute mutually recursive groups of tycons in topological order.
66 --
67 tyConGroups :: [TyCon] -> [TyConGroup]
68 tyConGroups tcs = map mk_grp (stronglyConnCompFromEdgedVertices edges)
69 where
70 edges = [((tc, ds), tc, uniqSetToList ds) | tc <- tcs
71 , let ds = tyConsOfTyCon tc]
72
73 mk_grp (AcyclicSCC (tc, ds)) = ([tc], ds)
74 mk_grp (CyclicSCC els) = (tcs, unionManyUniqSets dss)
75 where
76 (tcs, dss) = unzip els
77
78 -- |Collect the set of TyCons used by the representation of some data type.
79 --
80 tyConsOfTyCon :: TyCon -> UniqSet TyCon
81 tyConsOfTyCon = tyConsOfTypes . concatMap dataConRepArgTys . tyConDataCons
82
83 -- |Collect the set of TyCons that occur in these types.
84 --
85 tyConsOfTypes :: [Type] -> UniqSet TyCon
86 tyConsOfTypes = unionManyUniqSets . map tyConsOfType
87
88 -- |Collect the set of TyCons that occur in this type.
89 --
90 tyConsOfType :: Type -> UniqSet TyCon
91 tyConsOfType ty
92 | Just ty' <- coreView ty = tyConsOfType ty'
93 tyConsOfType (TyVarTy _) = emptyUniqSet
94 tyConsOfType (TyConApp tc tys) = extend (tyConsOfTypes tys)
95 where
96 extend | isUnLiftedTyCon tc
97 || isTupleTyCon tc = id
98
99 | otherwise = (`addOneToUniqSet` tc)
100
101 tyConsOfType (AppTy a b) = tyConsOfType a `unionUniqSets` tyConsOfType b
102 tyConsOfType (FunTy a b) = (tyConsOfType a `unionUniqSets` tyConsOfType b)
103 `addOneToUniqSet` funTyCon
104 tyConsOfType (ForAllTy _ ty) = tyConsOfType ty
105 tyConsOfType other = pprPanic "ClosureConv.tyConsOfType" $ ppr other