Strict maps

CHANGELOG.md

@@ -10,6 +10,8 @@
     `MinPQueue` work "incrementally", like the one for `MinQueue`.
     ([#92](http://github.com/lspitzner/pqueue/pull/92))
 
+  * Add strict maps and traversals.
+
 ## 1.4.3.0 -- 2022-10-30
 
   * Add instances for [indexed-traversable](https://hackage.haskell.org/package/indexed-traversable).

src/Data/PQueue/Prio/Internals.hs

@@ -24,6 +24,7 @@ module Data.PQueue.Prio.Internals (
   updateMinWithKeyA',
   minViewWithKey,
   mapWithKey,
+  mapWithKey',
   mapKeysMonotonic,
   mapMaybeWithKey,
   mapEitherWithKey,
@@ -43,8 +44,10 @@ module Data.PQueue.Prio.Internals (
   foldlWithKeyU,
   foldlWithKeyU',
   traverseWithKey,
+  traverseWithKey',
   mapMWithKey,
   traverseWithKeyU,
+  traverseWithKeyU',
   seqSpine,
   mapForest,
   unions
@@ -342,6 +345,12 @@ minViewWithKey (MinPQ n k a ts) = Just ((k, a), extractHeap n ts)
 mapWithKey :: (k -> a -> b) -> MinPQueue k a -> MinPQueue k b
 mapWithKey f = runIdentity . traverseWithKeyU (Identity .: f)
 
+-- | \(O(n)\). Map a function over all values in the queue, forcing the results.
+--
+-- @since 1.5.0
+mapWithKey' :: (k -> a -> b) -> MinPQueue k a -> MinPQueue k b
+mapWithKey' f = runIdentity . traverseWithKeyU' (Identity .: f)
+
 -- | \(O(n)\). @'mapKeysMonotonic' f q == 'mapKeys' f q@, but only works when @f@ is strictly
 -- monotonic. /The precondition is not checked./ This function has better performance than
 -- 'mapKeys'.
@@ -688,6 +697,15 @@ traverseWithKey f q = case minViewWithKey q of
   Nothing      -> pure empty
   Just ((k, a), q')  -> liftA2 (insertMin k) (f k a) (traverseWithKey f q')
 
+-- | A version of 'traverseWithKey' that forces all the results before
+-- installing them in a queue.
+--
+-- @since 1.5.0
+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 () = ⊥@.
@@ -709,6 +727,16 @@ traverseWithKeyU :: Applicative f => (k -> a -> f b) -> MinPQueue k a -> f (MinP
 traverseWithKeyU _ Empty = pure Empty
 traverseWithKeyU f (MinPQ n k a ts) = liftA2 (MinPQ n k) (f k a) (traverseForest f (const (pure Zero)) ts)
 
+-- | \(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. The results are forced before they are installed
+-- in the queue.
+--
+-- @since 1.5.0
+traverseWithKeyU' :: 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) (traverseForest' f (const (pure Zero)) ts)
+
 {-# SPECIALIZE traverseForest :: (k -> a -> Identity b) -> (rk k a -> Identity (rk k b)) -> BinomForest rk k a ->
   Identity (BinomForest rk k b) #-}
 traverseForest :: (Applicative f) => (k -> a -> f b) -> (rk k a -> f (rk k b)) -> BinomForest rk k a -> f (BinomForest rk k b)
@@ -721,6 +749,18 @@ traverseForest f fCh ts0 = case ts0 of
     fCh' (Succ (BinomTree k a ts) tss)
       = Succ <$> (BinomTree k <$> f k a <*> fCh ts) <*> fCh tss
 
+{-# SPECIALIZE traverseForest' :: (k -> a -> Identity b) -> (rk k a -> Identity (rk k b)) -> BinomForest rk k a ->
+  Identity (BinomForest rk k b) #-}
+traverseForest' :: Applicative f => (k -> a -> f b) -> (rk k a -> f (rk k b)) -> BinomForest rk k a -> f (BinomForest rk k b)
+traverseForest' f fCh ts0 = case ts0 of
+  Nil       -> pure Nil
+  Skip ts'  -> (Skip $!) <$> traverseForest f fCh' ts'
+  Cons (BinomTree k a ts) tss
+    -> liftA3 (\ !a' !ts' !tss' -> Cons (BinomTree k a' ts') tss') (f k a) (fCh ts) (traverseForest' f fCh' tss)
+  where
+    fCh' (Succ (BinomTree k a ts) tss)
+      = liftA3 (\ !a' !ts' !tss' -> Succ (BinomTree k a' ts') tss') (f k a) (fCh ts) (fCh tss)
+
 -- | Unordered right fold on a binomial forest.
 foldrWithKeyF_ :: (k -> a -> b -> b) -> (rk k a -> b -> b) -> BinomForest rk k a -> b -> b
 foldrWithKeyF_ f fCh ts0 z0 = case ts0 of

src/Data/PQueue/Prio/Max.hs

@@ -57,13 +57,15 @@ module Data.PQueue.Prio.Max (
   -- ** Map
   map,
   mapWithKey,
+  mapWithKey',
   mapKeys,
   mapKeysMonotonic,
   -- ** Fold
   foldrWithKey,
   foldlWithKey,
   -- ** Traverse
   traverseWithKey,
+  traverseWithKey',
   mapMWithKey,
   -- * Subsets
   -- ** Indexed
@@ -110,6 +112,7 @@ module Data.PQueue.Prio.Max (
   foldlWithKeyU',
   traverseU,
   traverseWithKeyU,
+  traverseWithKeyU',
   keysU,
   elemsU,
   assocsU,

src/Data/PQueue/Prio/Max/Internals.hs

@@ -43,13 +43,15 @@ module Data.PQueue.Prio.Max.Internals (
   -- ** Map
   map,
   mapWithKey,
+  mapWithKey',
   mapKeys,
   mapKeysMonotonic,
   -- ** Fold
   foldrWithKey,
   foldlWithKey,
   -- ** Traverse
   traverseWithKey,
+  traverseWithKey',
   mapMWithKey,
   -- * Subsets
   -- ** Indexed
@@ -96,6 +98,7 @@ module Data.PQueue.Prio.Max.Internals (
   foldlWithKeyU',
   traverseU,
   traverseWithKeyU,
+  traverseWithKeyU',
   keysU,
   elemsU,
   assocsU,
@@ -334,6 +337,13 @@ map = mapWithKey . const
 mapWithKey :: (k -> a -> b) -> MaxPQueue k a -> MaxPQueue k b
 mapWithKey f (MaxPQ q) = MaxPQ (Q.mapWithKey (f . unDown) q)
 
+-- | \(O(n)\). A version of 'mapWithKey' that forces all the elements before
+-- installing them in the result queue.
+--
+-- @since 1.5.0
+mapWithKey' :: (k -> a -> b) -> MaxPQueue k a -> MaxPQueue k b
+mapWithKey' f (MaxPQ q) = MaxPQ (Q.mapWithKey' (f . unDown) q)
+
 -- | \(O(n)\). Map a function over all values in the queue.
 mapKeys :: Ord k' => (k -> k') -> MaxPQueue k a -> MaxPQueue k' a
 mapKeys f (MaxPQ q) = MaxPQ (Q.mapKeys (fmap f) q)
@@ -367,6 +377,13 @@ foldlWithKey f z0 (MaxPQ q) = Q.foldlWithKey (\z -> f z . unDown) z0 q
 traverseWithKey :: (Ord k, Applicative f) => (k -> a -> f b) -> MaxPQueue k a -> f (MaxPQueue k b)
 traverseWithKey f (MaxPQ q) = MaxPQ <$> Q.traverseWithKey (f . unDown) q
 
+-- | A version of 'traverseWithKey' that forces each element before
+-- installing it in a result queue.
+--
+-- @since 1.5.0
+traverseWithKey' :: (Ord k, Applicative f) => (k -> a -> f b) -> MaxPQueue k a -> f (MaxPQueue k b)
+traverseWithKey' f (MaxPQ q) = MaxPQ <$> Q.traverseWithKey' (f . unDown) 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 () = ⊥@.
@@ -550,15 +567,22 @@ foldlWithKeyU' f z0 (MaxPQ q) = Q.foldlWithKeyU' (\z -> f z . unDown) z0 q
 -- | \(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.
-traverseU :: (Applicative f) => (a -> f b) -> MaxPQueue k a -> f (MaxPQueue k b)
+traverseU :: Applicative f => (a -> f b) -> MaxPQueue k a -> f (MaxPQueue k b)
 traverseU = traverseWithKeyU . const
 
 -- | \(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.
-traverseWithKeyU :: (Applicative f) => (k -> a -> f b) -> MaxPQueue k a -> f (MaxPQueue k b)
+traverseWithKeyU :: Applicative f => (k -> a -> f b) -> MaxPQueue k a -> f (MaxPQueue k b)
 traverseWithKeyU f (MaxPQ q) = MaxPQ <$> Q.traverseWithKeyU (f . unDown) q
 
+-- | A version of 'traverseWithKeyU' that forces each value before installing
+-- it in a result queue.
+--
+-- @since 1.5.0
+traverseWithKeyU' :: Applicative f => (k -> a -> f b) -> MaxPQueue k a -> f (MaxPQueue k b)
+traverseWithKeyU' f (MaxPQ q) = MaxPQ <$> Q.traverseWithKeyU' (f . unDown) q
+
 -- | \(O(n)\). Return all keys of the queue in no particular order.
 keysU :: MaxPQueue k a -> [k]
 keysU = fmap fst . toListU

src/Data/PQueue/Prio/Min.hs

@@ -67,13 +67,15 @@ module Data.PQueue.Prio.Min (
   -- ** Map
   map,
   mapWithKey,
+  mapWithKey',
   mapKeys,
   mapKeysMonotonic,
   -- ** Fold
   foldrWithKey,
   foldlWithKey,
   -- ** Traverse
   traverseWithKey,
+  traverseWithKey',
   mapMWithKey,
   -- * Subsets
   -- ** Indexed
@@ -120,6 +122,7 @@ module Data.PQueue.Prio.Min (
   foldlWithKeyU',
   traverseU,
   traverseWithKeyU,
+  traverseWithKeyU',
   keysU,
   elemsU,
   assocsU,