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

-- |
-- Copyright:   (c) 2024 Bodigrim
-- License:     BSD3
module Data.List.Infinite.Set (
  Set,
  empty,
  member,
  insert,
) where

-- | Okasaki red-black tree.
data Set a
  = Empty
  | Red !(Set a) !a !(Set a)
  | Black !(Set a) !a !(Set a)

empty :: Set a
empty :: forall a. Set a
empty = Set a
forall a. Set a
Empty

member :: (a -> a -> Ordering) -> a -> Set a -> Bool
member :: forall a. (a -> a -> Ordering) -> a -> Set a -> Bool
member a -> a -> Ordering
cmp = a -> Set a -> Bool
member'
  where
    member' :: a -> Set a -> Bool
member' !a
x = Set a -> Bool
go
      where
        go :: Set a -> Bool
go = \case
          Set a
Empty -> Bool
False
          Red Set a
left a
center Set a
right -> Set a -> a -> Set a -> Bool
whereToGo Set a
left a
center Set a
right
          Black Set a
left a
center Set a
right -> Set a -> a -> Set a -> Bool
whereToGo Set a
left a
center Set a
right

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

insert :: (a -> a -> Ordering) -> a -> Set a -> Set a
insert :: forall a. (a -> a -> Ordering) -> a -> Set a -> Set a
insert a -> a -> Ordering
cmp = a -> Set a -> Set a
insert'
  where
    insert' :: a -> Set a -> Set a
insert' !a
x = Set a -> Set a
forall {a}. Set a -> Set a
blacken (Set a -> Set a) -> (Set a -> Set a) -> Set a -> Set a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set a -> Set a
go
      where
        go :: Set a -> Set a
go Set a
node = case Set a
node of
          Set a
Empty -> Set a -> a -> Set a -> Set a
forall a. Set a -> a -> Set a -> Set a
Red Set a
forall a. Set a
Empty a
x Set a
forall a. Set a
Empty
          Red Set a
left a
center Set a
right -> case a
x a -> a -> Ordering
`cmp` a
center of
            Ordering
LT -> Set a -> a -> Set a -> Set a
forall a. Set a -> a -> Set a -> Set a
Red (Set a -> Set a
go Set a
left) a
center Set a
right
            Ordering
EQ -> Set a
node
            Ordering
GT -> Set a -> a -> Set a -> Set a
forall a. Set a -> a -> Set a -> Set a
Red Set a
left a
center (Set a -> Set a
go Set a
right)
          Black Set a
left a
center Set a
right -> case a
x a -> a -> Ordering
`cmp` a
center of
            Ordering
LT -> Set a -> a -> Set a -> Set a
forall a. Set a -> a -> Set a -> Set a
balanceLeft (Set a -> Set a
go Set a
left) a
center Set a
right
            Ordering
EQ -> Set a
node
            Ordering
GT -> Set a -> a -> Set a -> Set a
forall a. Set a -> a -> Set a -> Set a
balanceRight Set a
left a
center (Set a -> Set a
go Set a
right)

    blacken :: Set a -> Set a
blacken Set a
node = case Set a
node of
      Set a
Empty -> Set a
forall a. Set a
Empty
      Red Set a
left a
center Set a
right -> Set a -> a -> Set a -> Set a
forall a. Set a -> a -> Set a -> Set a
Black Set a
left a
center Set a
right
      Black {} -> Set a
node
{-# INLINE insert #-}

balanceLeft :: Set a -> a -> Set a -> Set a
balanceLeft :: forall a. Set a -> a -> Set a -> Set a
balanceLeft (Red (Red Set a
a a
b Set a
c) a
d Set a
e) a
f Set a
g =
  Set a -> a -> Set a -> Set a
forall a. Set a -> a -> Set a -> Set a
Red (Set a -> a -> Set a -> Set a
forall a. Set a -> a -> Set a -> Set a
Black Set a
a a
b Set a
c) a
d (Set a -> a -> Set a -> Set a
forall a. Set a -> a -> Set a -> Set a
Black Set a
e a
f Set a
g)
balanceLeft (Red Set a
a a
b (Red Set a
c a
d Set a
e)) a
f Set a
g =
  Set a -> a -> Set a -> Set a
forall a. Set a -> a -> Set a -> Set a
Red (Set a -> a -> Set a -> Set a
forall a. Set a -> a -> Set a -> Set a
Black Set a
a a
b Set a
c) a
d (Set a -> a -> Set a -> Set a
forall a. Set a -> a -> Set a -> Set a
Black Set a
e a
f Set a
g)
balanceLeft Set a
left a
center Set a
right =
  Set a -> a -> Set a -> Set a
forall a. Set a -> a -> Set a -> Set a
Black Set a
left a
center Set a
right
{-# INLINE balanceLeft #-}

balanceRight :: Set a -> a -> Set a -> Set a
balanceRight :: forall a. Set a -> a -> Set a -> Set a
balanceRight Set a
a a
b (Red (Red Set a
c a
d Set a
e) a
f Set a
g) =
  Set a -> a -> Set a -> Set a
forall a. Set a -> a -> Set a -> Set a
Red (Set a -> a -> Set a -> Set a
forall a. Set a -> a -> Set a -> Set a
Black Set a
a a
b Set a
c) a
d (Set a -> a -> Set a -> Set a
forall a. Set a -> a -> Set a -> Set a
Black Set a
e a
f Set a
g)
balanceRight Set a
a a
b (Red Set a
c a
d (Red Set a
e a
f Set a
g)) =
  Set a -> a -> Set a -> Set a
forall a. Set a -> a -> Set a -> Set a
Red (Set a -> a -> Set a -> Set a
forall a. Set a -> a -> Set a -> Set a
Black Set a
a a
b Set a
c) a
d (Set a -> a -> Set a -> Set a
forall a. Set a -> a -> Set a -> Set a
Black Set a
e a
f Set a
g)
balanceRight Set a
left a
center Set a
right =
  Set a -> a -> Set a -> Set a
forall a. Set a -> a -> Set a -> Set a
Black Set a
left a
center Set a
right
{-# INLINE balanceRight #-}