Add Map curry/uncurry utils

containers-tests/benchmarks/Map.hs

@@ -2,7 +2,7 @@
 {-# LANGUAGE BangPatterns #-}
 module Main where
 
-import Control.Applicative (Const(Const, getConst), pure)
+import Control.Applicative (Const(Const, getConst), liftA2, pure)
 import Control.DeepSeq (rnf)
 import Control.Exception (evaluate)
 import Test.Tasty.Bench (bench, bgroup, defaultMain, whnf, nf)
@@ -31,12 +31,16 @@ main = do
     let m = M.fromList elems :: M.Map Int Int
         m_even = M.fromList elems_even :: M.Map Int Int
         m_odd = M.fromList elems_odd :: M.Map Int Int
+        m_uncurried = M.fromList elems_uncurried :: M.Map (Int, Int) Int
+        m_curried = M.curry m_uncurried :: M.Map Int (M.Map Int Int)
         s_random = Set.fromList keys_random :: Set.Set Int
     evaluate $ rnf [m, m_even, m_odd]
     evaluate $ rnf [s_random]
     evaluate $ rnf
       [elems_distinct_asc, elems_distinct_desc, elems_asc, elems_desc]
     evaluate $ rnf [keys_random]
+    evaluate $ rnf [m_uncurried]
+    evaluate $ rnf [m_curried]
     defaultMain
         [ bench "lookup absent" $ whnf (lookup evens) m_odd
         , bench "lookup present" $ whnf (lookup evens) m_even
@@ -143,6 +147,8 @@ main = do
         , bench "Strict.fromSetA outer" $ whnf (MS.fromSetA (MkSolo . pred)) s_random
         , bench "Lazy.fromSetA inner" $ whnf (getSolo . M.fromSetA (MkSolo . pred)) s_random
         , bench "Strict.fromSetA inner" $ whnf (getSolo . MS.fromSetA (MkSolo . pred)) s_random
+        , bench "curry" $ whnf M.curry m_uncurried
+        , bench "uncurry" $ whnf M.uncurry m_curried
         , bench "minView" $ whnf (\m' -> case M.minViewWithKey m' of {Nothing -> 0; Just ((k,v),m'') -> k+v+M.size m''}) (M.fromAscList $ zip [1..10::Int] [100..110::Int])
         , bench "eq" $ whnf (\m' -> m' == m') m -- worst case, compares everything
         , bench "compare" $ whnf (\m' -> compare m' m') m -- worst case, compares everything
@@ -155,7 +161,8 @@ main = do
         , bench "mapKeysWith:desc" $ whnf (M.mapKeysWith (+) (negate . (`div` 2))) m
         ]
   where
-    bound = 2^14
+    magnitude = 14
+    bound = 2^magnitude
     elems = shuffle elems_distinct_asc
     elems_even = zip evens evens
     elems_odd = zip odds odds
@@ -172,6 +179,11 @@ main = do
     sumkv k v1 v2 = k + v1 + v2
     consPair k v xs = (k, v) : xs
     keys_random = take bound (randoms gen)
+    elems_uncurried = zip xs evens
+      where
+        left = magnitude `div` 2
+        right = magnitude - left
+        xs = shuffle $ liftA2 (,) [1..2^left] $ reverse [1..2^right]
 
 add3 :: Int -> Int -> Int -> Int
 add3 x y z = x + y + z

containers-tests/tests/map-properties.hs

@@ -34,7 +34,7 @@ import Data.Functor
 import qualified Data.Foldable as Foldable
 import qualified Data.Bifoldable as Bifoldable
 import Data.Proxy (Proxy(..))
-import Prelude hiding (lookup, null, map, filter, foldr, foldl, foldl', take, drop, splitAt)
+import Prelude hiding (lookup, null, map, filter, foldr, foldl, foldl', take, drop, splitAt, curry, uncurry)
 import qualified Prelude
 
 import Data.List (nub,sort)
@@ -278,6 +278,7 @@ main = defaultMain $ testGroup "map-properties"
          , testProperty "fromSetMaybe"         prop_fromSetMaybe
          , testProperty "fromSetMaybeA"        prop_fromSetMaybeA
          , testProperty "fromArgSet"           prop_fromArgSet
+         , testProperty "curry"                prop_curry
          , testProperty "takeWhileAntitone"    prop_takeWhileAntitone
          , testProperty "dropWhileAntitone"    prop_dropWhileAntitone
          , testProperty "spanAntitone"         prop_spanAntitone
@@ -1576,7 +1577,7 @@ prop_fromList :: [Int] -> Bool
 prop_fromList xs
   = case fromList (zip xs xs) of
       t -> t == fromAscList (zip sort_xs sort_xs) &&
-           t == List.foldr (uncurry insert) empty (zip xs xs)
+           t == List.foldr (Prelude.uncurry insert) empty (zip xs xs)
   where sort_xs = sort xs
 
 prop_fromAscList :: SortedOnFst Int A -> Property
@@ -1848,7 +1849,7 @@ prop_mapAssocsMonotonic
   :: MonotonicFun -> Fun (Int, A) B -> Map Int A -> Property
 prop_mapAssocsMonotonic f1 f2 m =
   valid m' .&&.
-  toList m' === fmap (uncurry f) (toList m)
+  toList m' === fmap (Prelude.uncurry f) (toList m)
   where
     m' = mapAssocsMonotonic f m
     f k x = (applyMonotonicFun f1 k, applyFun2 f2 k x)
@@ -1878,7 +1879,7 @@ prop_foldMap = \m -> Foldable.foldMap f m === Foldable.foldMap f (elems m)
     f v = [v]
 
 prop_foldMapWithKey :: Map Int A -> Property
-prop_foldMapWithKey = \m -> foldMapWithKey (curry f) m === Foldable.foldMap f (toList m)
+prop_foldMapWithKey = \m -> foldMapWithKey (Prelude.curry f) m === Foldable.foldMap f (toList m)
   where
     f kv = [kv]
 
@@ -1887,15 +1888,15 @@ prop_foldMapWithKey = \m -> foldMapWithKey (curry f) m === Foldable.foldMap f (t
 prop_foldr :: Fun (A, B) B -> B -> [(Int, A)] -> Property
 prop_foldr c n ys = foldr c' n m === Foldable.foldr c' n (snd <$> xs)
   where
-    c' = curry (apply c)
+    c' = Prelude.curry (apply c)
     xs = List.sortBy (comparing fst) (List.nubBy ((==) `on` fst) ys)
     m  = fromList xs
 
 
 -- toList is implemented in terms of foldrWithKey, so we don't want to rely on it
 -- when we're trying to test foldrWithKey.
 prop_foldrWithKey :: Fun (Int, A, B) B -> B -> [(Int, A)] -> Property
-prop_foldrWithKey c n ys = foldrWithKey c' n m === Foldable.foldr (uncurry c') n xs
+prop_foldrWithKey c n ys = foldrWithKey c' n m === Foldable.foldr (Prelude.uncurry c') n xs
   where
     c' k v acc = apply c (k, v, acc)
     xs = List.sortBy (comparing fst) (List.nubBy ((==) `on` fst) ys)
@@ -1904,30 +1905,30 @@ prop_foldrWithKey c n ys = foldrWithKey c' n m === Foldable.foldr (uncurry c') n
 prop_foldr' :: Fun (A, B) B -> B -> Map Int A -> Property
 prop_foldr' c n m = foldr' c' n m === Foldable.foldr' c' n (elems m)
   where
-    c' = curry (apply c)
+    c' = Prelude.curry (apply c)
 
 prop_foldrWithKey' :: Fun (Int, A, B) B -> B -> Map Int A -> Property
-prop_foldrWithKey' c n m = foldrWithKey' c' n m === Foldable.foldr' (uncurry c') n (toList m)
+prop_foldrWithKey' c n m = foldrWithKey' c' n m === Foldable.foldr' (Prelude.uncurry c') n (toList m)
   where
     c' k v acc = apply c (k, v, acc)
 
 prop_foldl :: Fun (B, A) B -> B -> Map Int A -> Property
 prop_foldl c n m = foldl c' n m === Foldable.foldl c' n (elems m)
   where
-    c' = curry (apply c)
+    c' = Prelude.curry (apply c)
 
 prop_foldlWithKey :: Fun (B, Int, A) B -> B -> Map Int A -> Property
-prop_foldlWithKey c n m = foldlWithKey c' n m === Foldable.foldl (uncurry . c') n (toList m)
+prop_foldlWithKey c n m = foldlWithKey c' n m === Foldable.foldl (Prelude.uncurry . c') n (toList m)
   where
     c' acc k v = apply c (acc, k, v)
 
 prop_foldl' :: Fun (B, A) B -> B -> Map Int A -> Property
 prop_foldl' c n m = foldl' c' n m === Foldable.foldl' c' n (elems m)
   where
-    c' = curry (apply c)
+    c' = Prelude.curry (apply c)
 
 prop_foldlWithKey' :: Fun (B, Int, A) B -> B -> Map Int A -> Property
-prop_foldlWithKey' c n m = foldlWithKey' c' n m === Foldable.foldl' (uncurry . c') n (toList m)
+prop_foldlWithKey' c n m = foldlWithKey' c' n m === Foldable.foldl' (Prelude.uncurry . c') n (toList m)
   where
     c' acc k v = apply c (acc, k, v)
 
@@ -1940,29 +1941,29 @@ prop_bifoldMap m = Bifoldable.bifoldMap (:[]) (:[]) m === Foldable.foldMap (\(k,
 prop_bifoldr :: Fun (Int, B) B -> Fun (A, B) B -> B -> Map Int A -> Property
 prop_bifoldr ck cv n m = Bifoldable.bifoldr ck' cv' n m === Foldable.foldr c' n (toList m)
   where
-    ck' = curry (apply ck)
-    cv' = curry (apply cv)
+    ck' = Prelude.curry (apply ck)
+    cv' = Prelude.curry (apply cv)
     (k,v) `c'` acc = k `ck'` (v `cv'` acc)
 
 prop_bifoldr' :: Fun (Int, B) B -> Fun (A, B) B -> B -> Map Int A -> Property
 prop_bifoldr' ck cv n m = Bifoldable.bifoldr' ck' cv' n m === Foldable.foldr' c' n (toList m)
   where
-    ck' = curry (apply ck)
-    cv' = curry (apply cv)
+    ck' = Prelude.curry (apply ck)
+    cv' = Prelude.curry (apply cv)
     (k,v) `c'` acc = k `ck'` (v `cv'` acc)
 
 prop_bifoldl :: Fun (B, Int) B -> Fun (B, A) B -> B -> Map Int A -> Property
 prop_bifoldl ck cv n m = Bifoldable.bifoldl ck' cv' n m === Foldable.foldl c' n (toList m)
   where
-    ck' = curry (apply ck)
-    cv' = curry (apply cv)
+    ck' = Prelude.curry (apply ck)
+    cv' = Prelude.curry (apply cv)
     acc `c'` (k,v) = (acc `ck'` k) `cv'` v
 
 prop_bifoldl' :: Fun (B, Int) B -> Fun (B, A) B -> B -> Map Int A -> Property
 prop_bifoldl' ck cv n m = Bifoldable.bifoldl' ck' cv' n m === Foldable.foldl' c' n (toList m)
   where
-    ck' = curry (apply ck)
-    cv' = curry (apply cv)
+    ck' = Prelude.curry (apply ck)
+    cv' = Prelude.curry (apply cv)
     acc `c'` (k,v) = (acc `ck'` k) `cv'` v
 
 prop_keysSet :: [OrdA] -> Property
@@ -1971,7 +1972,7 @@ prop_keysSet keys =
 
 prop_argSet :: [(OrdA, B)] -> Property
 prop_argSet xs =
-  argSet (fromList xs) === Set.fromList (List.map (uncurry Arg) xs)
+  argSet (fromList xs) === Set.fromList (List.map (Prelude.uncurry Arg) xs)
 
 prop_fromSet :: Set OrdA -> Fun OrdA B -> Property
 prop_fromSet keys funF =
@@ -2011,7 +2012,11 @@ prop_fromSetMaybeA keys f =
 
 prop_fromArgSet :: [(OrdA, B)] -> Property
 prop_fromArgSet ys =
-  fromArgSet (Set.fromList $ List.map (uncurry Arg) ys) === fromList ys
+  fromArgSet (Set.fromList $ List.map (Prelude.uncurry Arg) ys) === fromList ys
+
+prop_curry :: Map Int (Map Int A) -> Property
+prop_curry m = m' === Data.Map.uncurry (Data.Map.curry m')
+  where m' = Data.Map.uncurry m
 
 prop_eq :: Map Int A -> Map Int A -> Property
 prop_eq m1 m2 = (m1 == m2) === (toList m1 == toList m2)

containers/src/Data/Map/Internal.hs

@@ -268,6 +268,10 @@ module Data.Map.Internal (
     , fromSetMaybeA
     , fromArgSet
 
+    -- ** Maps
+    , curry
+    , uncurry
+
     -- ** Lists
     , toList
     , fromList
@@ -391,7 +395,7 @@ import Control.DeepSeq (NFData(rnf),NFData1(liftRnf),NFData2(liftRnf2))
 import qualified Data.Foldable as Foldable
 import Data.Bifoldable
 import Utils.Containers.Internal.Prelude hiding
-  (lookup, map, filter, foldr, foldl, foldl', null, splitAt, take, drop)
+  (lookup, map, filter, foldr, foldl, foldl', null, splitAt, take, drop, curry, uncurry)
 import Prelude ()
 
 import qualified Data.Set.Internal as Set
@@ -3427,6 +3431,35 @@ fromArgSet :: Set.Set (Arg k a) -> Map k a
 fromArgSet Set.Tip = Tip
 fromArgSet (Set.Bin sz (Arg x v) l r) = Bin sz x v (fromArgSet l) (fromArgSet r)
 
+{--------------------------------------------------------------------
+  Maps
+--------------------------------------------------------------------}
+-- | \(O(n)\). Group map entries by the first component.
+--
+-- > curry $ fromList [((1,2),12),((1,3),13)] == fromList [(1,fromList [(2,12),(3,13)])]
+--
+-- @since FIXME
+curry :: (Ord a, Ord b) => Map (a,b) c -> Map a (Map b c)
+curry m = fmap (fromDescList . ($ [])) $ fromAscListWith (.) $ fmap (\((a,b),c) -> (a, ((b,c):))) $ toAscList m
+
+-- | \(O(n)\). Flatten nested maps.
+--
+-- Note
+--
+-- > uncurry . curry = id
+--
+-- but not the other way around
+--
+-- > uncurry (fromList [(1, fromList [])]) == fromList []
+-- > uncurry (fromList [(1, fromList [(2,12),(3,13)])]) == fromList [((1,2),12),((1,3),13)]
+--
+-- @since FIXME
+uncurry :: (Ord a, Ord b) => Map a (Map b c) -> Map (a,b) c
+uncurry m = fromAscList $ do
+  (a,b2c) <- toAscList m
+  (b,c) <- toAscList b2c
+  pure ((a,b), c)
+
 {--------------------------------------------------------------------
   Lists
 --------------------------------------------------------------------}

containers/src/Data/Map/Lazy.hs

@@ -240,6 +240,10 @@ module Data.Map.Lazy (
     , keysSet
     , argSet
 
+    -- ** Maps
+    , curry
+    , uncurry
+
     -- ** Lists
     , toList
 

containers/src/Data/Map/Strict.hs

@@ -254,6 +254,10 @@ module Data.Map.Strict
     , keysSet
     , argSet
 
+    -- ** Maps
+    , curry
+    , uncurry
+
     -- ** Lists
     , toList
 

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

@@ -219,6 +219,10 @@ module Data.Map.Strict.Internal
     , fromSetMaybeA
     , fromArgSet
 
+    -- ** Maps
+    , curry
+    , uncurry
+
     -- ** Lists
     , toList
     , fromList
@@ -298,7 +302,7 @@ module Data.Map.Strict.Internal
     ) where
 
 import Utils.Containers.Internal.Prelude hiding
-  (lookup,map,filter,foldr,foldl,foldl',null,take,drop,splitAt)
+  (lookup,map,filter,foldr,foldl,foldl',null,take,drop,splitAt,curry,uncurry)
 import Prelude ()
 
 import Data.Map.Internal
@@ -1427,7 +1431,7 @@ fromSet f = runIdentity . fromSetA (pure . f)
 -- @since FIXME
 fromSetA :: Applicative f => (k -> f a) -> Set k -> f (Map k a)
 fromSetA _ Set.Tip = pure Tip
-fromSetA f (Set.Bin sz x l r) = 
+fromSetA f (Set.Bin sz x l r) =
   liftA3 (flip (Bin sz x $!)) (fromSetA f l) (f x) (fromSetA f r)
 {-# INLINABLE fromSetA #-}
 
@@ -1471,6 +1475,35 @@ fromArgSet :: Set.Set (Arg k a) -> Map k a
 fromArgSet Set.Tip = Tip
 fromArgSet (Set.Bin sz (Arg x v) l r) = v `seq` Bin sz x v (fromArgSet l) (fromArgSet r)
 
+{--------------------------------------------------------------------
+  Maps
+--------------------------------------------------------------------}
+-- | \(O(n)\). Group map entries by the first component.
+--
+-- > curry $ fromList [((1,2),12),((1,3),13)] == fromList [(1,fromList [(2,12),(3,13)])]
+--
+-- @since FIXME
+curry :: (Ord a, Ord b) => Map (a,b) c -> Map a (Map b c)
+curry m = fmap (fromDescList . ($ [])) $ fromAscListWith (.) $ fmap (\((a,b),c) -> (a, ((b,c):))) $ toAscList m
+
+-- | \(O(n)\). Flatten nested maps.
+--
+-- Note
+--
+-- > uncurry . curry = id
+--
+-- but not the other way around
+--
+-- > uncurry (fromList [(1, fromList [])]) == fromList []
+-- > uncurry (fromList [(1, fromList [(2,12),(3,13)])]) == fromList [((1,2),12),((1,3),13)]
+--
+-- @since FIXME
+uncurry :: (Ord a, Ord b) => Map a (Map b c) -> Map (a,b) c
+uncurry m = fromAscList $ do
+  (a,b2c) <- toAscList m
+  (b,c) <- toAscList b2c
+  pure ((a,b), c)
+
 {--------------------------------------------------------------------
   Lists
 --------------------------------------------------------------------}