Optimized IntMap's withoutKeys

haskell · Feb 8, 2017 · 0ec279b · 0ec279b
1 parent eeb3921
commit 0ec279b
Showing 1 changed file with 27 additions and 14 deletions.
diff --git a/Data/IntMap/Internal.hs b/Data/IntMap/Internal.hs
@@ -1046,26 +1046,39 @@ withoutKeys :: IntMap a -> IntSet.IntSet -> IntMap a
 withoutKeys = go
   where
     go t1@(Bin p1 m1 l1 r1) t2@(IntSet.Bin p2 m2 l2 r2)
-      | shorter m1 m2  = merge1
-      | shorter m2 m1  = merge2
+      | shorter m1 m2  = difference1
+      | shorter m2 m1  = difference2
       | p1 == p2       = bin p1 m1 (go l1 l2) (go r1 r2)
       | otherwise      = t1
       where
-        merge1 | nomatch p2 p1 m1  = t1
-               | zero p2 m1        = binCheckLeft p1 m1 (go l1 t2) r1
-               | otherwise         = binCheckRight p1 m1 l1 (go r1 t2)
-        merge2 | nomatch p1 p2 m2  = t1
-               | zero p1 m2        = bin p2 m2 (go t1 l2) Nil
-               | otherwise         = bin p2 m2 Nil (go t1 r2)
-
-    go t1'@(Bin _ _ _ _) t2'@(IntSet.Tip _ _) =
-      filterWithKey (\k _ -> k `IntSet.notMember` t2') t1'
+        difference1
+            | nomatch p2 p1 m1  = t1
+            | zero p2 m1        = binCheckLeft p1 m1 (go l1 t2) r1
+            | otherwise         = binCheckRight p1 m1 l1 (go r1 t2)
+        difference2
+            | nomatch p1 p2 m2  = t1
+            | zero p1 m2        = go t1 l2
+            | otherwise         = go t1 r2
+
+    -- TODO(wrengr): should we inline the top-level 'deleteBM' here?
+    go t1@(Bin _ _ _ _) (IntSet.Tip kx' bm') = deleteBM kx' bm' t1
+      where
+        deleteBM !kx !bm t@(Bin p m l r)
+          | nomatch kx p m = t
+          | zero kx m      = binCheckLeft p m (deleteBM kx bm l) r
+          | otherwise      = binCheckRight p m l (deleteBM kx bm r)
+        deleteBM kx bm t@(Tip ky _)
+          -- TODO(wrengr): should we inline 'IntSet.member' here?
+          | ky `IntSet.member` IntSet.Tip kx bm = Nil
+          | otherwise = t
+        deleteBM _ _ Nil = Nil
 
     go t1@(Bin _ _ _ _) IntSet.Nil = t1
 
-    go t1'@(Tip k1' _) t2'
-      | k1' `IntSet.member` t2' = Nil
-      | otherwise = t1'
+    go t1@(Tip k1 _) t2
+      | k1 `IntSet.member` t2 = Nil
+      | otherwise = t1
+
     go Nil _ = Nil