Add Map.mapAssocsMonotonic

containers-tests/tests/map-properties.hs

@@ -1,5 +1,7 @@
 {-# LANGUAGE CPP #-}
 {-# LANGUAGE TupleSections #-}
+{-# LANGUAGE DerivingStrategies #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
 
 #ifdef STRICT
 import Data.Map.Strict as Data.Map
@@ -246,6 +248,7 @@ main = defaultMain $ testGroup "map-properties"
          , testProperty "mapKeys"              prop_mapKeys
          , testProperty "mapKeysWith"          prop_mapKeysWith
          , testProperty "mapKeysMonotonic"     prop_mapKeysMonotonic
+         , testProperty "mapAssocsMonotonic"   prop_mapAssocsMonotonic
          , testProperty "split"                prop_splitModel
          , testProperty "fold"                 prop_fold
          , testProperty "foldMap"              prop_foldMap
@@ -1698,13 +1701,28 @@ prop_mapKeysWith f g m =
   where
     m' = mapKeysWith (applyFun2 f) (applyFun g) m
 
-prop_mapKeysMonotonic :: Positive Integer -> Integer -> Map Int A -> Property
-prop_mapKeysMonotonic (Positive p) q m =
+prop_mapKeysMonotonic :: MonotonicFun -> Map Int A -> Property
+prop_mapKeysMonotonic f m =
   valid m' .&&.
-  toList m' == [(f k, x) | (k,x) <- toList m]
+  toList m' === [(applyMonotonicFun f k, x) | (k,x) <- toList m]
   where
-    m' = mapKeysMonotonic f m
-    f x = fromIntegral x * p + q
+    m' = mapKeysMonotonic (applyMonotonicFun f) m
+
+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)
+  where
+    m' = mapAssocsMonotonic f m
+    f k x = (applyMonotonicFun f1 k, applyFun2 f2 k x)
+
+newtype MonotonicFun = MonotonicFun (Positive Integer, Integer)
+  deriving stock Show
+  deriving newtype Arbitrary
+
+applyMonotonicFun :: MonotonicFun -> Int -> Integer
+applyMonotonicFun (MonotonicFun (Positive p, q)) x = fromIntegral x * p + q
 
 prop_splitModel :: Int -> [(Int, Int)] -> Property
 prop_splitModel n ys = length ys > 0 ==>

containers/src/Data/Map/Internal.hs

@@ -241,6 +241,7 @@ module Data.Map.Internal (
     , mapKeys
     , mapKeysWith
     , mapKeysMonotonic
+    , mapAssocsMonotonic
 
     -- * Folds
     , foldr
@@ -3154,9 +3155,23 @@ mapKeysWith c f m =
 -- > valid (mapKeysMonotonic (\ _ -> 1)     (fromList [(5,"a"), (3,"b")])) == False
 
 mapKeysMonotonic :: (k1->k2) -> Map k1 a -> Map k2 a
-mapKeysMonotonic _ Tip = Tip
-mapKeysMonotonic f (Bin sz k x l r) =
-    Bin sz (f k) x (mapKeysMonotonic f l) (mapKeysMonotonic f r)
+mapKeysMonotonic f = mapAssocsMonotonic (\k x -> (f k, x))
+
+-- | \(O(n)\). Map over keys and values with a function @f@ that is
+-- monotonically strictly increasing in the keys. That is, for keys @kx@ and
+-- @ky@ and values @x@ and @y@, if @kx@ < @ky@ then
+-- @fst (f kx x)@ < @fst (f ky y)@.
+--
+-- __Warning__: This function should be used only if @f@ is monotonically
+-- strictly increasing in the key. This precondition is not checked.
+--
+-- @since FIXME
+mapAssocsMonotonic :: (k1 -> a1 -> (k2, a2)) -> Map k1 a1 -> Map k2 a2
+mapAssocsMonotonic f = go
+  where
+    go Tip = Tip
+    go (Bin sz k1 x1 l r) = case f k1 x1 of
+      (k2, x2) -> Bin sz k2 x2 (go l) (go r)
 
 {--------------------------------------------------------------------
   Folds

containers/src/Data/Map/Lazy.hs

@@ -214,6 +214,7 @@ module Data.Map.Lazy (
     , mapKeys
     , mapKeysWith
     , mapKeysMonotonic
+    , mapAssocsMonotonic
 
     -- * Folds
     , foldr

containers/src/Data/Map/Strict.hs

@@ -228,6 +228,7 @@ module Data.Map.Strict
     , mapKeys
     , mapKeysWith
     , mapKeysMonotonic
+    , mapAssocsMonotonic
 
     -- * Folds
     , foldr

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

@@ -192,6 +192,7 @@ module Data.Map.Strict.Internal
     , mapKeys
     , mapKeysWith
     , mapKeysMonotonic
+    , mapAssocsMonotonic
 
     -- * Folds
     , foldr
@@ -1376,6 +1377,22 @@ mapKeysWith c f m =
   finishB (foldlWithKey' (\b kx x -> insertWithB c (f kx) x b) emptyB m)
 {-# INLINABLE mapKeysWith #-}
 
+-- | \(O(n)\). Map over keys and values with a function @f@ that is
+-- monotonically strictly increasing in the keys. That is, for keys @kx@ and
+-- @ky@ and values @x@ and @y@, if @kx@ < @ky@ then
+-- @fst (f kx x)@ < @fst (f ky y)@.
+--
+-- __Warning__: This function should be used only if @f@ is monotonically
+-- strictly increasing in the key. This precondition is not checked.
+--
+-- @since FIXME
+mapAssocsMonotonic :: (k1 -> a1 -> (k2, a2)) -> Map k1 a1 -> Map k2 a2
+mapAssocsMonotonic f = go
+  where
+    go Tip = Tip
+    go (Bin sz k1 x1 l r) = case f k1 x1 of
+      (k2, !x2) -> Bin sz k2 x2 (go l) (go r)
+
 {--------------------------------------------------------------------
   Conversions
 --------------------------------------------------------------------}