1{-
2(c) The University of Glasgow 2006
3(c) The GRASP/AQUA Project, Glasgow University, 1992-1998
4-}
5
6{-# LANGUAGE RankNTypes #-}
7{-# LANGUAGE TypeFamilies #-}
8{-# LANGUAGE FlexibleContexts #-}
9{-# LANGUAGE TypeSynonymInstances #-}
10{-# LANGUAGE FlexibleInstances #-}
11{-# LANGUAGE UndecidableInstances #-}
12module TrieMap(
13   -- * Maps over 'Maybe' values
14   MaybeMap,
15   -- * Maps over 'List' values
16   ListMap,
17   -- * Maps over 'Literal's
18   LiteralMap,
19   -- * 'TrieMap' class
20   TrieMap(..), insertTM, deleteTM,
21
22   -- * Things helpful for adding additional Instances.
23   (>.>), (|>), (|>>), XT,
24   foldMaybe,
25   -- * Map for leaf compression
26   GenMap,
27   lkG, xtG, mapG, fdG,
28   xtList, lkList
29
30 ) where
31
32import GhcPrelude
33
34import Literal
35import UniqDFM
36import Unique( Unique )
37
38import qualified Data.Map    as Map
39import qualified Data.IntMap as IntMap
40import Outputable
41import Control.Monad( (>=>) )
42
43{-
44This module implements TrieMaps, which are finite mappings
45whose key is a structured value like a CoreExpr or Type.
46
47This file implements tries over general data structures.
48Implementation for tries over Core Expressions/Types are
49available in coreSyn/TrieMap.
50
51The regular pattern for handling TrieMaps on data structures was first
52described (to my knowledge) in Connelly and Morris's 1995 paper "A
53generalization of the Trie Data Structure"; there is also an accessible
54description of the idea in Okasaki's book "Purely Functional Data
55Structures", Section 10.3.2
56
57************************************************************************
58*                                                                      *
59                   The TrieMap class
60*                                                                      *
61************************************************************************
62-}
63
64type XT a = Maybe a -> Maybe a  -- How to alter a non-existent elt (Nothing)
65                                --               or an existing elt (Just)
66
67class TrieMap m where
68   type Key m :: *
69   emptyTM  :: m a
70   lookupTM :: forall b. Key m -> m b -> Maybe b
71   alterTM  :: forall b. Key m -> XT b -> m b -> m b
72   mapTM    :: (a->b) -> m a -> m b
73
74   foldTM   :: (a -> b -> b) -> m a -> b -> b
75      -- The unusual argument order here makes
76      -- it easy to compose calls to foldTM;
77      -- see for example fdE below
78
79insertTM :: TrieMap m => Key m -> a -> m a -> m a
80insertTM k v m = alterTM k (\_ -> Just v) m
81
82deleteTM :: TrieMap m => Key m -> m a -> m a
83deleteTM k m = alterTM k (\_ -> Nothing) m
84
85----------------------
86-- Recall that
87--   Control.Monad.(>=>) :: (a -> Maybe b) -> (b -> Maybe c) -> a -> Maybe c
88
89(>.>) :: (a -> b) -> (b -> c) -> a -> c
90-- Reverse function composition (do f first, then g)
91infixr 1 >.>
92(f >.> g) x = g (f x)
93infixr 1 |>, |>>
94
95(|>) :: a -> (a->b) -> b     -- Reverse application
96x |> f = f x
97
98----------------------
99(|>>) :: TrieMap m2
100      => (XT (m2 a) -> m1 (m2 a) -> m1 (m2 a))
101      -> (m2 a -> m2 a)
102      -> m1 (m2 a) -> m1 (m2 a)
103(|>>) f g = f (Just . g . deMaybe)
104
105deMaybe :: TrieMap m => Maybe (m a) -> m a
106deMaybe Nothing  = emptyTM
107deMaybe (Just m) = m
108
109{-
110************************************************************************
111*                                                                      *
112                   IntMaps
113*                                                                      *
114************************************************************************
115-}
116
117instance TrieMap IntMap.IntMap where
118  type Key IntMap.IntMap = Int
119  emptyTM = IntMap.empty
120  lookupTM k m = IntMap.lookup k m
121  alterTM = xtInt
122  foldTM k m z = IntMap.foldr k z m
123  mapTM f m = IntMap.map f m
124
125xtInt :: Int -> XT a -> IntMap.IntMap a -> IntMap.IntMap a
126xtInt k f m = IntMap.alter f k m
127
128instance Ord k => TrieMap (Map.Map k) where
129  type Key (Map.Map k) = k
130  emptyTM = Map.empty
131  lookupTM = Map.lookup
132  alterTM k f m = Map.alter f k m
133  foldTM k m z = Map.foldr k z m
134  mapTM f m = Map.map f m
135
136
137{-
138Note [foldTM determinism]
139~~~~~~~~~~~~~~~~~~~~~~~~~
140We want foldTM to be deterministic, which is why we have an instance of
141TrieMap for UniqDFM, but not for UniqFM. Here's an example of some things that
142go wrong if foldTM is nondeterministic. Consider:
143
144  f a b = return (a <> b)
145
146Depending on the order that the typechecker generates constraints you
147get either:
148
149  f :: (Monad m, Monoid a) => a -> a -> m a
150
151or:
152
153  f :: (Monoid a, Monad m) => a -> a -> m a
154
155The generated code will be different after desugaring as the dictionaries
156will be bound in different orders, leading to potential ABI incompatibility.
157
158One way to solve this would be to notice that the typeclasses could be
159sorted alphabetically.
160
161Unfortunately that doesn't quite work with this example:
162
163  f a b = let x = a <> a; y = b <> b in x
164
165where you infer:
166
167  f :: (Monoid m, Monoid m1) => m1 -> m -> m1
168
169or:
170
171  f :: (Monoid m1, Monoid m) => m1 -> m -> m1
172
173Here you could decide to take the order of the type variables in the type
174according to depth first traversal and use it to order the constraints.
175
176The real trouble starts when the user enables incoherent instances and
177the compiler has to make an arbitrary choice. Consider:
178
179  class T a b where
180    go :: a -> b -> String
181
182  instance (Show b) => T Int b where
183    go a b = show a ++ show b
184
185  instance (Show a) => T a Bool where
186    go a b = show a ++ show b
187
188  f = go 10 True
189
190GHC is free to choose either dictionary to implement f, but for the sake of
191determinism we'd like it to be consistent when compiling the same sources
192with the same flags.
193
194inert_dicts :: DictMap is implemented with a TrieMap. In getUnsolvedInerts it
195gets converted to a bag of (Wanted) Cts using a fold. Then in
196solve_simple_wanteds it's merged with other WantedConstraints. We want the
197conversion to a bag to be deterministic. For that purpose we use UniqDFM
198instead of UniqFM to implement the TrieMap.
199
200See Note [Deterministic UniqFM] in UniqDFM for more details on how it's made
201deterministic.
202-}
203
204instance TrieMap UniqDFM where
205  type Key UniqDFM = Unique
206  emptyTM = emptyUDFM
207  lookupTM k m = lookupUDFM m k
208  alterTM k f m = alterUDFM f m k
209  foldTM k m z = foldUDFM k z m
210  mapTM f m = mapUDFM f m
211
212{-
213************************************************************************
214*                                                                      *
215                   Maybes
216*                                                                      *
217************************************************************************
218
219If              m is a map from k -> val
220then (MaybeMap m) is a map from (Maybe k) -> val
221-}
222
223data MaybeMap m a = MM { mm_nothing  :: Maybe a, mm_just :: m a }
224
225instance TrieMap m => TrieMap (MaybeMap m) where
226   type Key (MaybeMap m) = Maybe (Key m)
227   emptyTM  = MM { mm_nothing = Nothing, mm_just = emptyTM }
228   lookupTM = lkMaybe lookupTM
229   alterTM  = xtMaybe alterTM
230   foldTM   = fdMaybe
231   mapTM    = mapMb
232
233mapMb :: TrieMap m => (a->b) -> MaybeMap m a -> MaybeMap m b
234mapMb f (MM { mm_nothing = mn, mm_just = mj })
235  = MM { mm_nothing = fmap f mn, mm_just = mapTM f mj }
236
237lkMaybe :: (forall b. k -> m b -> Maybe b)
238        -> Maybe k -> MaybeMap m a -> Maybe a
239lkMaybe _  Nothing  = mm_nothing
240lkMaybe lk (Just x) = mm_just >.> lk x
241
242xtMaybe :: (forall b. k -> XT b -> m b -> m b)
243        -> Maybe k -> XT a -> MaybeMap m a -> MaybeMap m a
244xtMaybe _  Nothing  f m = m { mm_nothing  = f (mm_nothing m) }
245xtMaybe tr (Just x) f m = m { mm_just = mm_just m |> tr x f }
246
247fdMaybe :: TrieMap m => (a -> b -> b) -> MaybeMap m a -> b -> b
248fdMaybe k m = foldMaybe k (mm_nothing m)
249            . foldTM k (mm_just m)
250
251{-
252************************************************************************
253*                                                                      *
254                   Lists
255*                                                                      *
256************************************************************************
257-}
258
259data ListMap m a
260  = LM { lm_nil  :: Maybe a
261       , lm_cons :: m (ListMap m a) }
262
263instance TrieMap m => TrieMap (ListMap m) where
264   type Key (ListMap m) = [Key m]
265   emptyTM  = LM { lm_nil = Nothing, lm_cons = emptyTM }
266   lookupTM = lkList lookupTM
267   alterTM  = xtList alterTM
268   foldTM   = fdList
269   mapTM    = mapList
270
271instance (TrieMap m, Outputable a) => Outputable (ListMap m a) where
272  ppr m = text "List elts" <+> ppr (foldTM (:) m [])
273
274mapList :: TrieMap m => (a->b) -> ListMap m a -> ListMap m b
275mapList f (LM { lm_nil = mnil, lm_cons = mcons })
276  = LM { lm_nil = fmap f mnil, lm_cons = mapTM (mapTM f) mcons }
277
278lkList :: TrieMap m => (forall b. k -> m b -> Maybe b)
279        -> [k] -> ListMap m a -> Maybe a
280lkList _  []     = lm_nil
281lkList lk (x:xs) = lm_cons >.> lk x >=> lkList lk xs
282
283xtList :: TrieMap m => (forall b. k -> XT b -> m b -> m b)
284        -> [k] -> XT a -> ListMap m a -> ListMap m a
285xtList _  []     f m = m { lm_nil  = f (lm_nil m) }
286xtList tr (x:xs) f m = m { lm_cons = lm_cons m |> tr x |>> xtList tr xs f }
287
288fdList :: forall m a b. TrieMap m
289       => (a -> b -> b) -> ListMap m a -> b -> b
290fdList k m = foldMaybe k          (lm_nil m)
291           . foldTM    (fdList k) (lm_cons m)
292
293foldMaybe :: (a -> b -> b) -> Maybe a -> b -> b
294foldMaybe _ Nothing  b = b
295foldMaybe k (Just a) b = k a b
296
297{-
298************************************************************************
299*                                                                      *
300                   Basic maps
301*                                                                      *
302************************************************************************
303-}
304
305type LiteralMap  a = Map.Map Literal a
306
307{-
308************************************************************************
309*                                                                      *
310                   GenMap
311*                                                                      *
312************************************************************************
313
314Note [Compressed TrieMap]
315~~~~~~~~~~~~~~~~~~~~~~~~~
316
317The GenMap constructor augments TrieMaps with leaf compression.  This helps
318solve the performance problem detailed in #9960: suppose we have a handful
319H of entries in a TrieMap, each with a very large key, size K. If you fold over
320such a TrieMap you'd expect time O(H). That would certainly be true of an
321association list! But with TrieMap we actually have to navigate down a long
322singleton structure to get to the elements, so it takes time O(K*H).  This
323can really hurt on many type-level computation benchmarks:
324see for example T9872d.
325
326The point of a TrieMap is that you need to navigate to the point where only one
327key remains, and then things should be fast.  So the point of a SingletonMap
328is that, once we are down to a single (key,value) pair, we stop and
329just use SingletonMap.
330
331'EmptyMap' provides an even more basic (but essential) optimization: if there is
332nothing in the map, don't bother building out the (possibly infinite) recursive
333TrieMap structure!
334
335Compressed triemaps are heavily used by CoreMap. So we have to mark some things
336as INLINEABLE to permit specialization.
337-}
338
339data GenMap m a
340   = EmptyMap
341   | SingletonMap (Key m) a
342   | MultiMap (m a)
343
344instance (Outputable a, Outputable (m a)) => Outputable (GenMap m a) where
345  ppr EmptyMap = text "Empty map"
346  ppr (SingletonMap _ v) = text "Singleton map" <+> ppr v
347  ppr (MultiMap m) = ppr m
348
349-- TODO undecidable instance
350instance (Eq (Key m), TrieMap m) => TrieMap (GenMap m) where
351   type Key (GenMap m) = Key m
352   emptyTM  = EmptyMap
353   lookupTM = lkG
354   alterTM  = xtG
355   foldTM   = fdG
356   mapTM    = mapG
357
358--We want to be able to specialize these functions when defining eg
359--tries over (GenMap CoreExpr) which requires INLINEABLE
360
361{-# INLINEABLE lkG #-}
362lkG :: (Eq (Key m), TrieMap m) => Key m -> GenMap m a -> Maybe a
363lkG _ EmptyMap                         = Nothing
364lkG k (SingletonMap k' v') | k == k'   = Just v'
365                           | otherwise = Nothing
366lkG k (MultiMap m)                     = lookupTM k m
367
368{-# INLINEABLE xtG #-}
369xtG :: (Eq (Key m), TrieMap m) => Key m -> XT a -> GenMap m a -> GenMap m a
370xtG k f EmptyMap
371    = case f Nothing of
372        Just v  -> SingletonMap k v
373        Nothing -> EmptyMap
374xtG k f m@(SingletonMap k' v')
375    | k' == k
376    -- The new key matches the (single) key already in the tree.  Hence,
377    -- apply @f@ to @Just v'@ and build a singleton or empty map depending
378    -- on the 'Just'/'Nothing' response respectively.
379    = case f (Just v') of
380        Just v'' -> SingletonMap k' v''
381        Nothing  -> EmptyMap
382    | otherwise
383    -- We've hit a singleton tree for a different key than the one we are
384    -- searching for. Hence apply @f@ to @Nothing@. If result is @Nothing@ then
385    -- we can just return the old map. If not, we need a map with *two*
386    -- entries. The easiest way to do that is to insert two items into an empty
387    -- map of type @m a@.
388    = case f Nothing of
389        Nothing  -> m
390        Just v   -> emptyTM |> alterTM k' (const (Just v'))
391                           >.> alterTM k  (const (Just v))
392                           >.> MultiMap
393xtG k f (MultiMap m) = MultiMap (alterTM k f m)
394
395{-# INLINEABLE mapG #-}
396mapG :: TrieMap m => (a -> b) -> GenMap m a -> GenMap m b
397mapG _ EmptyMap = EmptyMap
398mapG f (SingletonMap k v) = SingletonMap k (f v)
399mapG f (MultiMap m) = MultiMap (mapTM f m)
400
401{-# INLINEABLE fdG #-}
402fdG :: TrieMap m => (a -> b -> b) -> GenMap m a -> b -> b
403fdG _ EmptyMap = \z -> z
404fdG k (SingletonMap _ v) = \z -> k v z
405fdG k (MultiMap m) = foldTM k m
406