differenceWithKey: Use hashOfLeafOrCollision

Author: sjakobi

Data/HashMap/Internal.hs

@@ -1939,7 +1939,26 @@ differenceWithKey f = go_differenceWithKey 0
     go_differenceWithKey _s a@(Collision hA aryA) (Leaf hB (L kB vB))
       | hA == hB = updateCollision (\vA -> f kB vA vB) hA kB aryA a
       | otherwise = a
-    go_differenceWithKey s a@(BitmapIndexed bA aryA) b@(Leaf hB _)
+    go_differenceWithKey s a@(Collision hA _) (BitmapIndexed bB aryB)
+        | bB .&. m == 0 = a
+        | otherwise =
+          case A.index# aryB (sparseIndex bB m) of
+            (# stB #) -> go_differenceWithKey (nextShift s) a stB
+      where m = mask hA s
+    go_differenceWithKey s a@(Collision hA _) (Full aryB)
+      = case A.index# aryB (index hA s) of
+          (# stB #) -> go_differenceWithKey (nextShift s) a stB
+    go_differenceWithKey s a@(BitmapIndexed bA aryA) (BitmapIndexed bB aryB)
+      = differenceWithKey_Arrays s bA aryA a bB aryB
+    go_differenceWithKey s a@(Full aryA) (BitmapIndexed bB aryB)
+      = differenceWithKey_Arrays s fullBitmap aryA a bB aryB
+    go_differenceWithKey s a@(BitmapIndexed bA aryA) (Full aryB)
+      = differenceWithKey_Arrays s bA aryA a fullBitmap aryB
+    go_differenceWithKey s a@(Full aryA) (Full aryB)
+      = differenceWithKey_Arrays s fullBitmap aryA a fullBitmap aryB
+    go_differenceWithKey _s a@(Collision hA aryA) (Collision hB aryB)
+      = differenceWithKey_Collisions f hA aryA a hB aryB
+    go_differenceWithKey s a@(BitmapIndexed bA aryA) b
       | bA .&. m == 0 = a
       | otherwise = case A.index# aryA i of
           (# !stA #) -> case go_differenceWithKey (nextShift s) stA b of
@@ -1953,36 +1972,10 @@ differenceWithKey f = go_differenceWithKey 0
                  | stA `ptrEq` stA' -> a
                  | otherwise -> BitmapIndexed bA (A.update aryA i stA')
       where
+        hB = hashOfLeafOrCollision b
         m = mask hB s
         i = sparseIndex bA m
-    go_differenceWithKey s a@(BitmapIndexed bA aryA) b@(Collision hB _)
-        | bA .&. m == 0 = a
-        | otherwise =
-            case A.index# aryA i of
-              (# !st #) -> case go_differenceWithKey (nextShift s) st b of
-                Empty | A.length aryA == 2
-                      , (# l #) <- A.index# aryA (otherOfOneOrZero i)
-                      , isLeafOrCollision l
-                      -> l
-                      | otherwise
-                      -> BitmapIndexed (bA .&. complement m) (A.delete aryA i)
-                st' | isLeafOrCollision st' && A.length aryA == 1 -> st'
-                    | st `ptrEq` st' -> a
-                    | otherwise -> BitmapIndexed bA (A.update aryA i st')
-      where
-        m = mask hB s
-        i = sparseIndex bA m
-    go_differenceWithKey s a@(Full aryA) b@(Leaf hB _)
-      = case A.index# aryA i of
-          (# !stA #) -> case go_differenceWithKey (nextShift s) stA b of
-            Empty ->
-                let aryA' = A.delete aryA i
-                    bm    = fullBitmap .&. complement (1 `unsafeShiftL` i)
-                in BitmapIndexed bm aryA'
-            stA' | stA `ptrEq` stA' -> a
-                 | otherwise -> Full (updateFullArray aryA i stA')
-      where i = index hB s
-    go_differenceWithKey s a@(Full aryA) b@(Collision hB _)
+    go_differenceWithKey s a@(Full aryA) b
       = case A.index# aryA i of
           (# !stA #) -> case go_differenceWithKey (nextShift s) stA b of
             Empty ->
@@ -1991,26 +1984,9 @@ differenceWithKey f = go_differenceWithKey 0
                 in BitmapIndexed bm aryA'
             stA' | stA `ptrEq` stA' -> a
                  | otherwise -> Full (updateFullArray aryA i stA')
-      where i = index hB s
-    go_differenceWithKey s a@(Collision hA _) (BitmapIndexed bB aryB)
-        | bB .&. m == 0 = a
-        | otherwise =
-          case A.index# aryB (sparseIndex bB m) of
-            (# stB #) -> go_differenceWithKey (nextShift s) a stB
-      where m = mask hA s
-    go_differenceWithKey s a@(Collision hA _) (Full aryB)
-      = case A.index# aryB (index hA s) of
-          (# stB #) -> go_differenceWithKey (nextShift s) a stB
-    go_differenceWithKey s a@(BitmapIndexed bA aryA) (BitmapIndexed bB aryB)
-      = differenceWithKey_Arrays s bA aryA a bB aryB
-    go_differenceWithKey s a@(Full aryA) (BitmapIndexed bB aryB)
-      = differenceWithKey_Arrays s fullBitmap aryA a bB aryB
-    go_differenceWithKey s a@(BitmapIndexed bA aryA) (Full aryB)
-      = differenceWithKey_Arrays s bA aryA a fullBitmap aryB
-    go_differenceWithKey s a@(Full aryA) (Full aryB)
-      = differenceWithKey_Arrays s fullBitmap aryA a fullBitmap aryB
-    go_differenceWithKey _s a@(Collision hA aryA) (Collision hB aryB)
-      = differenceWithKey_Collisions f hA aryA a hB aryB
+      where
+        hB = hashOfLeafOrCollision b
+        i = index hB s
 
     differenceWithKey_Arrays !s !bA !aryA tA !bB !aryB
       | bA .&. bB == 0 = tA