Improve query functions of Map and Set.
authorMilan Straka <fox@ucw.cz>
Sat, 21 Apr 2012 19:41:22 +0000 (21:41 +0200)
committerMilan Straka <fox@ucw.cz>
Sat, 21 Apr 2012 19:41:22 +0000 (21:41 +0200)
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 Map
are manually inlined, so 'member' does not have to call 'lookup'
and heap-allocate 'Just a'.

Tests of query functions are much improved too.

Data/Map/Base.hs
Data/Map/Strict.hs
Data/Set.hs
tests/map-properties.hs
tests/set-properties.hs

index 41072f9..8d18cd8 100644 (file)
 -- entry of the outer method.
 
 
+-- [Note: Local 'go' functions and capturing]
+-- ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+-- Care must be taken when using 'go' function which captures an argument.
+-- Sometimes (for example when the argument is passed to a data constructor,
+-- as in insert), GHC heap-allocates more than necessary. Therefore C-- code
+-- must be checked for increased allocation when creating and modifying such
+-- functions.
+
+
 -- [Note: Order of constructors]
 -- ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 -- The order of constructors of Map matters when considering performance.
@@ -271,7 +280,7 @@ infixl 9 !,\\ --
 -- > fromList [(5,'a'), (3,'b')] ! 5 == 'a'
 
 (!) :: Ord k => Map k a -> k -> a
