mini-1.4.3.0: Mini/Data/Map.hs

-- | A structure mapping unique keys to values
module Mini.Data.Map (
  -- * Type
  Map,

  -- * Construction
  empty,
  fromList,
  fromListWith,
  fromListWithKey,
  singleton,

  -- * Combination
  difference,
  intersection,
  union,

  -- * Conversion
  toAscList,
  toDescList,

  -- * Fold
  foldlWithKey,
  foldrWithKey,

  -- * Modification
  adjust,
  delete,
  filter,
  filterWithKey,
  insert,
  insertWith,
  insertWithKey,
  update,

  -- * Query
  isSubmapOf,
  lookup,
  lookupGE,
  lookupGT,
  lookupLE,
  lookupLT,
  lookupMax,
  lookupMin,
  member,
  null,
  size,

  -- * Traversal
  traverseWithKey,

  -- * Validation
  valid,

  -- * Examples
  -- $examples
) where

import Control.Applicative (
  (<|>),
 )
import Control.Monad (
  liftM2,
 )
import Data.Bool (
  bool,
 )
import Prelude (
  Applicative,
  Bool (
    False,
    True
  ),
  Eq,
  Foldable,
  Functor,
  Int,
  Maybe (
    Just,
    Nothing
  ),
  Monoid,
  Ord,
  Ordering (
    EQ,
    GT,
    LT
  ),
  Semigroup,
  Show,
  Traversable,
  compare,
  const,
  error,
  flip,
  fmap,
  foldl,
  foldr,
  max,
  maybe,
  mempty,
  pure,
  show,
  traverse,
  uncurry,
  ($),
  (&&),
  (+),
  (-),
  (.),
  (<),
  (<$>),
  (<*>),
  (<>),
  (==),
  (>),
 )

{-
 - Type
 -}

-- | A map from keys /k/ to values /a/, internally structured as an AVL tree
data Map k a
  = -- | Empty bin
    E
  | -- | Left-heavy bin
    L (Map k a) k a (Map k a)
  | -- | Balanced bin
    B (Map k a) k a (Map k a)
  | -- | Right-heavy bin
    R (Map k a) k a (Map k a)

instance (Eq k, Eq a) => Eq (Map k a) where
  t1 == t2 = toAscList t1 == toAscList t2

instance (Ord k, Ord a) => Ord (Map k a) where
  compare t1 t2 = compare (toAscList t1) (toAscList t2)

instance (Show k, Show a) => Show (Map k a) where
  show = curl . map [] go go go
   where
    go _ k a _ recl recr = recl <> show (k, a) <> "," <> recr
    curl = wrap "{" "}" . removeTrailingComma
    wrap open close s = open <> s <> close
    removeTrailingComma s = case s of
      [] -> []
      [_] -> []
      (c : cs) -> c : removeTrailingComma cs

instance Functor (Map k) where
  fmap f = map E (go L) (go B) (go R)
   where
    go c _ k a _ recl = c recl k (f a)

instance Foldable (Map k) where
  foldr = foldrWithKey . const

instance Traversable (Map k) where
  traverse = traverseWithKey . const

instance (Ord k) => Semigroup (Map k a) where
  (<>) = union

instance (Ord k) => Monoid (Map k a) where
  mempty = empty

{-
 - Primitive recursion
 -}

-- | Primitive recursion on maps
map
  :: b
  -- ^ Empty bin
  -> (Map k a -> k -> a -> Map k a -> b -> b -> b)
  -- ^ Left-heavy bin
  -> (Map k a -> k -> a -> Map k a -> b -> b -> b)
  -- ^ Balanced bin
  -> (Map k a -> k -> a -> Map k a -> b -> b -> b)
  -- ^ Right-heavy bin
  -> Map k a
  -- ^ Map
  -> b
map e _ _ _ E = e
map e f g h (L l k a r) = f l k a r (map e f g h l) (map e f g h r)
map e f g h (B l k a r) = g l k a r (map e f g h l) (map e f g h r)
map e f g h (R l k a r) = h l k a r (map e f g h l) (map e f g h r)

{-
 - Construction
 -}

-- | /O(1)/ The empty map
empty :: Map k a
empty = E

-- | /O(n log n)/ Make a map from a tail-biased list of @(key, value)@ pairs
fromList :: (Ord k) => [(k, a)] -> Map k a
fromList = foldl (flip $ uncurry insert) empty

-- | /O(n log n)/ Make a map from a list of pairs, combining matching keys
fromListWith :: (Ord k) => (a -> a -> a) -> [(k, a)] -> Map k a
fromListWith = fromListWithKey . const

