From 02cb5ba84a1f19e812bac813e2d23e954ed6d965 Mon Sep 17 00:00:00 2001
From: David Feuer <David.Feuer@gmail.com>
Date: Sun, 26 Mar 2023 21:01:48 -0400
Subject: [PATCH 1/2] Strict maps

Closes #60
---
 CHANGELOG.md                          |  2 ++
 src/Data/PQueue/Prio/Internals.hs     | 40 +++++++++++++++++++++++++++
 src/Data/PQueue/Prio/Max.hs           |  3 ++
 src/Data/PQueue/Prio/Max/Internals.hs | 28 +++++++++++++++++--
 src/Data/PQueue/Prio/Min.hs           |  3 ++
 5 files changed, 74 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index a247452..66ce8e6 100644
--- a/CHANGELOG.md
+++ b/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).
diff --git a/src/Data/PQueue/Prio/Internals.hs b/src/Data/PQueue/Prio/Internals.hs
index 131ec9d..5a03eb4 100644
--- a/src/Data/PQueue/Prio/Internals.hs
+++ b/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 (\ !b !q' -> MinPQ n k b q') (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 (\ !p !q -> Cons (BinomTree k p q)) (f k a) (fCh ts) (traverseForest' f fCh' tss)
+  where
+    fCh' (Succ (BinomTree k a ts) tss)
+      = liftA3 (\ !p !q -> Succ (BinomTree k p q)) (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
diff --git a/src/Data/PQueue/Prio/Max.hs b/src/Data/PQueue/Prio/Max.hs
index 1bc39d6..04c0368 100644
--- a/src/Data/PQueue/Prio/Max.hs
+++ b/src/Data/PQueue/Prio/Max.hs
@@ -57,6 +57,7 @@ module Data.PQueue.Prio.Max (
   -- ** Map
   map,
   mapWithKey,
+  mapWithKey',
   mapKeys,
   mapKeysMonotonic,
   -- ** Fold
@@ -64,6 +65,7 @@ module Data.PQueue.Prio.Max (
   foldlWithKey,
   -- ** Traverse
   traverseWithKey,
+  traverseWithKey',
   mapMWithKey,
   -- * Subsets
   -- ** Indexed
@@ -110,6 +112,7 @@ module Data.PQueue.Prio.Max (
   foldlWithKeyU',
   traverseU,
   traverseWithKeyU,
+  traverseWithKeyU',
   keysU,
   elemsU,
   assocsU,
diff --git a/src/Data/PQueue/Prio/Max/Internals.hs b/src/Data/PQueue/Prio/Max/Internals.hs
index 701bb28..359b424 100644
--- a/src/Data/PQueue/Prio/Max/Internals.hs
+++ b/src/Data/PQueue/Prio/Max/Internals.hs
@@ -43,6 +43,7 @@ module Data.PQueue.Prio.Max.Internals (
   -- ** Map
   map,
   mapWithKey,
+  mapWithKey',
   mapKeys,
   mapKeysMonotonic,
   -- ** Fold
@@ -50,6 +51,7 @@ module Data.PQueue.Prio.Max.Internals (
   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
diff --git a/src/Data/PQueue/Prio/Min.hs b/src/Data/PQueue/Prio/Min.hs
index c383d1b..1dbfd68 100644
--- a/src/Data/PQueue/Prio/Min.hs
+++ b/src/Data/PQueue/Prio/Min.hs
@@ -67,6 +67,7 @@ module Data.PQueue.Prio.Min (
   -- ** Map
   map,
   mapWithKey,
+  mapWithKey',
   mapKeys,
   mapKeysMonotonic,
   -- ** Fold
@@ -74,6 +75,7 @@ module Data.PQueue.Prio.Min (
   foldlWithKey,
   -- ** Traverse
   traverseWithKey,
+  traverseWithKey',
   mapMWithKey,
   -- * Subsets
   -- ** Indexed
@@ -120,6 +122,7 @@ module Data.PQueue.Prio.Min (
   foldlWithKeyU',
   traverseU,
   traverseWithKeyU,
+  traverseWithKeyU',
   keysU,
   elemsU,
   assocsU,

From 27605ba549019e1940598af18abe28747a69d297 Mon Sep 17 00:00:00 2001
From: David Feuer <David.Feuer@gmail.com>
Date: Fri, 31 Mar 2023 22:09:41 -0400
Subject: [PATCH 2/2] Fix strictness

---
 src/Data/PQueue/Prio/Internals.hs | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/Data/PQueue/Prio/Internals.hs b/src/Data/PQueue/Prio/Internals.hs
index 5a03eb4..86a26ad 100644
--- a/src/Data/PQueue/Prio/Internals.hs
+++ b/src/Data/PQueue/Prio/Internals.hs
@@ -704,7 +704,7 @@ traverseWithKey f q = case minViewWithKey q of
 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')
+  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,
@@ -735,7 +735,7 @@ traverseWithKeyU f (MinPQ n k a ts) = liftA2 (MinPQ n k) (f k a) (traverseForest
 -- @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 (\ !b !q' -> MinPQ n k b q') (f k a) (traverseForest' f (const (pure Zero)) ts)
+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) #-}
@@ -756,10 +756,10 @@ traverseForest' f fCh ts0 = case ts0 of
   Nil       -> pure Nil
   Skip ts'  -> (Skip $!) <$> traverseForest f fCh' ts'
   Cons (BinomTree k a ts) tss
-    -> liftA3 (\ !p !q -> Cons (BinomTree k p q)) (f k a) (fCh ts) (traverseForest' f fCh' 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 (\ !p !q -> Succ (BinomTree k p q)) (f k a) (fCh ts) (fCh 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