NubOrd.hs

#!/usr/bin/env cabal
{- cabal:
build-depends: base, containers, extra, random, tasty-bench
-}

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE ScopedTypeVariables #-}

import qualified Data.Containers.ListUtils as Containers
import qualified Data.List.Extra as Extra
import System.Random
import Test.Tasty.Bench
import qualified Data.List as List
import GHC.Exts (inline)

main :: IO ()
main = defaultMain
  [ benchIt 10 10
  , benchIt 10 30
  , benchIt 10 100
  , benchIt 100 100
  , benchIt 100 300
  , benchIt 100 1000
  , benchIt 1000 1000
  , benchIt 1000 3000
  , benchIt 1000 10000
  ]

randomList :: Int -> Int -> [Int]
randomList len upperBound =
  fst $ uniformListR len (1, upperBound) (mkStdGen 42)

nubs :: [(String, [Int] -> [Int])]
nubs =
  [ ("extra", Extra.nubOrd)
  , ("containers", Containers.nubOrd)
  , ("base", nubOrd)
  ]

benchIt :: Int -> Int -> Benchmark
benchIt upperBound len =
  bgroup (show len ++ " numbers from 1 to " ++ show upperBound)
  [ bgroup "random" $
    map (\(name, func) -> if func xs == Containers.nubOrd xs then bench name (nf func xs) else error "mismatch") nubs
  , bgroup "ascending" $
    map (\(name, func) -> if func ys == Containers.nubOrd ys then bench name (nf func ys) else error "mismatch") nubs
  , bgroup "descending" $
    map (\(name, func) -> if func zs == Containers.nubOrd zs then bench name (nf func zs) else error "mismatch") nubs
  ]
  where
    xs = randomList len upperBound
    ys = take len $ cycle [1 .. upperBound]
    zs = take len $ cycle [upperBound, upperBound - 1 .. 1]

nubOrd :: Ord a => [a] -> [a]
nubOrd = nubOrdBy compare
{-# INLINE nubOrd #-}

nubOrdBy :: (a -> a -> Ordering) -> [a] -> [a]
nubOrdBy cmp xs = foldr
  (\x acc seen -> if member cmp x seen then acc seen else x : acc (insert cmp x seen))
  (const [])
  xs
  empty
{-# INLINE nubOrdBy #-}

-------------------------------------------------------------------------------

-- | Implemented as a red-black tree, a la Okasaki.
data NubOrdSet a
  = Empty
  | NodeRed !(NubOrdSet a) !a !(NubOrdSet a)
  | NodeBlack !(NubOrdSet a) !a !(NubOrdSet a)

empty :: NubOrdSet a
empty = Empty

member :: (a -> a -> Ordering) -> a -> NubOrdSet a -> Bool
member cmp = member'
  where
    member' !x = go
      where
        go = \case
          Empty -> False
          NodeRed left center right -> chooseWay left center right
          NodeBlack left center right -> chooseWay left center right

        chooseWay left center right = case cmp x center of
            LT -> go left
            EQ -> True
            GT -> go right
{-# INLINE member #-}

insert :: (a -> a -> Ordering) -> a -> NubOrdSet a -> NubOrdSet a
insert cmp = insert'
  where
    insert' !x = blacken . go
      where
        go node = case node of
          Empty -> NodeRed Empty x Empty
          NodeRed left center right -> case cmp x center of
            LT -> NodeRed (go left) center right
            EQ -> node
            GT -> NodeRed left center (go right)
          NodeBlack left center right -> case cmp x center of
            LT -> balanceBlackLeft (go left) center right
            EQ -> node
            GT -> balanceBlackRight left center (go right)

    blacken node = case node of
      Empty -> Empty
      NodeRed left center right -> NodeBlack left center right
      NodeBlack{} -> node
{-# INLINE insert #-}

balanceBlackLeft :: NubOrdSet a -> a -> NubOrdSet a -> NubOrdSet a
balanceBlackLeft (NodeRed (NodeRed a b c) d e) f g =
  NodeRed (NodeBlack a b c) d (NodeBlack e f g)
balanceBlackLeft (NodeRed a b (NodeRed c d e)) f g =
  NodeRed (NodeBlack a b c) d (NodeBlack e f g)
balanceBlackLeft left center right =
  NodeBlack left center right
{-# INLINE balanceBlackLeft #-}

balanceBlackRight :: NubOrdSet a -> a -> NubOrdSet a -> NubOrdSet a
balanceBlackRight a b (NodeRed (NodeRed c d e) f g) =
  NodeRed (NodeBlack a b c) d (NodeBlack e f g)
balanceBlackRight a b (NodeRed c d (NodeRed e f g)) =
  NodeRed (NodeBlack a b c) d (NodeBlack e f g)
balanceBlackRight left center right =
  NodeBlack left center right
{-# INLINE balanceBlackRight #-}