Merge pull request #224 from treeowl/more-foldable

Add more Foldable methods

Author: treeowl

CHANGES.md

@@ -1,6 +1,8 @@
 ## next
 
  * Add `HashMap.findWithDefault` (deprecates `HashMap.lookupDefault`)
+ 
+ * Add more folding functions and use them in `Foldable` instances.
 
  * Add `HashMap.!?`, a flipped version of `lookup`.
 

Data/HashMap/Array.hs

@@ -41,8 +41,11 @@ module Data.HashMap.Array
     , copyM
 
       -- * Folds
+    , foldl
     , foldl'
     , foldr
+    , foldr'
+    , foldMap
 
     , thaw
     , map
@@ -63,9 +66,9 @@ import GHC.ST (ST(..))
 import Control.Monad.ST (stToIO)
 
 #if __GLASGOW_HASKELL__ >= 709
-import Prelude hiding (filter, foldr, length, map, read, traverse)
+import Prelude hiding (filter, foldMap, foldr, foldl, length, map, read, traverse)
 #else
-import Prelude hiding (filter, foldr, length, map, read)
+import Prelude hiding (filter, foldr, foldl, length, map, read)
 #endif
 
 #if __GLASGOW_HASKELL__ >= 710
@@ -79,6 +82,7 @@ import GHC.Exts (Array#, newArray#, readArray#, writeArray#,
                  indexArray#, unsafeFreezeArray#, unsafeThawArray#,
                  MutableArray#, sizeofArray#, copyArray#, thawArray#,
                  sizeofMutableArray#, copyMutableArray#, cloneMutableArray#)
+import Data.Monoid (Monoid (..))
 #endif
 
 #if defined(ASSERTS)
@@ -418,6 +422,15 @@ foldl' f = \ z0 ary0 -> go ary0 (length ary0) 0 z0
             (# x #) -> go ary n (i+1) (f z x)
 {-# INLINE foldl' #-}
 
+foldr' :: (a -> b -> b) -> b -> Array a -> b
+foldr' f = \ z0 ary0 -> go ary0 (length ary0 - 1) z0
+  where
+    go !_ary (-1) z = z
+    go !ary i !z
+      | (# x #) <- index# ary i
+      = go ary (i - 1) (f x z)
+{-# INLINE foldr' #-}
+
 foldr :: (a -> b -> b) -> b -> Array a -> b
 foldr f = \ z0 ary0 -> go ary0 (length ary0) 0 z0
   where
@@ -428,6 +441,29 @@ foldr f = \ z0 ary0 -> go ary0 (length ary0) 0 z0
             (# x #) -> f x (go ary n (i+1) z)
 {-# INLINE foldr #-}
 
+foldl :: (b -> a -> b) -> b -> Array a -> b
+foldl f = \ z0 ary0 -> go ary0 (length ary0 - 1) z0
+  where
+    go _ary (-1) z = z
+    go ary i z
+      | (# x #) <- index# ary i
+      = f (go ary (i - 1) z) x
+{-# INLINE foldl #-}
+
+-- We go to a bit of trouble here to avoid appending an extra mempty.
+-- The below implementation is by Mateusz Kowalczyk, who indicates that
+-- benchmarks show it to be faster than one that avoids lifting out
+-- lst.
+foldMap :: Monoid m => (a -> m) -> Array a -> m
+foldMap f = \ary0 -> case length ary0 of
+  0 -> mempty
+  len ->
+    let !lst = len - 1
+        go i | (# x #) <- index# ary0 i, let fx = f x =
+          if i == lst then fx else fx `mappend` go (i + 1)
+    in go 0
+{-# INLINE foldMap #-}
+
 undefinedElem :: a
 undefinedElem = error "Data.HashMap.Array: Undefined element"
 {-# NOINLINE undefinedElem #-}

Data/HashMap/Base.hs

@@ -58,10 +58,15 @@ module Data.HashMap.Base
     , intersectionWithKey
 
       -- * Folds
+    , foldr'
     , foldl'
+    , foldrWithKey'
     , foldlWithKey'
     , foldr
+    , foldl
     , foldrWithKey
+    , foldlWithKey
+    , foldMapWithKey
 
       -- * Filter
     , mapMaybe
@@ -128,7 +133,7 @@ import Data.Data hiding (Typeable)
 import qualified Data.Foldable as Foldable
 import qualified Data.List as L
 import GHC.Exts ((==#), build, reallyUnsafePtrEquality#)
-import Prelude hiding (filter, foldr, lookup, map, null, pred)
+import Prelude hiding (filter, foldl, foldr, lookup, map, null, pred)
 import Text.Read hiding (step)
 
 import qualified Data.HashMap.Array as A
@@ -200,14 +205,20 @@ instance Functor (HashMap k) where
     fmap = map
 
 instance Foldable.Foldable (HashMap k) where
-    foldr = Data.HashMap.Base.foldr
+    foldMap f = foldMapWithKey (\ _k v -> f v)
+    {-# INLINE foldMap #-}
+    foldr = foldr
     {-# INLINE foldr #-}
-    foldl' = Data.HashMap.Base.foldl'
+    foldl = foldl
+    {-# INLINE foldl #-}
+    foldr' = foldr'
+    {-# INLINE foldr' #-}
+    foldl' = foldl'
     {-# INLINE foldl' #-}
 #if MIN_VERSION_base(4,8,0)
-    null = Data.HashMap.Base.null
+    null = null
     {-# INLINE null #-}
-    length = Data.HashMap.Base.size
+    length = size
     {-# INLINE length #-}
 #endif
 
@@ -1584,6 +1595,15 @@ foldl' :: (a -> v -> a) -> a -> HashMap k v -> a
 foldl' f = foldlWithKey' (\ z _ v -> f z v)
 {-# INLINE foldl' #-}
 
+-- | /O(n)/ Reduce this map by applying a binary operator to all
+-- elements, using the given starting value (typically the
+-- right-identity of the operator).  Each application of the operator
+-- is evaluated before using the result in the next application.
+-- This function is strict in the starting value.
+foldr' :: (v -> a -> a) -> a -> HashMap k v -> a
+foldr' f = foldrWithKey' (\ _ v z -> f v z)
+{-# INLINE foldr' #-}
+
 -- | /O(n)/ Reduce this map by applying a binary operator to all
 -- elements, using the given starting value (typically the
 -- left-identity of the operator).  Each application of the operator
@@ -1599,26 +1619,73 @@ foldlWithKey' f = go
     go z (Collision _ ary)     = A.foldl' (\ z' (L k v) -> f z' k v) z ary
 {-# INLINE foldlWithKey' #-}
 
+-- | /O(n)/ Reduce this map by applying a binary operator to all
+-- elements, using the given starting value (typically the
+-- right-identity of the operator).  Each application of the operator
+-- is evaluated before using the result in the next application.
+-- This function is strict in the starting value.
+foldrWithKey' :: (k -> v -> a -> a) -> a -> HashMap k v -> a
+foldrWithKey' f = flip go
+  where
+    go Empty z                 = z
+    go (Leaf _ (L k v)) !z     = f k v z
+    go (BitmapIndexed _ ary) !z = A.foldr' go z ary
+    go (Full ary) !z           = A.foldr' go z ary
+    go (Collision _ ary) !z    = A.foldr' (\ (L k v) z' -> f k v z') z ary
+{-# INLINE foldrWithKey' #-}
+
 -- | /O(n)/ Reduce this map by applying a binary operator to all
 -- elements, using the given starting value (typically the
 -- right-identity of the operator).
 foldr :: (v -> a -> a) -> a -> HashMap k v -> a
 foldr f = foldrWithKey (const f)
 {-# INLINE foldr #-}
 
+-- | /O(n)/ Reduce this map by applying a binary operator to all
+-- elements, using the given starting value (typically the
+-- left-identity of the operator).
+foldl :: (a -> v -> a) -> a -> HashMap k v -> a
+foldl f = foldlWithKey (\a _k v -> f a v)
+{-# INLINE foldl #-}
+
 -- | /O(n)/ Reduce this map by applying a binary operator to all
 -- elements, using the given starting value (typically the
 -- right-identity of the operator).
 foldrWithKey :: (k -> v -> a -> a) -> a -> HashMap k v -> a
-foldrWithKey f = go
+foldrWithKey f = flip go
   where
-    go z Empty                 = z
-    go z (Leaf _ (L k v))      = f k v z
-    go z (BitmapIndexed _ ary) = A.foldr (flip go) z ary
-    go z (Full ary)            = A.foldr (flip go) z ary
-    go z (Collision _ ary)     = A.foldr (\ (L k v) z' -> f k v z') z ary
+    go Empty z                 = z
+    go (Leaf _ (L k v)) z      = f k v z
+    go (BitmapIndexed _ ary) z = A.foldr go z ary
+    go (Full ary) z            = A.foldr go z ary
+    go (Collision _ ary) z     = A.foldr (\ (L k v) z' -> f k v z') z ary
 {-# INLINE foldrWithKey #-}
 
+-- | /O(n)/ Reduce this map by applying a binary operator to all
+-- elements, using the given starting value (typically the
+-- left-identity of the operator).
+foldlWithKey :: (a -> k -> v -> a) -> a -> HashMap k v -> a
+foldlWithKey f = go
+  where
+    go z Empty                 = z
+    go z (Leaf _ (L k v))      = f z k v
+    go z (BitmapIndexed _ ary) = A.foldl go z ary
+    go z (Full ary)            = A.foldl go z ary
+    go z (Collision _ ary)     = A.foldl (\ z' (L k v) -> f z' k v) z ary
+{-# INLINE foldlWithKey #-}
+
+-- | /O(n)/ Reduce the map by applying a function to each element
+-- and combining the results with a monoid operation.
+foldMapWithKey :: Monoid m => (k -> v -> m) -> HashMap k v -> m
+foldMapWithKey f = go
+  where
+    go Empty = mempty
+    go (Leaf _ (L k v)) = f k v
+    go (BitmapIndexed _ ary) = A.foldMap go ary
+    go (Full ary) = A.foldMap go ary
+    go (Collision _ ary) = A.foldMap (\ (L k v) -> f k v) ary
+{-# INLINE foldMapWithKey #-}
+
 ------------------------------------------------------------------------
 -- * Filter
 

Data/HashMap/Lazy.hs

@@ -70,10 +70,15 @@ module Data.HashMap.Lazy
     , intersectionWithKey
 
       -- * Folds
+    , foldMapWithKey
+    , foldr
+    , foldl
+    , foldr'
     , foldl'
+    , foldrWithKey'
     , foldlWithKey'
-    , foldr
     , foldrWithKey
+    , foldlWithKey
 
       -- * Filter
     , filter

Data/HashMap/Strict.hs

@@ -69,10 +69,15 @@ module Data.HashMap.Strict
     , intersectionWithKey
 
       -- * Folds
+    , foldMapWithKey
+    , foldr
+    , foldl
+    , foldr'
     , foldl'
+    , foldrWithKey'
     , foldlWithKey'
-    , foldr
     , foldrWithKey
+    , foldlWithKey
 
       -- * Filter
     , filter

Data/HashMap/Strict/Base.hs

@@ -71,10 +71,15 @@ module Data.HashMap.Strict.Base
     , intersectionWithKey
 
       -- * Folds
+    , foldMapWithKey
+    , foldr'
     , foldl'
+    , foldrWithKey'
     , foldlWithKey'
     , HM.foldr
+    , HM.foldl
     , foldrWithKey
+    , foldlWithKey
 
       -- * Filter
     , HM.filter

Data/HashSet/Base.hs

@@ -55,8 +55,10 @@ module Data.HashSet.Base
     , intersection
 
     -- * Folds
-    , foldl'
     , foldr
+    , foldr'
+    , foldl
+    , foldl'
 
     -- * Filter
     , filter
@@ -77,15 +79,17 @@ module Data.HashSet.Base
 
 import Control.DeepSeq (NFData(..))
 import Data.Data hiding (Typeable)
-import Data.HashMap.Base (HashMap, foldrWithKey, equalKeys, equalKeys1)
+import Data.HashMap.Base
+  ( HashMap, foldMapWithKey, foldlWithKey, foldrWithKey
+  , equalKeys, equalKeys1)
 import Data.Hashable (Hashable(hashWithSalt))
 #if __GLASGOW_HASKELL__ >= 711
 import Data.Semigroup (Semigroup(..))
 #elif __GLASGOW_HASKELL__ < 709
 import Data.Monoid (Monoid(..))
 #endif
 import GHC.Exts (build)
-import Prelude hiding (filter, foldr, map, null)
+import Prelude hiding (filter, foldr, foldl, map, null)
 import qualified Data.Foldable as Foldable
 import qualified Data.HashMap.Base as H
 import qualified Data.List as List
@@ -138,16 +142,21 @@ instance Ord1 HashSet where
 #endif
 
 instance Foldable.Foldable HashSet where
-    foldr = Data.HashSet.Base.foldr
+    foldMap f = foldMapWithKey (\a _ -> f a) . asMap
+    foldr = foldr
     {-# INLINE foldr #-}
-    foldl' = Data.HashSet.Base.foldl'
+    foldl = foldl
+    {-# INLINE foldl #-}
+    foldl' = foldl'
     {-# INLINE foldl' #-}
+    foldr' = foldr'
+    {-# INLINE foldr' #-}
 #if MIN_VERSION_base(4,8,0)
-    toList = Data.HashSet.Base.toList
+    toList = toList
     {-# INLINE toList #-}
-    null = Data.HashSet.Base.null
+    null = null
     {-# INLINE null #-}
-    length = Data.HashSet.Base.size
+    length = size
     {-# INLINE length #-}
 #endif
 
@@ -303,6 +312,16 @@ foldl' f z0 = H.foldlWithKey' g z0 . asMap
   where g z k _ = f z k
 {-# INLINE foldl' #-}
 
+-- | /O(n)/ Reduce this set by applying a binary operator to all
+-- elements, using the given starting value (typically the
+-- right-identity of the operator). Each application of the operator
+-- is evaluated before before using the result in the next
+-- application. This function is strict in the starting value.
+foldr' :: (b -> a -> a) -> a -> HashSet b -> a
+foldr' f z0 = H.foldrWithKey' g z0 . asMap
+  where g k _ z = f k z
+{-# INLINE foldr' #-}
+
 -- | /O(n)/ Reduce this set by applying a binary operator to all
 -- elements, using the given starting value (typically the
 -- right-identity of the operator).
@@ -311,6 +330,14 @@ foldr f z0 = foldrWithKey g z0 . asMap
   where g k _ z = f k z
 {-# INLINE foldr #-}
 
+-- | /O(n)/ Reduce this set by applying a binary operator to all
+-- elements, using the given starting value (typically the
+-- left-identity of the operator).
+foldl :: (a -> b -> a) -> a -> HashSet b -> a
+foldl f z0 = foldlWithKey g z0 . asMap
+  where g z k _ = f z k
+{-# INLINE foldl #-}
+
 -- | /O(n)/ Filter this set by retaining only elements satisfying a
 -- predicate.
 filter :: (a -> Bool) -> HashSet a -> HashSet a