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))