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