Improve query functions of IntMap and IntSet.

As documented in the Note: Local 'go' functions and capturing,
it is safe to use captured key argument in query functions.

Also, in order to decrease allocation, the query functions in IntMap
are manually inlined, so 'member' does not have to call 'lookup'
and heap-allocate 'Just a'.

Tests of query functions are much improved too.

Author: foxik

Data/IntMap/Base.hs

@@ -281,7 +281,7 @@ type Key    = Int
 -- > fromList [(5,'a'), (3,'b')] ! 5 == 'a'
 
 (!) :: IntMap a -> Key -> a
-m ! k    = find k m
+m ! k = find k m
 
 -- | Same as 'difference'.
 (\\) :: IntMap a -> IntMap b -> IntMap a
@@ -363,11 +363,15 @@ size t
 -- > member 5 (fromList [(5,'a'), (3,'b')]) == True
 -- > member 1 (fromList [(5,'a'), (3,'b')]) == False
 
+-- See Note: Local 'go' functions and capturing]
 member :: Key -> IntMap a -> Bool
-member k m
-  = case lookup k m of
-      Nothing -> False
-      Just _  -> True
+member k = k `seq` go
+  where
+    go (Bin p m l r) | nomatch k p m = False
+                     | zero k m  = go l
+                     | otherwise = go r
+    go (Tip kx _) = k == kx
+    go Nil = False
 
 -- | /O(min(n,W))/. Is the key not a member of the map?
 --
@@ -377,28 +381,32 @@ member k m
 notMember :: Key -> IntMap a -> Bool
 notMember k m = not $ member k m
 
--- The 'go' function in the lookup causes 10% speedup, but also an increased
--- memory allocation. It does not cause speedup with other methods like insert
--- and delete, so it is present only in lookup.
-
 -- | /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]
 lookup :: Key -> IntMap a -> Maybe a
 lookup k = k `seq` go
   where
-    go (Bin _ m l r)
-      | zero k m  = go l
-      | otherwise = go r
-    go (Tip kx x)
-      | k == kx   = Just x
-      | otherwise = Nothing
-    go Nil      = Nothing
+    go (Bin p m l r) | nomatch k p m = Nothing
+                     | 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 m
-  = case lookup k m of
-      Nothing -> error ("IntMap.find: key " ++ show k ++ " is not an element of the map")
-      Just x  -> x
+find k = k `seq` go
+  where
+    go (Bin p m l r) | nomatch k p m = not_found
+                     | zero k m  = go l
+                     | otherwise = go r
+    go (Tip kx x) | k == kx   = x
+                  | otherwise = not_found
+    go Nil = not_found
+
+    not_found = error ("IntMap.!: key" ++ show k ++ " is not an element of the map")
 
 -- | /O(min(n,W))/. The expression @('findWithDefault' def k map)@
 -- returns the value at key @k@ or returns @def@ when the key is not an
@@ -407,11 +415,16 @@ find k m
 -- > findWithDefault 'x' 1 (fromList [(5,'a'), (3,'b')]) == 'x'
 -- > findWithDefault 'x' 5 (fromList [(5,'a'), (3,'b')]) == 'a'
 
+-- See Note: Local 'go' functions and capturing]
 findWithDefault :: a -> Key -> IntMap a -> a
-findWithDefault def k m
-  = case lookup k m of
-      Nothing -> def
-      Just x  -> x
+findWithDefault def k = k `seq` go
+  where
+    go (Bin p m l r) | nomatch k p m = def
+                     | zero k m  = go l
+                     | otherwise = go r
+    go (Tip kx x) | k == kx   = x
+                  | otherwise = def
+    go Nil = def
 
 {--------------------------------------------------------------------
   Construction

Data/IntMap/Strict.hs

@@ -277,11 +277,16 @@ import Data.StrictPair
 -- > findWithDefault 'x' 1 (fromList [(5,'a'), (3,'b')]) == 'x'
 -- > findWithDefault 'x' 5 (fromList [(5,'a'), (3,'b')]) == 'a'
 
+-- See IntMap.Base.Note: Local 'go' functions and capturing]
 findWithDefault :: a -> Key -> IntMap a -> a
-findWithDefault def k m
-  = def `seq` case lookup k m of
-      Nothing -> def
-      Just x  -> x
+findWithDefault def k = def `seq` k `seq` go
+  where
+    go (Bin p m l r) | nomatch k p m = def
+                     | zero k m  = go l
+                     | otherwise = go r
+    go (Tip kx x) | k == kx   = x
+                  | otherwise = def
+    go Nil = def
 
 {--------------------------------------------------------------------
   Construction

Data/IntSet.hs

@@ -304,6 +304,7 @@ size t
 
 -- | /O(min(n,W))/. Is the value a member of the set?
 
+-- See Note: Local 'go' functions and capturing]
 member :: Int -> IntSet -> Bool
 member x = x `seq` go
   where

tests/intmap-properties.hs

