Worker/wrapper can lead to sharing failure

nofib turned up a serious performance bug in the implementation of insert in containers-0.5.10.1. The function was defined thus:

origInsert :: Ord k => k -> a -> Map k a -> Map k a
origInsert = go
  where
    go :: Ord k => k -> a -> Map k a -> Map k a
    go !kx x Tip = singleton kx x
    go !kx x t@(Bin sz ky y l r) =
        case compare kx ky of
            LT | l' `ptrEq` l -> t
               | otherwise -> balanceL ky y l' r
               where !l' = go kx x l
            GT | r' `ptrEq` r -> t
               | otherwise -> balanceR ky y l r'
               where !r' = go kx x r
            EQ | kx `ptrEq` ky && x `ptrEq` y -> t
               | otherwise -> Bin sz kx x l r

{-# INLINABLE origInsert #-}

When this specializes to Int keys (or any other "unboxable" ones, including tuples), worker/wrapper botches the job:

Rec {
-- RHS size: {terms: 102, types: 65, coercions: 0}
$w$sgo
  :: forall a_a7M6.
     Int# -> a_a7M6 -> Map Int a_a7M6 -> Map Int a_a7M6
$w$sgo =
  \ (@ a_a7M6)
    (ww_s8oI :: Int#)
    (w_s8oE :: a_a7M6)
    (w1_s8oF :: Map Int a_a7M6) ->
    let {
      kx_X7KQ :: Int
      kx_X7KQ = I# ww_s8oI } in
    case w1_s8oF of wild_Xg {

[...]

origInsertInt_$sgo
  :: forall a_a7M6. Int -> a_a7M6 -> Map Int a_a7M6 -> Map Int a_a7M6
origInsertInt_$sgo =
  \ (@ a_a7M6)
    (w_s8oD :: Int)
    (w1_s8oE :: a_a7M6)
    (w2_s8oF :: Map Int a_a7M6) ->
    case w_s8oD of _ { I# ww1_s8oI -> $w$sgo ww1_s8oI w1_s8oE w2_s8oF }

The wrapper opens the box, throws it away, and passes the contents to the worker, which immediately builds a new box with exactly the same contents. This prevents the pointer equality tests from succeeding for these types, and it also turns out to cause quite a lot of extra allocation for some types (leading to the severe nofib regression).

One could reasonably argue that the code above is a bit complicated, and that GHC could be forgiven for failing to realize that the box should be saved. Unfortunately, a straightforward change that would seem to make this clear does not in fact convince GHC:

myInsert :: Ord k => k -> a -> Map k a -> Map k a
myInsert kx0 = go kx0
  where
    go !kx x Tip = singleton kx0 x
    go !kx x t@(Bin sz ky y l r) =
        case compare kx ky of
            LT | l' `ptrEq` l -> t
               | otherwise -> balanceL ky y l' r
               where !l' = go kx x l
            GT | r' `ptrEq` r -> t
               | otherwise -> balanceR ky y l r'
               where !r' = go kx x r
            EQ | kx0 `ptrEq` ky && x `ptrEq` y -> t
               | otherwise -> Bin sz kx0 x l r

{-# INLINABLE myInsert #-}

does exactly the same thing. The only simple way I found to avoid that is to remove the bang patterns, which really ''shouldn't'' work, but does. This, however, is prohibited by the desired semantics—I believe we want to be strict in the key even if comparison is not. In any case, that really shouldn't be causing trouble and it is. The only fix I've found thus far is truly disgusting, and seems to work at least partly by mistake:

insert :: Ord k => k -> a -> Map k a -> Map k a
insert kx0 = go kx0 kx0
  where
    go :: Ord k => k -> k -> a -> Map k a -> Map k a
    go orig !kx x Tip = singleton (lazy orig) x
    go orig !kx x t@(Bin sz ky y l r) =
        case compare kx ky of
            LT | l' `ptrEq` l -> t
               | otherwise -> balanceL ky y l' r
               where !l' = go orig kx x l
            GT | r' `ptrEq` r -> t
               | otherwise -> balanceR ky y l r'
               where !r' = go orig kx x r
            EQ | x `ptrEq` y && (lazy orig `seq` (orig `ptrEq` ky)) -> t
               | otherwise -> Bin sz (lazy orig) x l r

{-# INLINABLE insert #-}

We would also like to be able to experiment with an implementation that uses CPS (recursive join points today!) rather than pointer equality tests for the internal nodes, leaving pointer equality to the leaves. But I have not found any way whatsoever to avoid this W/W problem in that version.