Make intersections much faster

[treeowl]

We currently calculate intersections completely naively:

intersectionWithKey f a b = foldlWithKey' go empty a
  where
    go m k v = case lookup k b of
                 Just w -> insert k (f k v w) m
                 _      -> m

Surely we can do much better by taking advantage of the structure of a HashMap. There are four easy cases:

intersectionWithKey :: (Eq k, Hashable k) => (k -> v1 -> v2 -> v3)
                    -> HashMap k v1 -> HashMap k v2 -> HashMap k v3
intersectionWithKey _ Empty !_ = Empty
intersectionWithKey _ _ Empty = Empty
intersectionWithKey f (Leaf h (L k v)) bs =
  case lookup' h k bs of
    Nothing -> Empty
    Just vb -> Leaf h (L k (f k v vb))
intersectionWithKey f as (Leaf h (L k v)) =
  case lookup' h k as of
    Nothing -> Empty
    Just va -> Leaf h (L k (f k va v))

We can deal with Collision naively, or fuss about a bit to make it efficient: perform a sort of hash-only lookup in the other map to obtain a Leaf or a Collision, then compare all the pairs.

For BitmapIndexed and Full, we can first combine the bitmaps to figure out which children we even have to bother with. Pair those up and take intersections recursively. Then filter out any empties and calculate the final bitmap (I imagine this will look a lot like the nasty mapMaybeWithKey code).

[treeowl]

For intersection, we can even try to preserve equal subtrees.

[treeowl]

I've started to dig into this a bit. If either argument is a Leaf or Empty, then we can do something pretty trivial (though it gets a tad annoying if we want to be sure to use the exact key from the first argument, rather than either of two that compare EQ). We can try to get a bit clever about Collisions, or assume they're small and rare and just perform a separate lookup for each key. That leaves cases where both trees are BitmapIndexed or Full. I believe we can handle these the same, treating Full as BitmapIndexed with fullNodeMask. We create a temporary array to hold the results of recursive calls; this array has an effective bitmap produced by applying .&. to the bitmaps of the given nodes, and has the popCount of that as its length. There'll be some bit fiddling to match up the two inputs with the output. Then we get to filter out the Emptys, deal with the special case that we end up with a single Leaf or Collision result, and otherwise call bitmapIndexedOrFull to create the new node.

[treeowl]

@emilypi, do you think you could give this a shot?

[treeowl]

Unless @sjakobi gets to it soon....

[sjakobi]

Yeah, I guess I'm currently a bit too spread out over multiple projects to focus on this issue. Can't promise when I'll return to this.

[sjakobi]

For BitmapIndexed and Full, we can first combine the bitmaps to figure out which children we even have to bother with. Pair those up and take intersections recursively. Then filter out any empties and calculate the final bitmap (I imagine this will look a lot like the nasty mapMaybeWithKey code).

As a start, we could even continue using folds for these, but be a bit smarter about inserting things into the result map. Ideally, we

• re-use the existing hashes instead of computing new ones.
• use an unsafe, modifying insertion (Speed up difference* and intersection* with unsafeInsert #372)
• avoid key equality checks in the style of insertNewKey

The result wouldn't be as fast as @treeowl's suggestion, but it would possibly be easier to implement and already much better than the current implementation,