Speed up folds on Sequences

Data/Sequence/Internal.hs

@@ -248,6 +248,7 @@ import qualified Data.Array
 import qualified GHC.Arr
 #endif
 
+import Utils.Containers.Internal.Coercions ((.#), (.^#))
 -- Coercion on GHC 7.8+
 #if __GLASGOW_HASKELL__ >= 708
 import Data.Coerce
@@ -385,20 +386,29 @@ fmapSeq f (Seq xs) = Seq (fmap (fmap f) xs)
  #-}
 #endif
 
+getSeq :: Seq a -> FingerTree (Elem a)
+getSeq (Seq xs) = xs
+
 instance Foldable Seq where
-    foldMap f (Seq xs) = foldMap (foldMap f) xs
-#if __GLASGOW_HASKELL__ >= 708
-    foldr f z (Seq xs) = foldr (coerce f) z xs
-    foldr' f z (Seq xs) = foldr' (coerce f) z xs
-#else
-    foldr f z (Seq xs) = foldr (flip (foldr f)) z xs
-#if MIN_VERSION_base(4,6,0)
-    foldr' f z (Seq xs) = foldr' (flip (foldr' f)) z xs
-#endif
+    foldMap f = foldMap (f .# getElem) .# getSeq
+    foldr f z = foldr (f .# getElem) z .# getSeq
+    foldl f z = foldl (f .^# getElem) z .# getSeq
+
+#if __GLASGOW_HASKELL__
+    {-# INLINABLE foldMap #-}
+    {-# INLINABLE foldr #-}
+    {-# INLINABLE foldl #-}
 #endif
-    foldl f z (Seq xs) = foldl (foldl f) z xs
+
 #if MIN_VERSION_base(4,6,0)
-    foldl' f z (Seq xs) = foldl' (foldl' f) z xs
+    foldr' f z = foldr' (f .# getElem) z .# getSeq
+    foldl' f z = foldl' (f .^# getElem) z .# getSeq
+
+#if __GLASGOW_HASKELL__
+    {-# INLINABLE foldr' #-}
+    {-# INLINABLE foldl' #-}
+#endif
+
 #endif
 
     foldr1 f (Seq xs) = getElem (foldr1 f' xs)
@@ -894,32 +904,126 @@ instance Sized a => Sized (FingerTree a) where
 
 instance Foldable FingerTree where
     foldMap _ EmptyT = mempty
-    foldMap f (Single x) = f x
-    foldMap f (Deep _ pr m sf) =
-        foldMap f pr <> foldMap (foldMap f) m <> foldMap f sf
+    foldMap f' (Single x') = f' x'
+    foldMap f' (Deep _ pr' m' sf') = 
+        foldMapDigit f' pr' <>
+        foldMapTree (foldMapNode f') m' <>
+        foldMapDigit f' sf'
+      where
+        foldMapTree :: Monoid m => (Node a -> m) -> FingerTree (Node a) -> m
+        foldMapTree _ EmptyT = mempty
+        foldMapTree f (Single x) = f x
+        foldMapTree f (Deep _ pr m sf) = 
+            foldMapDigitN f pr <>
+            foldMapTree (foldMapNodeN f) m <>
+            foldMapDigitN f sf
 
-    foldr _ z EmptyT = z
-    foldr f z (Single x) = x `f` z
-    foldr f z (Deep _ pr m sf) =
-        foldr f (foldr (flip (foldr f)) (foldr f z sf) m) pr
+        foldMapDigit :: Monoid m => (a -> m) -> Digit a -> m
+        foldMapDigit f t = foldDigit (<>) f t
+
+        foldMapDigitN :: Monoid m => (Node a -> m) -> Digit (Node a) -> m
+        foldMapDigitN f t = foldDigit (<>) f t
+
+        foldMapNode :: Monoid m => (a -> m) -> Node a -> m
+        foldMapNode f t = foldNode (<>) f t
+
+        foldMapNodeN :: Monoid m => (Node a -> m) -> Node (Node a) -> m
+        foldMapNodeN f t = foldNode (<>) f t
+#if __GLASGOW_HASKELL__
+    {-# INLINABLE foldMap #-}
+#endif
+
+    foldr _ z' EmptyT = z'
+    foldr f' z' (Single x') = x' `f'` z'
+    foldr f' z' (Deep _ pr' m' sf') =
+        foldrDigit f' (foldrTree (foldrNode f') (foldrDigit f' z' sf') m') pr'
+      where
+        foldrTree :: (Node a -> b -> b) -> b -> FingerTree (Node a) -> b
+        foldrTree _ z EmptyT = z
+        foldrTree f z (Single x) = x `f` z
+        foldrTree f z (Deep _ pr m sf) =
+            foldrDigitN f (foldrTree (foldrNodeN f) (foldrDigitN f z sf) m) pr
+
+        foldrDigit :: (a -> b -> b) -> b -> Digit a -> b
+        foldrDigit f z t = foldr f z t
+
+        foldrDigitN :: (Node a -> b -> b) -> b -> Digit (Node a) -> b
+        foldrDigitN f z t = foldr f z t
+
+        foldrNode :: (a -> b -> b) -> Node a -> b -> b
+        foldrNode f t z = foldr f z t
+
+        foldrNodeN :: (Node a -> b -> b) -> Node (Node a) -> b -> b
+        foldrNodeN f t z = foldr f z t
+    {-# INLINE foldr #-}
+
+
+    foldl _ z' EmptyT = z'
+    foldl f' z' (Single x') = z' `f'` x'
+    foldl f' z' (Deep _ pr' m' sf') =
+        foldlDigit f' (foldlTree (foldlNode f') (foldlDigit f' z' pr') m') sf'
+      where
+        foldlTree :: (b -> Node a -> b) -> b -> FingerTree (Node a) -> b
+        foldlTree _ z EmptyT = z
+        foldlTree f z (Single x) = z `f` x
+        foldlTree f z (Deep _ pr m sf) =
+            foldlDigitN f (foldlTree (foldlNodeN f) (foldlDigitN f z pr) m) sf
+
+        foldlDigit :: (b -> a -> b) -> b -> Digit a -> b
+        foldlDigit f z t = foldl f z t
+
+        foldlDigitN :: (b -> Node a -> b) -> b -> Digit (Node a) -> b
+        foldlDigitN f z t = foldl f z t
+
+        foldlNode :: (b -> a -> b) -> b -> Node a -> b
+        foldlNode f z t = foldl f z t
+
+        foldlNodeN :: (b -> Node a -> b) -> b -> Node (Node a) -> b
+        foldlNodeN f z t = foldl f z t
+    {-# INLINE foldl #-}
 
-    foldl _ z EmptyT = z
-    foldl f z (Single x) = z `f` x
-    foldl f z (Deep _ pr m sf) =
-        foldl f (foldl (foldl f) (foldl f z pr) m) sf
 
 #if MIN_VERSION_base(4,6,0)
-    foldr' _ z EmptyT = z
-    foldr' f z (Single x) = f x z
-    foldr' f z (Deep _ pr m sf) = foldr' f mres pr
-        where !sfRes = foldr' f z sf
-              !mres = foldr' (flip (foldr' f)) sfRes m
-
-    foldl' _ z EmptyT = z
-    foldl' f z (Single x) = z `f` x
-    foldl' f z (Deep _ pr m sf) = foldl' f mres sf
-        where !prRes = foldl' f z pr
-              !mres = foldl' (foldl' f) prRes m
+    foldr' _ z' EmptyT = z'
+    foldr' f' z' (Single x') = f' x' z'
+    foldr' f' z' (Deep _ pr' m' sf') =
+        (foldrDigit' f' $! (foldrTree' (foldrNode' f') $! (foldrDigit' f' z') sf') m') pr'
+      where
+        foldrTree' :: (Node a -> b -> b) -> b -> FingerTree (Node a) -> b
+        foldrTree' _ z EmptyT = z
+        foldrTree' f z (Single x) = f x $! z
+        foldrTree' f z (Deep _ pr m sf) =
+            (foldr' f $! (foldrTree' (foldrNodeN' f) $! (foldr' f $! z) sf) m) pr
+
+        foldrDigit' :: (a -> b -> b) -> b -> Digit a -> b
+        foldrDigit' f z t = foldr' f z t
+
+        foldrNode' :: (a -> b -> b) -> Node a -> b -> b
+        foldrNode' f t z = foldr' f z t
+
+        foldrNodeN' :: (Node a -> b -> b) -> Node (Node a) -> b -> b
+        foldrNodeN' f t z = foldr' f z t
+    {-# INLINE foldr' #-}
+
+    foldl' _ z' EmptyT = z'
+    foldl' f' z' (Single x') = f' z' x'
+    foldl' f' z' (Deep _ pr' m' sf') =
+        (foldlDigit' f' $!
+         (foldlTree' (foldlNode' f') $! (foldlDigit' f' z') pr') m')
+            sf'
+      where
+        foldlTree' :: (b -> Node a -> b) -> b -> FingerTree (Node a) -> b
+        foldlTree' _ z EmptyT = z
+        foldlTree' f z (Single xs) = f z xs
+        foldlTree' f z (Deep _ pr m sf) =
+            (foldl' f $! (foldlTree' (foldl' f) $! foldl' f z pr) m) sf
+
+        foldlDigit' :: (b -> a -> b) -> b -> Digit a -> b
+        foldlDigit' f z t = foldl' f z t
+
+        foldlNode' :: (b -> a -> b) -> b -> Node a -> b
+        foldlNode' f z t = foldl' f z t
+    {-# INLINE foldl' #-}
 #endif
 
     foldr1 _ EmptyT = error "foldr1: empty sequence"
@@ -991,22 +1095,26 @@ instance Foldable Digit where
     foldr f z (Two a b) = a `f` (b `f` z)
     foldr f z (Three a b c) = a `f` (b `f` (c `f` z))
     foldr f z (Four a b c d) = a `f` (b `f` (c `f` (d `f` z)))
+    {-# INLINE foldr #-}
 
     foldl f z (One a) = z `f` a
     foldl f z (Two a b) = (z `f` a) `f` b
     foldl f z (Three a b c) = ((z `f` a) `f` b) `f` c
     foldl f z (Four a b c d) = (((z `f` a) `f` b) `f` c) `f` d
+    {-# INLINE foldl #-}
 
 #if MIN_VERSION_base(4,6,0)
-    foldr' f z (One a) = a `f` z
+    foldr' f z (One a) = f a z
     foldr' f z (Two a b) = f a $! f b z
     foldr' f z (Three a b c) = f a $! f b $! f c z
     foldr' f z (Four a b c d) = f a $! f b $! f c $! f d z
+    {-# INLINE foldr' #-}
 
     foldl' f z (One a) = f z a
     foldl' f z (Two a b) = (f $! f z a) b
     foldl' f z (Three a b c) = (f $! (f $! f z a) b) c
     foldl' f z (Four a b c d) = (f $! (f $! (f $! f z a) b) c) d
+    {-# INLINE foldl' #-}
 #endif
 
     foldr1 _ (One a) = a
@@ -1078,16 +1186,20 @@ instance Foldable Node where
 
     foldr f z (Node2 _ a b) = a `f` (b `f` z)
     foldr f z (Node3 _ a b c) = a `f` (b `f` (c `f` z))
+    {-# INLINE foldr #-}
 
     foldl f z (Node2 _ a b) = (z `f` a) `f` b
     foldl f z (Node3 _ a b c) = ((z `f` a) `f` b) `f` c
+    {-# INLINE foldl #-}
 
 #if MIN_VERSION_base(4,6,0)
     foldr' f z (Node2 _ a b) = f a $! f b z
     foldr' f z (Node3 _ a b c) = f a $! f b $! f c z
+    {-# INLINE foldr' #-}
 
     foldl' f z (Node2 _ a b) = (f $! f z a) b
     foldl' f z (Node3 _ a b c) = (f $! (f $! f z a) b) c
+    {-# INLINE foldl' #-}
 #endif
 
 instance Functor Node where

Utils/Containers/Internal/Coercions.hs

@@ -0,0 +1,44 @@
+{-# LANGUAGE CPP #-}
+{-# OPTIONS_HADDOCK hide #-}
+
+#include "containers.h"
+
+module Utils.Containers.Internal.Coercions where
+
+#if __GLASGOW_HASKELL__ >= 708
+import Data.Coerce
+#endif
+
+infixl 8 .#
+#if __GLASGOW_HASKELL__ >= 708
+(.#) :: Coercible b a => (b -> c) -> (a -> b) -> a -> c
+(.#) f _ = coerce f
+#else
+(.#) :: (b -> c) -> (a -> b) -> a -> c
+(.#) = (.)
+#endif
+{-# INLINE (.#) #-}
+
+infix 9 .^#
+
+-- | Coerce the second argument of a function. Conceptually,
+-- can be thought of as:
+--
+-- @
+--   (f .^# g) x y = f x (g y)
+-- @
+--
+-- However it is most useful when coercing the arguments to
+-- 'foldl':
+--
+-- @
+--   foldl f b . fmap g = foldl (f .^# g) b
+-- @
+#if __GLASGOW_HASKELL__ >= 708
+(.^#) :: Coercible c b => (a -> c -> d) -> (b -> c) -> (a -> b -> d)
+(.^#) f _ = coerce f
+#else
+(.^#) :: (a -> c -> d) -> (b -> c) -> (a -> b -> d)
+(f .^# g) x y = f x (g y)
+#endif
+{-# INLINE (.^#) #-}

benchmarks/Sequence.hs

@@ -159,15 +159,15 @@ main = do
          ]
       , bgroup "unstableSortOn"
          [ bgroup "already sorted"
-            [ bench "10" $ nf S.unstableSortOn id s10
-            , bench "100" $ nf S.unstableSortOn id s100
-            , bench "1000" $ nf S.unstableSortOn id s1000
-            , bench "10000" $ nf S.unstableSortOn id s10000]
+            [ bench "10"    $ nf (S.unstableSortOn id) s10
+            , bench "100"   $ nf (S.unstableSortOn id) s100
+            , bench "1000"  $ nf (S.unstableSortOn id) s1000
+            , bench "10000" $ nf (S.unstableSortOn id) s10000]
          , bgroup "random"
-            [ bench "10" $ nf S.unstableSortOn id rs10
-            , bench "100" $ nf S.unstableSortOn id rs100
-            , bench "1000" $ nf S.unstableSortOn id rs1000
-            , bench "10000" $ nf S.unstableSortOn id rs10000]
+            [ bench "10"    $ nf (S.unstableSortOn id) rs10
+            , bench "100"   $ nf (S.unstableSortOn id) rs100
+            , bench "1000"  $ nf (S.unstableSortOn id) rs1000
+            , bench "10000" $ nf (S.unstableSortOn id) rs10000]
          ]
       ]
 

containers.cabal

@@ -84,6 +84,7 @@ Library
         Utils.Containers.Internal.StrictFold
         Utils.Containers.Internal.StrictMaybe
         Utils.Containers.Internal.PtrEquality
+        Utils.Containers.Internal.Coercions
         Data.Map.Internal.DeprecatedShowTree
         Data.IntMap.Internal.DeprecatedDebug
 

tests/seq-properties.hs

@@ -54,9 +54,11 @@ main = defaultMain
        , testProperty "(<$)" prop_constmap
        , testProperty "foldr" prop_foldr
        , testProperty "foldr'" prop_foldr'
+       , testProperty "lazy foldr'" prop_lazyfoldr'
        , testProperty "foldr1" prop_foldr1
        , testProperty "foldl" prop_foldl
        , testProperty "foldl'" prop_foldl'
+       , testProperty "lazy foldl'" prop_lazyfoldl'
        , testProperty "foldl1" prop_foldl1
        , testProperty "(==)" prop_equals
        , testProperty "compare" prop_compare
@@ -306,6 +308,16 @@ prop_foldr' xs =
     f = (:)
     z = []
 
+prop_lazyfoldr' :: Seq () -> Property
+prop_lazyfoldr' xs =
+    not (null xs) ==>
+    foldr'
+        (\e _ ->
+              e)
+        (error "Data.Sequence.foldr': should be lazy in initial accumulator")
+        xs ===
+    ()
+
 prop_foldr1 :: Seq Int -> Property
 prop_foldr1 xs =
     not (null xs) ==> foldr1 f xs == Data.List.foldr1 f (toList xs)
@@ -325,6 +337,16 @@ prop_foldl' xs =
     f = flip (:)
     z = []
 
+prop_lazyfoldl' :: Seq () -> Property
+prop_lazyfoldl' xs =
+    not (null xs) ==>
+    foldl'
+        (\_ e ->
+              e)
+        (error "Data.Sequence.foldl': should be lazy in initial accumulator")
+        xs ===
+    ()
+
 prop_foldl1 :: Seq Int -> Property
 prop_foldl1 xs =
     not (null xs) ==> foldl1 f xs == Data.List.foldl1 f (toList xs)