Change IntMap.lookup and add new IntMap.query function

containers-tests/benchmarks/IntMap.hs

@@ -11,10 +11,18 @@ import Data.Maybe (fromMaybe)
 import Prelude hiding (lookup)
 
 main = do
-    let m = M.fromAscList elems :: M.IntMap Int
-    evaluate $ rnf [m]
+    let m     = M.fromAscList elems_hits   :: M.IntMap Int
+    let m'    = M.fromAscList elems_mid    :: M.IntMap Int
+    let m''   = M.fromAscList elems_most   :: M.IntMap Int
+    let m'''  = M.fromAscList elems_misses :: M.IntMap Int
+    let m'''' = M.fromAscList elems_mixed  :: M.IntMap Int
+    evaluate $ rnf [m, m', m'', m''', m'''']
     defaultMain
-        [ bench "lookup" $ whnf (lookup keys) m
+        [ bench "lookup_hits" $ whnf (lookup keys) m
+        , bench "lookup_half" $ whnf (lookup keys) m'
+        , bench "lookup_most" $ whnf (lookup keys) m''
+        , bench "lookup_misses" $ whnf (lookup keys'') m'''
+        , bench "lookup_mixed" $ whnf (lookup keys) m''''
         , bench "insert" $ whnf (ins elems) M.empty
         , bench "insertWith empty" $ whnf (insWith elems) M.empty
         , bench "insertWith update" $ whnf (insWith elems) m
@@ -44,12 +52,23 @@ main = do
                     (M.fromList $ zip [1..10] [1..10])
         ]
   where
-    elems = zip keys values
+    elems = elems_hits
+    elems_hits   = zip keys values
+    elems_mid    = zip (map (+ (2^12 `div` 2)) keys) values
+    elems_most   = zip (map (+ (2^12 `div` 10)) keys) values
+    elems_misses = zip (map (\x-> x * 2 + 1) keys) values
+    elems_mixed = zip mixedKeys values
+    --------------------------------------------------------
     keys = [1..2^12]
+    keys' = fmap (+ 1000000) keys
+    keys'' = fmap (* 2) [1..2^12]
+    mixedKeys = interleave keys keys'
     values = [1..2^12]
+    --------------------------------------------------------
     sum k v1 v2 = k + v1 + v2
     consPair k v xs = (k, v) : xs
 
+------------------------------------------------------------
 add3 :: Int -> Int -> Int -> Int
 add3 x y z = x + y + z
 {-# INLINE add3 #-}
@@ -95,3 +114,8 @@ alt xs m = foldl' (\m k -> M.alter id k m) m xs
 maybeDel :: Int -> Maybe Int
 maybeDel n | n `mod` 3 == 0 = Nothing
            | otherwise      = Just n
+
+------------------------------------------------------------
+interleave :: [Int] -> [Int] -> [Int]
+interleave [] ys = ys
+interleave (x:xs) (y:ys) = x : y : interleave xs ys

containers-tests/tests/intmap-properties.hs

@@ -295,12 +295,17 @@ test_notMember = do
 
 test_lookup :: Assertion
 test_lookup = do
-    employeeCurrency 1 @?= Just 1
-    employeeCurrency 2 @?= Nothing
+    employeeCurrency 1      @?= Just 1
+    employeeCurrency 2      @?= Just 2
+    employeeCurrency 3      @?= Just 3
+    employeeCurrency 4      @?= Just 4
+    employeeCurrency 5      @?= Nothing
+    employeeCurrency (2^10) @?= Just 42
+    employeeCurrency 6      @?= Nothing
   where
-    employeeDept = fromList([(1,2), (3,1)])
-    deptCountry = fromList([(1,1), (2,2)])
-    countryCurrency = fromList([(1, 2), (2, 1)])
+    employeeDept    = fromList [(1,2), (2, 14), (3, 10), (4, 18), (2^10, 100)]
+    deptCountry     = fromList [(1,1), (14, 14), (10, 10), (18, 18), (100, 100), (2,2)]
+    countryCurrency = fromList [(1, 2), (2, 1), (14, 2), (10, 3), (18, 4), (100, 42)]
     employeeCurrency :: Int -> Maybe Int
     employeeCurrency name = do
         dept <- lookup name employeeDept

containers/src/Data/IntMap/Internal.hs

@@ -594,25 +594,22 @@ notMember k m = not $ member k m
 
 -- | /O(min(n,W))/. Lookup the value at a key in the map. See also 'Data.Map.lookup'.
 
--- See Note: Local 'go' functions and capturing]
+-- See Note: Local 'go' functions and capturing
 lookup :: Key -> IntMap a -> Maybe a
 lookup !k = go
   where
-    go (Bin p m l r) | nomatch k p m = Nothing
-                     | zero k m  = go l
-                     | otherwise = go r
+    go (Bin _p m l r) | zero k m  = go l
+                      | otherwise = go r
     go (Tip kx x) | k == kx   = Just x
                   | otherwise = Nothing
     go Nil = Nothing
 
-
 -- See Note: Local 'go' functions and capturing]
 find :: Key -> IntMap a -> a
 find !k = go
   where
-    go (Bin p m l r) | nomatch k p m = not_found
-                     | zero k m  = go l
-                     | otherwise = go r
+    go (Bin _p m l r) | zero k m  = go l
+                      | otherwise = go r
     go (Tip kx x) | k == kx   = x
                   | otherwise = not_found
     go Nil = not_found
@@ -943,8 +940,7 @@ adjust f k m
 -- > adjustWithKey f 7 empty                         == empty
 
 adjustWithKey ::  (Key -> a -> a) -> Key -> IntMap a -> IntMap a
-adjustWithKey f !k t@(Bin p m l r)
-  | nomatch k p m = t
+adjustWithKey f !k (Bin p m l r)
   | zero k m      = Bin p m (adjustWithKey f k l) r
   | otherwise     = Bin p m l (adjustWithKey f k r)
 adjustWithKey f k t@(Tip ky y)
@@ -976,8 +972,7 @@ update f
 -- > updateWithKey f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a"
 
 updateWithKey ::  (Key -> a -> Maybe a) -> Key -> IntMap a -> IntMap a
-updateWithKey f !k t@(Bin p m l r)
-  | nomatch k p m = t
+updateWithKey f !k (Bin p m l r)
   | zero k m      = binCheckLeft p m (updateWithKey f k l) r
   | otherwise     = binCheckRight p m l (updateWithKey f k r)
 updateWithKey f k t@(Tip ky y)
@@ -998,8 +993,7 @@ updateWithKey _ _ Nil = Nil
 -- > updateLookupWithKey f 3 (fromList [(5,"a"), (3,"b")]) == (Just "b", singleton 5 "a")
 
 updateLookupWithKey ::  (Key -> a -> Maybe a) -> Key -> IntMap a -> (Maybe a,IntMap a)
-updateLookupWithKey f !k t@(Bin p m l r)
-  | nomatch k p m = (Nothing,t)
+updateLookupWithKey f !k (Bin p m l r)
   | zero k m      = let !(found,l') = updateLookupWithKey f k l
                     in (found,binCheckLeft p m l' r)
   | otherwise     = let !(found,r') = updateLookupWithKey f k r