Add fromSet method.
authorMilan Straka <fox@ucw.cz>
Wed, 25 Apr 2012 16:54:04 +0000 (18:54 +0200)
committerMilan Straka <fox@ucw.cz>
Wed, 25 Apr 2012 16:54:04 +0000 (18:54 +0200)
Following a proposal on libraries@..., we add fromSet method to Map and
IntMap:
  Map.fromSet :: (k -> a) -> Set k -> Map k a
  IntMap.fromSet :: (Key -> a) -> IntSet -> IntMap a
It is implemented using exported Set and IntSet constructors.

Map.fromSet is trivial, as Map and Set have the same structure.

The IntMap.fromSet implementation is more complicated, as IntSet uses
dense representation of several last leves of the trie.

Data/IntMap/Base.hs
Data/IntMap/Lazy.hs
Data/IntMap/Strict.hs
Data/Map/Base.hs
Data/Map/Lazy.hs
Data/Map/Strict.hs
tests/intmap-properties.hs
tests/map-properties.hs

index fb0560e..1e90f6b 100644 (file)
@@ -132,8 +132,9 @@ module Data.IntMap.Base (
             -- * Conversion
             , elems
             , keys
-            , keysSet
             , assocs
+            , keysSet
+            , fromSet
 
             -- ** Lists
             , toList
@@ -194,6 +195,7 @@ module Data.IntMap.Base (
             , natFromInt
             , intFromNat
             , shiftRL
+            , shiftLL
             , join
             , bin
             , zero
@@ -226,7 +228,8 @@ import Data.Data (Data(..), mkNoRepType)
 #endif
 
 #if __GLASGOW_HASKELL__
-import GHC.Exts ( Word(..), Int(..), shiftRL#, build )
+import GHC.Exts ( Word(..), Int(..), build )
+import GHC.Prim ( uncheckedShiftL#, uncheckedShiftRL# )
 #else
 import Data.Word
 #endif
@@ -253,17 +256,20 @@ intFromNat :: Nat -> Key
 intFromNat = fromIntegral
 {-# INLINE intFromNat #-}
 
-shiftRL :: Nat -> Key -> Nat
+-- Right and left logical shifts.
+shiftRL, shiftLL :: Nat -> Key -> Nat
 #if __GLASGOW_HASKELL__
 {--------------------------------------------------------------------
   GHC: use unboxing to get @shiftRL@ inlined.
 --------------------------------------------------------------------}
-shiftRL (W# x) (I# i)
-  = W# (shiftRL# x i)
+shiftRL (W# x) (I# i) = W# (uncheckedShiftRL# x i)
+shiftLL (W# x) (I# i) = W# (uncheckedShiftL#  x i)
 #else
 shiftRL x i   = shiftR x i
-{-# INLINE shiftRL #-}
+shiftLL x i   = shiftL x i
 #endif
+{-# INLINE shiftRL #-}
+{-# INLINE shiftLL #-}
 
 {--------------------------------------------------------------------
   Types
@@ -1677,6 +1683,15 @@ elems = foldr (:) []
 keys  :: IntMap a -> [Key]
 keys = foldrWithKey (\k _ ks -> k : ks) []
 
+-- | /O(n)/. An alias for 'toAscList'. Returns all key\/value pairs in the
+-- map in ascending key order. Subject to list fusion.
+--
+-- > assocs (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")]
+-- > assocs empty == []
+
+assocs :: IntMap a -> [(Key,a)]
+assocs = toAscList
+
 -- | /O(n*min(n,W))/. The set of all keys of the map.
 --
 -- > keysSet (fromList [(5,"a"), (3,"b")]) == Data.IntSet.fromList [3,5]
@@ -1693,15 +1708,32 @@ keysSet (Bin p m 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.
---
--- > assocs (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")]
--- > assocs empty == []
-
-assocs :: IntMap a -> [(Key,a)]
-assocs = toAscList
-
+-- | /O(n)/. Build a map from a set of keys and a function which for each key
+-- computes its value.
+--
+-- > fromSet (\k -> replicate k 'a') (Data.IntSet.fromList [3, 5]) == fromList [(5,"aaaaa"), (3,"aaa")]
+-- > fromSet undefined Data.IntSet.empty == empty
+
+fromSet :: (Key -> a) -> IntSet.IntSet -> IntMap a
+fromSet _ IntSet.Nil = Nil
+fromSet f (IntSet.Bin p m l r) = Bin p m (fromSet f l) (fromSet f r)
+fromSet f (IntSet.Tip kx bm) = buildTree f kx bm (IntSet.suffixBitMask + 1)
+  where -- This is slightly complicated, as we to convert the dense
+        -- representation of IntSet into tree representation of IntMap.
+        --
+        -- We are given a nonzero bit mask 'bmask' of 'bits' bits with prefix 'prefix'.
+        -- We split bmask into halves corresponding to left and right subtree.
+        -- If they are both nonempty, we create a Bin node, otherwise exactly
+        -- one of them is nonempty and we construct the IntMap from that half.
+        buildTree g prefix bmask bits = prefix `seq` bmask `seq` case bits of
+          0 -> Tip prefix (g prefix)
+          _ -> case intFromNat ((natFromInt bits) `shiftRL` 1) of
+                 bits2 | bmask .&. ((1 `shiftLL` bits2) - 1) == 0 ->
+                           buildTree g (prefix + bits2) (bmask `shiftRL` bits2) bits2
+                       | (bmask `shiftRL` bits2) .&. ((1 `shiftLL` bits2) - 1) == 0 ->
+                           buildTree g prefix bmask bits2
+                       | otherwise ->
+                           Bin prefix bits2 (buildTree g prefix bmask bits2) (buildTree g (prefix + bits2) (bmask `shiftRL` bits2) bits2)
 
 {--------------------------------------------------------------------
   Lists
index cbfb34c..7e735a9 100644 (file)
@@ -142,8 +142,9 @@ module Data.IntMap.Lazy (
             -- * Conversion
             , elems
             , keys
-            , keysSet
             , assocs
+            , keysSet
+            , fromSet
 
             -- ** Lists
             , toList
index b6e911a..db472da 100644 (file)
@@ -148,8 +148,9 @@ module Data.IntMap.Strict (
             -- * Conversion
             , elems
             , keys
-            , keysSet
             , assocs
+            , keysSet
+            , fromSet
 
             -- ** Lists
             , toList
@@ -206,6 +207,7 @@ module Data.IntMap.Strict (
 
 import Prelude hiding (lookup,map,filter,foldr,foldl,null)
 
+import Data.Bits
 import Data.IntMap.Base hiding
     ( findWithDefault
     , singleton
@@ -241,6 +243,7 @@ import Data.IntMap.Base hiding
     , mapMaybeWithKey
     , mapEither
     , mapEitherWithKey
+    , fromSet
     , fromList
     , fromListWith
     , fromListWithKey
@@ -249,6 +252,7 @@ import Data.IntMap.Base hiding
     , fromAscListWithKey
     , fromDistinctAscList
     )
+import qualified Data.IntSet.Base as IntSet
 import Data.StrictPair
 
 -- $strictness
@@ -828,6 +832,36 @@ mapEitherWithKey f (Tip k x) = case f k x of
   Right z -> z `seq` (Nil, Tip k z)
 mapEitherWithKey _ Nil = (Nil, Nil)
 
+{--------------------------------------------------------------------
+  Conversions
+--------------------------------------------------------------------}
+
+-- | /O(n)/. Build a map from a set of keys and a function which for each key
+-- computes its value.
+--
+-- > fromSet (\k -> replicate k 'a') (Data.IntSet.fromList [3, 5]) == fromList [(5,"aaaaa"), (3,"aaa")]
+-- > fromSet undefined Data.IntSet.empty == empty
+
+fromSet :: (Key -> a) -> IntSet.IntSet -> IntMap a
+fromSet _ IntSet.Nil = Nil
+fromSet f (IntSet.Bin p m l r) = Bin p m (fromSet f l) (fromSet f r)
+fromSet f (IntSet.Tip kx bm) = buildTree f kx bm (IntSet.suffixBitMask + 1)
+  where -- This is slightly complicated, as we to convert the dense
+        -- representation of IntSet into tree representation of IntMap.
+        --
+        -- We are given a nonzero bit mask 'bmask' of 'bits' bits with prefix 'prefix'.
+        -- We split bmask into halves corresponding to left and right subtree.
+        -- If they are both nonempty, we create a Bin node, otherwise exactly
+        -- one of them is nonempty and we construct the IntMap from that half.
+        buildTree g prefix bmask bits = prefix `seq` bmask `seq` case bits of
+          0 -> Tip prefix $! g prefix
+          _ -> case intFromNat ((natFromInt bits) `shiftRL` 1) of
+                 bits2 | bmask .&. ((1 `shiftLL` bits2) - 1) == 0 ->
+                           buildTree g (prefix + bits2) (bmask `shiftRL` bits2) bits2
+                       | (bmask `shiftRL` bits2) .&. ((1 `shiftLL` bits2) - 1) == 0 ->
+                           buildTree g prefix bmask bits2
+                       | otherwise ->
+                           Bin prefix bits2 (buildTree g prefix bmask bits2) (buildTree g (prefix + bits2) (bmask `shiftRL` bits2) bits2)
 
 {--------------------------------------------------------------------
   Lists
index 2e72687..c20e7f5 100644 (file)
@@ -171,8 +171,9 @@ module Data.Map.Base (
             -- * Conversion
             , elems
             , keys
-            , keysSet
             , assocs
+            , keysSet
+            , fromSet
 
             -- ** Lists
             , toList
@@ -1873,6 +1874,16 @@ elems = foldr (:) []
 keys  :: Map k a -> [k]
 keys = foldrWithKey (\k _ ks -> k : ks) []
 
+-- | /O(n)/. An alias for 'toAscList'. Return all key\/value pairs in the map
+-- in ascending key order. Subject to list fusion.
+--
+-- > assocs (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")]
+-- > assocs empty == []
+
+assocs :: Map k a -> [(k,a)]
+assocs m
+  = toAscList m
+
 -- | /O(n)/. The set of all keys of the map.
 --
 -- > keysSet (fromList [(5,"a"), (3,"b")]) == Data.Set.fromList [3,5]
@@ -1882,15 +1893,15 @@ keysSet :: Map k a -> Set.Set k
 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.
+-- | /O(n)/. Build a map from a set of keys and a function which for each key
+-- computes its value.
 --
--- > assocs (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")]
--- > assocs empty == []
+-- > fromSet (\k -> replicate k 'a') (Data.Set.fromList [3, 5]) == fromList [(5,"aaaaa"), (3,"aaa")]
+-- > fromSet undefined Data.Set.empty == empty
 
-assocs :: Map k a -> [(k,a)]
-assocs m
-  = toAscList m
+fromSet :: (k -> a) -> Set.Set k -> Map k a
+fromSet _ Set.Tip = Tip
+fromSet f (Set.Bin sz x l r) = Bin sz x (f x) (fromSet f l) (fromSet f r)
 
 {--------------------------------------------------------------------
   Lists
index 253ea90..1a44933 100644 (file)
@@ -135,8 +135,9 @@ module Data.Map.Lazy (
             -- * Conversion
             , elems
             , keys
-            , keysSet
             , assocs
+            , keysSet
+            , fromSet
 
             -- ** Lists
             , toList
index 7f17b50..c0b7af3 100644 (file)
@@ -142,8 +142,9 @@ module Data.Map.Strict
     -- * Conversion
     , elems
     , keys
-    , keysSet
     , assocs
+    , keysSet
+    , fromSet
 
     -- ** Lists
     , toList
@@ -242,6 +243,7 @@ import Data.Map.Base hiding
     , mapAccumWithKey
     , mapAccumRWithKey
     , mapKeysWith
+    , fromSet
     , fromList
     , fromListWith
     , fromListWithKey
@@ -259,6 +261,7 @@ import Data.Map.Base hiding
     , updateMinWithKey
     , updateMaxWithKey
     )
+import qualified Data.Set.Base as Set
 import Data.StrictPair
 
 -- Use macros to define strictness of functions.  STRICT_x_OF_y
@@ -975,6 +978,20 @@ mapKeysWith c f = fromListWith c . foldrWithKey (\k x xs -> (f k, x) : xs) []
 #endif
 
 {--------------------------------------------------------------------
+  Conversions
+--------------------------------------------------------------------}
+
+-- | /O(n)/. Build a map from a set of keys and a function which for each key
+-- computes its value.
+--
+-- > fromSet (\k -> replicate k 'a') (Data.Set.fromList [3, 5]) == fromList [(5,"aaaaa"), (3,"aaa")]
+-- > fromSet undefined Data.Set.empty == empty
+
+fromSet :: (k -> a) -> Set.Set k -> Map k a
+fromSet _ Set.Tip = Tip
+fromSet f (Set.Bin sz x l r) = case f x of v -> v `seq` Bin sz x v (fromSet f l) (fromSet f r)
+
+{--------------------------------------------------------------------
   Lists
   use [foldlStrict] to reduce demand on the control-stack
 --------------------------------------------------------------------}
index d26d1fa..4f5d06e 100644 (file)
@@ -77,8 +77,9 @@ main = defaultMainWithOpts
              , testCase "mapKeysMonotonic" test_mapKeysMonotonic
              , testCase "elems" test_elems
              , testCase "keys" test_keys
+             , testCase "assocs" test_assocs
              , testCase "keysSet" test_keysSet
-             , testCase "associative" test_assocs
+             , testCase "keysSet" test_fromSet
              , testCase "toList" test_toList
              , testCase "fromList" test_fromList
              , testCase "fromListWith" test_fromListWith
@@ -163,6 +164,7 @@ main = defaultMainWithOpts
              , testProperty "foldl"                prop_foldl
              , testProperty "foldl'"               prop_foldl'
              , testProperty "keysSet"              prop_keysSet
+             , testProperty "fromSet"              prop_fromSet
              ] opts
 
   where
@@ -494,15 +496,20 @@ test_keys = do
     keys (fromList [(5,"a"), (3,"b")]) @?= [3,5]
     keys (empty :: UMap) @?= []
 
+test_assocs :: Assertion
+test_assocs = do
+    assocs (fromList [(5,"a"), (3,"b")]) @?= [(3,"b"), (5,"a")]
+    assocs (empty :: UMap) @?= []
+
 test_keysSet :: Assertion
 test_keysSet = do
     keysSet (fromList [(5,"a"), (3,"b")]) @?= Data.IntSet.fromList [3,5]
     keysSet (empty :: UMap) @?= Data.IntSet.empty
 
-test_assocs :: Assertion
-test_assocs = do
-    assocs (fromList [(5,"a"), (3,"b")]) @?= [(3,"b"), (5,"a")]
-    assocs (empty :: UMap) @?= []
+test_fromSet :: Assertion
+test_fromSet = do
+   fromSet (\k -> replicate k 'a') (Data.IntSet.fromList [3, 5]) @?= fromList [(5,"aaaaa"), (3,"aaa")]
+   fromSet undefined Data.IntSet.empty @?= (empty :: IMap)
 
 ----------------------------------------------------------------
 -- Lists
@@ -1027,3 +1034,8 @@ prop_foldl' n ys = length ys > 0 ==>
 prop_keysSet :: [(Int, Int)] -> Bool
 prop_keysSet xs =
   keysSet (fromList xs) == Data.IntSet.fromList (List.map fst xs)
+
+prop_fromSet :: [(Int, Int)] -> Bool
+prop_fromSet ys =
+  let xs = List.nubBy ((==) `on` fst) ys
+  in fromSet (\k -> fromJust $ List.lookup k xs) (Data.IntSet.fromList $ List.map fst xs) == fromList xs
index 068ee0a..77b2236 100644 (file)
@@ -76,8 +76,9 @@ main = defaultMainWithOpts
          , testCase "mapKeysMonotonic" test_mapKeysMonotonic
          , testCase "elems" test_elems
          , testCase "keys" test_keys
+         , testCase "assocs" test_assocs
          , testCase "keysSet" test_keysSet
-         , testCase "associative" test_assocs
+         , testCase "fromSet" test_fromSet
          , testCase "toList" test_toList
          , testCase "fromList" test_fromList
          , testCase "fromListWith" test_fromListWith
@@ -184,6 +185,7 @@ main = defaultMainWithOpts
          , testProperty "foldl"                prop_foldl
          , testProperty "foldl'"               prop_foldl'
          , testProperty "keysSet"              prop_keysSet
+         , testProperty "fromSet"              prop_fromSet
          ] opts
 
   where
@@ -524,15 +526,20 @@ test_keys = do
     keys (fromList [(5,"a"), (3,"b")]) @?= [3,5]
     keys (empty :: UMap) @?= []
 
+test_assocs :: Assertion
+test_assocs = do
+    assocs (fromList [(5,"a"), (3,"b")]) @?= [(3,"b"), (5,"a")]
+    assocs (empty :: UMap) @?= []
+
 test_keysSet :: Assertion
 test_keysSet = do
     keysSet (fromList [(5,"a"), (3,"b")]) @?= Data.Set.fromList [3,5]
     keysSet (empty :: UMap) @?= Data.Set.empty
 
-test_assocs :: Assertion
-test_assocs = do
-    assocs (fromList [(5,"a"), (3,"b")]) @?= [(3,"b"), (5,"a")]
-    assocs (empty :: UMap) @?= []
+test_fromSet :: Assertion
+test_fromSet = do
+   fromSet (\k -> replicate k 'a') (Data.Set.fromList [3, 5]) @?= fromList [(5,"aaaaa"), (3,"aaa")]
+   fromSet undefined Data.Set.empty @?= (empty :: IMap)
 
 ----------------------------------------------------------------
 -- Lists
@@ -1137,3 +1144,8 @@ prop_foldl' n ys = length ys > 0 ==>
 prop_keysSet :: [(Int, Int)] -> Bool
 prop_keysSet xs =
   keysSet (fromList xs) == Data.Set.fromList (List.map fst xs)
+
+prop_fromSet :: [(Int, Int)] -> Bool
+prop_fromSet ys =
+  let xs = List.nubBy ((==) `on` fst) ys
+  in fromSet (\k -> fromJust $ List.lookup k xs) (Data.Set.fromList $ List.map fst xs) == fromList xs