Introduce ‘partitionKeys’ that fuses ‘restrictKeys’ and ‘withoutKeys’ in one go

Author: sergv

containers-tests/benchmarks/Map.hs

@@ -5,7 +5,7 @@ module Main where
 import Control.Applicative (Const(Const, getConst), pure)
 import Control.DeepSeq (rnf)
 import Control.Exception (evaluate)
-import Test.Tasty.Bench (bench, defaultMain, whnf, nf)
+import Test.Tasty.Bench (bench, defaultMain, whnf, nf, bcompare)
 import Data.Functor.Identity (Identity(..))
 import Data.List (foldl')
 import qualified Data.Map as M
@@ -15,13 +15,16 @@ import Data.Maybe (fromMaybe)
 import Data.Functor ((<$))
 import Data.Coerce
 import Prelude hiding (lookup)
+import Utils.Containers.Internal.StrictPair
 
 main = do
     let m = M.fromAscList elems :: M.Map Int Int
         m_even = M.fromAscList elems_even :: M.Map Int Int
         m_odd = M.fromAscList elems_odd :: M.Map Int Int
+        m_odd_keys = M.keysSet m_odd
     evaluate $ rnf [m, m_even, m_odd]
     evaluate $ rnf elems_rev
+    evaluate $ rnf m_odd_keys
     defaultMain
         [ bench "lookup absent" $ whnf (lookup evens) m_odd
         , bench "lookup present" $ whnf (lookup evens) m_even
@@ -95,6 +98,12 @@ main = do
         , bench "fromDistinctDescList" $ whnf M.fromDistinctDescList elems_rev
         , bench "fromDistinctDescList:fusion" $ whnf (\n -> M.fromDistinctDescList [(i,i) | i <- [n,n-1..1]]) bound
         , 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 "restrictKeys+withoutKeys"
+        $ whnf (\ks -> M.restrictKeys m ks :*: M.withoutKeys m ks) m_odd_keys
+        , bcompare "/restrictKeys+withoutKeys/"
+        $ bench "partitionKeys"
+        $ whnf (M.partitionKeys m) m_odd_keys
         ]
   where
     bound = 2^12

containers-tests/tests/map-properties.hs

@@ -173,6 +173,7 @@ main = defaultMain $ testGroup "map-properties"
          , testProperty "withoutKeys"          prop_withoutKeys
          , testProperty "intersection"         prop_intersection
          , testProperty "restrictKeys"         prop_restrictKeys
+         , testProperty "partitionKeys"        prop_partitionKeys
          , testProperty "intersection model"   prop_intersectionModel
          , testProperty "intersectionWith"     prop_intersectionWith
          , testProperty "intersectionWithModel" prop_intersectionWithModel
@@ -1138,6 +1139,12 @@ prop_withoutKeys m s0 = valid reduced .&&. (m `withoutKeys` s === filterWithKey
     s = keysSet s0
     reduced = withoutKeys m s
 
+prop_partitionKeys :: IMap -> IMap -> Property
+prop_partitionKeys m s0 = valid with .&&. valid without .&&. (m `partitionKeys` s === (m `restrictKeys` s, m `withoutKeys` s))
+  where
+    s = keysSet s0
+    (with, without) = partitionKeys m s
+
 prop_intersection :: IMap -> IMap -> Bool
 prop_intersection t1 t2 = valid (intersection t1 t2)
 

containers/src/Data/Map/Internal.hs

@@ -7,6 +7,7 @@
 {-# LANGUAGE StandaloneDeriving #-}
 {-# LANGUAGE Trustworthy #-}
 {-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE ScopedTypeVariables #-}
 #endif
 #define USE_MAGIC_PROXY 1
 
@@ -296,6 +297,7 @@ module Data.Map.Internal (
 
     , restrictKeys
     , withoutKeys
+    , partitionKeys
     , partition
     , partitionWithKey
 
@@ -1943,6 +1945,48 @@ withoutKeys m (Set.Bin _ k ls rs) = case splitMember k m of
 {-# INLINABLE withoutKeys #-}
 #endif
 
+-- | \(O\bigl(m \log\bigl(\frac{n}{m}+1\bigr)\bigr), \; 0 < m \leq n\). Partition the map according to a set.
+-- The first map contains the input 'Map' restricted to those keys found in the 'Set',
+-- the second map contains the input 'Map' without all keys in the 'Set'.
+-- This is more efficient than using 'restrictKeys' and 'withoutKeys' together.
+--
+-- @
+-- m \`partitionKeys\` s = (m ``restrictKeys`` s, m ``withoutKeys`` s)
+-- @
+partitionKeys :: forall k a. Ord k => Map k a -> Set k -> (Map k a, Map k a)
+partitionKeys xs ys =
+  case go xs ys of
+    xs' :*: ys' -> (xs', ys')
+  where
+    go :: Map k a -> Set k -> StrictPair (Map k a) (Map k a)
+    go Tip                 _           = Tip :*: Tip
+    go m                   Set.Tip     = Tip :*: m
+    go m@(Bin _ k x lm rm) s@Set.Bin{} =
+      case b of
+        True  -> with :*: without
+          where
+            with    =
+              if lmWith `ptrEq` lm && rmWith `ptrEq` rm
+              then m
+              else link k x lmWith rmWith
+            without =
+              link2 lmWithout rmWithout
+        False -> with :*: without
+          where
+            with    = link2 lmWith rmWith
+            without =
+              if lmWithout `ptrEq` lm && rmWithout `ptrEq` rm
+              then m
+              else link k x lmWithout rmWithout
+        where
+          !(lmWith :*: lmWithout) = go lm ls'
+          !(rmWith :*: rmWithout) = go rm rs'
+
+          !(!ls', b, !rs') = Set.splitMember k s
+#if __GLASGOW_HASKELL__
+{-# INLINABLE partitionKeys #-}
+#endif
+
 -- | \(O(n+m)\). Difference with a combining function.
 -- When two equal keys are
 -- encountered, the combining function is applied to the values of these keys.

containers/src/Data/Map/Lazy.hs

@@ -222,6 +222,7 @@ module Data.Map.Lazy (
     , filterWithKey
     , restrictKeys
     , withoutKeys
+    , partitionKeys
     , partition
     , partitionWithKey
     , takeWhileAntitone

containers/src/Data/Map/Strict.hs

@@ -238,6 +238,7 @@ module Data.Map.Strict
     , filterWithKey
     , restrictKeys
     , withoutKeys
+    , partitionKeys
     , partition
     , partitionWithKey
 

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

@@ -253,6 +253,7 @@ module Data.Map.Strict.Internal
     , filterWithKey
     , restrictKeys
     , withoutKeys
+    , partitionKeys
     , partition
     , partitionWithKey
     , takeWhileAntitone
@@ -416,7 +417,8 @@ import Data.Map.Internal
   , toDescList
   , union
   , unions
-  , withoutKeys )
+  , withoutKeys
+  , partitionKeys )
 
 #if defined(__GLASGOW_HASKELL__)
 import Data.Map.Internal.DeprecatedShowTree (showTree, showTreeWith)