Foldable traversable

CHANGELOG.md

@@ -1,5 +1,16 @@
 # Revision history for pqueue
 
+## 1.4.2
+
+  * Overall performance has improved greatly, especially when there are many
+    insertions and/or merges in a row. Insertion, deletion, and merge are now
+    *worst case* logarithmic, while maintaining their previous amortized
+    bounds. (#26)
+
+  * New `mapMWithKey` functions optimized for working in strict monads. These
+    are used to implement the `mapM` and `sequence` methods of `Traversable`.
+    (#46)
+
 ## 1.4.1.4  -- 2021-12-04
 
   * Maintenance release for ghc-9.0 & ghc-9.2 support

pqueue.cabal

@@ -14,7 +14,6 @@ build-type:         Simple
 cabal-version:      >= 1.10
 tested-with:        GHC == 7.10.3, GHC == 8.0.2, GHC == 8.2.2, GHC == 8.4.4, GHC == 8.6.5, GHC == 8.8.4, GHC == 8.10.7, GHC == 9.0.1, GHC == 9.2.1
 extra-source-files:
-  include/Typeable.h
   CHANGELOG.md
 
 source-repository head
@@ -42,6 +41,10 @@ library
   if impl(ghc) {
     default-extensions: DeriveDataTypeable
   }
+  other-extensions:
+      BangPatterns
+    , CPP
+    , StandaloneDeriving
   ghc-options: {
     -- We currently need -fspec-constr to get GHC to compile conversions
     -- from lists well. We could (and probably should) write those a

src/Data/PQueue/Internals.hs

@@ -46,12 +46,6 @@ import Prelude hiding (null)
 
 -- | A priority queue with elements of type @a@. Supports extracting the minimum element.
 data MinQueue a = Empty | MinQueue {-# UNPACK #-} !Int !a !(BinomHeap a)
-#if __GLASGOW_HASKELL__>=707
-  deriving Typeable
-#else
-#include "Typeable.h"
-INSTANCE_TYPEABLE1(MinQueue,minQTC,"MinQueue")
-#endif
 
 #ifdef __GLASGOW_HASKELL__
 instance (Ord a, Data a) => Data (MinQueue a) where

src/Data/PQueue/Max.hs

@@ -108,7 +108,7 @@ build f = f (:) []
 -- Implemented as a wrapper around 'Min.MinQueue'.
 newtype MaxQueue a = MaxQ (Min.MinQueue (Down a))
 # if __GLASGOW_HASKELL__
-  deriving (Eq, Ord, Data, Typeable)
+  deriving (Eq, Ord, Data)
 # else
   deriving (Eq, Ord)
 # endif

src/Data/PQueue/Prio/Internals.hs

@@ -37,6 +37,7 @@ module Data.PQueue.Prio.Internals (
   ) where
 
 import Control.Applicative.Identity (Identity(Identity, runIdentity))
+import Control.Applicative (liftA2, liftA3)
 import Control.DeepSeq (NFData(rnf), deepseq)
 import Data.List (foldl')
 
@@ -74,9 +75,6 @@ infixr 8 .:
 -- | A priority queue where values of type @a@ are annotated with keys of type @k@.
 -- The queue supports extracting the element with minimum key.
 data MinPQueue k a = Empty | MinPQ {-# UNPACK #-} !Int !k a !(BinomHeap k a)
-#if __GLASGOW_HASKELL__
-  deriving (Typeable)
-#endif
 
 data BinomForest rk k a =
   Nil |
@@ -504,7 +502,7 @@ foldlWithKeyU f z0 (MinPQ _ k0 a0 ts) = foldlWithKeyF_ (\k a z -> f z k a) (cons
 -- priority queue will be perfectly valid.
 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) = MinPQ n k <$> f k a <*> traverseForest f (const (pure Zero)) ts
+traverseWithKeyU f (MinPQ n k a ts) = liftA2 (MinPQ n k) (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) #-}
@@ -513,7 +511,7 @@ traverseForest f fCh ts0 = case ts0 of
   Nil       -> pure Nil
   Skip ts'  -> Skip <$> traverseForest f fCh' ts'
   Cons (BinomTree k a ts) tss
-    -> Cons <$> (BinomTree k <$> f k a <*> fCh ts) <*> traverseForest f fCh' 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)
       = Succ <$> (BinomTree k <$> f k a <*> fCh ts) <*> fCh tss

src/Data/PQueue/Prio/Max.hs

@@ -1,3 +1,4 @@
+{-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE CPP #-}
 {-# OPTIONS_GHC -fno-warn-orphans #-}
 
@@ -63,6 +64,7 @@ module Data.PQueue.Prio.Max (
   foldlWithKey,
   -- ** Traverse
   traverseWithKey,
+  mapMWithKey,
   -- * Subsets
   -- ** Indexed
   take,
@@ -114,6 +116,7 @@ module Data.PQueue.Prio.Max (
   )
   where
 
+import Control.Applicative (liftA2)
 import Data.Maybe (fromMaybe)
 import Data.PQueue.Prio.Max.Internals
 
@@ -122,6 +125,7 @@ import Data.Semigroup (Semigroup((<>)))
 #endif
 
 import Prelude hiding (map, filter, break, span, takeWhile, dropWhile, splitAt, take, drop, (!!), null)
+import qualified Data.Foldable as F
 
 import qualified Data.PQueue.Prio.Min as Q
 
@@ -172,8 +176,15 @@ instance Ord k => Foldable (MaxPQueue k) where
   foldr f z (MaxPQ q) = foldr f z q
   foldl f z (MaxPQ q) = foldl f z q
 
+  length = size
+  null = null
+
+-- | Traverses in descending order. 'mapM' is strictly accumulating like
+-- 'mapMWithKey'.
 instance Ord k => Traversable (MaxPQueue k) where
   traverse f (MaxPQ q) = MaxPQ <$> traverse f q
+  mapM = mapMWithKey . const
+  sequence = mapM id
 
 -- | /O(1)/. Returns the empty priority queue.
 empty :: MaxPQueue k a
@@ -298,9 +309,25 @@ foldlWithKey f z0 (MaxPQ q) = Q.foldlWithKey (\z -> f z . unDown) z0 q
 -- (@'traverseWithKey' f q == 'fromDescList' <$> 'traverse' ('uncurry' f) ('toDescList' 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) -> 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 () = ⊥@.
+mapMWithKey :: (Ord k, Monad m) => (k -> a -> m b) -> MaxPQueue k a -> m (MaxPQueue k b)
+mapMWithKey f = go empty
+  where
+    go !acc q =
+      case maxViewWithKey q of
+        Nothing           -> pure acc
+        Just ((k, a), q') -> do
+          b <- f k a
+          let !acc' = insertMin' k b acc
+          go acc' q'
+
 -- | /O(k log n)/. Takes the first @k@ (key, value) pairs in the queue, or the first @n@ if @k >= n@.
 -- (@'take' k q == 'List.take' k ('toDescList' q)@)
 take :: Ord k => Int -> MaxPQueue k a -> [(k, a)]

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

@@ -4,24 +4,26 @@ module Data.PQueue.Prio.Max.Internals where
 
 import Control.DeepSeq (NFData(rnf))
 
-# if __GLASGOW_HASKELL__
-import Data.Data (Data, Typeable)
-# endif
+#if __GLASGOW_HASKELL__
+import Data.Data (Data)
+#endif
 
 import Data.PQueue.Prio.Internals (MinPQueue)
+import qualified Data.PQueue.Prio.Internals as Internals
 
 newtype Down a = Down { unDown :: a }
 # if __GLASGOW_HASKELL__
-  deriving (Eq, Data, Typeable)
+  deriving (Eq, Data)
 # else
   deriving (Eq)
 # endif
 
+
 -- | A priority queue where values of type @a@ are annotated with keys of type @k@.
 -- The queue supports extracting the element with maximum key.
 newtype MaxPQueue k a = MaxPQ (MinPQueue (Down k) a)
 # if __GLASGOW_HASKELL__
-  deriving (Eq, Ord, Data, Typeable)
+  deriving (Eq, Ord, Data)
 # else
   deriving (Eq, Ord)
 # endif
@@ -45,3 +47,6 @@ instance Foldable Down where
 
 instance Traversable Down where
   traverse f (Down a) = Down <$> f a
+
+insertMin' :: k -> a -> MaxPQueue k a -> MaxPQueue k a
+insertMin' k a (MaxPQ q) = MaxPQ (Internals.insertMax' (Down k) a q)

src/Data/PQueue/Prio/Min.hs

@@ -1,3 +1,4 @@
+{-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE CPP #-}
 {-# OPTIONS_GHC -fno-warn-orphans #-}
 
@@ -63,6 +64,7 @@ module Data.PQueue.Prio.Min (
   foldlWithKey,
   -- ** Traverse
   traverseWithKey,
+  mapMWithKey,
   -- * Subsets
   -- ** Indexed
   take,
@@ -114,6 +116,7 @@ module Data.PQueue.Prio.Min (
   )
   where
 
+import Control.Applicative (liftA2)
 import qualified Data.List as List
 import Data.Maybe (fromMaybe)
 
@@ -216,10 +219,26 @@ mapKeys f q = fromList [(f k, a) | (k, a) <- toListU q]
 -- (@'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')  -> 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,
+-- 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'
 
 -- | /O(n)/. Map values and collect the 'Just' results.
 mapMaybe :: Ord k => (a -> Maybe b) -> MinPQueue k a -> MinPQueue k b
@@ -400,6 +419,12 @@ instance Functor (MinPQueue k) where
 instance Ord k => Foldable (MinPQueue k) where
   foldr   = foldrWithKey . const
   foldl f = foldlWithKey (const . f)
+  length = size
+  null = null
 
+-- | 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

tests/PQueueTests.hs

@@ -1,12 +1,13 @@
 module Main (main) where
 
-import qualified Data.PQueue.Prio.Max as PMax ()
+import qualified Data.PQueue.Prio.Max as PMax
 import qualified Data.PQueue.Prio.Min as PMin
 import qualified Data.PQueue.Max as Max ()
 import qualified Data.PQueue.Min as Min
 
 import Test.QuickCheck
-import Test.QuickCheck.Poly (OrdA)
+import Test.QuickCheck.Poly (OrdA, B, C)
+import Test.QuickCheck.Function (Fun, applyFun2)
 
 import System.Exit
 
@@ -125,6 +126,26 @@ validFoldrU :: [Int] -> Bool
 validFoldrU xs = Min.foldrU (+) 0 q == List.sum xs
   where q = Min.fromList xs
 
+validTraverseWithKey :: Fun (OrdA, B) [C] -> Short (OrdA, B) -> Property
+validTraverseWithKey f (Short xs) =
+  PMin.traverseWithKey (applyFun2 f) q === PMin.mapMWithKey (applyFun2 f) q
+  where
+    q = PMin.fromList xs
+
+validTraverseWithKeyMax :: Fun (OrdA, B) [C] -> Short (OrdA, B) -> Property
+validTraverseWithKeyMax f (Short xs) =
+  PMax.traverseWithKey (applyFun2 f) q === PMax.mapMWithKey (applyFun2 f) q
+  where
+    q = PMax.fromList xs
+
+-- Lists of between 0 and 3 elements
+newtype Short a = Short [a]
+  deriving Show
+instance Arbitrary a => Arbitrary (Short a) where
+  arbitrary = do
+    n <- frequency [(1, pure 0), (5, pure 1), (4, pure 2), (3, pure 3)]
+    Short <$> vectorOf n arbitrary
+
 main :: IO ()
 main = do
   check validMinToAscList
@@ -153,6 +174,8 @@ main = do
   check validFoldl
   check validFoldlU
   check validFoldrU
+  check validTraverseWithKey
+  check validTraverseWithKeyMax
   putStrLn "all tests passed"
 
 isPass :: Result -> Bool