-- | /O(n log n)/ Make a map from a list of pairs, combining matching keys
fromListWithKey :: (Ord k) => (k -> a -> a -> a) -> [(k, a)] -> Map k a
fromListWithKey f = foldl (flip . uncurry $ insertWithKey f) empty

-- | /O(1)/ Make a map with a single bin
singleton :: k -> a -> Map k a
singleton k a = B E k a E

{-
 - Combination
 -}

-- | /O(m log n)/ Subtract a map by another via key matching
difference :: (Ord k) => Map k a -> Map k b -> Map k a
difference = foldrWithKey (\k _ b -> delete k b)

-- | /O(n log m)/ Intersect a map with another via left-biased key matching
intersection :: (Ord k) => Map k a -> Map k b -> Map k a
intersection t1 t2 =
  foldrWithKey
    (\k a b -> bool b (insert k a b) $ k `member` t2)
    empty
    t1

-- | /O(m log n)/ Unite a map with another via left-biased key matching
union :: (Ord k) => Map k a -> Map k a -> Map k a
union t = foldrWithKey (\k a b -> bool (insert k a b) b $ k `member` t) t

{-
 - Conversion
 -}

-- | /O(n)/ Turn a map into a list of @(key, value)@ pairs in ascending order
toAscList :: Map k a -> [(k, a)]
toAscList = foldlWithKey (\b k a -> (k, a) : b) []

-- | /O(n)/ Turn a map into a list of @(key, value)@ pairs in descending order
toDescList :: Map k a -> [(k, a)]
toDescList = foldrWithKey (\k a b -> (k, a) : b) []

{-
 - Fold
 -}

-- | /O(n)/ Reduce a map with a left-associative operation and an accumulator
foldlWithKey :: (b -> k -> a -> b) -> b -> Map k a -> b
foldlWithKey f b = map b go go go
 where
  go l k a _ _ recr = foldlWithKey f (f recr k a) l

-- | /O(n)/ Reduce a map with a right-associative operation and an accumulator
foldrWithKey :: (k -> a -> b -> b) -> b -> Map k a -> b
foldrWithKey f b = map b go go go
 where
  go _ k a r recl _ = foldrWithKey f (f k a recl) r

{-
 - Modification
 -}

-- | /O(log n)/ Adjust with an operation the value of a key in a map
adjust :: (Ord k) => (a -> a) -> k -> Map k a -> Map k a
adjust f k0 = map E (go L) (go B) (go R)
 where
  go c l k a r recl recr = case compare k0 k of
    LT -> c recl k a r
    EQ -> c l k (f a) r
    GT -> c l k a recr

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

-- | /O(n log n)/ Keep the bins whose values satisfy a predicate
filter :: (Ord k) => (a -> Bool) -> Map k a -> Map k a
filter p = foldrWithKey (\k a b -> bool b (insert k a b) $ p a) empty

-- | /O(n log n)/ Keep the bins whose keys and values satisfy a predicate
filterWithKey :: (Ord k) => (k -> a -> Bool) -> Map k a -> Map k a
filterWithKey p = foldrWithKey (\k a b -> bool b (insert k a b) $ p k a) empty

