Write custom strict folds (#281)
authorDavid Feuer <David.Feuer@gmail.com>
Thu, 2 Jun 2016 02:20:50 +0000 (22:20 -0400)
committerDavid Feuer <David.Feuer@gmail.com>
Thu, 2 Jun 2016 02:20:50 +0000 (22:20 -0400)
Writing `foldl'` and `foldr'` by hand, instead of using the
default definitions, makes them about twice as fast.

Fix completely bogus definition of `length` for `ViewR`.

Data/Sequence.hs
benchmarks/Sequence.hs
changelog.md
tests/seq-properties.hs

index af06317..82afd61 100644 (file)
@@ -209,10 +209,12 @@ import Control.DeepSeq (NFData(rnf))
 import Control.Monad (MonadPlus(..), ap)
 import Data.Monoid (Monoid(..))
 import Data.Functor (Functor(..))
+#if MIN_VERSION_base(4,6,0)
+import Data.Foldable (Foldable(foldl, foldl1, foldr, foldr1, foldMap, foldl', foldr'), toList)
+#else
 import Data.Foldable (Foldable(foldl, foldl1, foldr, foldr1, foldMap), foldl', toList)
-#if MIN_VERSION_base(4,8,0)
-import Data.Foldable (foldr')
 #endif
+
 #if MIN_VERSION_base(4,9,0)
 import qualified Data.Semigroup as Semigroup
 #endif
@@ -375,8 +377,19 @@ fmapSeq f (Seq xs) = Seq (fmap (fmap f) 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
+#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
+#endif
 
     foldr1 f (Seq xs) = getElem (foldr1 f' xs)
       where f' (Elem x) (Elem y) = Elem (f x y)
@@ -763,6 +776,20 @@ instance Foldable FingerTree where
     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
+#endif
+
     foldr1 _ EmptyT = error "foldr1: empty sequence"
     foldr1 _ (Single x) = x
     foldr1 f (Deep _ pr m sf) =
@@ -834,6 +861,18 @@ instance Foldable Digit where
     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
 
+#if MIN_VERSION_base(4,6,0)
+    foldr' f z (One a) = a `f` 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
+
+    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
+#endif
+
     foldr1 _ (One a) = a
     foldr1 f (Two a b) = a `f` b
     foldr1 f (Three a b c) = a `f` (b `f` c)
@@ -921,6 +960,14 @@ instance Foldable Node where
     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
 
+#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
+
+    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
+#endif
+
 instance Functor Node where
     {-# INLINE fmap #-}
     fmap f (Node2 v a b) = Node2 v (f a) (f b)
@@ -970,9 +1017,18 @@ instance Functor Elem where
 #endif
 
 instance Foldable Elem where
-    foldMap f (Elem x) = f x
     foldr f z (Elem x) = f x z
+#if __GLASGOW_HASKELL__ >= 708
+    foldMap = coerce
+    foldl = coerce
+    foldl' = coerce
+#else
+    foldMap f (Elem x) = f x
     foldl f z (Elem x) = f z x
+#if MIN_VERSION_base(4,6,0)
+    foldl' f z (Elem x) = f z x
+#endif
+#endif
 
 instance Traversable Elem where
     traverse f (Elem x) = Elem <$> f x
@@ -1510,6 +1566,14 @@ instance Foldable ViewL where
     foldl1 _ EmptyL = error "foldl1: empty view"
     foldl1 f (x :< xs) = foldl f x xs
 
+#if MIN_VERSION_base(4,8,0)
+    null EmptyL = True
+    null (_ :< _) = False
+
+    length EmptyL = 0
+    length (_ :< xs) = 1 + length xs
+#endif
+
 instance Traversable ViewL where
     traverse _ EmptyL       = pure EmptyL
     traverse f (x :< xs)    = (:<) <$> f x <*> traverse f xs
@@ -1564,12 +1628,11 @@ instance Foldable ViewR where
     foldr1 _ EmptyR = error "foldr1: empty view"
     foldr1 f (xs :> x) = foldr f x xs
 #if MIN_VERSION_base(4,8,0)
-    -- The default definitions are sensible for ViewL, but not so much for
-    -- ViewR.
     null EmptyR = True
     null (_ :> _) = False
 
-    length = foldr' (\_ k -> k+1) 0
+    length EmptyR = 0
+    length (xs :> _) = length xs + 1
 #endif
 
 instance Traversable ViewR where
index 9b93fe8..f8437e2 100644 (file)
@@ -5,7 +5,7 @@ import Control.DeepSeq (rnf)
 import Control.Exception (evaluate)
 import Control.Monad.Trans.State.Strict
 import Criterion.Main (bench, bgroup, defaultMain, nf)
-import Data.List (foldl')
+import Data.Foldable (foldl', foldr')
 import qualified Data.Sequence as S
 import qualified Data.Foldable
 import Data.Traversable (traverse)
@@ -35,6 +35,18 @@ main = do
          , bench "100" $ nf (shuffle r100) s100
          , bench "1000" $ nf (shuffle r1000) s1000
          ]
+      , bgroup "foldl'"
+         [ bench "10" $ nf (foldl' (+) 0) s10
+         , bench "100" $ nf (foldl' (+) 0) s100
+         , bench "1000" $ nf (foldl' (+) 0) s1000
+         , bench "10000" $ nf (foldl' (+) 0) s10000
+         ]
+      , bgroup "foldr'"
+         [ bench "10" $ nf (foldr' (+) 0) s10
+         , bench "100" $ nf (foldr' (+) 0) s100
+         , bench "1000" $ nf (foldr' (+) 0) s1000
+         , bench "10000" $ nf (foldr' (+) 0) s10000
+         ]
       , bgroup "update"
          [ bench "10" $ nf (updatePoints r10 10) s10
          , bench "100" $ nf (updatePoints r100 10) s100
index e51dbe2..9200dae 100644 (file)
@@ -28,6 +28,8 @@
   * Make `splitAt` in `Data.Sequence` strict in its arguments. Previously,
     it returned a lazy pair.
 
+  * Fix completely erroneous definition of `length` for `ViewR`.
+
   * Derive `Generic` and `Generic1` for `Data.Tree`.
 
   * Add `foldTree` for `Data.Tree`.
@@ -40,6 +42,9 @@
     improvements are greatest for small sequences, but large even for long
     ones. Reimplement `take` and `drop` to avoid building trees only to discard them.
 
+  * Roughly double the speeds of `foldl'` and `foldr'` for `Data.Sequence`
+    by writing custom definitions instead of using the defaults.
+
   * Add rewrite rules to fuse `fmap` with `reverse` for `Data.Sequence`.
 
   * Speed up `adjust` for `Data.Map`.
index e76cbc6..fa1e567 100644 (file)
@@ -4,7 +4,7 @@ import Control.Applicative (Applicative(..))
 import Control.Arrow ((***))
 import Control.Monad.Trans.State.Strict
 import Data.Array (listArray)
-import Data.Foldable (Foldable(foldl, foldl1, foldr, foldr1, foldMap, fold), toList, all, sum)
+import Data.Foldable (Foldable(foldl, foldl1, foldr, foldr1, foldMap, fold), toList, all, sum, foldl', foldr')
 import Data.Functor ((<$>), (<$))
 import Data.Maybe
 import Data.Monoid (Monoid(..), All(..), Endo(..), Dual(..))
@@ -28,8 +28,10 @@ main = defaultMain
        [ testProperty "fmap" prop_fmap
        , testProperty "(<$)" prop_constmap
        , testProperty "foldr" prop_foldr
+       , testProperty "foldr'" prop_foldr'
        , testProperty "foldr1" prop_foldr1
        , testProperty "foldl" prop_foldl
+       , testProperty "foldl'" prop_foldl'
        , testProperty "foldl1" prop_foldl1
        , testProperty "(==)" prop_equals
        , testProperty "compare" prop_compare
@@ -233,9 +235,16 @@ prop_constmap :: A -> Seq A -> Bool
 prop_constmap x xs =
     toList' (x <$ xs) ~= map (const x) (toList xs)
 
-prop_foldr :: Seq A -> Bool
+prop_foldr :: Seq A -> Property
 prop_foldr xs =
-    foldr f z xs == Prelude.foldr f z (toList xs)
+    foldr f z xs === Prelude.foldr f z (toList xs)
+  where
+    f = (:)
+    z = []
+
+prop_foldr' :: Seq A -> Property
+prop_foldr' xs =
+    foldr' f z xs === foldr' f z (toList xs)
   where
     f = (:)
     z = []
@@ -245,9 +254,16 @@ prop_foldr1 xs =
     not (null xs) ==> foldr1 f xs == Data.List.foldr1 f (toList xs)
   where f = (-)
 
-prop_foldl :: Seq A -> Bool
+prop_foldl :: Seq A -> Property
 prop_foldl xs =
-    foldl f z xs == Prelude.foldl f z (toList xs)
+    foldl f z xs === Prelude.foldl f z (toList xs)
+  where
+    f = flip (:)
+    z = []
+
+prop_foldl' :: Seq A -> Property
+prop_foldl' xs =
+    foldl' f z xs === foldl' f z (toList xs)
   where
     f = flip (:)
     z = []