Change IntMap.lookup and add new IntMap.query function

  - `IntMap.lookup` no longer checks for short circuit failure.
  - Add a new function `IntMap.query` with the old fast-fail behaviour.

Author: Boarders

containers-tests/benchmarks/IntMap.hs

@@ -11,10 +11,23 @@ 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 "query_hits" $ whnf (query keys) m
+        , bench "query_half" $ whnf (query keys) m'
+        , bench "query_most" $ whnf (query keys) m''
+        , bench "query_misses" $ whnf (query keys'') m'''
+        , bench "query_mixed" $ whnf (query 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,19 +57,33 @@ 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 #-}
 
 lookup :: [Int] -> M.IntMap Int -> Int
 lookup xs m = foldl' (\n k -> fromMaybe n (M.lookup k m)) 0 xs
 
+query :: [Int] -> M.IntMap Int -> Int
+query xs m = foldl' (\n k -> fromMaybe n (M.query k m)) 0 xs
+
 ins :: [(Int, Int)] -> M.IntMap Int -> M.IntMap Int
 ins xs m = foldl' (\m (k, v) -> M.insert k v m) m xs
 
@@ -95,3 +122,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

@@ -45,6 +45,7 @@ main = defaultMain $ testGroup "intmap-properties"
              , testCase "member"     test_member
              , testCase "notMember"  test_notMember
              , testCase "lookup"     test_lookup
+             , testCase "query "     test_query
              , testCase "findWithDefault"     test_findWithDefault
              , testCase "lookupLT"   test_lookupLT
              , testCase "lookupGT"   test_lookupGT
@@ -295,18 +296,42 @@ 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
         country <- lookup dept deptCountry
         lookup country countryCurrency
 
+test_query :: Assertion
+test_query = do
+    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), (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 <- query name employeeDept
+        country <- query dept deptCountry
+        query country countryCurrency
+
 test_findWithDefault :: Assertion
 test_findWithDefault = do
     findWithDefault 'x' 1 (fromList [(5,'a'), (3,'b')]) @?= 'x'

containers/src/Data/IntMap/Internal.hs

@@ -86,6 +86,7 @@ module Data.IntMap.Internal (
     , member
     , notMember
     , lookup
+    , query
     , findWithDefault
     , lookupLT
     , lookupGT
@@ -594,9 +595,23 @@ 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) | zero k m  = go l
+                      | otherwise = go r
+    go (Tip kx x) | k == kx   = Just x
+                  | otherwise = Nothing
+    go Nil = Nothing
+
+-- | /O(min(n,W))/. 'query' has identical behaviour to 'Data.IntMap.Internal.lookup', but
+-- will fail faster if the key is not present. When the key is (likely) present, it is
+-- preferable to use 'Data.IntMap.Internal.lookup' instead.
+
+-- See Note: Local 'go' functions and capturing
+query :: Key -> IntMap a -> Maybe a
+query !k = go
   where
     go (Bin p m l r) | nomatch k p m = Nothing
                      | zero k m  = go l

containers/src/Data/IntMap/Lazy.hs

@@ -108,6 +108,7 @@ module Data.IntMap.Lazy (
     -- * Query
     -- ** Lookup
     , IM.lookup
+    , IM.query
     , (!?)
     , (!)
     , findWithDefault

containers/src/Data/IntMap/Strict.hs

@@ -127,6 +127,7 @@ module Data.IntMap.Strict (
     -- * Query
     -- ** Lookup
     , lookup
+    , query
     , (!?)
     , (!)
     , findWithDefault

containers/src/Data/IntMap/Strict/Internal.hs

@@ -127,6 +127,7 @@ module Data.IntMap.Strict.Internal (
     -- * Query
     -- ** Lookup
     , lookup
+    , query
     , (!?)
     , (!)
     , findWithDefault
@@ -327,6 +328,7 @@ import Data.IntMap.Internal
   , null
   , partition
   , partitionWithKey
+  , query
   , restrictKeys
   , size
   , split