[commit: containers] master: Improve {Map, IntMap}.keysSet. (ab8460f)

[Paolo Capriotti]

Repository : ssh://darcs.haskell.org//srv/darcs/packages/containers

On branch  : master

http://hackage.haskell.org/trac/ghc/changeset/ab8460f8017c1defb649df8413407153d6750d1c

>---------------------------------------------------------------

commit ab8460f8017c1defb649df8413407153d6750d1c
Author: Milan Straka <fox at ucw.cz>
Date:   Tue Apr 24 17:32:04 2012 +0200

    Improve {Map, IntMap}.keysSet.
    
    The keysSet method is now implemented using the exported constructors of
    Data.{Set, IntSet}.Base.
    
    The implementation of Map.keysSet is trivial, as Set and Map use same
    tree structure.
    
    The implementation of IntMap.keysSet is slightly complicated, because of
    the dense representation of IntSet, where several last levels of the
    tree are flatten into a bitmap.

>---------------------------------------------------------------

 Data/IntMap/Base.hs        |   13 ++++++++++---
 Data/IntSet/Base.hs        |    3 +++
 Data/Map/Base.hs           |    5 +++--
 tests/intmap-properties.hs |    5 +++++
 tests/map-properties.hs    |    5 +++++
 5 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/Data/IntMap/Base.hs b/Data/IntMap/Base.hs
index 0f8a11e..a2f35ed 100644
--- a/Data/IntMap/Base.hs
+++ b/Data/IntMap/Base.hs
@@ -210,7 +210,7 @@ module Data.IntMap.Base (
 import Data.Bits
 
 import Prelude hiding (lookup,map,filter,foldr,foldl,null)
-import qualified Data.IntSet as IntSet
+import qualified Data.IntSet.Base as IntSet
 import Data.Monoid (Monoid(..))
 import Data.Maybe (fromMaybe)
 import Data.Typeable
@@ -1683,8 +1683,15 @@ keys = foldrWithKey (\k _ ks -> k : ks) []
 -- > keysSet empty == Data.IntSet.empty
 
 keysSet :: IntMap a -> IntSet.IntSet
-keysSet m = IntSet.fromDistinctAscList (keys m)
-
+keysSet Nil = IntSet.Nil
+keysSet (Tip kx _) = IntSet.singleton kx
+keysSet (Bin p m l r)
+  | m .&. IntSet.suffixBitMask == 0 = IntSet.Bin p m (keysSet l) (keysSet r)
+  | otherwise = IntSet.Tip (p .&. IntSet.prefixBitMask) (computeBm (computeBm 0 l) r)
+  where STRICT_1_OF_2(computeBm)
+        computeBm acc (Bin _ _ l r) = computeBm (computeBm acc l) r
+        computeBm acc (Tip kx _) = acc .|. IntSet.bitmapOf kx
+        computeBm _   Nil = error "Data.IntSet.keysSet: Nil"
 
 -- | /O(n)/. An alias for 'toAscList'. Returns all key\/value pairs in the
 -- map in ascending key order. Subject to list fusion.
diff --git a/Data/IntSet/Base.hs b/Data/IntSet/Base.hs
index ed77b7c..056aa52 100644
--- a/Data/IntSet/Base.hs
+++ b/Data/IntSet/Base.hs
@@ -153,6 +153,9 @@ module Data.IntSet.Base (
 
             -- * Internals
             , match
+            , suffixBitMask
+            , prefixBitMask
+            , bitmapOf
             ) where
 
 
diff --git a/Data/Map/Base.hs b/Data/Map/Base.hs
index e5c33e0..2e72687 100644
--- a/Data/Map/Base.hs
+++ b/Data/Map/Base.hs
@@ -254,7 +254,7 @@ module Data.Map.Base (
             ) where
 
 import Prelude hiding (lookup,map,filter,foldr,foldl,null)
-import qualified Data.Set as Set
+import qualified Data.Set.Base as Set
 import Data.Monoid (Monoid(..))
 import Control.Applicative (Applicative(..), (<$>))
 import Data.Traversable (Traversable(traverse))
@@ -1879,7 +1879,8 @@ keys = foldrWithKey (\k _ ks -> k : ks) []
 -- > keysSet empty == Data.Set.empty
 
 keysSet :: Map k a -> Set.Set k
-keysSet m = Set.fromDistinctAscList (keys m)
+keysSet Tip = Set.Tip
+keysSet (Bin sz kx _ l r) = Set.Bin sz kx (keysSet l) (keysSet r)
 
 -- | /O(n)/. An alias for 'toAscList'. Return all key\/value pairs in the map
 -- in ascending key order. Subject to list fusion.
diff --git a/tests/intmap-properties.hs b/tests/intmap-properties.hs
index bbc068b..d26d1fa 100644
--- a/tests/intmap-properties.hs
+++ b/tests/intmap-properties.hs
@@ -162,6 +162,7 @@ main = defaultMainWithOpts
              , testProperty "foldr'"               prop_foldr'
              , testProperty "foldl"                prop_foldl
              , testProperty "foldl'"               prop_foldl'
+             , testProperty "keysSet"              prop_keysSet
              ] opts
 
   where
@@ -1022,3 +1023,7 @@ prop_foldl' n ys = length ys > 0 ==>
       foldlWithKey' (\b _ a -> a + b) n m == List.foldr (+) n (List.map snd xs) &&
       foldlWithKey' (\b k _ -> k + b) n m == List.foldr (+) n (List.map fst xs) &&
       foldlWithKey' (\xs k x -> (k,x):xs) [] m == reverse (List.sort xs)
+
+prop_keysSet :: [(Int, Int)] -> Bool
+prop_keysSet xs =
+  keysSet (fromList xs) == Data.IntSet.fromList (List.map fst xs)
diff --git a/tests/map-properties.hs b/tests/map-properties.hs
index 4b4efa6..068ee0a 100644
--- a/tests/map-properties.hs
+++ b/tests/map-properties.hs
@@ -183,6 +183,7 @@ main = defaultMainWithOpts
          , testProperty "foldr'"               prop_foldr'
          , testProperty "foldl"                prop_foldl
          , testProperty "foldl'"               prop_foldl'
+         , testProperty "keysSet"              prop_keysSet
          ] opts
 
   where
@@ -1132,3 +1133,7 @@ prop_foldl' n ys = length ys > 0 ==>
       foldlWithKey' (\b _ a -> a + b) n m == List.foldr (+) n (List.map snd xs) &&
       foldlWithKey' (\b k _ -> k + b) n m == List.foldr (+) n (List.map fst xs) &&
       foldlWithKey' (\xs k x -> (k,x):xs) [] m == reverse (List.sort xs)
+
+prop_keysSet :: [(Int, Int)] -> Bool
+prop_keysSet xs =
+  keysSet (fromList xs) == Data.Set.fromList (List.map fst xs)