Faster IntSet and IntMap split functions (#879)

* Add property test for IntMap.splitLookup

* Add split, splitMember, splitLookup benchmarks

* Use bin instead of union

In IntMap and Intset split, splitLookup, splitMember

* Move a few things around for clarity, no effect on performance

* Fix missed strictness in IntSet splitMember

Author: meooow25

containers-tests/benchmarks/IntMap.hs

@@ -51,6 +51,8 @@ main = do
         , bench "minView" $ whnf (maybe 0 (\((k,v), m) -> k+v+M.size m) . M.minViewWithKey)
                     (M.fromList $ zip [1..10] [1..10])
         , bench "spanAntitone" $ whnf (M.spanAntitone (<key_mid)) m
+        , bench "split" $ whnf (M.split key_mid) m
+        , bench "splitLookup" $ whnf (M.splitLookup key_mid) m
         ]
   where
     elems = elems_hits

containers-tests/benchmarks/IntSet.hs

@@ -50,14 +50,18 @@ main = do
           $ whnf (num_transitions . det 2 0) $ hard_nfa 1111 16
         , bench "spanAntitone:dense" $ whnf (IS.spanAntitone (<elem_mid)) s
         , bench "spanAntitone:sparse" $ whnf (IS.spanAntitone (<elem_sparse_mid)) s_sparse
+        , bench "split:dense" $ whnf (IS.split elem_mid) s
+        , bench "split:sparse" $ whnf (IS.split elem_sparse_mid) s_sparse
+        , bench "splitMember:dense" $ whnf (IS.splitMember elem_mid) s
+        , bench "splitMember:sparse" $ whnf (IS.splitMember elem_sparse_mid) s_sparse
         ]
   where
     elems = [1..2^12]
     elems_even = [2,4..2^12]
     elems_odd = [1,3..2^12]
-    elem_mid = 2^11
+    elem_mid = 2^11 + 31 -- falls in the middle of a packed Tip bitmask (assuming 64-bit words)
     elems_sparse = map (*64) elems -- when built into a map, each Tip is a singleton
-    elem_sparse_mid = 64 * elem_mid
+    elem_sparse_mid = 2^11 * 64
 
 member :: [Int] -> IS.IntSet -> Int
 member xs s = foldl' (\n x -> if IS.member x s then n + 1 else n) 0 xs

containers-tests/tests/intmap-properties.hs

@@ -187,6 +187,7 @@ main = defaultMain $ testGroup "intmap-properties"
              , testProperty "fmap"                 prop_fmap
              , testProperty "mapkeys"              prop_mapkeys
              , testProperty "split"                prop_splitModel
+             , testProperty "splitLookup"          prop_splitLookup
              , testProperty "splitRoot"            prop_splitRoot
              , testProperty "foldr"                prop_foldr
              , testProperty "foldr'"               prop_foldr'
@@ -1519,6 +1520,16 @@ prop_splitModel n ys = length ys > 0 ==>
       toAscList l === sort [(k, v) | (k,v) <- xs, k < n] .&&.
       toAscList r === sort [(k, v) | (k,v) <- xs, k > n]
 
+prop_splitLookup :: Int -> [(Int, Int)] -> Property
+prop_splitLookup n ys =
+    let xs = List.nubBy ((==) `on` fst) ys
+        (l, x, r) = splitLookup n (fromList xs)
+    in  valid l .&&.
+        valid r .&&.
+        x === List.lookup n xs .&&.
+        toAscList l === sort [(k, v) | (k,v) <- xs, k < n] .&&.
+        toAscList r === sort [(k, v) | (k,v) <- xs, k > n]
+
 prop_splitRoot :: IMap -> Bool
 prop_splitRoot s = loop ls && (s == unions ls)
  where

containers/src/Data/IntMap/Internal.hs

@@ -2785,26 +2785,26 @@ mapEitherWithKey f0 t0 = toPair $ go f0 t0
 split :: Key -> IntMap a -> (IntMap a, IntMap a)
 split k t =
   case t of
-    Bin _ m l r
+    Bin p m l r
       | m < 0 ->
         if k >= 0 -- handle negative numbers.
         then
           case go k l of
             (lt :*: gt) ->
-              let !lt' = union r lt
+              let !lt' = bin p m lt r
               in (lt', gt)
         else
           case go k r of
             (lt :*: gt) ->
-              let !gt' = union gt l
+              let !gt' = bin p m l gt
               in (lt, gt')
     _ -> case go k t of
           (lt :*: gt) -> (lt, gt)
   where
     go k' t'@(Bin p m l r)
       | nomatch k' p m = if k' > p then t' :*: Nil else Nil :*: t'
-      | zero k' m = case go k' l of (lt :*: gt) -> lt :*: union gt r
-      | otherwise = case go k' r of (lt :*: gt) -> union l lt :*: gt
+      | zero k' m = case go k' l of (lt :*: gt) -> lt :*: bin p m gt r
+      | otherwise = case go k' r of (lt :*: gt) -> bin p m l lt :*: gt
     go k' t'@(Tip ky _)
       | k' > ky   = (t' :*: Nil)
       | k' < ky   = (Nil :*: t')
