mini-1.5.2.0: Mini/Data/Set.hs

{-# LANGUAGE TupleSections #-}

-- | A structure containing unique elements
module Mini.Data.Set (
  -- * Type
  Set,

  -- * Construction
  empty,
  fromList,
  singleton,

  -- * Combination
  difference,
  intersection,
  union,
  unions,

  -- * Conversion
  toAscList,
  toDescList,

  -- * Modification
  delete,
  deleteMax,
  deleteMin,
  filter,
  insert,

  -- * Partition
  partition,
  split,
  splitMax,
  splitMin,

  -- * Query
  disjoint,
  isSubsetOf,
  lookupGE,
  lookupGT,
  lookupLE,
  lookupLT,
  lookupMax,
  lookupMin,
  member,
  null,
  size,

  -- * Validation
  valid,
) where

import Control.Applicative (
  liftA2,
  (<|>),
 )
import Data.Bifunctor (
  first,
 )
import Data.Bool (
  bool,
 )
import Prelude (
  Bool (
    False,
    True
  ),
  Eq,
  Foldable,
  Int,
  Maybe (
    Just,
    Nothing
  ),
  Monoid,
  Ord,
  Ordering (
    EQ,
    GT,
    LT
  ),
  Semigroup,
  Show,
  any,
  compare,
  error,
  flip,
  foldl,
  foldr,
  max,
  maybe,
  mempty,
  not,
  show,
  uncurry,
  ($),
  (&&),
  (+),
  (-),
  (.),
  (<),
  (<$>),
  (<>),
  (==),
  (>),
 )

{-
 - Type
 -}

-- | A set containing elements of type /a/, internally structured as an AVL tree
data Set a
  = -- | Empty node
    E
  | -- | Left-heavy node
    L (Set a) a (Set a)
  | -- | Balanced node
    B (Set a) a (Set a)
  | -- | Right-heavy node
    R (Set a) a (Set a)

instance (Eq a) => Eq (Set a) where
  t1 == t2 = toAscList t1 == toAscList t2

instance (Ord a) => Ord (Set a) where
  compare t1 t2 = compare (toAscList t1) (toAscList t2)

instance (Show a) => Show (Set a) where
  show = show . toAscList

instance Foldable Set where
  foldr f b = set b go go go
   where
    go l a _ _ recr = foldr f (f a recr) l

instance (Ord a) => Semigroup (Set a) where
  (<>) = union

instance (Ord a) => Monoid (Set a) where
  mempty = empty

{-
 - Primitive recursion
 -}

-- | Primitive recursion on sets
set
  :: b
  -- ^ Empty node
  -> (Set a -> a -> Set a -> b -> b -> b)
  -- ^ Left-heavy node
  -> (Set a -> a -> Set a -> b -> b -> b)
  -- ^ Balanced node
  -> (Set a -> a -> Set a -> b -> b -> b)
  -- ^ Right-heavy node
  -> Set a
  -- ^ Set
  -> b
set e _ _ _ E = e
set e f g h (L l a r) = f l a r (set e f g h l) (set e f g h r)
set e f g h (B l a r) = g l a r (set e f g h l) (set e f g h r)
set e f g h (R l a r) = h l a r (set e f g h l) (set e f g h r)

{-
 - Construction
 -}

-- | /O(1)/ The empty set
empty :: Set a
empty = E

-- | /O(n log n)/ Make a set from a list of elements
fromList :: (Ord a) => [a] -> Set a
fromList = foldl (flip insert) empty

-- | /O(1)/ Make a set with a single element
singleton :: a -> Set a
singleton a = B E a E

{-
 - Combination
 -}

-- | /O(m log n)/ Subtract a set by another
difference :: (Ord a) => Set a -> Set a -> Set a
difference = foldr delete

-- | /O(m log n)/ Intersect a set with another
intersection :: (Ord a) => Set a -> Set a -> Set a
intersection t1 t2 =
  foldr
    (\a b -> bool b (insert a b) (a `member` t2))
    empty
    t1

-- | /O(m log n)/ Unite a set with another
union :: (Ord a) => Set a -> Set a -> Set a
union = flip $ foldr insert

-- | Unite a collection of sets
unions :: (Foldable t, Ord a) => t (Set a) -> Set a
unions = foldr union empty

{-
 - Conversion
 -}

-- | /O(n)/ Turn a set into a list of elements in ascending order
toAscList :: Set a -> [a]
toAscList = foldr (:) []

-- | /O(n)/ Turn a set into a list of elements in descending order
toDescList :: Set a -> [a]
toDescList = foldl (flip (:)) []

{-
 - Modification
 -}

-- | /O(log n)/ Delete an element from a set
delete :: (Ord a) => a -> Set a -> Set a
delete a0 t = bool t (go t) (a0 `member` t)
 where
  go =
    set
      (error "Set.delete: L0")
      ( \l a r _ _ ->
          case compare a0 a of
            LT -> deleteLl l a r
            EQ -> substituteL l r
            GT -> deleteLr l a r
      )
      ( \l a r _ _ ->
          case compare a0 a of
            LT -> deleteBl l a r
            EQ -> substituteBr l r
            GT -> deleteBr l a r
      )
      ( \l a r _ _ ->
          case compare a0 a of
            LT -> deleteRl l a r
            EQ -> substituteR l r
            GT -> deleteRr l a r
      )
  deleteRl l a r =
    set
      (error "Set.delete: L1")
      ( \ll la lr _ _ ->
          case compare a0 la of
            LT -> checkLeftR (deleteLl ll la lr) a r
            EQ -> checkLeftR (substituteL ll lr) a r
            GT -> checkLeftR (deleteLr ll la lr) a r
      )
      ( \ll la lr _ _ ->
          case compare a0 la of
            LT -> R (deleteBl ll la lr) a r
            EQ -> checkLeftR' (substituteBr ll lr) a r
            GT -> R (deleteBr ll la lr) a r
      )
      ( \ll la lr _ _ ->
          case compare a0 la of
            LT -> checkLeftR (deleteRl ll la lr) a r
            EQ -> checkLeftR (substituteR ll lr) a r
            GT -> checkLeftR (deleteRr ll la lr) a r
      )
      l
  deleteRr l a =
    set
      (error "Set.delete: L2")
      ( \rl ra rr _ _ ->
          case compare a0 ra of
            LT -> checkRightR l a (deleteLl rl ra rr)
            EQ -> checkRightR l a (substituteL rl rr)
            GT -> checkRightR l a (deleteLr rl ra rr)
      )
      ( \rl ra rr _ _ ->
          case compare a0 ra of
            LT -> R l a (deleteBl rl ra rr)
            EQ -> checkRightR' l a (substituteBl rl rr)
            GT -> R l a (deleteBr rl ra rr)
      )
      ( \rl ra rr _ _ ->
          case compare a0 ra of
            LT -> checkRightR l a (deleteRl rl ra rr)
            EQ -> checkRightR l a (substituteR rl rr)
            GT -> checkRightR l a (deleteRr rl ra rr)
      )
  deleteBl l a r =
    set
      (error "Set.delete: L3")
      ( \ll la lr _ _ ->
          case compare a0 la of
            LT -> checkLeftB (deleteLl ll la lr) a r
            EQ -> checkLeftB (substituteL ll lr) a r
            GT -> checkLeftB (deleteLr ll la lr) a r
      )
      ( \ll la lr _ _ ->
          case compare a0 la of
            LT -> B (deleteBl ll la lr) a r
            EQ -> checkLeftB' (substituteBr ll lr) a r
            GT -> B (deleteBr ll la lr) a r
      )
      ( \ll la lr _ _ ->
          case compare a0 la of
            LT -> checkLeftB (deleteRl ll la lr) a r
            EQ -> checkLeftB (substituteR ll lr) a r
            GT -> checkLeftB (deleteRr ll la lr) a r
      )
      l
  deleteBr l a =
    set
      (error "Set.delete: L4")
      ( \rl ra rr _ _ ->
          case compare a0 ra of
            LT -> checkRightB l a (deleteLl rl ra rr)
            EQ -> checkRightB l a (substituteL rl rr)
            GT -> checkRightB l a (deleteLr rl ra rr)
      )
      ( \rl ra rr _ _ ->
          case compare a0 ra of
            LT -> B l a (deleteBl rl ra rr)
            EQ -> checkRightB' l a (substituteBl rl rr)
            GT -> B l a (deleteBr rl ra rr)
      )
      ( \rl ra rr _ _ ->
          case compare a0 ra of
            LT -> checkRightB l a (deleteRl rl ra rr)
            EQ -> checkRightB l a (substituteR rl rr)
            GT -> checkRightB l a (deleteRr rl ra rr)
      )
  deleteLl l a r =
    set
      (error "Set.delete: L5")
      ( \ll la lr _ _ ->
          case compare a0 la of
            LT -> checkLeftL (deleteLl ll la lr) a r
            EQ -> checkLeftL (substituteL ll lr) a r
            GT -> checkLeftL (deleteLr ll la lr) a r
      )
      ( \ll la lr _ _ ->
          case compare a0 la of
            LT -> L (deleteBl ll la lr) a r
            EQ -> checkLeftL' (substituteBr ll lr) a r
            GT -> L (deleteBr ll la lr) a r
      )
      ( \ll la lr _ _ ->
          case compare a0 la of
            LT -> checkLeftL (deleteRl ll la lr) a r
            EQ -> checkLeftL (substituteR ll lr) a r
            GT -> checkLeftL (deleteRr ll la lr) a r
      )
      l
  deleteLr l a =
    set
      (error "Set.delete: L6")
      ( \rl ra rr _ _ ->
          case compare a0 ra of
            LT -> checkRightL l a (deleteLl rl ra rr)
            EQ -> checkRightL l a (substituteL rl rr)
            GT -> checkRightL l a (deleteLr rl ra rr)
      )
      ( \rl ra rr _ _ ->
          case compare a0 ra of
            LT -> L l a (deleteBl rl ra rr)
            EQ -> checkRightL' l a (substituteBl rl rr)
            GT -> L l a (deleteBr rl ra rr)
      )
      ( \rl ra rr _ _ ->
          case compare a0 ra of
            LT -> checkRightL l a (deleteRl rl ra rr)
            EQ -> checkRightL l a (substituteR rl rr)
            GT -> checkRightL l a (deleteRr rl ra rr)
      )
  rebalanceR l a =
    set
      (error "Set.delete: L7")
      ( \rl ra rr _ _ ->
          set
            (error "Set.delete: L8")
            (\rll rla rlr _ _ -> B (B l a rll) rla (R rlr ra rr))
            (\rll rla rlr _ _ -> B (B l a rll) rla (B rlr ra rr))
            (\rll rla rlr _ _ -> B (L l a rll) rla (B rlr ra rr))
            rl
      )
      (\rl ra rr _ _ -> L (R l a rl) ra rr)
      (\rl ra rr _ _ -> B (B l a rl) ra rr)
  rebalanceL l a r =
    set
      (error "Set.delete: L9")
      (\ll la lr _ _ -> B ll la (B lr a r))
      (\ll la lr _ _ -> R ll la (L lr a r))
      ( \ll la lr _ _ ->
          set
            (error "Set.delete: L10")
            (\lrl lra lrr _ _ -> B (B ll la lrl) lra (R lrr a r))
            (\lrl lra lrr _ _ -> B (B ll la lrl) lra (B lrr a r))
            (\lrl lra lrr _ _ -> B (L ll la lrl) lra (B lrr a r))
            lr
      )
      l
  checkLeftR l a r =
    set
      (error "Set.delete: L11")
      (\_ _ _ _ _ -> R l a r)
      (\_ _ _ _ _ -> rebalanceR l a r)
      (\_ _ _ _ _ -> R l a r)
      l
  checkLeftB l a r =
    set
      (error "Set.delete: L12")
      (\_ _ _ _ _ -> B l a r)
      (\_ _ _ _ _ -> R l a r)
      (\_ _ _ _ _ -> B l a r)
      l
  checkLeftL l a r =
    set
      (error "Set.delete: L13")
      (\_ _ _ _ _ -> L l a r)
      (\_ _ _ _ _ -> B l a r)
      (\_ _ _ _ _ -> L l a r)
      l
  checkRightR l a r =
    set
      (error "Set.delete: L14")
      (\_ _ _ _ _ -> R l a r)
      (\_ _ _ _ _ -> B l a r)
      (\_ _ _ _ _ -> R l a r)
      r
  checkRightB l a r =
    set
      (error "Set.delete: L15")
      (\_ _ _ _ _ -> B l a r)
      (\_ _ _ _ _ -> L l a r)
      (\_ _ _ _ _ -> B l a r)
      r
  checkRightL l a r =
    set
      (error "Set.delete: L16")
      (\_ _ _ _ _ -> L l a r)
      (\_ _ _ _ _ -> rebalanceL l a r)
      (\_ _ _ _ _ -> L l a r)
      r
  substituteR l =
    set
      (error "Set.delete: L17")
      (\rl ra rr _ _ -> uncurry (checkRightR l) $ popLeftL rl ra rr)
      (\rl ra rr _ _ -> uncurry (checkRightR' l) $ popLeftB rl ra rr)
      (\rl ra rr _ _ -> uncurry (checkRightR l) $ popLeftR rl ra rr)
  substituteBr l =
    set
      E
      (\rl ra rr _ _ -> uncurry (checkRightB l) $ popLeftL rl ra rr)
      (\rl ra rr _ _ -> uncurry (checkRightB' l) $ popLeftB rl ra rr)
      (\rl ra rr _ _ -> uncurry (checkRightB l) $ popLeftR rl ra rr)
  substituteBl l r =
    set
      E
      (\ll la lr _ _ -> (\(l', a) -> checkLeftB l' a r) $ popRightL ll la lr)
      (\ll la lr _ _ -> (\(l', a) -> checkLeftB' l' a r) $ popRightB ll la lr)
      (\ll la lr _ _ -> (\(l', a) -> checkLeftB l' a r) $ popRightR ll la lr)
      l
  substituteL l r =
    set
      (error "Set.delete: L18")
      (\ll la lr _ _ -> (\(l', a) -> checkLeftL l' a r) $ popRightL ll la lr)
      (\ll la lr _ _ -> (\(l', a) -> checkLeftL' l' a r) $ popRightB ll la lr)
      (\ll la lr _ _ -> (\(l', a) -> checkLeftL l' a r) $ popRightR ll la lr)
      l
  checkLeftR' l a r =
    set
      (rebalanceR l a r)
      (\_ _ _ _ _ -> R l a r)
      (\_ _ _ _ _ -> R l a r)
      (\_ _ _ _ _ -> R l a r)
      l
  checkLeftB' l a r =
    set
      (R l a r)
      (\_ _ _ _ _ -> B l a r)
      (\_ _ _ _ _ -> B l a r)
      (\_ _ _ _ _ -> B l a r)
      l
  checkLeftL' l a r =
    set
      (B l a r)
      (\_ _ _ _ _ -> L l a r)
      (\_ _ _ _ _ -> L l a r)
      (\_ _ _ _ _ -> L l a r)
      l
  checkRightR' l a r =
    set
      (B l a r)
      (\_ _ _ _ _ -> R l a r)
      (\_ _ _ _ _ -> R l a r)
      (\_ _ _ _ _ -> R l a r)
      r
  checkRightB' l a r =
    set
      (L l a r)
      (\_ _ _ _ _ -> B l a r)
      (\_ _ _ _ _ -> B l a r)
      (\_ _ _ _ _ -> B l a r)
      r
  checkRightL' l a r =
    set
      (rebalanceL l a r)
      (\_ _ _ _ _ -> L l a r)
      (\_ _ _ _ _ -> L l a r)
      (\_ _ _ _ _ -> L l a r)
      r
  popLeftR l a r =
    set
      (a, r)
      ( \ll la lr _ _ ->
          (\(a', l') -> (a', checkLeftR l' a r)) $
            popLeftL ll la lr
      )
      (\ll la lr _ _ -> popLeftRB ll la lr a r)
      ( \ll la lr _ _ ->
          (\(a', l') -> (a', checkLeftR l' a r)) $
            popLeftR ll la lr
      )
      l
  popLeftB l a r =
    set
      (a, E)
      (\ll la lr _ _ -> popLeftBL ll la lr a r)
      (\ll la lr _ _ -> popLeftBB ll la lr a r)
      (\ll la lr _ _ -> popLeftBR ll la lr a r)
      l
  popLeftL l a r =
    set
      (error "Set.delete: L19")
      ( \ll la lr _ _ ->
          (\(a', l') -> (a', checkLeftL l' a r)) $
            popLeftL ll la lr
      )
      (\ll la lr _ _ -> popLeftLB ll la lr a r)
      ( \ll la lr _ _ ->
          (\(a', l') -> (a', checkLeftL l' a r)) $
            popLeftR ll la lr
      )
      l
  popLeftRB ll la lr a r =
    set
      (la, rebalanceR E a r)
      ( \lll lla llr _ _ ->
          (\(a', l) -> (a', R l a r)) $
            popLeftBL lll lla llr la lr
      )
      ( \lll lla llr _ _ ->
          (\(a', l) -> (a', R l a r)) $
            popLeftBB lll lla llr la lr
      )
      ( \lll lla llr _ _ ->
          (\(a', l) -> (a', R l a r)) $
            popLeftBR lll lla llr la lr
      )
      ll
  popLeftBB ll la lr a r =
    set
      (la, R E a r)
      ( \lll lla llr _ _ ->
          (\(a', l) -> (a', B l a r)) $
            popLeftBL lll lla llr la lr
      )
      ( \lll lla llr _ _ ->
          (\(a', l) -> (a', B l a r)) $
            popLeftBB lll lla llr la lr
      )
      ( \lll lla llr _ _ ->
          (\(a', l) -> (a', B l a r)) $
            popLeftBR lll lla llr la lr
      )
      ll
  popLeftLB ll la lr a r =
    set
      (la, B E a E)
      ( \lll lla llr _ _ ->
          (\(a', l) -> (a', L l a r)) $
            popLeftBL lll lla llr la lr
      )
      ( \lll lla llr _ _ ->
          (\(a', l) -> (a', L l a r)) $
            popLeftBB lll lla llr la lr
      )
      ( \lll lla llr _ _ ->
          (\(a', l) -> (a', L l a r)) $
            popLeftBR lll lla llr la lr
      )
      ll
  popLeftBR ll la lr a r =
    (\(a', l) -> (a', checkLeftB l a r)) $
      popLeftR ll la lr
  popLeftBL ll la lr a r =
    (\(a', l) -> (a', checkLeftB l a r)) $
      popLeftL ll la lr
  popRightR l a =
    set
      (error "Set.delete: L20")
      (\rl ra rr _ _ -> first (checkRightR l a) $ popRightL rl ra rr)
      (\rl ra rr _ _ -> popRightRB l a rl ra rr)
      (\rl ra rr _ _ -> first (checkRightR l a) $ popRightR rl ra rr)
  popRightB l a =
    set
      (E, a)
      (\rl ra rr _ _ -> popRightBL l a rl ra rr)
      (\rl ra rr _ _ -> popRightBB l a rl ra rr)
      (\rl ra rr _ _ -> popRightBR l a rl ra rr)
  popRightL l a =
    set
      (l, a)
      (\rl ra rr _ _ -> first (checkRightL l a) $ popRightL rl ra rr)
      (\rl ra rr _ _ -> popRightLB l a rl ra rr)
      (\rl ra rr _ _ -> first (checkRightL l a) $ popRightR rl ra rr)
  popRightRB l a rl ra =
    set
      (B E a E, ra)
      (\rrl rra rrr _ _ -> first (R l a) $ popRightBL rl ra rrl rra rrr)
      (\rrl rra rrr _ _ -> first (R l a) $ popRightBB rl ra rrl rra rrr)
      (\rrl rra rrr _ _ -> first (R l a) $ popRightBR rl ra rrl rra rrr)
  popRightBB l a rl ra =
    set
      (L l a E, ra)
      (\rrl rra rrr _ _ -> first (B l a) $ popRightBL rl ra rrl rra rrr)
      (\rrl rra rrr _ _ -> first (B l a) $ popRightBB rl ra rrl rra rrr)
      (\rrl rra rrr _ _ -> first (B l a) $ popRightBR rl ra rrl rra rrr)
  popRightLB l a rl ra =
    set
      (rebalanceL l a E, ra)
      (\rrl rra rrr _ _ -> first (L l a) $ popRightBL rl ra rrl rra rrr)
      (\rrl rra rrr _ _ -> first (L l a) $ popRightBB rl ra rrl rra rrr)
      (\rrl rra rrr _ _ -> first (L l a) $ popRightBR rl ra rrl rra rrr)
  popRightBR l a rl ra rr = first (checkRightB l a) $ popRightR rl ra rr
  popRightBL l a rl ra rr = first (checkRightB l a) $ popRightL rl ra rr

-- | /O(log n)/ Delete the maximum element from a set
deleteMax :: (Ord a) => Set a -> Set a
deleteMax t = maybe t (`delete` t) $ lookupMax t

-- | /O(log n)/ Delete the minimum element from a set
deleteMin :: (Ord a) => Set a -> Set a
deleteMin t = maybe t (`delete` t) $ lookupMin t

-- | /O(n log n)/ Keep the elements satisfying a predicate
filter :: (Ord a) => (a -> Bool) -> Set a -> Set a
filter p = foldr (\a b -> bool b (insert a b) $ p a) empty

-- | /O(log n)/ Insert an element into a set
insert :: (Ord a) => a -> Set a -> Set a
insert a0 =
  set
    (B E a0 E)
    (\l a r _ _ -> insertL l a r)
    (\l a r _ _ -> insertB l a r)
    (\l a r _ _ -> insertR l a r)
 where
  insertR l a r =
    case compare a0 a of
      LT -> insertRl l a r
      EQ -> R l a0 r
      GT -> insertRr l a r
  insertB l a r =
    case compare a0 a of
      LT -> insertBl l a r
      EQ -> B l a0 r
      GT -> insertBr l a r
  insertL l a r =
    case compare a0 a of
      LT -> insertLl l a r
      EQ -> L l a0 r
      GT -> insertLr l a r
  insertRl l a r =
    set
      (B (B E a0 E) a r)
      (\ll la lr _ _ -> R (insertL ll la lr) a r)
      ( \ll la lr _ _ ->
          let l' = insertB ll la lr
           in set
                (error "Set.insert: L0")
                (\_ _ _ _ _ -> B l' a r)
                (\_ _ _ _ _ -> R l' a r)
                (\_ _ _ _ _ -> B l' a r)
                l'
      )
      (\ll la lr _ _ -> R (insertR ll la lr) a r)
      l
  insertBl l a r =
    set
      (L (B E a0 E) a r)
      (\ll la lr _ _ -> B (insertL ll la lr) a r)
      ( \ll la lr _ _ ->
          let l' = insertB ll la lr
           in set
                (error "Set.insert: L1")
                (\_ _ _ _ _ -> L l' a r)
                (\_ _ _ _ _ -> B l' a r)
                (\_ _ _ _ _ -> L l' a r)
                l'
      )
      (\ll la lr _ _ -> B (insertR ll la lr) a r)
      l
  insertBr l a =
    set
      (R l a (B E a0 E))
      (\rl ra rr _ _ -> B l a (insertL rl ra rr))
      ( \rl ra rr _ _ ->
          let r = insertB rl ra rr
           in set
                (error "Set.insert: L2")
                (\_ _ _ _ _ -> R l a r)
                (\_ _ _ _ _ -> B l a r)
                (\_ _ _ _ _ -> R l a r)
                r
      )
      (\rl ra rr _ _ -> B l a (insertR rl ra rr))
  insertLr l a =
    set
      (B l a (B E a0 E))
      (\rl ra rr _ _ -> L l a (insertL rl ra rr))
      ( \rl ra rr _ _ ->
          let r = insertB rl ra rr
           in set
                (error "Set.insert: L3")
                (\_ _ _ _ _ -> B l a r)
                (\_ _ _ _ _ -> L l a r)
                (\_ _ _ _ _ -> B l a r)
                r
      )
      (\rl ra rr _ _ -> L l a (insertR rl ra rr))
  insertRr l a =
    set
      (error "Set.insert: L4")
      (\rl ra rr _ _ -> R l a (insertL rl ra rr))
      ( \rl ra rr _ _ ->
          case compare a0 ra of
            LT -> insertRrl l a rl ra rr
            EQ -> R l a (B rl a0 rr)
            GT -> insertRrr l a rl ra rr
      )
      (\rl ra rr _ _ -> R l a (insertR rl ra rr))
  insertLl l a r =
    set
      (error "Set.insert: L5")
      (\ll la lr _ _ -> L (insertL ll la lr) a r)
      ( \ll la lr _ _ ->
          case compare a0 la of
            LT -> insertLll ll la lr a r
            EQ -> L (B ll a0 lr) a r
            GT -> insertLlr ll la lr a r
      )
      (\ll la lr _ _ -> L (insertR ll la lr) a r)
      l
  insertRrr l a rl ra =
    set
      (B (B l a rl) ra (B E a0 E))
      (\rrl rra rrr _ _ -> R l a (B rl ra (insertL rrl rra rrr)))
      ( \rrl rra rrr _ _ ->
          let rr = insertB rrl rra rrr
           in set
                (error "Set.insert: L6")
                (\_ _ _ _ _ -> B (B l a rl) ra rr)
                (\_ _ _ _ _ -> R l a (B rl ra rr))
                (\_ _ _ _ _ -> B (B l a rl) ra rr)
                rr
      )
      (\rrl rra rrr _ _ -> R l a (B rl ra (insertR rrl rra rrr)))
  insertLll ll la lr a r =
    set
      (B (B E a0 E) la (B lr a r))
      (\lll lla llr _ _ -> L (B (insertL lll lla llr) la lr) a r)
      ( \lll lla llr _ _ ->
          let ll' = insertB lll lla llr
           in set
                (error "Set.insert: L7")
                (\_ _ _ _ _ -> B ll' la (B lr a r))
                (\_ _ _ _ _ -> L (B ll' la lr) a r)
                (\_ _ _ _ _ -> B ll' la (B lr a r))
                ll'
      )
      (\lll lla llr _ _ -> L (B (insertR lll lla llr) la lr) a r)
      ll
  insertRrl l a rl ra rr =
    set
      (B (B l a E) a0 (B E ra rr))
      (\rll rla rlr _ _ -> R l a (B (insertL rll rla rlr) ra rr))
      ( \rll rla rlr _ _ ->
          let rl' = insertB rll rla rlr
           in set
                (error "Set.insert: L8")
                (\rll' rla' rlr' _ _ -> B (B l a rll') rla' (R rlr' ra rr))
                (\_ _ _ _ _ -> R l a (B rl' ra rr))
                (\rll' rla' rlr' _ _ -> B (L l a rll') rla' (B rlr' ra rr))
                rl'
      )
      (\rll rla rlr _ _ -> R l a (B (insertR rll rla rlr) ra rr))
      rl
  insertLlr ll la lr a r =
    set
      (B (B ll la E) a0 (B E a r))
      (\lrl lra lrr _ _ -> L (B ll la (insertL lrl lra lrr)) a r)
      ( \lrl lra lrr _ _ ->
          let lr' = insertB lrl lra lrr
           in set
                (error "Set.insert: L9")
                (\lrl' lra' lrr' _ _ -> B (B ll la lrl') lra' (R lrr' a r))
                (\_ _ _ _ _ -> L (B ll la lr') a r)
                (\lrl' lra' lrr' _ _ -> B (L ll la lrl') lra' (B lrr' a r))
                lr'
      )
      (\lrl lra lrr _ _ -> L (B ll la (insertR lrl lra lrr)) a r)
      lr

{-
 - Partition
 -}

-- | /O(n log n)/ Partition a set with a predicate into @(true, false)@ subsets
partition :: (Ord a) => (a -> Bool) -> Set a -> (Set a, Set a)
partition p =
  foldr
    (\a (t1, t2) -> bool (t1, insert a t2) (insert a t1, t2) $ p a)
    (empty, empty)

-- | /O(n log n)/ Split a set by an element into @(lt, eq, gt)@ subsets
split :: (Ord a) => a -> Set a -> (Set a, Bool, Set a)
split a0 =
  foldr
    ( \a (t1, a', t2) -> case compare a a0 of
        LT -> (insert a t1, a', t2)
        EQ -> (t1, True, t2)
        GT -> (t1, a', insert a t2)
    )
    (empty, False, empty)

-- | /O(log n)/ Split a set by its maximum element
splitMax :: (Ord a) => Set a -> Maybe (a, Set a)
splitMax t = (,deleteMax t) <$> lookupMax t

-- | /O(log n)/ Split a set by its minimum element
splitMin :: (Ord a) => Set a -> Maybe (a, Set a)
splitMin t = (,deleteMin t) <$> lookupMin t

{-
 - Query
 -}

-- | /O(m log n)/ Check whether two sets have no elements in common
disjoint :: (Ord a) => Set a -> Set a -> Bool
disjoint t1 = not . any (`member` t1)

-- | /O(n log m)/ Check whether the elements of a set exist in the other
isSubsetOf :: (Ord a) => Set a -> Set a -> Bool
isSubsetOf t1 t2 = foldr (\a b -> a `member` t2 && b) True t1

-- | /O(log n)/ Fetch the least element greater than or equal to the given one
lookupGE :: (Ord a) => a -> Set a -> Maybe a
lookupGE a0 = set Nothing go go go
 where
  go _ a _ recl recr = case compare a a0 of
    LT -> recr
    EQ -> Just a
    GT -> recl <|> Just a

-- | /O(log n)/ Fetch the least element strictly greater than the given one
lookupGT :: (Ord a) => a -> Set a -> Maybe a
lookupGT a0 = set Nothing go go go
 where
  go _ a _ recl recr = case compare a a0 of
    LT -> recr
    EQ -> recr
    GT -> recl <|> Just a

-- | /O(log n)/ Fetch the greatest element less than or equal to the given one
lookupLE :: (Ord a) => a -> Set a -> Maybe a
lookupLE a0 = set Nothing go go go
 where
  go _ a _ recl recr = case compare a a0 of
    LT -> recr <|> Just a
    EQ -> Just a
    GT -> recl

-- | /O(log n)/ Fetch the greatest element strictly less than the given one
lookupLT :: (Ord a) => a -> Set a -> Maybe a
lookupLT a0 = set Nothing go go go
 where
  go _ a _ recl recr = case compare a a0 of
    LT -> recr <|> Just a
    EQ -> recl
    GT -> recl

-- | /O(log n)/ Fetch the maximum element, or 'Nothing' if the set is empty
lookupMax :: Set a -> Maybe a
lookupMax = set Nothing go go go
 where
  go _ a r _ recr = set (Just a) go' go' go' r
   where
    go' _ _ _ _ _ = recr

-- | /O(log n)/ Fetch the minimum element, or 'Nothing' if the set is empty
lookupMin :: Set a -> Maybe a
lookupMin = set Nothing go go go
 where
  go l a _ recl _ = set (Just a) go' go' go' l
   where
    go' _ _ _ _ _ = recl

-- | /O(log n)/ Check whether an element is in a set
member :: (Ord a) => a -> Set a -> Bool
member a0 = set False go go go
 where
  go _ a _ recl recr = case compare a0 a of
    LT -> recl
    EQ -> True
    GT -> recr

-- | /O(1)/ Check whether a set is empty
null :: Set a -> Bool
null = set True go go go where go _ _ _ _ _ = False

-- | /O(n)/ Get the size of a set
size :: Set a -> Int
size = set 0 go go go where go _ _ _ recl recr = 1 + recl + recr

{-
 - Validation
 -}

-- | /O(n)/ Check whether a set is internally height-balanced and ordered
valid :: (Ord a) => Set a -> Bool
valid = liftA2 (&&) balanced ordered
 where
  balanced =
    set
      True
      (\l _ r recl recr -> levels l - levels r == 1 && recl && recr)
      (\l _ r recl recr -> levels l - levels r == 0 && recl && recr)
      (\l _ r recl recr -> levels r - levels l == 1 && recl && recr)
  levels = set 0 go go go
   where
    go _ _ _ recl recr = 1 + max recl recr :: Int
  ordered = set True go go go
   where
    go l a r recl recr =
      set True lt lt lt l
        && set True gt gt gt r
     where
      lt _ la _ _ _ = la < a && recl && recr
      gt _ ra _ _ _ = ra > a && recl && recr