-m ! k    = find k m
+m ! k = find k m
 #if __GLASGOW_HASKELL__ >= 700
 {-# INLINABLE (!) #-}
 #endif
@@ -371,32 +380,30 @@ size (Bin sz _ _ _ _) = sz
 -- >   John's currency: Just "Euro"
 -- >   Pete's currency: Nothing
 
+-- See Note: Local 'go' functions and capturing
 lookup :: Ord k => k -> Map k a -> Maybe a
-lookup = go
+lookup k = k `seq` go
   where
-    STRICT_1_OF_2(go)
-    go _ Tip = Nothing
-    go k (Bin _ kx x l r) =
-        case compare k kx of
-            LT -> go k l
-            GT -> go k r
-            EQ -> Just x
+    go Tip = Nothing
+    go (Bin _ kx x l r) = case compare k kx of
+      LT -> go l
+      GT -> go r
+      EQ -> Just x
 #if __GLASGOW_HASKELL__ >= 700
 {-# INLINABLE lookup #-}
 #else
 {-# INLINE lookup #-}
 #endif
 
+-- See Note: Local 'go' functions and capturing
 lookupAssoc :: Ord k => k -> Map k a -> Maybe (k,a)
-lookupAssoc = go
+lookupAssoc k = k `seq` go
   where
-    STRICT_1_OF_2(go)
-    go _ Tip = Nothing
-    go k (Bin _ kx x l r) =
-        case compare k kx of
-            LT -> go k l
-            GT -> go k r
-            EQ -> Just (kx,x)
+    go Tip = Nothing
+    go (Bin _ kx x l r) = case compare k kx of
+      LT -> go l
+      GT -> go r
+      EQ -> Just (kx,x)
 #if __GLASGOW_HASKELL__ >= 700
 {-# INLINABLE lookupAssoc #-}
 #else
@@ -408,10 +415,15 @@ lookupAssoc = go
 -- > member 5 (fromList [(5,'a'), (3,'b')]) == True
 -- > member 1 (fromList [(5,'a'), (3,'b')]) == False
 
+-- See Note: Local 'go' functions and capturing
 member :: Ord k => k -> Map k a -> Bool
-member k m = case lookup k m of
-    Nothing -> False
-    Just _  -> True
+member k = k `seq` go
+  where
+    go Tip = False
+    go (Bin _ kx x l r) = case compare k kx of
+      LT -> go l
+      GT -> go r
+      EQ -> True
 #if __GLASGOW_HASKELL__ >= 700
 {-# INLINABLE member #-}
 #else
@@ -433,11 +445,16 @@ notMember k m = not $ member k m
 
 -- | /O(log n)/. Find the value at a key.
 -- Calls 'error' when the element can not be found.
--- Consider using 'lookup' when elements may not be present.
+
+-- See Note: Local 'go' functions and capturing
 find :: Ord k => k -> Map k a -> a
-find k m = case lookup k m of
-    Nothing -> error "Map.find: element not in the map"
-    Just x  -> x
+find k = k `seq` go
+  where
+    go Tip = error "Map.!: given key is not an element in the map"
+    go (Bin _ kx x l r) = case compare k kx of
+      LT -> go l
+      GT -> go r
+      EQ -> x
 #if __GLASGOW_HASKELL__ >= 700
 {-# INLINABLE find #-}
 #else
@@ -451,10 +468,15 @@ find k m = case lookup k m of
 -- > 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 :: Ord k => a -> k -> Map k a -> a
-findWithDefault def k m = case lookup k m of
-    Nothing -> def
-    Just x  -> x
+findWithDefault def k = k `seq` go
+  where
+    go Tip = def
+    go (Bin _ kx x l r) = case compare k kx of
+      LT -> go l
+      GT -> go r
+      EQ -> x
 #if __GLASGOW_HASKELL__ >= 700
 {-# INLINABLE findWithDefault #-}
 #else
index 4d722ff..e67b04f 100644 (file)
@@ -297,10 +297,15 @@ import Data.StrictPair
 -- > 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 :: Ord k => a -> k -> Map k 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 Tip = def
+    go (Bin _ kx x l r) = case compare k kx of
+      LT -> go l
+      GT -> go r
+      EQ -> x
 #if __GLASGOW_HASKELL__ >= 700
 {-# INLINABLE findWithDefault #-}
 #else
index 04e6278..73bca08 100644 (file)
 -- entry of the outer method.
 
 
+-- [Note: Local 'go' functions and capturing]
+-- ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+-- Care must be taken when using 'go' function which captures an argument.
+-- Sometimes (for example when the argument is passed to a data constructor,
+-- as in insert), GHC heap-allocates more than necessary. Therefore C-- code
+-- must be checked for increased allocation when creating and modifying such
+-- functions.
+
+
 -- [Note: Order of constructors]
 -- ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 -- The order of constructors of Set matters when considering performance.
@@ -261,15 +270,16 @@ size (Bin sz _ _ _) = sz
 {-# INLINE size #-}
 
 -- | /O(log n)/. Is the element in the set?
+
+-- See Note: Local 'go' functions and capturing
 member :: Ord a => a -> Set a -> Bool
-member = go
+member x = x `seq` go
   where
-    STRICT_1_OF_2(go)
-    go _ Tip = False
-    go x (Bin _ y l r) = case compare x y of
-          LT -> go x l
-          GT -> go x r
-          EQ -> True
+    go Tip = False
+    go (Bin _ y l r) = case compare x y of
+      LT -> go l
+      GT -> go r
+      EQ -> True
 #if __GLASGOW_HASKELL__ >= 700
 {-# INLINABLE member #-}
 #else
index 8b9a058..1275462 100644 (file)
@@ -123,7 +123,7 @@ main = defaultMainWithOpts
          , testProperty "fromList"             prop_fromList
          , testProperty "insert to singleton"  prop_singleton
          , testProperty "insert"               prop_insert
-         , testProperty "insert then lookup"   prop_lookup
+         , testProperty "insert then lookup"   prop_insertLookup
          , testProperty "insert then delete"   prop_insertDelete
          , testProperty "insert then delete2"  prop_insertDelete2
          , testProperty "delete non member"    prop_deleteNonMember
@@ -156,6 +156,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 "findIndex"            prop_findIndex
          , testProperty "lookupIndex"          prop_lookupIndex
@@ -798,8 +800,8 @@ prop_singleton k x = insert k x empty == singleton k x
 prop_insert :: Int -> UMap -> Bool
 prop_insert k t = valid $ insert k () t
 
-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 -> Bool
 prop_insertDelete k t = valid $ delete k (insert k () t)
@@ -938,18 +940,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_findIndex :: [(Int, Int)] -> Property
 prop_findIndex ys = length ys > 0 ==>
@@ -1013,7 +1027,7 @@ prop_mapkeys :: (Int -> Int) -> [(Int, Int)] -> Property
 prop_mapkeys f ys = length ys > 0 ==>
   let xs = List.nubBy ((==) `on` fst) ys
       m  = fromList xs
-  in  Data.Map.mapKeys f m == (fromList $ List.nubBy ((==) `on` fst) $ reverse [ (f a, b) | (a,b) <- sort xs])
+  in  mapKeys f m == (fromList $ List.nubBy ((==) `on` fst) $ reverse [ (f a, b) | (a,b) <- sort xs])
 
 prop_splitModel :: Int -> [(Int, Int)] -> Property
 prop_splitModel n ys = length ys > 0 ==>
index 19a678a..87ffa2b 100644 (file)
@@ -11,6 +11,8 @@ import Test.Framework.Providers.QuickCheck2
 main :: IO ()
 main = defaultMainWithOpts [ testProperty "prop_Valid" prop_Valid
                            , testProperty "prop_Single" prop_Single
+                           , testProperty "prop_Member" prop_Member
+                           , testProperty "prop_NotMember" prop_NotMember
                            , testProperty "prop_InsertValid" prop_InsertValid
                            , testProperty "prop_InsertDelete" prop_InsertDelete
                            , testProperty "prop_DeleteValid" prop_DeleteValid
@@ -98,11 +100,21 @@ prop_Valid :: Property
 prop_Valid = forValidUnitTree $ \t -> valid t
 
 {--------------------------------------------------------------------
-  Single, Insert, Delete
+  Single, Member, Insert, Delete
 --------------------------------------------------------------------}
 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_InsertValid :: Int -> Property
 prop_InsertValid k = forValidUnitTree $ \t -> valid (insert k t)