{-# LANGUAGE LambdaCase #-}

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

data Color = Red | Black
  deriving (Int -> Color -> ShowS
[Color] -> ShowS
Color -> String
(Int -> Color -> ShowS)
-> (Color -> String) -> ([Color] -> ShowS) -> Show Color
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Color -> ShowS
showsPrec :: Int -> Color -> ShowS
$cshow :: Color -> String
show :: Color -> String
$cshowList :: [Color] -> ShowS
showList :: [Color] -> ShowS
Show)

-- | Okasaki red-black tree.
data Set a = Empty | Node !Color !(Set a) !a !(Set a)
  deriving (Int -> Set a -> ShowS
[Set a] -> ShowS
Set a -> String
(Int -> Set a -> ShowS)
-> (Set a -> String) -> ([Set a] -> ShowS) -> Show (Set a)
forall a. Show a => Int -> Set a -> ShowS
forall a. Show a => [Set a] -> ShowS
forall a. Show a => Set a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> Set a -> ShowS
showsPrec :: Int -> Set a -> ShowS
$cshow :: forall a. Show a => Set a -> String
show :: Set a -> String
$cshowList :: forall a. Show a => [Set a] -> ShowS
showList :: [Set a] -> ShowS
Show)

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
          Node Color
_ Set a
left a
center Set a
right -> case a -> a -> Ordering
cmp a
x 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

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 = \case
          Set a
Empty -> Color -> Set a -> a -> Set a -> Set a
forall a. Color -> Set a -> a -> Set a -> Set a
Node Color
Red Set a
forall a. Set a
Empty a
x Set a
forall a. Set a
Empty
          Node Color
color Set a
left a
center Set a
right -> case a -> a -> Ordering
cmp a
x a
center of
            Ordering
LT -> Color -> Set a -> a -> Set a -> Set a
forall a. Color -> Set a -> a -> Set a -> Set a
balance Color
color (Set a -> Set a
go Set a
left) a
center Set a
right
            Ordering
EQ -> Color -> Set a -> a -> Set a -> Set a
forall a. Color -> Set a -> a -> Set a -> Set a
Node Color
color Set a
left a
center Set a
right
            Ordering
GT -> Color -> Set a -> a -> Set a -> Set a
forall a. Color -> Set a -> a -> Set a -> Set a
balance Color
color Set a
left a
center (Set a -> Set a
go Set a
right)

    blacken :: Set a -> Set a
blacken = \case
      Set a
Empty -> Set a
forall a. Set a
Empty
      Node Color
_ Set a
left a
center Set a
right -> Color -> Set a -> a -> Set a -> Set a
forall a. Color -> Set a -> a -> Set a -> Set a
Node Color
Black Set a
left a
center Set a
right

balance :: Color -> Set a -> a -> Set a -> Set a
balance :: forall a. Color -> Set a -> a -> Set a -> Set a
balance Color
Black (Node Color
Red (Node Color
Red Set a
a a
b Set a
c) a
d Set a
e) a
f Set a
g =
  Color -> Set a -> a -> Set a -> Set a
forall a. Color -> Set a -> a -> Set a -> Set a
Node Color
Red (Color -> Set a -> a -> Set a -> Set a
forall a. Color -> Set a -> a -> Set a -> Set a
Node Color
Black Set a
a a
b Set a
c) a
d (Color -> Set a -> a -> Set a -> Set a
forall a. Color -> Set a -> a -> Set a -> Set a
Node Color
Black Set a
e a
f Set a
g)
balance Color
Black (Node Color
Red Set a
a a
b (Node Color
Red Set a
c a
d Set a
e)) a
f Set a
g =
  Color -> Set a -> a -> Set a -> Set a
forall a. Color -> Set a -> a -> Set a -> Set a
Node Color
Red (Color -> Set a -> a -> Set a -> Set a
forall a. Color -> Set a -> a -> Set a -> Set a
Node Color
Black Set a
a a
b Set a
c) a
d (Color -> Set a -> a -> Set a -> Set a
forall a. Color -> Set a -> a -> Set a -> Set a
Node Color
Black Set a
e a
f Set a
g)
balance Color
Black Set a
a a
b (Node Color
Red (Node Color
Red Set a
c a
d Set a
e) a
f Set a
g) =
  Color -> Set a -> a -> Set a -> Set a
forall a. Color -> Set a -> a -> Set a -> Set a
Node Color
Red (Color -> Set a -> a -> Set a -> Set a
forall a. Color -> Set a -> a -> Set a -> Set a
Node Color
Black Set a
a a
b Set a
c) a
d (Color -> Set a -> a -> Set a -> Set a
forall a. Color -> Set a -> a -> Set a -> Set a
Node Color
Black Set a
e a
f Set a
g)
balance Color
Black Set a
a a
b (Node Color
Red Set a
c a
d (Node Color
Red Set a
e a
f Set a
g)) =
  Color -> Set a -> a -> Set a -> Set a
forall a. Color -> Set a -> a -> Set a -> Set a
Node Color
Red (Color -> Set a -> a -> Set a -> Set a
forall a. Color -> Set a -> a -> Set a -> Set a
Node Color
Black Set a
a a
b Set a
c) a
d (Color -> Set a -> a -> Set a -> Set a
forall a. Color -> Set a -> a -> Set a -> Set a
Node Color
Black Set a
e a
f Set a
g)
balance Color
color Set a
left a
center Set a
right =
  Color -> Set a -> a -> Set a -> Set a
forall a. Color -> Set a -> a -> Set a -> Set a
Node Color
color Set a
left a
center Set a
right