Set Unions – add light weight pointer equality check

[HuwCampbell]

It's not unusual when building compilers to do transformations which require collecting sets over an AST – but you'll often need to do branching and unions across very similar sets, for example, one side of an if branch added 2 free variables into the set of 1000. So you're doing a union where most of the Set is actually identical, but this is currently quite slow.

This pass in my language was actually half the compilation time, just doing similar set unions (PR is a work around – but it's easy to imagine a compiler pass where this work around would not be correct).

Adding a light weight pointer equality check can actually help quite a lot in these cases.

The reason I think this is a good idea is that when we add a few items to a large set, it's very likely that subtrees will remain the same. So if we add to the left the right subtree will still match. Even if there are some rebalance operations, we'll probably still hit identical trees a lot, just under one or two more levels of indirection.

If I add some small additional benchmarks to union as such:

defaultMain
      [ bench "union_same" $ whnf (S.union s_even) s_even
      , bench "union_same_but_different" $ whnf (S.union s_even) (S.map (\x -> x - 1) $ S.map (+1) s_even) // shuffling around so it's not shared. Actual benchmark put this above with a whnf.
      , bench "union_minor_diff" $ whnf (S.union s_even) (S.insert 1 s_even)
      , bench "union" $ whnf (S.union s_even) s_odd
      ]

I'll see these results:

Before:
  union_same:       OK
    11.4 μs ± 863 ns
  union_same_but_different:   OK
    29.0 μs ± 2.0 μs
  union_minor_diff: OK
    11.3 μs ± 850 ns
  union:            OK
    77.9 μs ± 7.0 μs

After:
  union_same:       OK
    2.02 ns ± 138 ps
  union_same_but_different:   OK
    40.0 μs ± 3.4 μs
  union_minor_diff: OK
    137  ns ± 7.2 ns
  union:            OK
    77.4 μs ± 6.9 μs

After adding the initial case I believe the older double ptr equality tests are pretty much redundant apart from very similar sets with no sharing. In the best case they don't stop full spine traversals, and only really ever prevented the rebalance in link, but there even just size checks would probably be more appropriate (see one of my earlier PRs).
It think it actually just made things on the whole slower in most cases.

It was in the margins, but here's the results from only adding a case, not including the second commit:

Middle:
  union_same:       OK
    2.02 ns ± 110 ps
  union_minor_diff: OK
    149  ns ±  13 ns
  union:            OK
    78.1 μs ± 7.4 μs

[HuwCampbell]

There's a pretty reasonable question to be asked about whether this optimisation should also be applied to maps. The insert operation there does quite a dance around pointer equalities, and it's kind of just for this case.

[treeowl]

I'm having a bit of trouble understanding the exact difference here. At first glance, it looks like a simplification of some sort, but it's not immediately clear if it catches more or fewer cases. Could you walk me through it?

[meooow25]

I'm not convinced that this is a good idea. You're proposing a time optimization that applies when the two maps being union-ed happen to have identical (by pointer) subtrees. This might be happening commonly for your case, but there is no reason to expect this in general and the price to pay is a linear overhead for everybody else.

This is also completely orthogonal to the pointer checks you have removed. Those are a memory optimization, to avoid allocating a result Bin when it would be equal to the input Bin.

[HuwCampbell]

I don't particularly mind putting back the second pointer equality checks. I just thought that with the proposed change they're less likely to be hit.

But let's talk it through. The existing code and checks are actually very good at dealing with trees of different balance; and works if the first set is a super set of the second; so two sets which are Eq equal should turn out with the balance (tree structure) of the first argument, and indeed, in the end, share a pointer. The same for the left being a proper superset of the second.

So what we had before will lower allocations and be faster for union supersets, or branches which themselves are supersets.

-- These two have the same elements, but are structured differently
s1 = Set.fromList [1,4,3,2,5]
-- 3
-- +--1
-- |  +--|
-- |  +--2
-- +--4
--    +--|
--    +--5
s2 = Set.fromList [1,5,2,3,4]
-- 2
-- +--1
-- +--4
--    +--3
--    +--5
s3 = Set.union s1 s2
-- ptrEq s1 s3 should be True

But where do these come from? In what scenarios do we create two sets with the same elements and then union them? In my case, it was because I was working with two sets which share a common history and have been rejoined, and I hypothesised that for large sets it's probably the main way this will happen.

So my thought was that we should be able to see if this is actually test if this is the case and just short-circuit where we see it.

For performance in cases with no sharing, the cost is very modest; 77.9 to 78.1 microseconds. The cost of the 2 pointer checks also seems similar if not slightly higher.

But, there are occasional significant wins with 100x seen above when we can effectively short circuit out all but one path to a new element.