Use Set.alterF in nubOrd (#1228)

Compared to the previous implementation of `member` followed by
`insert`, `alterF` uses `memberIndex` and `insertAt`. This avoids all
the element comparisons in `insert`, at the cost of `memberIndex` being
a little more expensive than `member`. I expect this change to be an
improvement in typical use cases.

Author: meooow25

containers-tests/benchmarks/ListUtils.hs

@@ -10,21 +10,26 @@ main :: IO ()
 main = do
   evaluate $ rnf [xs_distinct, xs_repeat]
   evaluate $ rnf [xss_distinct, xss_repeat]
+  evaluate $ rnf [strings_distinct, strings_repeat]
   defaultMain
     [ bgroup "nubOrd"
-      [ bench "no_fusion_distinct" $
+      [ bench "noFusion_distinct" $
           whnf (consumeNoFusion . LU.nubOrd) xs_distinct
-      , bench "no_fusion_repeat" $
+      , bench "noFusion_repeat" $
           whnf (consumeNoFusion . LU.nubOrd) xs_repeat
       , bench "issue1202_distinct" $
           whnf (consumeNoFusion . collectFrameworksDirs) xss_distinct
       , bench "issue1202_repeat" $
           whnf (consumeNoFusion . collectFrameworksDirs) xss_repeat
+      , bench "strings_noFusion_distinct" $
+          whnf (consumeNoFusion . LU.nubOrd) strings_distinct
+      , bench "strings_noFusion_repeat" $
+          whnf (consumeNoFusion . LU.nubOrd) strings_repeat
       ]
     , bgroup "nubInt"
-      [ bench "no_fusion_distinct" $
+      [ bench "noFusion_distinct" $
           whnf (consumeNoFusion . LU.nubInt) xs_distinct
-      , bench "no_fusion_repeat" $
+      , bench "noFusion_repeat" $
           whnf (consumeNoFusion . LU.nubInt) xs_repeat
       , bench "issue1202_distinct" $
           whnf (consumeNoFusion . collectFrameworksDirs_nubInt) xss_distinct
@@ -33,12 +38,17 @@ main = do
       ]
     ]
   where
-    bound = 1000 :: Int
+    !bound = 1000 :: Int
+
     xs_distinct = [1..bound]
     xs_repeat = replicate bound 1 :: [Int]
     xss_distinct = [[i] | i <- [1..bound]]
     xss_repeat = replicate bound [1] :: [[Int]]
 
+    -- Use strings as an example of expensive Ord comparisons.
+    strings_distinct = map show xs_distinct
+    strings_repeat = map show xs_repeat
+
 -- Simple version of the case reported in
 -- https://github.com/haskell/containers/issues/1202
 collectFrameworksDirs :: [[Int]] -> [Int]

containers-tests/containers-tests.cabal

@@ -286,7 +286,7 @@ benchmark set-operations-set
   build-depends:
       benchmark-utils
 
-benchmark listutils
+benchmark listutils-benchmarks
   import: benchmark-deps, warnings
   default-language: Haskell2010
   type:           exitcode-stdio-1.0

containers/src/Data/Containers/ListUtils.hs

@@ -82,11 +82,15 @@ nubOrdOnExcluding :: Ord b => (a -> b) -> Set b -> [a] -> [a]
 nubOrdOnExcluding f = go
   where
     go _ [] = []
-    go s (x:xs)
-      | fx `Set.member` s = go s xs
-      | otherwise = x : go (Set.insert fx s) xs
+    go s (x:xs) = case tryInsertSet fx s of
+      Nothing -> go s xs
+      Just s' -> x : go s' xs
       where !fx = f x
 
+tryInsertSet :: Ord a => a -> Set a -> Maybe (Set a)
+tryInsertSet = Set.alterF (\found -> if found then Nothing else Just True)
+{-# INLINE tryInsertSet #-}
+
 #ifdef __GLASGOW_HASKELL__
 -- We want this inlinable to specialize to the necessary Ord instance.
 {-# INLINABLE [1] nubOrdOnExcluding #-}
@@ -112,9 +116,9 @@ nubOrdOnFB :: Ord b
 nubOrdOnFB f c =  -- Inline with 2 args
   \x r -> oneShot (\s ->
     let !y = f x
-    in if y `Set.member` s
-       then r s
-       else x `c` r (Set.insert y s))
+    in case tryInsertSet y s of
+         Nothing -> r s
+         Just s' -> x `c` r s')
 {-# INLINE [0] nubOrdOnFB #-}
 
 constNubOn :: a -> b -> a