-- This is an internal module with a naive set implementation,
-- solely for the purposes of `Data.List.{,NonEmpty.}nubOrd{,By}`.

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

module Data.List.NubOrdSet (
  NubOrdSet,
  empty,
  member,
  insert,
) where

import GHC.Internal.Data.Bool (Bool(..))
import GHC.Internal.Data.Function ((.))
import GHC.Internal.Data.Ord (Ordering(..))

-- | 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 :: forall a. NubOrdSet a
empty = NubOrdSet a
forall a. NubOrdSet a
Empty

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

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

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

    blacken :: NubOrdSet a -> NubOrdSet a
blacken NubOrdSet a
node = case NubOrdSet a
node of
      NubOrdSet a
Empty -> NubOrdSet a
forall a. NubOrdSet a
Empty
      NodeRed NubOrdSet a
left a
center NubOrdSet a
right -> NubOrdSet a -> a -> NubOrdSet a -> NubOrdSet a
forall a. NubOrdSet a -> a -> NubOrdSet a -> NubOrdSet a
NodeBlack NubOrdSet a
left a
center NubOrdSet a
right
      NodeBlack{} -> NubOrdSet a
node
{-# INLINE insert #-}

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

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