Add Map curry/uncurry utils

Author: alexfmpe

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

@@ -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
@@ -2013,6 +2014,10 @@ prop_fromArgSet :: [(OrdA, B)] -> Property
 prop_fromArgSet ys =
   fromArgSet (Set.fromList $ List.map (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
 --------------------------------------------------------------------}