-- | /O(log n)/ Insert a key and its value into a map, overwriting if present
insert :: (Ord k) => k -> a -> Map k a -> Map k a
insert k0 a0 =
  map
    (B E k0 a0 E)
    (\l k a r _ _ -> insertL l k a r)
    (\l k a r _ _ -> insertB l k a r)
    (\l k a r _ _ -> insertR l k a r)
 where
  insertR l k a r =
    case compare k0 k of
      LT -> insertRl l k a r
      EQ -> R l k0 a0 r
      GT -> insertRr l k a r
  insertB l k a r =
    case compare k0 k of
      LT -> insertBl l k a r
      EQ -> B l k0 a0 r
      GT -> insertBr l k a r
  insertL l k a r =
    case compare k0 k of
      LT -> insertLl l k a r
      EQ -> L l k0 a0 r
      GT -> insertLr l k a r
  insertRl l k a r =
    map
      (B (B E k0 a0 E) k a r)
      (\ll lk la lr _ _ -> R (insertL ll lk la lr) k a r)
      ( \ll lk la lr _ _ ->
          let l' = insertB ll lk la lr
           in map
                (error "Map.insert: L0")
                (\_ _ _ _ _ _ -> B l' k a r)
                (\_ _ _ _ _ _ -> R l' k a r)
                (\_ _ _ _ _ _ -> B l' k a r)
                l'
      )
      (\ll lk la lr _ _ -> R (insertR ll lk la lr) k a r)
      l
  insertBl l k a r =
    map
      (L (B E k0 a0 E) k a r)
      (\ll lk la lr _ _ -> B (insertL ll lk la lr) k a r)
      ( \ll lk la lr _ _ ->
          let l' = insertB ll lk la lr
           in map
                (error "Map.insert: L1")
                (\_ _ _ _ _ _ -> L l' k a r)
                (\_ _ _ _ _ _ -> B l' k a r)
                (\_ _ _ _ _ _ -> L l' k a r)
                l'
      )
      (\ll lk la lr _ _ -> B (insertR ll lk la lr) k a r)
      l
  insertBr l k a =
    map
      (R l k a (B E k0 a0 E))
      (\rl rk ra rr _ _ -> B l k a (insertL rl rk ra rr))
      ( \rl rk ra rr _ _ ->
          let r = insertB rl rk ra rr
           in map
                (error "Map.insert: L2")
                (\_ _ _ _ _ _ -> R l k a r)
                (\_ _ _ _ _ _ -> B l k a r)
                (\_ _ _ _ _ _ -> R l k a r)
                r
      )
      (\rl rk ra rr _ _ -> B l k a (insertR rl rk ra rr))
  insertLr l k a =
    map
      (B l k a (B E k0 a0 E))
      (\rl rk ra rr _ _ -> L l k a (insertL rl rk ra rr))
      ( \rl rk ra rr _ _ ->
          let r = insertB rl rk ra rr
           in map
                (error "Map.insert: L3")
                (\_ _ _ _ _ _ -> B l k a r)
                (\_ _ _ _ _ _ -> L l k a r)
                (\_ _ _ _ _ _ -> B l k a r)
                r
      )
      (\rl rk ra rr _ _ -> L l k a (insertR rl rk ra rr))
  insertRr l k a =
    map
      (error "Map.insert: L4")
      (\rl rk ra rr _ _ -> R l k a (insertL rl rk ra rr))
      ( \rl rk ra rr _ _ ->
          case compare k0 rk of
            LT -> insertRrl l k a rl rk ra rr
            EQ -> R l k a (B rl k0 a0 rr)
            GT -> insertRrr l k a rl rk ra rr
      )
      (\rl rk ra rr _ _ -> R l k a (insertR rl rk ra rr))
  insertLl l k a r =
    map
      (error "Map.insert: L5")
      (\ll lk la lr _ _ -> L (insertL ll lk la lr) k a r)
      ( \ll lk la lr _ _ ->
          case compare k0 lk of
            LT -> insertLll ll lk la lr k a r
            EQ -> L (B ll k0 a0 lr) k a r
            GT -> insertLlr ll lk la lr k a r
      )
      (\ll lk la lr _ _ -> L (insertR ll lk la lr) k a r)
      l
  insertRrr l k a rl rk ra =
    map
      (B (B l k a rl) rk ra (B E k0 a0 E))
      (\rrl rrk rra rrr _ _ -> R l k a (B rl rk ra (insertL rrl rrk rra rrr)))
      ( \rrl rrk rra rrr _ _ ->
          let rr = insertB rrl rrk rra rrr
           in map
                (error "Map.insert: L6")
                (\_ _ _ _ _ _ -> B (B l k a rl) rk ra rr)
                (\_ _ _ _ _ _ -> R l k a (B rl rk ra rr))
                (\_ _ _ _ _ _ -> B (B l k a rl) rk ra rr)
                rr
      )
      (\rrl rrk rra rrr _ _ -> R l k a (B rl rk ra (insertR rrl rrk rra rrr)))
  insertLll ll lk la lr k a r =
    map
      (B (B E k0 a0 E) lk la (B lr k a r))
      (\lll llk lla llr _ _ -> L (B (insertL lll llk lla llr) lk la lr) k a r)
      ( \lll llk lla llr _ _ ->
          let ll' = insertB lll llk lla llr
           in map
                (error "Map.insert: L7")
                (\_ _ _ _ _ _ -> B ll' lk la (B lr k a r))
                (\_ _ _ _ _ _ -> L (B ll' lk la lr) k a r)
                (\_ _ _ _ _ _ -> B ll' lk la (B lr k a r))
                ll'
      )
      (\lll llk lla llr _ _ -> L (B (insertR lll llk lla llr) lk la lr) k a r)
      ll
  insertRrl l k a rl rk ra rr =
    map
      (B (B l k a E) k0 a0 (B E rk ra rr))
      (\rll rlk rla rlr _ _ -> R l k a (B (insertL rll rlk rla rlr) rk ra rr))
      ( \rll rlk rla rlr _ _ ->
          let rl' = insertB rll rlk rla rlr
           in map
                (error "Map.insert: L8")
                ( \rll' rlk' rla' rlr' _ _ ->
                    B
                      (B l k a rll')
                      rlk'
                      rla'
                      (R rlr' rk ra rr)
                )
                (\_ _ _ _ _ _ -> R l k a (B rl' rk ra rr))
                ( \rll' rlk' rla' rlr' _ _ ->
                    B
                      (L l k a rll')
                      rlk'
                      rla'
                      (B rlr' rk ra rr)
                )
                rl'
      )
      (\rll rlk rla rlr _ _ -> R l k a (B (insertR rll rlk rla rlr) rk ra rr))
      rl
  insertLlr ll lk la lr k a r =
    map
      (B (B ll lk la E) k0 a0 (B E k a r))
      (\lrl lrk lra lrr _ _ -> L (B ll lk la (insertL lrl lrk lra lrr)) k a r)
      ( \lrl lrk lra lrr _ _ ->
          let lr' = insertB lrl lrk lra lrr
           in map
                (error "Map.insert: L9")
                ( \lrl' lrk' lra' lrr' _ _ ->
                    B
                      (B ll lk la lrl')
                      lrk'
                      lra'
                      (R lrr' k a r)
                )
                (\_ _ _ _ _ _ -> L (B ll lk la lr') k a r)
                ( \lrl' lrk' lra' lrr' _ _ ->
                    B
                      (L ll lk la lrl')
                      lrk'
                      lra'
                      (B lrr' k a r)
                )
                lr'
      )
      (\lrl lrk lra lrr _ _ -> L (B ll lk la (insertR lrl lrk lra lrr)) k a r)
      lr

-- | /O(log n)/ Insert a key and its value, combining new and old if present
insertWith :: (Ord k) => (a -> a -> a) -> k -> a -> Map k a -> Map k a
insertWith = insertWithKey . const

-- | /O(log n)/ Insert a key and its value, combining new and old if present
insertWithKey :: (Ord k) => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a
insertWithKey f k a t = bool (insert k a t) (adjust (f k a) k t) $ k `member` t

-- | /O(log n)/ Modify the value of a key or delete its bin with an operation
update :: (Ord k) => (a -> Maybe a) -> k -> Map k a -> Map k a
update f k t =
  maybe
    t
    ( maybe
        (delete k t)
        (\a -> insert k a t)
        . f
    )
    $ lookup k t

{-
 - Query
 -}

-- | /O(n log m)/ Check whether the bins of one map exist in the other
isSubmapOf :: (Ord k, Eq a) => Map k a -> Map k a -> Bool
isSubmapOf t1 t2 =
  foldrWithKey
    (\k a b -> maybe False ((&& b) . (== a)) $ lookup k t2)
    True
    t1

-- | /O(log n)/ Fetch the value of a key in a map, or 'Nothing' if absent
lookup :: (Ord k) => k -> Map k a -> Maybe a
lookup k = map Nothing go go go
 where
  go _ k' a _ recl recr = case compare k k' of
    LT -> recl
    EQ -> Just a
    GT -> recr

-- | /O(log n)/ Fetch the least bin greater than or equal to a key
lookupGE :: (Ord k) => k -> Map k a -> Maybe (k, a)
lookupGE k0 = map Nothing go go go
 where
  go _ k a _ recl recr = case compare k k0 of
    LT -> recr
    EQ -> Just (k, a)
    GT -> recl <|> Just (k, a)

-- | /O(log n)/ Fetch the least bin strictly greater than a key
lookupGT :: (Ord k) => k -> Map k a -> Maybe (k, a)
lookupGT k0 = map Nothing go go go
 where
  go _ k a _ recl recr = case compare k k0 of
    LT -> recr
    EQ -> recr
    GT -> recl <|> Just (k, a)

-- | /O(log n)/ Fetch the greatest bin less than or equal to a key
lookupLE :: (Ord k) => k -> Map k a -> Maybe (k, a)
lookupLE k0 = map Nothing go go go
 where
  go _ k a _ recl recr = case compare k k0 of
    LT -> recr <|> Just (k, a)
    EQ -> Just (k, a)
    GT -> recl

-- | /O(log n)/ Fetch the greatest bin strictly less than a key
lookupLT :: (Ord k) => k -> Map k a -> Maybe (k, a)
lookupLT k0 = map Nothing go go go
 where
  go _ k a _ recl recr = case compare k k0 of
    LT -> recr <|> Just (k, a)
    EQ -> recl
    GT -> recl

-- | /O(log n)/ Fetch the bin with the maximum key, or 'Nothing' if empty
lookupMax :: Map k a -> Maybe (k, a)
lookupMax = map Nothing go go go
 where
  go _ k a r _ recr = map (Just (k, a)) go' go' go' r
   where
    go' _ _ _ _ _ _ = recr

-- | /O(log n)/ Fetch the bin with the minimum key, or 'Nothing' if empty
lookupMin :: Map k a -> Maybe (k, a)
lookupMin = map Nothing go go go
 where
  go l k a _ recl _ = map (Just (k, a)) go' go' go' l
   where
    go' _ _ _ _ _ _ = recl

-- | /O(log n)/ Check whether a key is in a map
member :: (Ord k) => k -> Map k a -> Bool
member k = map False go go go
 where
  go _ k' _ _ recl recr = case compare k k' of
    LT -> recl
    EQ -> True
    GT -> recr

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

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

{-
 - Traversal
 -}

-- | /O(n)/ Lift a map with a lifting operation on keys and values
traverseWithKey :: (Applicative f) => (k -> a -> f b) -> Map k a -> f (Map k b)
traverseWithKey f = map (pure E) (go L) (go B) (go R)
 where
  go c _ k a _ recl recr = c <$> recl <*> pure k <*> f k a <*> recr

{-
 - Validation
 -}

-- | /O(n)/ Check whether a map is internally height-balanced and ordered
valid :: (Ord k) => Map k a -> Bool
valid = liftM2 (&&) balanced ordered
 where
  balanced =
    map
      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 = map 0 go go go
   where
    go _ _ _ _ recl recr = 1 + max recl recr :: Int
  ordered = map True go go go
   where
    go l k _ r recl recr =
      map True lt lt lt l
        && map True gt gt gt r
     where
      lt _ lk _ _ _ _ = lk < k && recl && recr
      gt _ rk _ _ _ _ = rk > k && recl && recr

{-
 - Examples
 -}

{- $examples
'fromList': /tail-biased/ means that if a list of @(key, value)@ pairs contains
pairs with identical keys, the rightmost one is kept.

>>> fromList [('a',1),('b',2),('c',3),('b',4),('a',5)]
{('a',5),('b',4),('c',3)}

'fromListWith', 'fromListWithKey': If a list of @(key, value)@ pairs contains
pairs with identical keys, the leftmost one is inserted as is and the subsequent
ones adjust the value with the combining function left-associatively. The
combining function takes the new value as the left operand, and the existing
value as the right operand.

>>> fromListWith (<>) [(1,"a"),(2,"b"),(1,"c"),(1,"d")]
{(1,"dca"),(2,"b")}
>>> let f k new old = old <> ", " <> show k <> new
>>> fromListWithKey f [(1,"a"),(2,"b"),(1,"c"),(1,"d")]
{(1,"a, 1c, 1d"),(2,"b")}

'intersection', 'union': /left-biased/ means that if the operands contain bins
with identical keys, the bins from the /left/ operand is kept.

>>> fromList [('a',1),('b',2)] `intersection` fromList [('c',3),('b',4),('a',5)]
{('a',1),('b',2)}
>>> fromList [('a',1),('b',2)] `union` fromList [('c',3),('b',4),('a',5)]
{('a',1),('b',2),('c',3)}

'insertWith', 'insertWithKey': If the key does not exist in the map, it is
inserted with the given value as is. Otherwise, the existing value is adjusted
with the combining function, which takes the given value as the left operand and
the existing value as the right operand.

>>> insertWith (<>) 1 "foo" $ fromList [(2,"bar"),(3,"baz")]
{(1,"foo"),(2,"bar"),(3,"baz")}
>>> insertWith (<>) 2 "foo" $ fromList [(2,"bar"),(3,"baz")]
{(2,"foobar"),(3,"baz")}
>>> let f k new old = k + new - old
>>> insertWithKey f 1 2 $ fromList [(2,3),(3,5)]
{(1,2),(2,3),(3,5)}
>>> insertWithKey f 2 7 $ fromList [(2,3),(3,5)]
{(2,6),(3,5)}

'update': If the key does not exist, the map is unchanged. If the key exists and
the result of the operation is @Just x@, the value of the corresponding bin is
updated to @x@. If the key exists and the result of the operation is @Nothing@,
the corresponding bin is removed.

>>> f a = if a == 2 then Just 9 else Nothing
>>> update f 'c' $ fromList [('a',1),('b',2)]
{('a',1),('b',2)}
>>> update f 'b' $ fromList [('a',1),('b',2)]
{('a',1),('b',9)}
>>> update f 'a' $ fromList [('a',1),('b',2)]
{('b',2)}
-}