{-# 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, insertBehind, 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 import Control.Applicative (liftA2, liftA3, Const (..)) import Control.DeepSeq (NFData(rnf), deepseq) import Data.Coerce (coerce) import Data.Functor.Identity (Identity(Identity, runIdentity)) import qualified Data.List as List #if MIN_VERSION_base(4,9,0) import Data.Semigroup (Semigroup(..), stimesMonoid, Endo (..), Dual (..)) #else import Data.Monoid ((<>), Endo (..), Dual (..)) #endif 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 #if MIN_VERSION_base(4,9,0) instance Ord k => Semigroup (MinPQueue k a) where (<>) = union stimes = stimesMonoid {-# INLINABLE stimes #-} #endif 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 "fromAscList " . shows (toAscList xs) instance (Read k, Read a) => Read (MinPQueue k a) where #ifdef __GLASGOW_HASKELL__ readPrec = parens $ prec 10 $ do Ident "fromAscList" <- lexP xs <- readPrec return (fromAscList xs) readListPrec = readListPrecDefault #else readsPrec p = readParen (p > 10) $ \r -> do ("fromAscList",s) <- lex r (xs,t) <- reads s return (fromAscList 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) -- | \(O(n)\) (an earlier implementation had \(O(1)\) but was buggy). -- Insert an element with the specified key into the priority queue, -- putting it behind elements whose key compares equal to the -- inserted one. {-# DEPRECATED insertBehind "This function is not reliable." #-} insertBehind :: Ord k => k -> a -> MinPQueue k a -> MinPQueue k a insertBehind k v q = let (smaller, larger) = spanKey (<= k) q in foldr (uncurry insert) (insert k v larger) smaller spanKey :: Ord k => (k -> Bool) -> MinPQueue k a -> ([(k, a)], MinPQueue k a) spanKey p q = case minViewWithKey q of Just (t@(k, _), q') | p k -> let (kas, q'') = spanKey p q' in (t : kas, q'') _ -> ([], q) -- | 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 == 'mapKeys' f q@, but only works when -- @f@ is (weakly) monotonic. /The precondition is not checked./ This function -- has better performance than '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