Speedup 'nubBy' by doing more things to intermediate array in-place

bench/Data/Array.purs

@@ -3,20 +3,25 @@ module Bench.Data.Array where
 import Prelude
 
 import Data.Array as Array
+import Data.Traversable (sequence_)
 import Data.Maybe (Maybe(..))
 import Effect (Effect)
 import Effect.Console (log)
 import Performance.Minibench (benchWith)
 
 benchArray :: Effect Unit
 benchArray = do
-  log "mapMaybe"
-  log "---------------"
-  benchMapMaybe
+  sequence_ $ Array.intersperse (log "")
+    [ benchMapMaybe
+    , benchNubBy
+    ]
+
 
   where
 
   benchMapMaybe = do
+    log "mapMaybe"
+    log "---------------"
     let shortNats = Array.range 0 100
         longNats = Array.range 0 10000
         onlyEven x = if x `mod` 2 == 0 then Just x else Nothing
@@ -26,3 +31,16 @@ benchArray = do
 
     log $ "mapMaybe (" <> show (Array.length longNats) <> ")"
     benchWith 100 \_ -> Array.mapMaybe onlyEven longNats
+
+  benchNubBy = do
+    log "nubBy"
+    log "---------------"
+    let shortNats = Array.range 0 100
+        longNats = Array.range 0 10000
+        mod3Cmp x y = compare (x `mod` 3) (y `mod` 3)
+
+    log $ "nubBy (" <> show (Array.length shortNats) <> ")"
+    benchWith 1000 \_ -> Array.nubBy mod3Cmp shortNats
+
+    log $ "nubBy (" <> show (Array.length longNats) <> ")"
+    benchWith 100 \_ -> Array.nubBy mod3Cmp longNats

src/Data/Array.purs

@@ -1031,17 +1031,20 @@ nubEq = nubByEq eq
 nubBy :: forall a. (a -> a -> Ordering) -> Array a -> Array a
 nubBy comp xs = case head indexedAndSorted of
   Nothing -> []
-  Just x -> map snd $ sortWith fst $ ST.run do
+  Just x -> map snd $ ST.run do
      -- TODO: use NonEmptyArrays here to avoid partial functions
      result <- STA.unsafeThaw $ singleton x
      ST.foreach indexedAndSorted \pair@(Tuple i x') -> do
        lst <- snd <<< unsafePartial (fromJust <<< last) <$> STA.unsafeFreeze result
        when (comp lst x' /= EQ) $ void $ STA.push pair result
+     _ <- STA.sortWith fst result
      STA.unsafeFreeze result
   where
   indexedAndSorted :: Array (Tuple Int a)
-  indexedAndSorted = sortBy (\x y -> comp (snd x) (snd y))
-                            (mapWithIndex Tuple xs)
+  indexedAndSorted = ST.run do
+    arr <- STA.unsafeThaw (mapWithIndex Tuple xs)
+    _ <- STA.sortBy (\x y -> comp (snd x) (snd y)) arr
+    STA.unsafeFreeze arr
 
 -- | Remove the duplicates from an array, where element equality is determined
 -- | by the specified equivalence relation, creating a new array.