    {-# 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 #-}