Fixed the bugs in restrictKeys/withoutKeys

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