Map.hs

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

  -- * Construction
  empty,
  fromList,
  singleton,

  -- * Combination
  difference,
  intersection,
  union,

  -- * Conversion
  toAscList,
  toDescList,

  -- * Fold
  foldlWithKey,
  foldrWithKey,

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

  -- * Query
  isSubmapOf,
  lookup,
  lookupMax,
  lookupMin,
  member,
  null,
  size,

  -- * Traversal
  traverseWithKey,

  -- * Validation
  valid,

  -- * Examples
  -- $examples
) where

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,
  not,
  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)
  deriving (Eq, Ord)

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(1)/ Make a map with a single bin
singleton :: k -> a -> Map k a
singleton k a = B E k a E

{-
 - Combination
 -}

-- | /O(n 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 n)/ 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(n 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 b (insert k a b) . not $ 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)/ 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)/ 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)/ 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 n)/ Check whether the bins of one map exist in the other
isSubmapOf :: (Ord k, Eq a) => Map k a -> Map k a -> Bool
isSubmapOf p q =
  foldrWithKey
    (\k a b -> maybe False ((&& b) . (== a)) $ lookup k q)
    True
    p

-- | /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 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 one closest to the end of the list is kept.

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

'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)}

'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)}
-}