From 12cb12387010c0592b16260922bc6c6c0ca4ce8c Mon Sep 17 00:00:00 2001 From: Alexandre Esteves Date: Thu, 20 Aug 2026 05:09:00 +0100 Subject: [PATCH] Add Map curry/uncurry utils --- containers-tests/benchmarks/Map.hs | 16 ++++++- containers-tests/tests/map-properties.hs | 49 ++++++++++++---------- containers/src/Data/Map/Internal.hs | 35 +++++++++++++++- containers/src/Data/Map/Lazy.hs | 4 ++ containers/src/Data/Map/Strict.hs | 4 ++ containers/src/Data/Map/Strict/Internal.hs | 37 +++++++++++++++- 6 files changed, 118 insertions(+), 27 deletions(-) diff --git a/containers-tests/benchmarks/Map.hs b/containers-tests/benchmarks/Map.hs index c3e057a78..7bd014a2c 100644 --- a/containers-tests/benchmarks/Map.hs +++ b/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 diff --git a/containers-tests/tests/map-properties.hs b/containers-tests/tests/map-properties.hs index 62b2a4787..7768f57b8 100644 --- a/containers-tests/tests/map-properties.hs +++ b/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,7 +1888,7 @@ 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 @@ -1895,7 +1896,7 @@ prop_foldr c n ys = foldr c' n m === Foldable.foldr c' n (snd <$> 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) diff --git a/containers/src/Data/Map/Internal.hs b/containers/src/Data/Map/Internal.hs index e53aa616e..89a42440d 100644 --- a/containers/src/Data/Map/Internal.hs +++ b/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 --------------------------------------------------------------------} diff --git a/containers/src/Data/Map/Lazy.hs b/containers/src/Data/Map/Lazy.hs index b700c1805..a0d613327 100644 --- a/containers/src/Data/Map/Lazy.hs +++ b/containers/src/Data/Map/Lazy.hs @@ -240,6 +240,10 @@ module Data.Map.Lazy ( , keysSet , argSet + -- ** Maps + , curry + , uncurry + -- ** Lists , toList diff --git a/containers/src/Data/Map/Strict.hs b/containers/src/Data/Map/Strict.hs index c30e2900d..c3d9fe5de 100644 --- a/containers/src/Data/Map/Strict.hs +++ b/containers/src/Data/Map/Strict.hs @@ -254,6 +254,10 @@ module Data.Map.Strict , keysSet , argSet + -- ** Maps + , curry + , uncurry + -- ** Lists , toList diff --git a/containers/src/Data/Map/Strict/Internal.hs b/containers/src/Data/Map/Strict/Internal.hs index 12e24450b..4cba28f5d 100644 --- a/containers/src/Data/Map/Strict/Internal.hs +++ b/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 --------------------------------------------------------------------}