@@ -2835,11 +2835,11 @@ splitLookup :: Key -> IntMap a -> (IntMap a, Maybe a, IntMap a)
 splitLookup k t =
   case
     case t of
-      Bin _ m l r
+      Bin p m l r
         | m < 0 ->
           if k >= 0 -- handle negative numbers.
-          then mapLT (union r) (go k l)
-          else mapGT (`union` l) (go k r)
+          then mapLT (flip (bin p m) r) (go k l)
+          else mapGT (bin p m l) (go k r)
       _ -> go k t
   of SplitLookup lt fnd gt -> (lt, fnd, gt)
   where
@@ -2848,8 +2848,8 @@ splitLookup k t =
           if k' > p
           then SplitLookup t' Nothing Nil
           else SplitLookup Nil Nothing t'
-      | zero k' m = mapGT (`union` r) (go k' l)
-      | otherwise = mapLT (union l) (go k' r)
+      | zero k' m = mapGT (flip (bin p m) r) (go k' l)
+      | otherwise = mapLT (bin p m l) (go k' r)
     go k' t'@(Tip ky y)
       | k' > ky   = SplitLookup t'  Nothing  Nil
       | k' < ky   = SplitLookup Nil Nothing  t'

containers/src/Data/IntSet/Internal.hs

@@ -884,23 +884,26 @@ spanAntitone predicate t =
 split :: Key -> IntSet -> (IntSet,IntSet)
 split x t =
   case t of
-      Bin _ m l r
-          | m < 0 -> if x >= 0  -- handle negative numbers.
-                     then case go x l of (lt :*: gt) -> let !lt' = union lt r
-                                                        in (lt', gt)
-                     else case go x r of (lt :*: gt) -> let !gt' = union gt l
-                                                        in (lt, gt')
-      _ -> case go x t of
+    Bin p m l r
+      | m < 0 ->
+        if x >= 0  -- handle negative numbers.
+        then
+          case go x l of
+            (lt :*: gt) ->
+              let !lt' = bin p m lt r
+              in (lt', gt)
+        else
+          case go x r of
+            (lt :*: gt) ->
+              let !gt' = bin p m l gt
+              in (lt, gt')
+    _ -> case go x t of
           (lt :*: gt) -> (lt, gt)
   where
     go !x' t'@(Bin p m l r)
-        | match x' p m = if zero x' m
-                         then case go x' l of
-                             (lt :*: gt) -> lt :*: union gt r
-                         else case go x' r of
-                             (lt :*: gt) -> union lt l :*: gt
-        | otherwise   = if x' < p then (Nil :*: t')
-                        else (t' :*: Nil)
+        | nomatch x' p m = if x' < p then (Nil :*: t') else (t' :*: Nil)
+        | zero x' m      = case go x' l of (lt :*: gt) -> lt :*: bin p m gt r
+        | otherwise      = case go x' r of (lt :*: gt) -> bin p m l lt :*: gt
     go x' t'@(Tip kx' bm)
         | kx' > x'          = (Nil :*: t')
           -- equivalent to kx' > prefixOf x'
@@ -915,22 +918,33 @@ split x t =
 splitMember :: Key -> IntSet -> (IntSet,Bool,IntSet)
 splitMember x t =
   case t of
-      Bin _ m l r | m < 0 -> if x >= 0
-                             then case go x l of
-                                 (lt, fnd, gt) -> let !lt' = union lt r
-                                                  in (lt', fnd, gt)
-                             else case go x r of
-                                 (lt, fnd, gt) -> let !gt' = union gt l
-                                                  in (lt, fnd, gt')
-      _ -> go x t
+    Bin p m l r
+      | m < 0 ->
+        if x >= 0 -- handle negative numbers.
+        then
+          case go x l of
+            (lt, fnd, gt) ->
+              let !lt' = bin p m lt r
+              in (lt', fnd, gt)
+        else
+          case go x r of
+            (lt, fnd, gt) ->
+              let !gt' = bin p m l gt
+              in (lt, fnd, gt')
+    _ -> go x t
   where
     go x' t'@(Bin p m l r)
-        | match x' p m = if zero x' m
-                         then case go x' l of
-                             (lt, fnd, gt) -> (lt, fnd, union gt r)
-                         else case go x' r of
-                             (lt, fnd, gt) -> (union lt l, fnd, gt)
-        | otherwise   = if x' < p then (Nil, False, t') else (t', False, Nil)
+        | nomatch x' p m = if x' < p then (Nil, False, t') else (t', False, Nil)
+        | zero x' m =
+          case go x' l of
+            (lt, fnd, gt) ->
+              let !gt' = bin p m gt r
+              in (lt, fnd, gt')
+        | otherwise =
+          case go x' r of
+            (lt, fnd, gt) ->
+              let !lt' = bin p m l lt
+              in (lt', fnd, gt)
     go x' t'@(Tip kx' bm)
         | kx' > x'          = (Nil, False, t')
           -- equivalent to kx' > prefixOf x'