base-4.3.0.0: System/Event/IntMap.hs
{-# LANGUAGE CPP, MagicHash, NoImplicitPrelude #-}
-----------------------------------------------------------------------------
-- |
-- Module : System.Event.IntMap
-- Copyright : (c) Daan Leijen 2002
-- (c) Andriy Palamarchuk 2008
-- License : BSD-style
-- Maintainer : libraries@haskell.org
-- Stability : provisional
-- Portability : portable
--
-- An efficient implementation of maps from integer keys to values.
--
-- Since many function names (but not the type name) clash with
-- "Prelude" names, this module is usually imported @qualified@, e.g.
--
-- > import Data.IntMap (IntMap)
-- > import qualified Data.IntMap as IntMap
--
-- The implementation is based on /big-endian patricia trees/. This data
-- structure performs especially well on binary operations like 'union'
-- and 'intersection'. However, my benchmarks show that it is also
-- (much) faster on insertions and deletions when compared to a generic
-- size-balanced map implementation (see "Data.Map").
--
-- * Chris Okasaki and Andy Gill, \"/Fast Mergeable Integer Maps/\",
-- Workshop on ML, September 1998, pages 77-86,
-- <http://citeseer.ist.psu.edu/okasaki98fast.html>
--
-- * D.R. Morrison, \"/PATRICIA -- Practical Algorithm To Retrieve
-- Information Coded In Alphanumeric/\", Journal of the ACM, 15(4),
-- October 1968, pages 514-534.
--
-- Operation comments contain the operation time complexity in
-- the Big-O notation <http://en.wikipedia.org/wiki/Big_O_notation>.
-- Many operations have a worst-case complexity of /O(min(n,W))/.
-- This means that the operation can become linear in the number of
-- elements with a maximum of /W/ -- the number of bits in an 'Int'
-- (32 or 64).
-----------------------------------------------------------------------------
module System.Event.IntMap
(
-- * Map type
IntMap
, Key
-- * Query
, lookup
, member
-- * Construction
, empty
-- * Insertion
, insertWith
-- * Delete\/Update
, delete
, updateWith
-- * Traversal
-- ** Fold
, foldWithKey
-- * Conversion
, keys
) where
import Data.Bits
import Data.Maybe (Maybe(..))
import GHC.Base hiding (foldr)
import GHC.Num (Num(..))
import GHC.Real (fromIntegral)
import GHC.Show (Show(showsPrec), showParen, shows, showString)
#if __GLASGOW_HASKELL__
import GHC.Word (Word(..))
#else
import Data.Word
#endif
-- | A @Nat@ is a natural machine word (an unsigned Int)
type Nat = Word
natFromInt :: Key -> Nat
natFromInt i = fromIntegral i
intFromNat :: Nat -> Key
intFromNat w = fromIntegral w
shiftRL :: Nat -> Key -> Nat
#if __GLASGOW_HASKELL__
-- GHC: use unboxing to get @shiftRL@ inlined.
shiftRL (W# x) (I# i) = W# (shiftRL# x i)
#else
shiftRL x i = shiftR x i
#endif
------------------------------------------------------------------------
-- Types
-- | A map of integers to values @a@.
data IntMap a = Nil
| Tip {-# UNPACK #-} !Key !a
| Bin {-# UNPACK #-} !Prefix
{-# UNPACK #-} !Mask
!(IntMap a)
!(IntMap a)
type Prefix = Int
type Mask = Int
type Key = Int
------------------------------------------------------------------------
-- Query
-- | /O(min(n,W))/ Lookup the value at a key in the map. See also
-- 'Data.Map.lookup'.
lookup :: Key -> IntMap a -> Maybe a
lookup k t = let nk = natFromInt k in seq nk (lookupN nk t)
lookupN :: Nat -> IntMap a -> Maybe a
lookupN k t
= case t of
Bin _ m l r
| zeroN k (natFromInt m) -> lookupN k l
| otherwise -> lookupN k r
Tip kx x
| (k == natFromInt kx) -> Just x
| otherwise -> Nothing
Nil -> Nothing
-- | /O(min(n,W))/. Is the key a member of the map?
--
-- > member 5 (fromList [(5,'a'), (3,'b')]) == True
-- > member 1 (fromList [(5,'a'), (3,'b')]) == False
member :: Key -> IntMap a -> Bool
member k m
= case lookup k m of
Nothing -> False
Just _ -> True
------------------------------------------------------------------------
-- Construction
-- | /O(1)/ The empty map.
--
-- > empty == fromList []
-- > size empty == 0
empty :: IntMap a
empty = Nil
------------------------------------------------------------------------
-- Insert
-- | /O(min(n,W))/ Insert with a function, combining new value and old
-- value. @insertWith f key value mp@ will insert the pair (key,
-- value) into @mp@ if key does not exist in the map. If the key does
-- exist, the function will insert the pair (key, f new_value
-- old_value). The result is a pair where the first element is the
-- old value, if one was present, and the second is the modified map.
insertWith :: (a -> a -> a) -> Key -> a -> IntMap a -> (Maybe a, IntMap a)
insertWith f k x t = case t of
Bin p m l r
| nomatch k p m -> (Nothing, join k (Tip k x) p t)
| zero k m -> let (found, l') = insertWith f k x l
in (found, Bin p m l' r)
| otherwise -> let (found, r') = insertWith f k x r
in (found, Bin p m l r')
Tip ky y
| k == ky -> (Just y, Tip k (f x y))
| otherwise -> (Nothing, join k (Tip k x) ky t)
Nil -> (Nothing, Tip k x)
------------------------------------------------------------------------
-- Delete/Update
-- | /O(min(n,W))/. Delete a key and its value from the map. When the
-- key is not a member of the map, the original map is returned. The
-- result is a pair where the first element is the value associated
-- with the deleted key, if one existed, and the second element is the
-- modified map.
delete :: Key -> IntMap a -> (Maybe a, IntMap a)
delete k t = case t of
Bin p m l r
| nomatch k p m -> (Nothing, t)
| zero k m -> let (found, l') = delete k l
in (found, bin p m l' r)
| otherwise -> let (found, r') = delete k r
in (found, bin p m l r')
Tip ky y
| k == ky -> (Just y, Nil)
| otherwise -> (Nothing, t)
Nil -> (Nothing, Nil)
updateWith :: (a -> Maybe a) -> Key -> IntMap a -> (Maybe a, IntMap a)
updateWith f k t = case t of
Bin p m l r
| nomatch k p m -> (Nothing, t)
| zero k m -> let (found, l') = updateWith f k l
in (found, bin p m l' r)
| otherwise -> let (found, r') = updateWith f k r
in (found, bin p m l r')
Tip ky y
| k == ky -> case (f y) of
Just y' -> (Just y, Tip ky y')
Nothing -> (Just y, Nil)
| otherwise -> (Nothing, t)
Nil -> (Nothing, Nil)
-- | /O(n)/. Fold the keys and values in the map, such that
-- @'foldWithKey' f z == 'Prelude.foldr' ('uncurry' f) z . 'toAscList'@.
-- For example,
--
-- > keys map = foldWithKey (\k x ks -> k:ks) [] map
--
-- > let f k a result = result ++ "(" ++ (show k) ++ ":" ++ a ++ ")"
-- > foldWithKey f "Map: " (fromList [(5,"a"), (3,"b")]) == "Map: (5:a)(3:b)"
foldWithKey :: (Key -> a -> b -> b) -> b -> IntMap a -> b
foldWithKey f z t
= foldr f z t
-- | /O(n)/. Convert the map to a list of key\/value pairs.
--
-- > toList (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")]
-- > toList empty == []
toList :: IntMap a -> [(Key,a)]
toList t
= foldWithKey (\k x xs -> (k,x):xs) [] t
foldr :: (Key -> a -> b -> b) -> b -> IntMap a -> b
foldr f z t
= case t of
Bin 0 m l r | m < 0 -> foldr' f (foldr' f z l) r -- put negative numbers before.
Bin _ _ _ _ -> foldr' f z t
Tip k x -> f k x z
Nil -> z
foldr' :: (Key -> a -> b -> b) -> b -> IntMap a -> b
foldr' f z t
= case t of
Bin _ _ l r -> foldr' f (foldr' f z r) l
Tip k x -> f k x z
Nil -> z
-- | /O(n)/. Return all keys of the map in ascending order.
--
-- > keys (fromList [(5,"a"), (3,"b")]) == [3,5]
-- > keys empty == []
keys :: IntMap a -> [Key]
keys m
= foldWithKey (\k _ ks -> k:ks) [] m
------------------------------------------------------------------------
-- Eq
instance Eq a => Eq (IntMap a) where
t1 == t2 = equal t1 t2
t1 /= t2 = nequal t1 t2
equal :: Eq a => IntMap a -> IntMap a -> Bool
equal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
= (m1 == m2) && (p1 == p2) && (equal l1 l2) && (equal r1 r2)
equal (Tip kx x) (Tip ky y)
= (kx == ky) && (x==y)
equal Nil Nil = True
equal _ _ = False
nequal :: Eq a => IntMap a -> IntMap a -> Bool
nequal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2)
= (m1 /= m2) || (p1 /= p2) || (nequal l1 l2) || (nequal r1 r2)
nequal (Tip kx x) (Tip ky y)
= (kx /= ky) || (x/=y)
nequal Nil Nil = False
nequal _ _ = True
instance Show a => Show (IntMap a) where
showsPrec d m = showParen (d > 10) $
showString "fromList " . shows (toList m)
------------------------------------------------------------------------
-- Utility functions
join :: Prefix -> IntMap a -> Prefix -> IntMap a -> IntMap a
join p1 t1 p2 t2
| zero p1 m = Bin p m t1 t2
| otherwise = Bin p m t2 t1
where
m = branchMask p1 p2
p = mask p1 m
-- | @bin@ assures that we never have empty trees within a tree.
bin :: Prefix -> Mask -> IntMap a -> IntMap a -> IntMap a
bin _ _ l Nil = l
bin _ _ Nil r = r
bin p m l r = Bin p m l r
------------------------------------------------------------------------
-- Endian independent bit twiddling
zero :: Key -> Mask -> Bool
zero i m = (natFromInt i) .&. (natFromInt m) == 0
nomatch :: Key -> Prefix -> Mask -> Bool
nomatch i p m = (mask i m) /= p
mask :: Key -> Mask -> Prefix
mask i m = maskW (natFromInt i) (natFromInt m)
zeroN :: Nat -> Nat -> Bool
zeroN i m = (i .&. m) == 0
------------------------------------------------------------------------
-- Big endian operations
maskW :: Nat -> Nat -> Prefix
maskW i m = intFromNat (i .&. (complement (m-1) `xor` m))
branchMask :: Prefix -> Prefix -> Mask
branchMask p1 p2
= intFromNat (highestBitMask (natFromInt p1 `xor` natFromInt p2))
{-
Finding the highest bit mask in a word [x] can be done efficiently in
three ways:
* convert to a floating point value and the mantissa tells us the
[log2(x)] that corresponds with the highest bit position. The mantissa
is retrieved either via the standard C function [frexp] or by some bit
twiddling on IEEE compatible numbers (float). Note that one needs to
use at least [double] precision for an accurate mantissa of 32 bit
numbers.
* use bit twiddling, a logarithmic sequence of bitwise or's and shifts (bit).
* use processor specific assembler instruction (asm).
The most portable way would be [bit], but is it efficient enough?
I have measured the cycle counts of the different methods on an AMD
Athlon-XP 1800 (~ Pentium III 1.8Ghz) using the RDTSC instruction:
highestBitMask: method cycles
--------------
frexp 200
float 33
bit 11
asm 12
Wow, the bit twiddling is on today's RISC like machines even faster
than a single CISC instruction (BSR)!
-}
-- | @highestBitMask@ returns a word where only the highest bit is
-- set. It is found by first setting all bits in lower positions than
-- the highest bit and than taking an exclusive or with the original
-- value. Allthough the function may look expensive, GHC compiles
-- this into excellent C code that subsequently compiled into highly
-- efficient machine code. The algorithm is derived from Jorg Arndt's
-- FXT library.
highestBitMask :: Nat -> Nat
highestBitMask x0
= case (x0 .|. shiftRL x0 1) of
x1 -> case (x1 .|. shiftRL x1 2) of
x2 -> case (x2 .|. shiftRL x2 4) of
x3 -> case (x3 .|. shiftRL x3 8) of
x4 -> case (x4 .|. shiftRL x4 16) of
x5 -> case (x5 .|. shiftRL x5 32) of -- for 64 bit platforms
x6 -> (x6 `xor` (shiftRL x6 1))