e85377910abb743b5c0b297ab749240be4382dac
[packages/base.git] / GHC / Event / IntMap.hs
1 {-# LANGUAGE Unsafe #-}
2 {-# LANGUAGE CPP, MagicHash, NoImplicitPrelude #-}
3
4 -----------------------------------------------------------------------------
5 -- |
6 -- Module : GHC.Event.IntMap
7 -- Copyright : (c) Daan Leijen 2002
8 -- (c) Andriy Palamarchuk 2008
9 -- License : BSD-style
10 -- Maintainer : libraries@haskell.org
11 -- Stability : provisional
12 -- Portability : portable
13 --
14 -- An efficient implementation of maps from integer keys to values.
15 --
16 -- Since many function names (but not the type name) clash with
17 -- "Prelude" names, this module is usually imported @qualified@, e.g.
18 --
19 -- > import Data.IntMap (IntMap)
20 -- > import qualified Data.IntMap as IntMap
21 --
22 -- The implementation is based on /big-endian patricia trees/. This data
23 -- structure performs especially well on binary operations like 'union'
24 -- and 'intersection'. However, my benchmarks show that it is also
25 -- (much) faster on insertions and deletions when compared to a generic
26 -- size-balanced map implementation (see "Data.Map").
27 --
28 -- * Chris Okasaki and Andy Gill, \"/Fast Mergeable Integer Maps/\",
29 -- Workshop on ML, September 1998, pages 77-86,
30 -- <http://citeseer.ist.psu.edu/okasaki98fast.html>
31 --
32 -- * D.R. Morrison, \"/PATRICIA -- Practical Algorithm To Retrieve
33 -- Information Coded In Alphanumeric/\", Journal of the ACM, 15(4),
34 -- October 1968, pages 514-534.
35 --
36 -- Operation comments contain the operation time complexity in
37 -- the Big-O notation <http://en.wikipedia.org/wiki/Big_O_notation>.
38 -- Many operations have a worst-case complexity of /O(min(n,W))/.
39 -- This means that the operation can become linear in the number of
40 -- elements with a maximum of /W/ -- the number of bits in an 'Int'
41 -- (32 or 64).
42 -----------------------------------------------------------------------------
43
44 module GHC.Event.IntMap
45 (
46 -- * Map type
47 IntMap
48 , Key
49
50 -- * Query
51 , lookup
52 , member
53
54 -- * Construction
55 , empty
56
57 -- * Insertion
58 , insertWith
59
60 -- * Delete\/Update
61 , delete
62 , updateWith
63
64 -- * Traversal
65 -- ** Fold
66 , foldWithKey
67
68 -- * Conversion
69 , keys
70 ) where
71
72 import Data.Bits
73
74 import Data.Maybe (Maybe(..))
75 import GHC.Base hiding (foldr)
76 import GHC.Num (Num(..))
77 import GHC.Real (fromIntegral)
78 import GHC.Show (Show(showsPrec), showParen, shows, showString)
79
80 #if __GLASGOW_HASKELL__
81 import GHC.Word (Word(..))
82 #else
83 import Data.Word
84 #endif
85
86 -- | A @Nat@ is a natural machine word (an unsigned Int)
87 type Nat = Word
88
89 natFromInt :: Key -> Nat
90 natFromInt i = fromIntegral i
91
92 intFromNat :: Nat -> Key
93 intFromNat w = fromIntegral w
94
95 shiftRL :: Nat -> Key -> Nat
96 #if __GLASGOW_HASKELL__
97 -- GHC: use unboxing to get @shiftRL@ inlined.
98 shiftRL (W# x) (I# i) = W# (shiftRL# x i)
99 #else
100 shiftRL x i = shiftR x i
101 #endif
102
103 ------------------------------------------------------------------------
104 -- Types
105
106 -- | A map of integers to values @a@.
107 data IntMap a = Nil
108 | Tip {-# UNPACK #-} !Key !a
109 | Bin {-# UNPACK #-} !Prefix
110 {-# UNPACK #-} !Mask
111 !(IntMap a)
112 !(IntMap a)
113
114 type Prefix = Int
115 type Mask = Int
116 type Key = Int
117
118 ------------------------------------------------------------------------
119 -- Query
120
121 -- | /O(min(n,W))/ Lookup the value at a key in the map. See also
122 -- 'Data.Map.lookup'.
123 lookup :: Key -> IntMap a -> Maybe a
124 lookup k t = let nk = natFromInt k in seq nk (lookupN nk t)
125
126 lookupN :: Nat -> IntMap a -> Maybe a
127 lookupN k t
128 = case t of
129 Bin _ m l r
130 | zeroN k (natFromInt m) -> lookupN k l
131 | otherwise -> lookupN k r
132 Tip kx x
133 | (k == natFromInt kx) -> Just x
134 | otherwise -> Nothing
135 Nil -> Nothing
136
137 -- | /O(min(n,W))/. Is the key a member of the map?
138 --
139 -- > member 5 (fromList [(5,'a'), (3,'b')]) == True
140 -- > member 1 (fromList [(5,'a'), (3,'b')]) == False
141
142 member :: Key -> IntMap a -> Bool
143 member k m
144 = case lookup k m of
145 Nothing -> False
146 Just _ -> True
147
148 ------------------------------------------------------------------------
149 -- Construction
150
151 -- | /O(1)/ The empty map.
152 --
153 -- > empty == fromList []
154 -- > size empty == 0
155 empty :: IntMap a
156 empty = Nil
157
158 ------------------------------------------------------------------------
159 -- Insert
160
161 -- | /O(min(n,W))/ Insert with a function, combining new value and old
162 -- value. @insertWith f key value mp@ will insert the pair (key,
163 -- value) into @mp@ if key does not exist in the map. If the key does
164 -- exist, the function will insert the pair (key, f new_value
165 -- old_value). The result is a pair where the first element is the
166 -- old value, if one was present, and the second is the modified map.
167 insertWith :: (a -> a -> a) -> Key -> a -> IntMap a -> (Maybe a, IntMap a)
168 insertWith f k x t = case t of
169 Bin p m l r
170 | nomatch k p m -> (Nothing, join k (Tip k x) p t)
171 | zero k m -> let (found, l') = insertWith f k x l
172 in (found, Bin p m l' r)
173 | otherwise -> let (found, r') = insertWith f k x r
174 in (found, Bin p m l r')
175 Tip ky y
176 | k == ky -> (Just y, Tip k (f x y))
177 | otherwise -> (Nothing, join k (Tip k x) ky t)
178 Nil -> (Nothing, Tip k x)
179
180
181 ------------------------------------------------------------------------
182 -- Delete/Update
183
184 -- | /O(min(n,W))/. Delete a key and its value from the map. When the
185 -- key is not a member of the map, the original map is returned. The
186 -- result is a pair where the first element is the value associated
187 -- with the deleted key, if one existed, and the second element is the
188 -- modified map.
189 delete :: Key -> IntMap a -> (Maybe a, IntMap a)
190 delete k t = case t of
191 Bin p m l r
192 | nomatch k p m -> (Nothing, t)
193 | zero k m -> let (found, l') = delete k l
194 in (found, bin p m l' r)
195 | otherwise -> let (found, r') = delete k r
196 in (found, bin p m l r')
197 Tip ky y
198 | k == ky -> (Just y, Nil)
199 | otherwise -> (Nothing, t)
200 Nil -> (Nothing, Nil)
201
202 updateWith :: (a -> Maybe a) -> Key -> IntMap a -> (Maybe a, IntMap a)
203 updateWith f k t = case t of
204 Bin p m l r
205 | nomatch k p m -> (Nothing, t)
206 | zero k m -> let (found, l') = updateWith f k l
207 in (found, bin p m l' r)
208 | otherwise -> let (found, r') = updateWith f k r
209 in (found, bin p m l r')
210 Tip ky y
211 | k == ky -> case (f y) of
212 Just y' -> (Just y, Tip ky y')
213 Nothing -> (Just y, Nil)
214 | otherwise -> (Nothing, t)
215 Nil -> (Nothing, Nil)
216 -- | /O(n)/. Fold the keys and values in the map, such that
217 -- @'foldWithKey' f z == 'Prelude.foldr' ('uncurry' f) z . 'toAscList'@.
218 -- For example,
219 --
220 -- > keys map = foldWithKey (\k x ks -> k:ks) [] map
221 --
222 -- > let f k a result = result ++ "(" ++ (show k) ++ ":" ++ a ++ ")"
223 -- > foldWithKey f "Map: " (fromList [(5,"a"), (3,"b")]) == "Map: (5:a)(3:b)"
224
225 foldWithKey :: (Key -> a -> b -> b) -> b -> IntMap a -> b
226 foldWithKey f z t
227 = foldr f z t
228
229 -- | /O(n)/. Convert the map to a list of key\/value pairs.
230 --
231 -- > toList (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")]
232 -- > toList empty == []
233
234 toList :: IntMap a -> [(Key,a)]
235 toList t
236 = foldWithKey (\k x xs -> (k,x):xs) [] t
237
238 foldr :: (Key -> a -> b -> b) -> b -> IntMap a -> b
239 foldr f z t
240 = case t of
241 Bin 0 m l r | m < 0 -> foldr' f (foldr' f z l) r -- put negative numbers before.
242 Bin _ _ _ _ -> foldr' f z t
243 Tip k x -> f k x z
244 Nil -> z
245
246 foldr' :: (Key -> a -> b -> b) -> b -> IntMap a -> b
247 foldr' f z t
248 = case t of
249 Bin _ _ l r -> foldr' f (foldr' f z r) l
250 Tip k x -> f k x z
251 Nil -> z
252
253 -- | /O(n)/. Return all keys of the map in ascending order.
254 --
255 -- > keys (fromList [(5,"a"), (3,"b")]) == [3,5]
256 -- > keys empty == []
257
258 keys :: IntMap a -> [Key]
259 keys m
260 = foldWithKey (\k _ ks -> k:ks) [] m
261
262 ------------------------------------------------------------------------
263 -- Eq
264
265 instance Eq a => Eq (IntMap a) where
266 t1 == t2 = equal t1 t2
267 t1 /= t2 = nequal t1 t2
268
269 equal :: Eq a => IntMap a -> IntMap a -> Bool
270 equal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
271 = (m1 == m2) && (p1 == p2) && (equal l1 l2) && (equal r1 r2)
272 equal (Tip kx x) (Tip ky y)
273 = (kx == ky) && (x==y)
274 equal Nil Nil = True
275 equal _ _ = False
276
277 nequal :: Eq a => IntMap a -> IntMap a -> Bool
278 nequal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
279 = (m1 /= m2) || (p1 /= p2) || (nequal l1 l2) || (nequal r1 r2)
280 nequal (Tip kx x) (Tip ky y)
281 = (kx /= ky) || (x/=y)
282 nequal Nil Nil = False
283 nequal _ _ = True
284
285 instance Show a => Show (IntMap a) where
286 showsPrec d m = showParen (d > 10) $
287 showString "fromList " . shows (toList m)
288
289 ------------------------------------------------------------------------
290 -- Utility functions
291
292 join :: Prefix -> IntMap a -> Prefix -> IntMap a -> IntMap a
293 join p1 t1 p2 t2
294 | zero p1 m = Bin p m t1 t2
295 | otherwise = Bin p m t2 t1
296 where
297 m = branchMask p1 p2
298 p = mask p1 m
299
300 -- | @bin@ assures that we never have empty trees within a tree.
301 bin :: Prefix -> Mask -> IntMap a -> IntMap a -> IntMap a
302 bin _ _ l Nil = l
303 bin _ _ Nil r = r
304 bin p m l r = Bin p m l r
305
306 ------------------------------------------------------------------------
307 -- Endian independent bit twiddling
308
309 zero :: Key -> Mask -> Bool
310 zero i m = (natFromInt i) .&. (natFromInt m) == 0
311
312 nomatch :: Key -> Prefix -> Mask -> Bool
313 nomatch i p m = (mask i m) /= p
314
315 mask :: Key -> Mask -> Prefix
316 mask i m = maskW (natFromInt i) (natFromInt m)
317
318 zeroN :: Nat -> Nat -> Bool
319 zeroN i m = (i .&. m) == 0
320
321 ------------------------------------------------------------------------
322 -- Big endian operations
323
324 maskW :: Nat -> Nat -> Prefix
325 maskW i m = intFromNat (i .&. (complement (m-1) `xor` m))
326
327 branchMask :: Prefix -> Prefix -> Mask
328 branchMask p1 p2
329 = intFromNat (highestBitMask (natFromInt p1 `xor` natFromInt p2))
330
331 {-
332 Finding the highest bit mask in a word [x] can be done efficiently in
333 three ways:
334
335 * convert to a floating point value and the mantissa tells us the
336 [log2(x)] that corresponds with the highest bit position. The mantissa
337 is retrieved either via the standard C function [frexp] or by some bit
338 twiddling on IEEE compatible numbers (float). Note that one needs to
339 use at least [double] precision for an accurate mantissa of 32 bit
340 numbers.
341
342 * use bit twiddling, a logarithmic sequence of bitwise or's and shifts (bit).
343
344 * use processor specific assembler instruction (asm).
345
346 The most portable way would be [bit], but is it efficient enough?
347 I have measured the cycle counts of the different methods on an AMD
348 Athlon-XP 1800 (~ Pentium III 1.8Ghz) using the RDTSC instruction:
349
350 highestBitMask: method cycles
351 --------------
352 frexp 200
353 float 33
354 bit 11
355 asm 12
356
357 Wow, the bit twiddling is on today's RISC like machines even faster
358 than a single CISC instruction (BSR)!
359 -}
360
361 -- | @highestBitMask@ returns a word where only the highest bit is
362 -- set. It is found by first setting all bits in lower positions than
363 -- the highest bit and than taking an exclusive or with the original
364 -- value. Allthough the function may look expensive, GHC compiles
365 -- this into excellent C code that subsequently compiled into highly
366 -- efficient machine code. The algorithm is derived from Jorg Arndt's
367 -- FXT library.
368 highestBitMask :: Nat -> Nat
369 highestBitMask x0
370 = case (x0 .|. shiftRL x0 1) of
371 x1 -> case (x1 .|. shiftRL x1 2) of
372 x2 -> case (x2 .|. shiftRL x2 4) of
373 x3 -> case (x3 .|. shiftRL x3 8) of
374 x4 -> case (x4 .|. shiftRL x4 16) of
375 x5 -> case (x5 .|. shiftRL x5 32) of -- for 64 bit platforms
376 x6 -> (x6 `xor` (shiftRL x6 1))