Internals.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Data.PQueue.Prio.Internals (
  MinPQueue(..),
  BinomForest(..),
  BinomHeap,
  BinomTree(..),
  Zero(..),
  Succ(..),
  empty,
  null,
  size,
  singleton,
  insert,
  insertEager,
  union,
  getMin,
  adjustMinWithKey,
  adjustMinWithKeyA',
  updateMinWithKey,
  updateMinWithKeyA',
  minViewWithKey,
  mapWithKey,
  mapKeysMonotonic,
  mapMaybeWithKey,
  mapEitherWithKey,
  foldrWithKey,
  foldlWithKey,
  foldrU,
  toAscList,
  toDescList,
  toListU,
  insertMin,
  insertMin',
  insertMax',
  fromList,
  fromAscList,
  foldrWithKeyU,
  foldMapWithKeyU,
  foldlWithKeyU,
  foldlWithKeyU',
  traverseWithKey,
  mapMWithKey,
  traverseWithKeyU,
  seqSpine,
  unions
  ) where

#if MIN_VERSION_base(4,18,0)
import Control.Applicative (Const (..))
#else
import Control.Applicative (liftA2, Const (..))
#endif
import Control.DeepSeq (NFData(rnf), deepseq)
import Data.Coerce (coerce)
import Data.Functor.Identity (Identity(Identity, runIdentity))
import qualified Data.List as List

import Data.Semigroup (Semigroup(..), stimesMonoid, Endo (..), Dual (..))

import Prelude hiding (null, map)
#ifdef __GLASGOW_HASKELL__
import Data.Data
import GHC.Exts (build, inline)
import Text.Read (Lexeme(Ident), lexP, parens, prec,
  readPrec, readListPrec, readListPrecDefault)
#endif

import Data.Functor.WithIndex
import Data.Foldable.WithIndex
import Data.Traversable.WithIndex
import Nattish (Nattish (..))

#ifndef __GLASGOW_HASKELL__
build :: ((a -> [a] -> [a]) -> [a] -> [a]) -> [a]
build f = f (:) []
#endif

#if __GLASGOW_HASKELL__

