Improve performance of nub{Ord,Int}On after fusion (#1206)

INLINE the FB functions and use oneShot.

Improves benchmark times on GHC 9.14:

nubInt.issue1202_distinct  -45%
nubInt.issue1202_repeat    -76%
nubOrd.issue1202_distinct   -8%
nubOrd.issue1202_repeat    -60%

Author: meooow25

containers-tests/benchmarks/ListUtils.hs

@@ -0,0 +1,56 @@
+{-# LANGUAGE BangPatterns #-}
+
+import Control.DeepSeq (rnf)
+import Control.Exception (evaluate)
+import qualified Data.Containers.ListUtils as LU
+
+import Test.Tasty.Bench (bench, bgroup, defaultMain, whnf)
+
+main :: IO ()
+main = do
+  evaluate $ rnf [xs_distinct, xs_repeat]
+  evaluate $ rnf [xss_distinct, xss_repeat]
+  defaultMain
+    [ bgroup "nubOrd"
+      [ bench "no_fusion_distinct" $
+          whnf (consumeNoFusion . LU.nubOrd) xs_distinct
+      , bench "no_fusion_repeat" $
+          whnf (consumeNoFusion . LU.nubOrd) xs_repeat
+      , bench "issue1202_distinct" $
+          whnf (consumeNoFusion . collectFrameworksDirs) xss_distinct
+      , bench "issue1202_repeat" $
+          whnf (consumeNoFusion . collectFrameworksDirs) xss_repeat
+      ]
+    , bgroup "nubInt"
+      [ bench "no_fusion_distinct" $
+          whnf (consumeNoFusion . LU.nubInt) xs_distinct
+      , bench "no_fusion_repeat" $
+          whnf (consumeNoFusion . LU.nubInt) xs_repeat
+      , bench "issue1202_distinct" $
+          whnf (consumeNoFusion . collectFrameworksDirs_nubInt) xss_distinct
+      , bench "issue1202_repeat" $
+          whnf (consumeNoFusion . collectFrameworksDirs_nubInt) xss_repeat
+      ]
+    ]
+  where
+    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]]
+
+-- Simple version of the case reported in
+-- https://github.com/haskell/containers/issues/1202
+collectFrameworksDirs :: [[Int]] -> [Int]
+collectFrameworksDirs =
+  map (*2) . LU.nubOrd . filter (/=0) . concatMap id
+{-# NOINLINE collectFrameworksDirs #-}
+
+collectFrameworksDirs_nubInt :: [[Int]] -> [Int]
+collectFrameworksDirs_nubInt =
+  map (*2) . LU.nubInt . filter (/=0) . concatMap id
+{-# NOINLINE collectFrameworksDirs_nubInt #-}
+
+consumeNoFusion :: [a] -> ()
+consumeNoFusion = foldr (\_ z -> z) ()
+{-# NOINLINE consumeNoFusion #-}

containers-tests/containers-tests.cabal

@@ -286,6 +286,14 @@ benchmark set-operations-set
   build-depends:
       benchmark-utils
 
+benchmark listutils
+  import: benchmark-deps, warnings
+  default-language: Haskell2010
+  type:           exitcode-stdio-1.0
+  hs-source-dirs: benchmarks
+  main-is:        ListUtils.hs
+  ghc-options:    -O2
+
 benchmark lookupge-intmap
   import: benchmark-deps, warnings
   default-language: Haskell2010

containers/src/Data/Containers/ListUtils.hs

@@ -33,7 +33,7 @@ import qualified Data.Set as Set
 import qualified Data.IntSet as IntSet
 import Data.IntSet (IntSet)
 #ifdef __GLASGOW_HASKELL__
-import GHC.Exts ( build )
+import GHC.Exts (build, oneShot)
 #endif
 
 -- *** Ord-based nubbing ***
@@ -72,9 +72,8 @@ nubOrd = nubOrdOn id
 --
 -- @since 0.6.0.1
 nubOrdOn :: Ord b => (a -> b) -> [a] -> [a]
--- For some reason we need to write an explicit lambda here to allow this
--- to inline when only applied to a function.
-nubOrdOn f = \xs -> nubOrdOnExcluding f Set.empty xs
+nubOrdOn f =  -- Inline with 1 arg
+  \xs -> nubOrdOnExcluding f Set.empty xs
 {-# INLINE nubOrdOn #-}
 
 -- Splitting nubOrdOn like this means that we don't have to worry about
@@ -110,11 +109,13 @@ nubOrdOnFB :: Ord b
            -> (Set b -> r)
            -> Set b
            -> r
-nubOrdOnFB f c x r s
-  | fx `Set.member` s = r s
-  | otherwise = x `c` r (Set.insert fx s)
-  where !fx = f x
-{-# INLINABLE [0] nubOrdOnFB #-}
+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))
+{-# INLINE [0] nubOrdOnFB #-}
 
 constNubOn :: a -> b -> a
 constNubOn x _ = x
@@ -153,9 +154,8 @@ nubInt = nubIntOn id
 --
 -- @since 0.6.0.1
 nubIntOn :: (a -> Int) -> [a] -> [a]
--- For some reason we need to write an explicit lambda here to allow this
--- to inline when only applied to a function.
-nubIntOn f = \xs -> nubIntOnExcluding f IntSet.empty xs
+nubIntOn f =  -- Inline with 1 arg
+  \xs -> nubIntOnExcluding f IntSet.empty xs
 {-# INLINE nubIntOn #-}
 
 -- Splitting nubIntOn like this means that we don't have to worry about
@@ -189,9 +189,11 @@ nubIntOnFB :: (a -> Int)
            -> (IntSet -> r)
            -> IntSet
            -> r
-nubIntOnFB f c x r s
-  | fx `IntSet.member` s = r s
-  | otherwise = x `c` r (IntSet.insert fx s)
-  where !fx = f x
-{-# INLINABLE [0] nubIntOnFB #-}
+nubIntOnFB f c =  -- Inline with 2 args
+  \x r -> oneShot (\s ->
+    let !y = f x
+    in if y `IntSet.member` s
+       then r s
+       else x `c` r (IntSet.insert y s))
+{-# INLINE [0] nubIntOnFB #-}
 #endif