@@ -115,7 +115,7 @@ main = defaultMainWithOpts
              , testCase "minViewWithKey" test_minViewWithKey
              , testCase "maxViewWithKey" test_maxViewWithKey
              , testProperty "insert to singleton"  prop_singleton
-             , testProperty "insert then lookup"   prop_lookup
+             , testProperty "insert then lookup"   prop_insertLookup
              , testProperty "insert then delete"   prop_insertDelete
              , testProperty "delete non member"    prop_deleteNonMember
              , testProperty "union model"          prop_unionModel
@@ -137,6 +137,8 @@ main = defaultMainWithOpts
              , testProperty "null"                 prop_null
              , testProperty "member"               prop_member
              , testProperty "notmember"            prop_notmember
+             , testProperty "lookup"               prop_lookup
+             , testProperty "find"                 prop_find
              , testProperty "findWithDefault"      prop_findWithDefault
              , testProperty "findMin"              prop_findMin
              , testProperty "findMax"              prop_findMax
@@ -704,8 +706,8 @@ test_maxViewWithKey = do
 prop_singleton :: Int -> Int -> Bool
 prop_singleton k x = insert k x empty == singleton k x
 
-prop_lookup :: Int -> UMap -> Bool
-prop_lookup k t = lookup k (insert k () t) /= Nothing
+prop_insertLookup :: Int -> UMap -> Bool
+prop_insertLookup k t = lookup k (insert k () t) /= Nothing
 
 prop_insertDelete :: Int -> UMap -> Property
 prop_insertDelete k t = (lookup k t == Nothing) ==> (delete k (insert k () t) == t)
@@ -833,18 +835,30 @@ prop_null m = null m == (size m == 0)
 prop_member :: [Int] -> Int -> Bool
 prop_member xs n =
   let m  = fromList (zip xs xs)
-  in  (n `elem` xs) == (n `member` m)
+  in all (\k -> k `member` m == (k `elem` xs)) (n : xs)
 
 prop_notmember :: [Int] -> Int -> Bool
 prop_notmember xs n =
   let m  = fromList (zip xs xs)
-  in  (n `notElem` xs) == (n `notMember` m)
-
-prop_findWithDefault :: [(Int, Int)] -> Property
-prop_findWithDefault ys = length ys > 0 ==>
-  let xs = List.nubBy ((==) `on` fst) ys
-      m  = fromList xs
-  in  and [ findWithDefault 0 i m == j | (i,j) <- xs ]
+  in all (\k -> k `notMember` m == (k `notElem` xs)) (n : xs)
+
+prop_lookup :: [(Int, Int)] -> Int -> Bool
+prop_lookup xs n =
+  let xs' = List.nubBy ((==) `on` fst) xs
+      m = fromList xs'
+  in all (\k -> lookup k m == List.lookup k xs') (n : List.map fst xs')
+
+prop_find :: [(Int, Int)] -> Bool
+prop_find xs =
+  let xs' = List.nubBy ((==) `on` fst) xs
+      m = fromList xs'
+  in all (\(k, v) -> m ! k == v) xs'
+
+prop_findWithDefault :: [(Int, Int)] -> Int -> Int -> Bool
+prop_findWithDefault xs n x =
+  let xs' = List.nubBy ((==) `on` fst) xs
+      m = fromList xs'
+  in all (\k -> findWithDefault x k m == maybe x id (List.lookup k xs')) (n : List.map fst xs')
 
 prop_findMin :: [(Int, Int)] -> Property
 prop_findMin ys = length ys > 0 ==>

tests/intset-properties.hs

@@ -11,6 +11,8 @@ import Test.Framework.Providers.QuickCheck2
 
 main :: IO ()
 main = defaultMainWithOpts [ testProperty "prop_Single" prop_Single
+                           , testProperty "prop_Member" prop_Member
+                           , testProperty "prop_NotMember" prop_NotMember
                            , testProperty "prop_InsertDelete" prop_InsertDelete
                            , testProperty "prop_MemberFromList" prop_MemberFromList
                            , testProperty "prop_UnionInsert" prop_UnionInsert
@@ -63,12 +65,22 @@ instance Arbitrary IntSet where
 
 
 {--------------------------------------------------------------------
-  Single, Insert, Delete, Member, FromList
+  Single, Member, Insert, Delete, Member, FromList
 --------------------------------------------------------------------}
 prop_Single :: Int -> Bool
 prop_Single x
   = (insert x empty == singleton x)
 
+prop_Member :: [Int] -> Int -> Bool
+prop_Member xs n =
+  let m  = fromList xs
+  in all (\k -> k `member` m == (k `elem` xs)) (n : xs)
+
+prop_NotMember :: [Int] -> Int -> Bool
+prop_NotMember xs n =
+  let m  = fromList xs
+  in all (\k -> k `notMember` m == (k `notElem` xs)) (n : xs)
+
 prop_InsertDelete :: Int -> IntSet -> Property
 prop_InsertDelete k t
   = not (member k t) ==> delete k (insert k t) == t