Streamline and expand folds

pqueue.cabal

@@ -37,6 +37,7 @@ library
     Data.PQueue.Prio.Internals
     Data.PQueue.Internals
     Data.PQueue.Internals.Down
+    Data.PQueue.Internals.Foldable
     Data.PQueue.Prio.Max.Internals
     Control.Applicative.Identity
   if impl(ghc) {

src/Data/PQueue/Internals.hs

@@ -1,4 +1,6 @@
-{-# LANGUAGE CPP, StandaloneDeriving #-}
+{-# LANGUAGE CPP #-}
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE StandaloneDeriving #-}
 
 module Data.PQueue.Internals (
   MinQueue (..),
@@ -34,8 +36,10 @@ module Data.PQueue.Internals (
   fromList,
   mapU,
   fromAscList,
+  foldMapU,
   foldrU,
   foldlU,
+  foldlU',
 --   traverseU,
   seqSpine,
   unions
@@ -47,6 +51,7 @@ import Data.Foldable (foldl')
 import Data.Semigroup (Semigroup(..), stimesMonoid)
 #endif
 
+import Data.PQueue.Internals.Foldable
 #ifdef __GLASGOW_HASKELL__
 import Data.Data
 import Text.Read (Lexeme(Ident), lexP, parens, prec,
@@ -603,6 +608,63 @@ instance Functor rk => Functor (BinomForest rk) where
   fmap f (Skip ts) = Skip (fmap f ts)
   fmap f (Cons t ts) = Cons (fmap f t) (fmap f ts)
 
+instance Foldr Zero where
+  foldr_ _ z ~Zero = z
+
+instance Foldl Zero where
+  foldl_ _ z ~Zero = z
+
+instance Foldl' Zero where
+  foldl'_ _ z ~Zero = z
+
+instance FoldMap Zero where
+  foldMap_ _ ~Zero = mempty
+
+instance Foldr rk => Foldr (Succ rk) where
+  foldr_ f z (Succ t ts) = foldr_ f (foldr_ f z ts) t
+
+instance Foldl rk => Foldl (Succ rk) where
+  foldl_ f z (Succ t ts) = foldl_ f (foldl_ f z t) ts
+
+instance Foldl' rk => Foldl' (Succ rk) where
+  foldl'_ f !z (Succ t ts) = foldl'_ f (foldl'_ f z t) ts
+
+instance FoldMap rk => FoldMap (Succ rk) where
+  foldMap_ f (Succ t ts) = foldMap_ f t `mappend` foldMap_ f ts
+
+instance Foldr rk => Foldr (BinomTree rk) where
+  foldr_ f z (BinomTree x ts) = x `f` foldr_ f z ts
+
+instance Foldl rk => Foldl (BinomTree rk) where
+  foldl_ f z (BinomTree x ts) = foldl_ f (z `f` x) ts
+
+instance Foldl' rk => Foldl' (BinomTree rk) where
+  foldl'_ f !z (BinomTree x ts) = foldl'_ f (z `f` x) ts
+
+instance FoldMap rk => FoldMap (BinomTree rk) where
+  foldMap_ f (BinomTree x ts) = f x `mappend` foldMap_ f ts
+
+instance Foldr rk => Foldr (BinomForest rk) where
+  foldr_ _ z Nil          = z
+  foldr_ f z (Skip tss)   = foldr_ f z tss
+  foldr_ f z (Cons t tss) = foldr_ f (foldr_ f z tss) t
+
+instance Foldl rk => Foldl (BinomForest rk) where
+  foldl_ _ z Nil          = z
+  foldl_ f z (Skip tss)   = foldl_ f z tss
+  foldl_ f z (Cons t tss) = foldl_ f (foldl_ f z t) tss
+
+instance Foldl' rk => Foldl' (BinomForest rk) where
+  foldl'_ _ !z Nil          = z
+  foldl'_ f !z (Skip tss)   = foldl'_ f z tss
+  foldl'_ f !z (Cons t tss) = foldl'_ f (foldl'_ f z t) tss
+
+instance FoldMap rk => FoldMap (BinomForest rk) where
+  foldMap_ _ Nil = mempty
+  foldMap_ f (Skip tss)   = foldMap_ f tss
+  foldMap_ f (Cons t tss) = foldMap_ f t `mappend` foldMap_ f tss
+
+{-
 instance Foldable Zero where
   foldr _ z ~Zero = z
   foldl _ z ~Zero = z
@@ -622,6 +684,7 @@ instance Foldable rk => Foldable (BinomForest rk) where
   foldl _ z Nil          = z
   foldl f z (Skip tss)   = foldl f z tss
   foldl f z (Cons t tss) = foldl f (foldl f z t) tss
+-}
 
 -- instance Traversable Zero where
 --   traverse _ _ = pure Zero
@@ -645,12 +708,28 @@ mapU f (MinQueue n x ts) = MinQueue n (f x) (f <$> ts)
 -- | /O(n)/. Unordered right fold on a priority queue.
 foldrU :: (a -> b -> b) -> b -> MinQueue a -> b
 foldrU _ z Empty = z
-foldrU f z (MinQueue _ x ts) = x `f` foldr f z ts
+foldrU f z (MinQueue _ x ts) = x `f` foldr_ f z ts
 
--- | /O(n)/. Unordered left fold on a priority queue.
+-- | /O(n)/. Unordered left fold on a priority queue. This is rarely
+-- what you want; 'foldrU' and 'foldlU'' are more likely to perform
+-- well.
 foldlU :: (b -> a -> b) -> b -> MinQueue a -> b
 foldlU _ z Empty = z
-foldlU f z (MinQueue _ x ts) = foldl f (z `f` x) ts
+foldlU f z (MinQueue _ x ts) = foldl_ f (z `f` x) ts
+
+-- | /O(n)/. Unordered strict left fold on a priority queue.
+--
+-- @since 1.4.2
+foldlU' :: (b -> a -> b) -> b -> MinQueue a -> b
+foldlU' _ z Empty = z
+foldlU' f z (MinQueue _ x ts) = foldl'_ f (z `f` x) ts
+
+-- | /O(n)/. Unordered monoidal fold on a priority queue.
+--
+-- @since 1.4.2
+foldMapU :: Monoid m => (a -> m) -> MinQueue a -> m
+foldMapU _ Empty = mempty
+foldMapU f (MinQueue _ x ts) = f x `mappend` foldMap_ f ts
 
 {-# NOINLINE toListU #-}
 -- | /O(n)/. Returns the elements of the queue, in no particular order.
@@ -660,7 +739,7 @@ toListU q = foldrU (:) [] q
 {-# NOINLINE toListUApp #-}
 toListUApp :: MinQueue a -> [a] -> [a]
 toListUApp Empty app = app
-toListUApp (MinQueue _ x ts) app = x : foldr (:) app ts
+toListUApp (MinQueue _ x ts) app = x : foldr_ (:) app ts
 
 {-# RULES
 "toListU/build" [~1] forall q. toListU q = build (\c n -> foldrU c n q)

src/Data/PQueue/Internals/Down.hs

@@ -1,8 +1,10 @@
 {-# LANGUAGE CPP #-}
+{-# LANGUAGE BangPatterns #-}
 
 module Data.PQueue.Internals.Down where
 
 import Control.DeepSeq (NFData(rnf))
+import Data.Foldable (Foldable (..))
 
 #if __GLASGOW_HASKELL__
 import Data.Data (Data)
@@ -29,3 +31,4 @@ instance Functor Down where
 instance Foldable Down where
   foldr f z (Down a) = a `f` z
   foldl f z (Down a) = z `f` a
+  foldl' f !z (Down a) = z `f` a

src/Data/PQueue/Internals/Foldable.hs

@@ -0,0 +1,38 @@
+-- | Writing 'Foldable' instances for non-regular (AKA, nested) types in the
+-- natural manner leads to full `Foldable` dictionaries being constructed on
+-- each recursive call. This is pretty inefficient. It's better to construct
+-- exactly what we need instead.
+module Data.PQueue.Internals.Foldable
+  ( Foldr (..)
+  , Foldl (..)
+  , FoldMap (..)
+  , Foldl' (..)
+  , IFoldr (..)
+  , IFoldl (..)
+  , IFoldMap (..)
+  , IFoldl' (..)
+  ) where
+
+class Foldr t where
+  foldr_ :: (a -> b -> b) -> b -> t a -> b
+
+class IFoldr t where
+  foldrWithKey_ :: (k -> a -> b -> b) -> b -> t k a -> b
+
+class Foldl t where
+  foldl_ :: (b -> a -> b) -> b -> t a -> b
+
+class IFoldl t where
+  foldlWithKey_ :: (b -> k -> a -> b) -> b -> t k a -> b
+
+class FoldMap t where
+  foldMap_ :: Monoid m => (a -> m) -> t a -> m
+
+class IFoldMap t where
+  foldMapWithKey_ :: Monoid m => (k -> a -> m) -> t k a -> m
+
+class Foldl' t where
+  foldl'_ :: (b -> a -> b) -> b -> t a -> b
+
+class IFoldl' t where
+  foldlWithKey'_ :: (b -> k -> a -> b) -> b -> t k a -> b

src/Data/PQueue/Max.hs

@@ -74,6 +74,8 @@ module Data.PQueue.Max (
   mapU,
   foldrU,
   foldlU,
+  foldlU',
+  foldMapU,
   elemsU,
   toListU,
   -- * Miscellaneous operations
@@ -88,6 +90,8 @@ import Data.Maybe (fromMaybe)
 import Data.Semigroup (Semigroup(..), stimesMonoid)
 #endif
 
+import Data.Foldable (foldl')
+
 import qualified Data.PQueue.Min as Min
 import qualified Data.PQueue.Prio.Max.Internals as Prio
 import Data.PQueue.Internals.Down (Down(..))
@@ -273,10 +277,18 @@ mapU f (MaxQ q) = MaxQ (Min.mapU (\(Down a) -> Down (f a)) q)
 foldrU :: (a -> b -> b) -> b -> MaxQueue a -> b
 foldrU f z (MaxQ q) = Min.foldrU (flip (foldr f)) z q
 
+-- | /O(n)/. Unordered monoidal fold on a priority queue.
+foldMapU :: Monoid m => (a -> m) -> MaxQueue a -> m
+foldMapU f (MaxQ q) = Min.foldMapU (f . unDown) q
+
 -- | /O(n)/. Unordered left fold on a priority queue.
 foldlU :: (b -> a -> b) -> b -> MaxQueue a -> b
 foldlU f z (MaxQ q) = Min.foldlU (foldl f) z q
 
+-- | /O(n)/. Unordered strict left fold on a priority queue.
+foldlU' :: (b -> a -> b) -> b -> MaxQueue a -> b
+foldlU' f z (MaxQ q) = Min.foldlU' (foldl' f) z q
+
 {-# INLINE elemsU #-}
 -- | Equivalent to 'toListU'.
 elemsU :: MaxQueue a -> [a]

src/Data/PQueue/Min.hs

@@ -73,6 +73,8 @@ module Data.PQueue.Min (
   mapU,
   foldrU,
   foldlU,
+  foldlU',
+  foldMapU,
   elemsU,
   toListU,
   -- * Miscellaneous operations

src/Data/PQueue/Prio/Internals.hs

@@ -36,7 +36,9 @@ module Data.PQueue.Prio.Internals (
   fromList,
   fromAscList,
   foldrWithKeyU,
+  foldMapWithKeyU,
   foldlWithKeyU,
+  foldlWithKeyU',
   traverseWithKey,
   mapMWithKey,
   traverseWithKeyU,
@@ -49,6 +51,7 @@ import Control.Applicative.Identity (Identity(Identity, runIdentity))
 import Control.Applicative (liftA2, liftA3)
 import Control.DeepSeq (NFData(rnf), deepseq)
 import qualified Data.List as List
+import Data.PQueue.Internals.Foldable
 
 #if MIN_VERSION_base(4,9,0)
 import Data.Semigroup (Semigroup(..), stimesMonoid)
@@ -150,6 +153,40 @@ data BinomTree rk k a = BinomTree !k a !(rk k a)
 data Zero k a = Zero
 data Succ rk k a = Succ {-# UNPACK #-} !(BinomTree rk k a) !(rk k a)
 
+instance IFoldl' Zero where
+  foldlWithKey'_ _ z ~Zero = z
+
+instance IFoldMap Zero where
+  foldMapWithKey_ _ ~Zero = mempty
+
+instance IFoldl' t => IFoldl' (Succ t) where
+  foldlWithKey'_ f z (Succ t rk) = foldlWithKey'_ f z' rk
+    where
+      !z' = foldlWithKey'_ f z t
+
+instance IFoldMap t => IFoldMap (Succ t) where
+  foldMapWithKey_ f (Succ t rk) = foldMapWithKey_ f t `mappend` foldMapWithKey_ f rk
+
+instance IFoldl' rk => IFoldl' (BinomTree rk) where
+  foldlWithKey'_ f !z (BinomTree k a rk) = foldlWithKey'_ f ft rk
+    where
+      !ft = f z k a
+
+instance IFoldMap rk => IFoldMap (BinomTree rk) where
+  foldMapWithKey_ f (BinomTree k a rk) = f k a `mappend` foldMapWithKey_ f rk
+
+instance IFoldl' t => IFoldl' (BinomForest t) where
+  foldlWithKey'_ _f z Nil = z
+  foldlWithKey'_ f !z (Skip ts) = foldlWithKey'_ f z ts
+  foldlWithKey'_ f !z (Cons t ts) = foldlWithKey'_ f ft ts
+    where
+      !ft = foldlWithKey'_ f z t
+
+instance IFoldMap t => IFoldMap (BinomForest t) where
+  foldMapWithKey_ _f Nil = mempty
+  foldMapWithKey_ f (Skip ts) = foldMapWithKey_ f ts
+  foldMapWithKey_ f (Cons t ts) = foldMapWithKey_ f t `mappend` foldMapWithKey_ f ts
+
 type CompF a = a -> a -> Bool
 
 instance (Ord k, Eq a) => Eq (MinPQueue k a) where
@@ -586,11 +623,23 @@ foldrWithKeyU :: (k -> a -> b -> b) -> b -> MinPQueue k a -> b
 foldrWithKeyU _ z Empty            = z
 foldrWithKeyU f z (MinPQ _ k a ts) = f k a (foldrWithKeyF_ f (const id) ts z)
 
--- | /O(n)/. An unordered left fold over the elements of the queue, in no particular order.
+-- | /O(n)/. An unordered monoidal fold over the elements of the queue, in no particular order.
+foldMapWithKeyU :: Monoid m => (k -> a -> m) -> MinPQueue k a -> m
+foldMapWithKeyU _ Empty            = mempty
+foldMapWithKeyU f (MinPQ _ k a ts) = f k a `mappend` foldMapWithKey_ f ts
+
+-- | /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 _ z Empty = z
 foldlWithKeyU f z0 (MinPQ _ k0 a0 ts) = foldlWithKeyF_ (\k a z -> f z k a) (const id) ts (f z0 k0 a0)
 
+-- | /O(n)/. An unordered strict left fold over the elements of the queue, in no particular order.
+foldlWithKeyU' :: (b -> k -> a -> b) -> b -> MinPQueue k a -> b
+foldlWithKeyU' _ z Empty = z
+foldlWithKeyU' f !z0 (MinPQ _ k0 a0 ts) = foldlWithKey'_ f (f z0 k0 a0) ts
+
 -- | /O(n)/. Map a function over all values in the queue.
 map :: (a -> b) -> MinPQueue k a -> MinPQueue k b
 map = mapWithKey . const

src/Data/PQueue/Prio/Max.hs

@@ -99,8 +99,10 @@ module Data.PQueue.Prio.Max (
   -- * Unordered operations
   foldrU,
   foldrWithKeyU,
+  foldMapWithKeyU,
   foldlU,
   foldlWithKeyU,
+  foldlWithKeyU',
   traverseU,
   traverseWithKeyU,
   keysU,

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

@@ -82,9 +82,12 @@ module Data.PQueue.Prio.Max.Internals (
   toList,
   -- * Unordered operations
   foldrU,
+  foldMapWithKeyU,
   foldrWithKeyU,
   foldlU,
+  foldlU',
   foldlWithKeyU,
+  foldlWithKeyU',
   traverseU,
   traverseWithKeyU,
   keysU,
@@ -464,14 +467,31 @@ foldrU = foldrWithKeyU . const
 foldrWithKeyU :: (k -> a -> b -> b) -> b -> MaxPQueue k a -> b
 foldrWithKeyU f z (MaxPQ q) = Q.foldrWithKeyU (f . unDown) z q
 
--- | /O(n)/. An unordered left fold over the elements of the queue, in no particular order.
+-- | /O(n)/. An unordered monoidal fold over the elements of the queue, in no particular order.
+foldMapWithKeyU :: Monoid m => (k -> a -> m) -> MaxPQueue k a -> m
+foldMapWithKeyU f (MaxPQ q) = Q.foldMapWithKeyU (f . unDown) q
+
+-- | /O(n)/. An unordered left fold over the elements of the queue, in no
+-- particular order.  This is rarely what you want; 'foldrU' and 'foldlU'' are
+-- more likely to perform well.
 foldlU :: (b -> a -> b) -> b -> MaxPQueue k a -> b
 foldlU f = foldlWithKeyU (const . f)
 
--- | /O(n)/. An unordered left fold over the elements of the queue, in no particular order.
+-- | /O(n)/. An unordered strict left fold over the elements of the queue, in no
+-- particular order.
+foldlU' :: (b -> a -> b) -> b -> MaxPQueue k a -> b
+foldlU' f = foldlWithKeyU' (const . f)
+
+-- | /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 -> MaxPQueue k a -> b
 foldlWithKeyU f z0 (MaxPQ q) = Q.foldlWithKeyU (\z -> f z . unDown) z0 q
 
+-- | /O(n)/. An unordered left fold over the elements of the queue, in no particular order.
+foldlWithKeyU' :: (b -> k -> a -> b) -> b -> MaxPQueue k a -> b
+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.