[commit: packages/containers] master, revert-408-bugfix_394: Fixed the bugs in restrictKeys/withoutKeys (acc1581)

[git at git.haskell.org]

Repository : ssh://git@git.haskell.org/containers

On branches: master,revert-408-bugfix_394
Link       : http://git.haskell.org/packages/containers.git/commitdiff/acc158123dde87ca32603d67a4c9a348fa94eb29

>---------------------------------------------------------------

commit acc158123dde87ca32603d67a4c9a348fa94eb29
Author: wren romano <wren at community.haskell.org>
Date:   Wed Feb 8 19:04:28 2017 -0800

    Fixed the bugs in restrictKeys/withoutKeys


>---------------------------------------------------------------

acc158123dde87ca32603d67a4c9a348fa94eb29
 Data/IntMap/Internal.hs | 56 ++++++++++++++++++++++++++++++-------------------
 1 file changed, 34 insertions(+), 22 deletions(-)

diff --git a/Data/IntMap/Internal.hs b/Data/IntMap/Internal.hs
index 9f2e35d..262b1a9 100644
--- a/Data/IntMap/Internal.hs
+++ b/Data/IntMap/Internal.hs
@@ -1057,18 +1057,23 @@ withoutKeys t1@(Bin p1 m1 l1 r1) t2@(IntSet.Bin p2 m2 l2 r2)
         | nomatch p1 p2 m2  = t1
         | zero p1 m2        = withoutKeys t1 l2
         | otherwise         = withoutKeys t1 r2
--- TODO(wrengr): should we inline the top-level 'withoutBM' here?
-withoutKeys t1@(Bin _ _ _ _) (IntSet.Tip kx' bm') = withoutBM kx' bm' t1
+withoutKeys t1@(Bin _ _ _ _) (IntSet.Tip kx' bm') =
+    withoutBM t1 kx' bm' (IntSet.suffixBitMask + 1)
     where
-    withoutBM !kx !bm t@(Bin p m l r)
-        | nomatch kx p m = t
-        | zero kx m      = binCheckLeft p m (withoutBM kx bm l) r
-        | otherwise      = binCheckRight p m l (withoutBM kx bm r)
-    withoutBM kx bm t@(Tip ky _)
-        -- TODO(wrengr): should we inline 'IntSet.member' here?
-        | ky `IntSet.member` IntSet.Tip kx bm = Nil
-        | otherwise = t
-    withoutBM _ _ Nil = Nil
+    -- TODO(wrengr): this is still pretty naive. It could be improved by restricting @t@ on the recursive calls, so that the 'delete' in the basis case is faster. As is, this is linear in the size of the IntSet (as opposed to the previous version which was linear in the size of the IntMap; we want /O(n+m)/ at worst, just like for 'intersection').
+    withoutBM t !prefix !_ 0 = delete prefix t
+    withoutBM t prefix bmask bits =
+        case intFromNat (natFromInt bits `shiftRL` 1) of
+        bits2
+          | bmask .&. (shiftLL 1 bits2 - 1) == 0 ->
+              withoutBM t (prefix + bits2) (shiftRL bmask bits2) bits2
+          | shiftRL bmask bits2 .&. (shiftLL 1 bits2 - 1) == 0 ->
+              withoutBM t prefix bmask bits2
+          | otherwise ->
+              -- withoutKeys t (bin prefix bits2 _ _)
+              withoutBM
+                (withoutBM t (prefix + bits2) (shiftRL bmask bits2) bits2)
+                prefix bmask bits2
 withoutKeys t1@(Bin _ _ _ _) IntSet.Nil = t1
 withoutKeys t1@(Tip k1 _) t2
     | k1 `IntSet.member` t2 = Nil
@@ -1109,18 +1114,25 @@ restrictKeys t1@(Bin p1 m1 l1 r1) t2@(IntSet.Bin p2 m2 l2 r2)
         | nomatch p1 p2 m2  = Nil
         | zero p1 m2        = restrictKeys t1 l2
         | otherwise         = restrictKeys t1 r2
--- TODO(wrengr): should we inline the top-level 'restrictBM' here?
-restrictKeys t1@(Bin _ _ _ _) (IntSet.Tip kx' bm') = restrictBM kx' bm' t1
+restrictKeys t1@(Bin _ _ _ _) (IntSet.Tip kx' bm') =
+    restrictBM t1 kx' bm' (IntSet.suffixBitMask + 1)
     where
-    restrictBM !kx !bm (Bin p1 m1 l1 r1)
-        | nomatch kx p1 m1 = Nil
-        | zero kx m1       = restrictBM kx bm l1
-        | otherwise        = restrictBM kx bm r1
-    restrictBM kx bm t@(Tip ky _)
-        -- TODO(wrengr): should we inline 'IntSet.member' here?
-        | ky `IntSet.member` IntSet.Tip kx bm = t
-        | otherwise = Nil
-    restrictBM _ _ Nil = Nil
+    -- TODO(wrengr): this is still pretty naive. It could be improved by restricting @t@ on the recursive calls, so that the 'lookup' in the basis case is faster. As is, this is linear in the size of the IntSet (as opposed to the previous version which was linear in the size of the IntMap; we want /O(n+m)/ at worst, just like for 'intersection').
+    restrictBM t !prefix !_ 0 =
+        case lookup prefix t of
+        Nothing -> Nil
+        Just x -> Tip prefix x
+    restrictBM t prefix bmask bits =
+        case intFromNat (natFromInt bits `shiftRL` 1) of
+        bits2
+          | bmask .&. (shiftLL 1 bits2 - 1) == 0 ->
+              restrictBM t (prefix + bits2) (shiftRL bmask bits2) bits2
+          | shiftRL bmask bits2 .&. (shiftLL 1 bits2 - 1) == 0 ->
+              restrictBM t prefix bmask bits2
+          | otherwise ->
+              bin prefix bits2
+                (restrictBM t prefix bmask bits2)
+                (restrictBM t (prefix + bits2) (shiftRL bmask bits2) bits2)
 restrictKeys (Bin _ _ _ _) IntSet.Nil = Nil
 restrictKeys t1@(Tip k1 _) t2
     | k1 `IntSet.member` t2 = t1