-- | Treats the priority queue as an empty queue or a minimal
-- key-value pair and a priority queue. The constructors, conceptually,
-- are 'Data.PQueue.Prio.Min.Empty' and '(Data.PQueue.Prio.Min.:<)'.
--
-- 'gfoldl' is nondeterministic; any minimal pair may be chosen as
-- the first. All constructed queues maintain the queue invariants.
instance (Ord k, Data k, Data a) => Data (MinPQueue k a) where
  gfoldl f z q = case minViewWithKey q of
    Nothing      -> z Empty
    Just (x, q') -> z (\(k, a) -> insert k a) `f` x `f` q'

  gunfold k z c = case constrIndex c of
    1 -> z Empty
    2 -> k (k (z (\(key, val) -> insert key val)))
    _ -> error "gunfold: invalid constructor for MinPQueue"

  toConstr q
    | null q = emptyConstr
    | otherwise = consConstr

  dataTypeOf _ = queueDataType
  dataCast1 f  = gcast1 f
  dataCast2 f  = gcast2 f

queueDataType :: DataType
queueDataType = mkDataType "Data.PQueue.Prio.Min.MinPQueue" [emptyConstr, consConstr]

emptyConstr, consConstr :: Constr
emptyConstr = mkConstr queueDataType "Empty" [] Prefix
consConstr  = mkConstr queueDataType ":<" [] Infix
#endif

instance Ord k => Semigroup (MinPQueue k a) where
  (<>) = union
  stimes = stimesMonoid
  {-# INLINABLE stimes #-}

instance Ord k => Monoid (MinPQueue k a) where
  mempty = empty
#if !MIN_VERSION_base(4,11,0)
  mappend = union
#endif
  mconcat = unions

instance (Ord k, Show k, Show a) => Show (MinPQueue k a) where
  showsPrec p xs = showParen (p > 10) $
    showString "fromList " . shows (toAscList xs)

instance (Ord k, Read k, Read a) => Read (MinPQueue k a) where
#ifdef __GLASGOW_HASKELL__
  readPrec = parens $ prec 10 $ do
    Ident "fromList" <- lexP
    xs <- readPrec
    return (fromList xs)

  readListPrec = readListPrecDefault
#else
  readsPrec p = readParen (p > 10) $ \r -> do
    ("fromList",s) <- lex r
    (xs,t) <- reads s
    return (fromList xs,t)
#endif

-- | The union of a list of queues: (@'unions' == 'List.foldl' 'union' 'empty'@).
unions :: Ord k => [MinPQueue k a] -> MinPQueue k a
unions = List.foldl' union empty


(.:) :: (c -> d) -> (a -> b -> c) -> a -> b -> d
(f .: g) x y = f (g x y)

infixr 8 .:

-- | A priority queue where keys of type @k@ are annotated with values of type
-- @a@.  The queue supports extracting the key-value pair with minimum key.
data MinPQueue k a = Empty | MinPQ {-# UNPACK #-} !Int !k a !(BinomHeap k a)

data BinomForest rk k a =
  Nil |
  Skip (BinomForest (Succ rk) k a) |
  Cons {-# UNPACK #-} !(BinomTree rk k a) (BinomForest (Succ rk) k a)
type BinomHeap = BinomForest Zero

data BinomTree rk k a = BinomTree !k (rk k a)
newtype Zero k a = Zero a
data Succ rk k a = Succ {-# UNPACK #-} !(BinomTree rk k a) (rk k a)

instance (Ord k, Eq a) => Eq (MinPQueue k a) where
  MinPQ n1 k1 a1 ts1 == MinPQ n2 k2 a2 ts2 =
    n1 == n2 && eqExtract k1 a1 ts1 k2 a2 ts2
  Empty == Empty = True
  _     == _     = False

eqExtract :: (Ord k, Eq a) => k -> a -> BinomHeap k a -> k -> a -> BinomHeap k a -> Bool
eqExtract k10 a10 ts10 k20 a20 ts20 =
  k10 == k20 && a10 == a20 &&
  case (extract ts10, extract ts20) of
    (Yes (Extract k1 (Zero a1) ts1'), Yes (Extract k2 (Zero a2) ts2'))
             -> eqExtract k1 a1 ts1' k2 a2 ts2'
    (No, No) -> True
    _        -> False

instance (Ord k, Ord a) => Ord (MinPQueue k a) where
  MinPQ _n1 k10 a10 ts10 `compare` MinPQ _n2 k20 a20 ts20 =
    cmpExtract k10 a10 ts10 k20 a20 ts20
  Empty `compare` Empty   = EQ
  Empty `compare` MinPQ{} = LT
  MinPQ{} `compare` Empty = GT

cmpExtract :: (Ord k, Ord a) => k -> a -> BinomHeap k a -> k -> a -> BinomHeap k a -> Ordering
cmpExtract k10 a10 ts10 k20 a20 ts20 =
  k10 `compare` k20 <> a10 `compare` a20 <>
  case (extract ts10, extract ts20) of
    (Yes (Extract k1 (Zero a1) ts1'), Yes (Extract k2 (Zero a2) ts2'))
                -> cmpExtract k1 a1 ts1' k2 a2 ts2'
    (No, Yes{}) -> LT
    (Yes{}, No) -> GT
    (No, No)    -> EQ

-- | \(O(1)\). Returns the empty priority queue.
empty :: MinPQueue k a
empty = Empty

-- | \(O(1)\). Checks if this priority queue is empty.
null :: MinPQueue k a -> Bool
null Empty = True
null _     = False

-- | \(O(1)\). Returns the size of this priority queue.
size :: MinPQueue k a -> Int
size Empty           = 0
size (MinPQ n _ _ _) = n

-- | \(O(1)\). Constructs a singleton priority queue.
singleton :: k -> a -> MinPQueue k a
singleton k a = MinPQ 1 k a Nil

-- | Amortized \(O(1)\), worst-case \(O(\log n)\). Inserts
-- an element with the specified key into the queue.
insert :: Ord k => k -> a -> MinPQueue k a -> MinPQueue k a
insert k a Empty = singleton k a
insert k a (MinPQ n k' a' ts)
  | k <= k' = MinPQ (n + 1) k  a  (incrMin (tip k' a') ts)
  | otherwise = MinPQ (n + 1) k' a' (incr (tip k  a ) ts)

insertEager :: Ord k => k -> a -> MinPQueue k a -> MinPQueue k a
insertEager k a Empty = singleton k a
insertEager k a (MinPQ n k' a' ts)
  | k <= k' = MinPQ (n + 1) k a  (insertEagerHeap k' a' ts)
  | otherwise = MinPQ (n + 1) k' a' (insertEagerHeap k a ts)

-- | Amortized \(O(\log \min(n_1,n_2))\), worst-case \(O(\log \max(n_1,n_2))\). Returns the union
-- of the two specified queues.
union :: Ord k => MinPQueue k a -> MinPQueue k a -> MinPQueue k a
union (MinPQ n1 k1 a1 ts1) (MinPQ n2 k2 a2 ts2)
  | k1 <= k2 = MinPQ (n1 + n2) k1 a1 (insMerge k2 a2)
  | otherwise  = MinPQ (n1 + n2) k2 a2 (insMerge k1 a1)
  where  insMerge k a = carryForest (tip k a) ts1 ts2
union Empty q2 = q2
union q1 Empty = q1

-- | \(O(1)\). The minimal (key, element) in the queue, if the queue is nonempty.
getMin :: MinPQueue k a -> Maybe (k, a)
getMin (MinPQ _ k a _) = Just (k, a)
getMin _               = Nothing

-- | \(O(1)\). Alter the value at the minimum key. If the queue is empty, does nothing.
adjustMinWithKey :: (k -> a -> a) -> MinPQueue k a -> MinPQueue k a
adjustMinWithKey _ Empty = Empty
adjustMinWithKey f (MinPQ n k a ts) = MinPQ n k (f k a) ts

-- | \(O(1)\) per operation. Alter the value at the minimum key in an 'Applicative' context. If the
-- queue is empty, does nothing.
adjustMinWithKeyA' :: Applicative f => (MinPQueue k a -> r) -> (k -> a -> f a) -> MinPQueue k a -> f r
adjustMinWithKeyA' g _ Empty = pure (g Empty)
adjustMinWithKeyA' g f (MinPQ n k a ts) = fmap (\a' -> g (MinPQ n k a' ts)) (f k a)

-- | \(O(\log n)\). (Actually \(O(1)\) if there's no deletion.) Update the value at the minimum key.
-- If the queue is empty, does nothing.
updateMinWithKey :: Ord k => (k -> a -> Maybe a) -> MinPQueue k a -> MinPQueue k a
updateMinWithKey _ Empty = Empty
updateMinWithKey f (MinPQ n k a ts) = case f k a of
  Nothing  -> extractHeap n ts
  Just a'  -> MinPQ n k a' ts

-- | \(O(\log n)\) per operation. (Actually \(O(1)\) if there's no deletion.) Update
-- the value at the minimum key in an 'Applicative' context. If the queue is
-- empty, does nothing.
updateMinWithKeyA'
  :: (Applicative f, Ord k)
  => (MinPQueue k a -> r)
  -> (k -> a -> f (Maybe a))
  -> MinPQueue k a
  -> f r
updateMinWithKeyA' g _ Empty = pure (g Empty)
updateMinWithKeyA' g f (MinPQ n k a ts) = fmap (g . tweak) (f k a)
  where
    tweak Nothing = extractHeap n ts
    tweak (Just a') = MinPQ n k a' ts

-- | \(O(\log n)\). Retrieves the minimal (key, value) pair of the map, and the map stripped of that
-- element, or 'Nothing' if passed an empty map.
minViewWithKey :: Ord k => MinPQueue k a -> Maybe ((k, a), MinPQueue k a)
minViewWithKey Empty            = Nothing
minViewWithKey (MinPQ n k a ts) = Just ((k, a), extractHeap n ts)

-- | \(O(n)\). Map a function over all values in the queue.
mapWithKey :: (k -> a -> b) -> MinPQueue k a -> MinPQueue k b
mapWithKey f = runIdentity . traverseWithKeyU (coerce f)

-- | \(O(n)\). @'mapKeysMonotonic' f q == 'Data.PQueue.Prio.Min.mapKeys' f q@,
-- but only works when @f@ is (weakly) monotonic (meaning that @x <= y@ implies
-- @f x <= f y@). /The precondition is not checked./ This function has better
-- performance than 'Data.PQueue.Prio.Min.mapKeys'.
--
-- Note: if the given function returns bottom for any of the keys in the queue, then the
-- portion of the queue which is bottom is /unspecified/.
mapKeysMonotonic :: (k -> k') -> MinPQueue k a -> MinPQueue k' a
mapKeysMonotonic _ Empty = Empty
mapKeysMonotonic f (MinPQ n k a ts) = MinPQ n (f k) a $! mapKeysMonoHeap f ts

mapKeysMonoHeap :: forall k k' a. (k -> k') -> BinomHeap k a -> BinomHeap k' a
mapKeysMonoHeap f = mapKeysMonoForest Zeroy
  where
    mapKeysMonoForest :: Ranky rk -> BinomForest rk k a -> BinomForest rk k' a
    mapKeysMonoForest !_rky Nil = Nil
    mapKeysMonoForest !rky (Skip rest) = Skip $! mapKeysMonoForest (Succy rky) rest
    mapKeysMonoForest !rky (Cons t rest) = Cons (mapKeysMonoTree rky t) $! mapKeysMonoForest (Succy rky) rest

    {-# INLINE mapKeysMonoTree #-}
    mapKeysMonoTree :: Ranky rk -> BinomTree rk k a -> BinomTree rk k' a
    mapKeysMonoTree Zeroy (BinomTree k (Zero a)) =
      -- We've reached a value, which we must not force.
      BinomTree (f k) (Zero a)
      -- We're not at a value; we force the result.
    mapKeysMonoTree (Succy rky) (BinomTree k ts) = BinomTree (f k) $! mapKeysMonoTrees rky ts

    mapKeysMonoTrees :: Ranky rk -> Succ rk k a -> Succ rk k' a
    mapKeysMonoTrees Zeroy (Succ t (Zero a)) =
      -- Don't force the value!
      Succ (mapKeysMonoTree Zeroy t) (Zero a)
    mapKeysMonoTrees (Succy rky) (Succ t ts) =
      -- Whew, no values; force the trees.
      Succ (mapKeysMonoTree (Succy rky) t) $! mapKeysMonoTrees rky ts

-- | \(O(n)\). Map values and collect the 'Just' results.
mapMaybeWithKey :: Ord k => (k -> a -> Maybe b) -> MinPQueue k a -> MinPQueue k b
mapMaybeWithKey f = fromBare .
  foldlWithKeyU'
    (\q k a -> case f k a of
        Nothing -> q
        Just b -> insertEagerHeap k b q)
    Nil
{-# INLINABLE mapMaybeWithKey #-}

-- | \(O(n)\). Map values and separate the 'Left' and 'Right' results.
mapEitherWithKey :: Ord k => (k -> a -> Either b c) -> MinPQueue k a -> (MinPQueue k b, MinPQueue k c)
mapEitherWithKey f q
  | (l, r) <- mapEitherHeap f q
  , let
      !l' = fromBare l
      !r' = fromBare r
  = (l', r')
{-# INLINABLE mapEitherWithKey #-}

data Partition k a b = Partition !(BinomHeap k a) !(BinomHeap k b)

fromPartition :: Partition k a b -> (BinomHeap k a, BinomHeap k b)
fromPartition (Partition p q) = (p, q)

mapEitherHeap :: Ord k => (k -> a -> Either b c) -> MinPQueue k a -> (BinomHeap k b, BinomHeap k c)
mapEitherHeap f = fromPartition .
  foldlWithKeyU'
    (\(Partition ls rs) k a ->
         case f k a of
           Left b -> Partition (insertEagerHeap k b ls) rs
           Right b -> Partition ls (insertEagerHeap k b rs))
    (Partition Nil Nil)

insertEagerHeap :: Ord k => k -> a -> BinomHeap k a -> BinomHeap k a
insertEagerHeap k a h = incr' (tip k a) h
{-# INLINE insertEagerHeap #-}

-- | \(O(n \log n)\). Fold the keys and values in the map, such that
-- @'foldrWithKey' f z q == 'List.foldr' ('uncurry' f) z ('toAscList' q)@.
--
-- If you do not care about the traversal order, consider using 'foldrWithKeyU'.
foldrWithKey :: Ord k => (k -> a -> b -> b) -> b -> MinPQueue k a -> b
foldrWithKey _ z Empty = z
foldrWithKey f z (MinPQ _ k0 a0 ts0) = f k0 a0 (foldF ts0) where
  foldF ts = case extract ts of
    Yes (Extract k (Zero a) ts') -> f k a (foldF ts')
    No                           -> z

-- | \(O(n \log n)\). Fold the keys and values in the map, such that
-- @'foldlWithKey' f z q == 'List.foldl' ('uncurry' . f) z ('toAscList' q)@.
--
-- If you do not care about the traversal order, consider using 'foldlWithKeyU'.
foldlWithKey :: Ord k => (b -> k -> a -> b) -> b -> MinPQueue k a -> b
foldlWithKey _ z Empty = z
foldlWithKey f z0 (MinPQ _ k0 a0 ts0) = foldF (f z0 k0 a0) ts0 where
  foldF z ts = case extract ts of
    Yes (Extract k (Zero a) ts') -> foldF (f z k a) ts'
    No                           -> z

{-# INLINABLE [1] toAscList #-}
-- | \(O(n \log n)\). Return all (key, value) pairs in ascending order by key.
toAscList :: Ord k => MinPQueue k a -> [(k, a)]
toAscList = foldrWithKey (curry (:)) []

{-# INLINABLE [1] toDescList #-}
-- | \(O(n \log n)\). Return all (key, value) pairs in descending order by key.
toDescList :: Ord k => MinPQueue k a -> [(k, a)]
toDescList = foldlWithKey (\z k a -> (k, a) : z) []

-- | \(O(n)\). Build a priority queue from an ascending list of (key, value) pairs. /The precondition is not checked./
fromAscList :: [(k, a)] -> MinPQueue k a
{-# INLINE fromAscList #-}
fromAscList xs = List.foldl' (\q (k, a) -> insertMax' k a q) empty xs

{-# RULES
  "toAscList" toAscList = \q -> build (\c n -> foldrWithKey (curry c) n q);
  "toDescList" toDescList = \q -> build (\c n -> foldlWithKey (\z k a -> (k, a) `c` z) n q);
  "toListU" toListU = \q -> build (\c n -> foldrWithKeyU (curry c) n q);
  #-}

{-# NOINLINE toListU #-}
-- | \(O(n)\). Returns all (key, value) pairs in the queue in no particular order.
toListU :: MinPQueue k a -> [(k, a)]
toListU = foldrWithKeyU (curry (:)) []

-- | \(O(n)\). An unordered right fold over the elements of the queue, in no particular order.
foldrU :: (a -> b -> b) -> b -> MinPQueue k a -> b
foldrU = foldrWithKeyU . const

-- | Equivalent to 'insert', save the assumption that this key is @<=@
-- every other key in the map. /The precondition is not checked./
insertMin :: k -> a -> MinPQueue k a -> MinPQueue k a
insertMin k a Empty = MinPQ 1 k a Nil
insertMin k a (MinPQ n k' a' ts) = MinPQ (n + 1) k a (incrMin (tip k' a') ts)

-- | Equivalent to 'insert', save the assumption that this key is @<=@
-- every other key in the map. /The precondition is not checked./ Additionally,
-- this eagerly constructs the new portion of the spine.
insertMin' :: k -> a -> MinPQueue k a -> MinPQueue k a
insertMin' k a Empty = MinPQ 1 k a Nil
insertMin' k a (MinPQ n k' a' ts) = MinPQ (n + 1) k a (incrMin' (tip k' a') ts)

-- | Inserts an entry with key @>=@ every key in the map. Assumes and preserves
-- an extra invariant: the roots of the binomial trees are decreasing along
-- the spine.
insertMax' :: k -> a -> MinPQueue k a -> MinPQueue k a
insertMax' k a Empty = MinPQ 1 k a Nil
insertMax' k a (MinPQ n k' a' ts) = MinPQ (n + 1) k' a' (incrMax' (tip k a) ts)

{-# INLINE fromList #-}
-- | \(O(n)\). Constructs a priority queue from an unordered list.
fromList :: Ord k => [(k, a)] -> MinPQueue k a
-- We build a forest first and then extract its minimum at the end.  Why not
-- just build the 'MinQueue' directly? This way typically saves us one
-- comparison per element, which roughly halves comparisons.
fromList xs = fromBare (fromListHeap xs)

fromBare :: Ord k => BinomHeap k a -> MinPQueue k a
fromBare xs = case extract xs of
  No -> Empty
  -- Should we track the size as we go instead? That saves O(log n)
  -- at the end, but it needs an extra register all along the way.
  -- The nodes should probably all be in L1 cache already thanks to the
  -- extractHeap.
  Yes (Extract k (Zero v) f) -> MinPQ (sizeHeap f + 1) k v f

{-# INLINE fromListHeap #-}
fromListHeap :: Ord k => [(k, a)] -> BinomHeap k a
fromListHeap xs = List.foldl' go Nil xs
  where
    go fr (k, a) = insertEagerHeap k a fr

sizeHeap :: BinomHeap k a -> Int
sizeHeap = go 0 1
  where
    go :: Int -> Int -> BinomForest rk k a -> Int
    go acc rk Nil = rk `seq` acc
    go acc rk (Skip f) = go acc (2 * rk) f
    go acc rk (Cons _t f) = go (acc + rk) (2 * rk) f

-- | \(O(1)\). Returns a binomial tree of rank zero containing this
-- key and value.
tip :: k -> a -> BinomTree Zero k a
tip k a = BinomTree k (Zero a)

-- | \(O(1)\). Takes the union of two binomial trees of the same rank.
meld :: Ord k => BinomTree rk k a -> BinomTree rk k a -> BinomTree (Succ rk) k a
meld t1@(BinomTree k1 ts1) t2@(BinomTree k2 ts2)
  | k1 <= k2 = BinomTree k1 (Succ t2 ts1)
  | otherwise  = BinomTree k2 (Succ t1 ts2)

-- | Takes the union of two binomial forests, starting at the same rank. Analogous to binary addition.
mergeForest :: Ord k => BinomForest rk k a -> BinomForest rk k a -> BinomForest rk k a
mergeForest f1 f2 = case (f1, f2) of
  (Skip ts1, Skip ts2)       -> Skip $! mergeForest ts1 ts2
  (Skip ts1, Cons t2 ts2)    -> Cons t2 $! mergeForest ts1 ts2
  (Cons t1 ts1, Skip ts2)    -> Cons t1 $! mergeForest ts1 ts2
  (Cons t1 ts1, Cons t2 ts2) -> Skip $! carryForest (meld t1 t2) ts1 ts2
  (Nil, _)                   -> f2
  (_, Nil)                   -> f1

-- | Takes the union of two binomial forests, starting at the same rank, with an additional tree.
-- Analogous to binary addition when a digit has been carried.
carryForest :: Ord k => BinomTree rk k a -> BinomForest rk k a -> BinomForest rk k a -> BinomForest rk k a
carryForest t0 f1 f2 = t0 `seq` case (f1, f2) of
  (Cons t1 ts1, Cons t2 ts2) -> Cons t0 $! carryMeld t1 t2 ts1 ts2
  (Cons t1 ts1, Skip ts2)    -> Skip $! carryMeld t0 t1 ts1 ts2
  (Skip ts1, Cons t2 ts2)    -> Skip $! carryMeld t0 t2 ts1 ts2
  (Skip ts1, Skip ts2)       -> Cons t0 $! mergeForest ts1 ts2
  -- Why do these use incr and not incr'? We want the merge to take
  -- O(log(min(|f1|, |f2|))) amortized time. If we performed this final
  -- increment eagerly, that would degrade to O(log(max(|f1|, |f2|))) time.
  (Nil, _)                   -> incr t0 f2
  (_, Nil)                   -> incr t0 f1
  where  carryMeld = carryForest .: meld

-- | Inserts a binomial tree into a binomial forest. Analogous to binary incrementation.
incr :: Ord k => BinomTree rk k a -> BinomForest rk k a -> BinomForest rk k a
incr t ts = t `seq` case ts of
  Nil         -> Cons t Nil
  Skip ts'    -> Cons t ts'
  Cons t' ts' -> ts' `seq` Skip (incr (meld t t') ts')

-- | Inserts a binomial tree into a binomial forest. Analogous to binary incrementation.
-- Forces the rebuilt portion of the spine.
incr' :: Ord k => BinomTree rk k a -> BinomForest rk k a -> BinomForest rk k a
incr' t ts = t `seq` case ts of
  Nil         -> Cons t Nil
  Skip ts'    -> Cons t ts'
  Cons t' ts' -> Skip $! incr' (meld t t') ts'

-- | Inserts a binomial tree into a binomial forest. Assumes that the root of this tree
-- is less than all other roots. Analogous to binary incrementation. Equivalent to
-- @'incr' (\_ _ -> True)@.
incrMin :: BinomTree rk k a -> BinomForest rk k a -> BinomForest rk k a
incrMin t@(BinomTree k ts) tss = case tss of
  Nil          -> Cons t Nil
  Skip tss'    -> Cons t tss'
  Cons t' tss' -> tss' `seq` Skip (incrMin (BinomTree k (Succ t' ts)) tss')

-- | Inserts a binomial tree into a binomial forest. Assumes that the root of this tree
-- is less than all other roots. Analogous to binary incrementation. Equivalent to
-- @'incr'' (\_ _ -> True)@. Forces the rebuilt portion of the spine.
incrMin' :: BinomTree rk k a -> BinomForest rk k a -> BinomForest rk k a
incrMin' t@(BinomTree k ts) tss = case tss of
  Nil          -> Cons t Nil
  Skip tss'    -> Cons t tss'
  Cons t' tss' -> Skip $! incrMin' (BinomTree k (Succ t' ts)) tss'

-- | See 'insertMax'' for invariant info.
incrMax' :: BinomTree rk k a -> BinomForest rk k a -> BinomForest rk k a
incrMax' t tss = t `seq` case tss of
  Nil          -> Cons t Nil
  Skip tss'    -> Cons t tss'
  Cons (BinomTree k ts) tss' -> Skip $! incrMax' (BinomTree k (Succ t ts)) tss'

extractHeap :: Ord k => Int -> BinomHeap k a -> MinPQueue k a
extractHeap n ts = n `seq` case extract ts of
  No                      -> Empty
  Yes (Extract k (Zero a) ts') -> MinPQ (n - 1) k a ts'

-- | A specialized type intended to organize the return of extract-min queries
-- from a binomial forest. We walk all the way through the forest, and then
-- walk backwards. @Extract rk a@ is the result type of an extract-min
-- operation that has walked as far backwards of rank @rk@ -- that is, it
-- has visited every root of rank @>= rk@.
--
-- The interpretation of @Extract minKey minVal children forest@ is
--
--   * @minKey@ is the key of the minimum root visited so far. It may have
--     any rank @>= rk@. We will denote the root corresponding to
--     @minKey@ as @minRoot@.
--
--   * @minVal@ is the value corresponding to @minKey@.
--
--   * @children@ is those children of @minRoot@ which have not yet been
--     merged with the rest of the forest. Specifically, these are
--     the children with rank @< rk@.
--
--   * @forest@ is an accumulating parameter that maintains the partial
--     reconstruction of the binomial forest without @minRoot@. It is
--     the union of all old roots with rank @>= rk@ (except @minRoot@),
--     with the set of all children of @minRoot@ with rank @>= rk@.
--     Note that @forest@ is lazy, so if we discover a smaller key
--     than @minKey@ later, we haven't wasted significant work.

data Extract rk k a = Extract !k (rk k a) !(BinomForest rk k a)
data MExtract rk k a = No | Yes {-# UNPACK #-} !(Extract rk k a)

incrExtract :: Extract (Succ rk) k a -> Extract rk k a
incrExtract (Extract minKey (Succ kChild kChildren) ts)
  = Extract minKey kChildren (Cons kChild ts)

incrExtract' :: Ord k => BinomTree rk k a -> Extract (Succ rk) k a -> Extract rk k a
incrExtract' t (Extract minKey (Succ kChild kChildren) ts)
  = Extract minKey kChildren (Skip $! incr' (t `meld` kChild) ts)

-- | Walks backward from the biggest key in the forest, as far as rank @rk@.
-- Returns its progress. Each successive application of @extractBin@ takes
-- amortized \(O(1)\) time, so applying it from the beginning takes \(O(\log n)\) time.
extract :: Ord k => BinomForest rk k a -> MExtract rk k a
extract = start
  where
    start :: Ord k => BinomForest rk k a -> MExtract rk k a
    start Nil = No
    start (Skip f) = case start f of
      No     -> No
      Yes ex -> Yes (incrExtract ex)
    start (Cons t@(BinomTree k ts) f) = Yes $ case go k f of
      No -> Extract k ts (skip f)
      Yes ex -> incrExtract' t ex

    go :: Ord k => k -> BinomForest rk k a -> MExtract rk k a
    go _min_above Nil = _min_above `seq` No
    go min_above (Skip f) = case go min_above f of
      No -> No
      Yes ex -> Yes (incrExtract ex)
    go min_above (Cons t@(BinomTree k ts) f)
      | min_above <= k = case go min_above f of
          No -> No
          Yes ex -> Yes (incrExtract' t ex)
      | otherwise = case go k f of
          No -> Yes (Extract k ts (skip f))
          Yes ex -> Yes (incrExtract' t ex)

skip :: BinomForest (Succ rk) k a -> BinomForest rk k a
skip Nil = Nil
skip f = Skip f
{-# INLINE skip #-}

-- | \(O(n)\). An unordered right fold over the elements of the queue, in no particular order.
foldrWithKeyU :: (k -> a -> b -> b) -> b -> MinPQueue k a -> b
foldrWithKeyU c n = flip appEndo n . inline foldMapWithKeyU (coerce c)

-- | \(O(n)\). An unordered monoidal fold over the elements of the queue, in no particular order.
--
-- @since 1.4.2
foldMapWithKeyU :: forall m k a. Monoid m => (k -> a -> m) -> MinPQueue k a -> m
foldMapWithKeyU = coerce
  (inline traverseWithKeyU :: (k -> a -> Const m ()) -> MinPQueue k a -> Const m (MinPQueue k ()))

-- | \(O(n)\). An unordered left fold over the elements of the queue, in no
-- particular order. This is rarely what you want; 'foldrWithKeyU' and
-- 'foldlWithKeyU'' are more likely to perform well.
foldlWithKeyU :: (b -> k -> a -> b) -> b -> MinPQueue k a -> b
foldlWithKeyU f b = flip appEndo b . getDual .
  foldMapWithKeyU (\k a -> Dual $ Endo $ \r -> f r k a)

-- | \(O(n)\). An unordered strict left fold over the elements of the queue, in no particular order.
--
-- @since 1.4.2
foldlWithKeyU' :: (b -> k -> a -> b) -> b -> MinPQueue k a -> b
foldlWithKeyU' f !b q =
  case q of
    Empty -> b
    MinPQ _n k a ts -> foldlHeapU' f (f b k a) ts

foldlHeapU' :: forall k a b. (b -> k -> a -> b) -> b -> BinomHeap k a -> b
foldlHeapU' f = \b -> foldlForest' Zeroy b
  where
    foldlForest' :: Ranky rk -> b -> BinomForest rk k a -> b
    foldlForest' !_rky !acc Nil = acc
    foldlForest' !rky !acc (Skip rest) = foldlForest' (Succy rky) acc rest
    foldlForest' !rky !acc (Cons t rest) =
      foldlForest' (Succy rky) (foldlTree' rky acc t) rest

    {-# INLINE foldlTree' #-}
    foldlTree' :: Ranky rk -> b -> BinomTree rk k a -> b
    foldlTree' !rky !acc (BinomTree k ts) = foldlTrees' rky acc k ts

    foldlTrees' :: Ranky rk -> b -> k -> rk k a -> b
    foldlTrees' Zeroy !acc !k (Zero a) = f acc k a
    foldlTrees' (Succy rky) !acc !k (Succ t ts) =
      foldlTrees' rky (foldlTree' rky acc t) k ts

-- | \(O(n \log n)\). Traverses the elements of the queue in ascending order by key.
-- (@'traverseWithKey' f q == 'fromAscList' <$> 'traverse' ('uncurry' f) ('toAscList' q)@)
--
-- If you do not care about the /order/ of the traversal, consider using 'traverseWithKeyU'.
--
-- If you are working in a strict monad, consider using 'mapMWithKey'.
traverseWithKey :: (Ord k, Applicative f) => (k -> a -> f b) -> MinPQueue k a -> f (MinPQueue k b)
traverseWithKey f q = case minViewWithKey q of
  Nothing      -> pure empty
  Just ((k, a), q')  -> liftA2 (insertMin k) (f k a) (traverseWithKey f q')

-- | A strictly accumulating version of 'traverseWithKey'. This works well in
-- 'IO' and strict @State@, and is likely what you want for other "strict" monads,
-- where @⊥ >>= pure () = ⊥@.
mapMWithKey :: (Ord k, Monad m) => (k -> a -> m b) -> MinPQueue k a -> m (MinPQueue k b)
mapMWithKey f = go empty
  where
    go !acc q =
      case minViewWithKey q of
        Nothing           -> pure acc
        Just ((k, a), q') -> do
          b <- f k a
          let !acc' = insertMax' k b acc
          go acc' q'

-- | Natural numbers revealing whether something is 'Zero' or 'Succ'.
type Ranky = Nattish Zero Succ

-- | \(O(n)\). An unordered traversal over a priority queue, in no particular order.
-- While there is no guarantee in which order the elements are traversed, the resulting
-- priority queue will be perfectly valid.
{-# INLINABLE traverseWithKeyU #-}
traverseWithKeyU :: forall f k a b. Applicative f => (k -> a -> f b) -> MinPQueue k a -> f (MinPQueue k b)
traverseWithKeyU _ Empty = pure Empty
traverseWithKeyU f (MinPQ n k a ts) = liftA2 (\a' !ts' -> MinPQ n k a' ts') (f k a) (traverseHeapU f ts)

{-# INLINABLE traverseHeapU #-}
traverseHeapU :: forall f k a b. Applicative f => (k -> a -> f b) -> BinomHeap k a -> f (BinomHeap k b)
traverseHeapU f = traverseForest Zeroy
  where
    traverseForest :: Ranky rk -> BinomForest rk k a -> f (BinomForest rk k b)
    traverseForest !_rky Nil = pure Nil
    traverseForest !rky (Skip rest) = (Skip $!) <$> traverseForest (Succy rky) rest
    traverseForest !rky (Cons t rest) =
      liftA2 (\ !t' !rest' -> Cons t' rest') (traverseTree rky t) (traverseForest (Succy rky) rest)

    {-# INLINE traverseTree #-}
    traverseTree :: Ranky rk -> BinomTree rk k a -> f (BinomTree rk k b)
    traverseTree Zeroy (BinomTree k (Zero a)) =
      -- We've reached a value, so we don't force the result.
      BinomTree k . Zero <$> f k a
    traverseTree (Succy rky) (BinomTree k ts) =
      -- We're not at a value, so we force the tree list.
      (BinomTree k $!) <$> traverseTrees rky k ts

    traverseTrees :: Ranky rk -> k -> Succ rk k a -> f (Succ rk k b)
    traverseTrees Zeroy !k2 (Succ (BinomTree k1 (Zero a1)) (Zero a2)) =
      -- The right subtree is a value, so we don't force it.
      liftA2 (\b1 b2 -> Succ (BinomTree k1 (Zero b1)) (Zero b2)) (f k1 a1) (f k2 a2)
    traverseTrees (Succy rky) !k (Succ t ts) =
      -- Whew; no values. We're safe to force.
      liftA2 (\ !t' !ts' -> Succ t' ts') (traverseTree (Succy rky) t) (traverseTrees rky k ts)

-- | \(O(\log n)\). @seqSpine q r@ forces the spine of @q@ and returns @r@.
--
-- Note: The spine of a 'MinPQueue' is stored somewhat lazily. In earlier
-- versions of this package, some operations could produce chains of thunks
-- along the spine, occasionally necessitating manual forcing. Now, all
-- operations are careful to force enough to avoid this problem.
{-# DEPRECATED seqSpine "This function is no longer necessary or useful." #-}
seqSpine :: MinPQueue k a -> b -> b
seqSpine Empty z0 = z0
seqSpine (MinPQ _ _ _ ts0) z0 = ts0 `seqSpineF` z0 where
  seqSpineF :: BinomForest rk k a -> b -> b
  seqSpineF ts z = case ts of
    Nil        -> z
    Skip ts'   -> seqSpineF ts' z
    Cons _ ts' -> seqSpineF ts' z

class NFRank rk where
  rnfRk :: (NFData k, NFData a) => rk k a -> ()

instance NFRank Zero where
  rnfRk (Zero a) = rnf a

instance NFRank rk => NFRank (Succ rk) where
  rnfRk (Succ t ts) = t `deepseq` rnfRk ts

instance (NFData k, NFData a, NFRank rk) => NFData (BinomTree rk k a) where
  rnf (BinomTree k ts) = k `deepseq` rnfRk ts

instance (NFData k, NFData a, NFRank rk) => NFData (BinomForest rk k a) where
  rnf Nil = ()
  rnf (Skip tss) = rnf tss
  rnf (Cons t tss) = t `deepseq` rnf tss

instance (NFData k, NFData a) => NFData (MinPQueue k a) where
  rnf Empty = ()
  rnf (MinPQ _ k a ts) = k `deepseq` a `deepseq` rnf ts

instance Functor (MinPQueue k) where
  fmap = imap . const

instance FunctorWithIndex k (MinPQueue k) where
  imap = coerce
    (traverseWithKeyU :: (k -> a -> Identity b) -> MinPQueue k a -> Identity (MinPQueue k b))

instance Ord k => Foldable (MinPQueue k) where
  foldr   = foldrWithKey . const
  foldl f = foldlWithKey (const . f)
  length = size
  null = null

instance Ord k => FoldableWithIndex k (MinPQueue k) where
  ifoldr   = foldrWithKey
  ifoldl f = foldlWithKey (flip f)

-- | Traverses in ascending order. 'mapM' is strictly accumulating like
-- 'mapMWithKey'.
instance Ord k => Traversable (MinPQueue k) where
  traverse = traverseWithKey . const
  mapM = mapMWithKey . const
  sequence = mapM id

instance Ord k => TraversableWithIndex k (MinPQueue k) where
  itraverse = traverseWithKey