From 65dd6e5633690fb0498165d391bf2c71e0ac6c74 Mon Sep 17 00:00:00 2001
From: ekmett <ekmett@gmail.com>
Date: Tue, 29 Jun 2010 15:37:21 -0700
Subject: [PATCH] fixed folds

---
 CHANGELOG.markdown                         |  5 ++
 Control/Concurrent/Speculation/Internal.hs | 25 ++++++
 Data/Foldable/Speculation.hs               | 38 +++------
 Data/Traversable/Speculation.hs            | 95 ++++++++++++++++++++--
 speculation.cabal                          |  9 +-
 5 files changed, 138 insertions(+), 34 deletions(-)

diff --git a/CHANGELOG.markdown b/CHANGELOG.markdown
index c571e87..814035a 100644
--- a/CHANGELOG.markdown
+++ b/CHANGELOG.markdown
@@ -1,3 +1,8 @@
+/0.8.1/:
+ * Added Data.List.Foldable
+ * Added Data.Traversable.Foldable
+ * Fixed an off-by-one error in the arguments to the speculative fold estimators
+
 /0.8.0.2/:
 
  * changed tests and benchmarks to not build by default to work around corruption in the hackage db
diff --git a/Control/Concurrent/Speculation/Internal.hs b/Control/Concurrent/Speculation/Internal.hs
index 05e4934..48e79f5 100644
--- a/Control/Concurrent/Speculation/Internal.hs
+++ b/Control/Concurrent/Speculation/Internal.hs
@@ -7,14 +7,39 @@ module Control.Concurrent.Speculation.Internal
     , returning
     ) where
 
+import Data.Foldable
+import Data.Traversable
+import Control.Applicative
+
+-- comonad!
 data Acc a = Acc {-# UNPACK #-} !Int a
 
+instance Functor Acc where
+    fmap f (Acc n a) = Acc n (f a)
+
+instance Foldable Acc where
+    foldMap = foldMapDefault
+
+instance Traversable Acc where
+    traverse f (Acc n a) = Acc n <$> f a
+
 extractAcc :: Acc a -> a
 extractAcc (Acc _ a) = a
 {-# INLINE extractAcc #-}
 
 data MaybeAcc a = JustAcc {-# UNPACK #-} !Int a | NothingAcc
 
+instance Functor MaybeAcc where
+    fmap f (JustAcc n a) = JustAcc n (f a)
+    fmap _ NothingAcc = NothingAcc
+
+instance Foldable MaybeAcc where
+    foldMap = foldMapDefault
+
+instance Traversable MaybeAcc where
+    traverse f (JustAcc n a) = JustAcc n <$> f a
+    traverse _ NothingAcc    = pure NothingAcc
+
 fromMaybeAcc :: a -> MaybeAcc a -> a
 fromMaybeAcc _ (JustAcc _ a) = a
 fromMaybeAcc a _ = a
diff --git a/Data/Foldable/Speculation.hs b/Data/Foldable/Speculation.hs
index 7c2014d..a95701d 100644
--- a/Data/Foldable/Speculation.hs
+++ b/Data/Foldable/Speculation.hs
@@ -105,7 +105,7 @@ foldr = foldrBy (==)
 foldrBy :: Foldable f => (b -> b -> Bool) -> (Int -> b) -> (a -> b -> b) -> b -> f a -> b
 foldrBy cmp g f z = extractAcc . Foldable.foldr mf (Acc 0 z)
   where 
-    mf a (Acc n b) = let n' = n + 1 in Acc n' (specBy' cmp (g n') (f a) b)
+    mf a (Acc n b) = Acc (n + 1) (specBy' cmp (g n) (f a) b)
 {-# INLINE foldrBy #-}
 
 foldlM :: (Foldable f, Monad m, Eq (m b)) => (Int -> m b) -> (b -> a -> m b) -> m b -> f a -> m b
@@ -117,9 +117,8 @@ foldlByM  cmp g f mz = liftM extractAcc . Foldable.foldl go (liftM (Acc 0) mz)
   where
     go mia b = do
       Acc n a <- mia
-      let !n' = n + 1
-      a' <- specBy' cmp (g n') (>>= (`f` b)) (return a)
-      return (Acc n' a')
+      a' <- specBy' cmp (g n) (>>= (`f` b)) (return a)
+      return (Acc (n + 1) a')
 {-# INLINE foldlByM #-}
 
 foldrM :: (Foldable f, Monad m, Eq (m b)) => (Int -> m b) -> (a -> b -> m b) -> m b -> f a -> m b
@@ -131,9 +130,8 @@ foldrByM cmp g f mz = liftM extractAcc . Foldable.foldr go (liftM (Acc 0) mz)
   where
     go a mib = do
       Acc n b <- mib
-      let !n' = n + 1
-      b' <- specBy' cmp (g n') (>>= f a) (return b)
-      return (Acc n' b')
+      b' <- specBy' cmp (g n) (>>= f a) (return b)
+      return (Acc (n + 1) b')
 {-# INLINE foldrByM #-}
 
 foldlSTM :: (Foldable f, Eq a) => (Int -> STM a) -> (a -> b -> STM a) -> STM a -> f b -> STM a
@@ -145,9 +143,8 @@ foldlBySTM cmp g f mz = liftM extractAcc . Foldable.foldl go (liftM (Acc 0) mz)
   where
     go mia b = do
       Acc n a <- mia
-      let !n' = n + 1
-      a' <- specBySTM' cmp (g n') (`f` b) a
-      return (Acc n' a')
+      a' <- specBySTM' cmp (g n) (`f` b) a
+      return (Acc (n + 1) a')
 {-# INLINE foldlBySTM #-}
 
 foldrSTM :: (Foldable f, Eq b) => (Int -> STM b) -> (a -> b -> STM b) -> STM b -> f a -> STM b
@@ -159,21 +156,10 @@ foldrBySTM cmp g f mz = liftM extractAcc . Foldable.foldr go (liftM (Acc 0) mz)
   where
     go a mib = do
       Acc n b <- mib
-      let !n' = n + 1
-      b' <- specBySTM' cmp (g n') (f a) b
-      return (Acc n' b')
+      b' <- specBySTM' cmp (g n) (f a) b
+      return (Acc (n + 1) b')
 {-# INLINE foldrBySTM #-}
 
-{-
-foldrSTMBy cmp g f z xs = liftM extractAcc . Foldable.foldl mf return xs (Acc 0 z)
-  where
-    mf h t = do
-        Acc n t' <- t
-        let !n' = n + 1 
-        specSTMBy' cmp (g n') (flip f h >=> t) 
-        ...
--}
-
 -- | Given a valid estimator @g@, @'foldl' g f z xs@ yields the same answer as @'foldl'' f z xs@.
 --
 -- @g n@ should supply an estimate of the value returned from folding over the first @n@ elements of the container.
@@ -188,7 +174,7 @@ foldl = foldlBy (==)
 foldlBy  :: Foldable f => (b -> b -> Bool) -> (Int -> b) -> (b -> a -> b) -> b -> f a -> b
 foldlBy cmp g f z = extractAcc . Foldable.foldl mf (Acc 0 z)
   where
-    mf (Acc n a) b = let n' = n + 1 in Acc n' (specBy' cmp (g n') (`f` b) a)
+    mf (Acc n a) b = Acc (n + 1) (specBy' cmp (g n) (`f` b) a)
 {-# INLINE foldlBy #-}
 
 foldr1 :: (Foldable f, Eq a) => (Int -> a) -> (a -> a -> a) -> f a -> a
@@ -199,7 +185,7 @@ foldr1By :: Foldable f => (a -> a -> Bool) -> (Int -> a) -> (a -> a -> a) -> f a
 foldr1By cmp g f xs = fromMaybeAcc (errorEmptyStructure "foldr1")
                                    (Foldable.foldr mf NothingAcc xs)
   where
-    mf a (JustAcc n b) = let n' = n + 1 in JustAcc n' (specBy' cmp (g n') (f a) b)
+    mf a (JustAcc n b) = JustAcc (n + 1) (specBy' cmp (g n) (f a) b)
     mf a NothingAcc = JustAcc 1 a
 {-# INLINE foldr1By #-}
 
@@ -211,7 +197,7 @@ foldl1By :: Foldable f => (a -> a -> Bool) -> (Int -> a) -> (a -> a -> a) -> f a
 foldl1By cmp g f xs = fromMaybeAcc (errorEmptyStructure "foldl1")
                                (Foldable.foldl mf NothingAcc xs)
   where
-    mf (JustAcc n a) b = let n' = n + 1 in JustAcc n' (specBy' cmp (g n') (`f` b) a)
+    mf (JustAcc n a) b = JustAcc (n + 1) (specBy' cmp (g n) (`f` b) a)
     mf NothingAcc b    = JustAcc 1 b
 {-# INLINE foldl1By #-}
 
diff --git a/Data/Traversable/Speculation.hs b/Data/Traversable/Speculation.hs
index 92575c2..ba17f15 100644
--- a/Data/Traversable/Speculation.hs
+++ b/Data/Traversable/Speculation.hs
@@ -1,17 +1,98 @@
+{-# LANGUAGE MagicHash, Rank2Types, UnboxedTuples #-}
 module Data.Traversable.Speculation
-    ( 
+    ( mapAccumL, mapAccumLBy
+    , mapAccumR, mapAccumRBy
     ) where
 
-{-
+import GHC.Prim
+import GHC.Types
 import Data.Traversable (Traversable)
 import qualified Data.Traversable as Traversable
-import Control.Concurrent.STM
-import Control.Concurrent.Speculation
-import Control.Concurrent.Speculation.Internal
 import Control.Applicative
-import Control.Monad hiding (mapM_, msum, forM_, sequence_)
+import Control.Concurrent.Speculation
 
+data IntAccumL s a = IntAccumL (Int# -> s -> (# Int#, s, a #))
+
+runIntAccumL :: IntAccumL s a -> Int -> s -> (s, a)
+runIntAccumL (IntAccumL m) (I# i) s = case m i s of
+    (# _, s1, a #) -> (s1, a)
+
+instance Functor (IntAccumL s) where
+    fmap f (IntAccumL m) = IntAccumL  (\i s -> case m i s of
+        (# i1, s1, a #) -> (# i1, s1, f a #))
+
+instance Applicative (IntAccumL s) where
+    pure a = IntAccumL (\i s -> (# i, s, a #))
+    IntAccumL mf <*> IntAccumL ma = IntAccumL (\i s ->
+        case mf i s of 
+            (# i1, s1, f #) ->
+                case ma i1 s1 of 
+                    (# i2, s2, a #) -> (# i2, s2, f a #))
+    
+mapAccumL :: (Traversable t, Eq a, Eq c) => (Int -> a) -> (a -> b -> (a, c)) -> a -> t b -> (a, t c)
+mapAccumL = mapAccumLBy (==)
+
+mapAccumLBy :: Traversable t => (a -> a -> Bool) -> (Int -> a) -> (a -> b -> (a, c)) -> a -> t b -> (a, t c)
+mapAccumLBy cmp g f z xs = runIntAccumL (Traversable.traverse go xs) 0 z
+  where
+    go b = IntAccumL (\n a -> 
+            let ~(a', c) = specBy' cmp (g (I# n)) (`f` b) a
+            in (# n +# 1#, a', c #))
+
+data IntAccumR s a = IntAccumR (Int# -> s -> (# Int#, s, a #))
+
+runIntAccumR :: IntAccumR s a -> Int -> s -> (s, a)
+runIntAccumR (IntAccumR m) (I# i) s = case m i s of
+    (# _, s1, a #) -> (s1, a)
+
+instance Functor (IntAccumR s) where
+    fmap f (IntAccumR m) = IntAccumR  (\i s -> case m i s of
+        (# i1, s1, a #) -> (# i1, s1, f a #))
+
+instance Applicative (IntAccumR s) where
+    pure a = IntAccumR (\i s -> (# i, s, a #))
+    IntAccumR mf <*> IntAccumR ma = IntAccumR (\i s ->
+        case ma i s of 
+            (# i1, s1, a #) ->
+                case mf i1 s1 of 
+                    (# i2, s2, f #) -> (# i2, s2, f a #))
+
+mapAccumR :: (Traversable t, Eq a, Eq c) => (Int -> a) -> (a -> b -> (a, c)) -> a -> t b -> (a, t c)
+mapAccumR = mapAccumRBy (==)
+
+mapAccumRBy :: Traversable t => (a -> a -> Bool) -> (Int -> a) -> (a -> b -> (a, c)) -> a -> t b -> (a, t c)
+mapAccumRBy cmp g f z xs = runIntAccumR (Traversable.traverse go xs) 0 z
+  where
+    go b = IntAccumR (\n a -> 
+            let ~(a', c) = specBy' cmp (g (I# n)) (`f` b) a
+            in (# n +# 1#, a', c #))
+
+{-
 traverse :: (Traversable t, Applicative f, Eq (f b)) => (Int -> f b) -> (a -> f b) -> t a -> f (t b)
+traverse = traverseBy (==)
+
+traverseBy :: (Traversable t, Applicative f) => (Int -> f b) -> (a -> f b) -> t a -> f (t b)
+-}
+
+-- note applicative composition doesn't give StateT
+-- There is a difference between StateT s m and State s (m a)
+-- 
+{-
+newtype AccT m a = AccT (Int# -> m (Acc a))
+
+instance Functor f => Applicative (AccT s f) where
+    fmap f (AccT m) = AccT (fmap (fmap f) . m)
+
+instance Applicative m => Applicative (AccT s m) where
+    pure a = AccT (\i -> return (Acc i a))
+    AccT mf <*> AccT ma = AccT (\i -> 
+        let maccf = mf i 
+    
+
+        m (Acc (a -> b)) -> m (Acc a)
+-}
+
+{-
 traverseBy :: (Traversable t, Applicative f) => (f b -> f b -> Bool) -> (Int -> f b) -> (a -> f b) -> t a -> f (t b)
 sequence   :: (Traversable t, Monad m, Eq (m a)) => (Int -> m a) -> t (m a) -> m (t a)
 sequenceBy :: (Traversable t, Monad m) => (m a -> m a -> Bool) -> (Int -> m a) -> t (m a) -> m (t a)
@@ -25,7 +106,6 @@ mapSTM :: (Traversable t, Eq b) => (Int -> STM b) -> (a -> STM b) -> t a -> STM
 mapBySTM :: Traversable t => (b -> b -> STM Bool) -> (Int -> STM b) -> (a -> STM b) -> t a -> STM (t b)
 mapAccumR :: (Traversable t, Eq a, Eq c) => (Int -> (a, c)) -> (a -> b -> (a, c)) -> a -> t b -> (a, t c)
 mapAccumRBy :: Traversable t => ((a, c) -> (a, c) -> Bool) -> (Int -> (a, c)) -> (a -> b -> (a, c)) -> a -> t b -> (a, t c)
-mapAccumL :: (Traversable t, Eq a, Eq c) => (Int -> (a, c)) -> (a -> b -> (a, c)) -> a -> t b -> (a, t c)
 mapAccumLBy :: Traversable t => ((a, c) -> (a, c) -> Bool) -> (Int -> (a, c)) -> (a -> b -> (a, c)) -> a -> t b -> (a, t c)
 for :: (Traversable t, Applicative f, Eq (f b)) => (Int -> f b) -> t a -> (a -> f b) -> f (t b)
 forBy :: (Traversable t, Applicative f) => (f b -> f b -> Bool) -> (Int -> f b) -> t a -> (a -> f b) -> f (t b)
@@ -34,3 +114,4 @@ forByM :: (Traversable t, Monad m) => (m b -> m b -> Bool) -> (Int -> m b) -> t
 forSTM :: (Traversable t, Eq b) => (Int -> STM b) -> t a -> (a -> STM b) -> STM (t b)
 forBySTM :: Traversable t => (b -> b -> STM Bool) -> (Int -> STM b) -> t a -> (a -> STM b) -> STM (t b)
 -}
+
diff --git a/speculation.cabal b/speculation.cabal
index 3d57dcf..b46049c 100644
--- a/speculation.cabal
+++ b/speculation.cabal
@@ -77,7 +77,6 @@ flag hpc
   description: Use HPC for tests
   default:     True
 
-
 library
   if !flag(lib)
     buildable: False
@@ -88,12 +87,14 @@ library
 
     build-depends:
       base >= 4 && < 6,
+      ghc-prim >= 0.2 && < 0.3,
       parallel >= 2.2 && < 2.3,
       stm >= 2.1 && < 2.2
 
     exposed-modules:
       Control.Concurrent.Speculation
       Data.Foldable.Speculation
+      Data.Traversable.Speculation
       Data.List.Speculation
     other-modules:
       Control.Concurrent.Speculation.Internal
@@ -109,6 +110,8 @@ executable test-speculation
     ghc-options: -Wall
     build-depends:
       base >= 4 && < 6, 
+      ghc-prim >= 0.2 && < 0.3,
+      parallel >= 2.2 && < 2.3,
       stm >= 2.1 && < 2.2,
       containers >= 0.3.0 && < 0.4,
       test-framework >= 0.2.4 && < 0.3,
@@ -120,6 +123,7 @@ executable test-speculation
       Control.Concurrent.Speculation.Internal
       Control.Concurrent.Speculation
       Data.Foldable.Speculation
+      Data.Traversable.Speculation
       Data.List.Speculation
 
 executable benchmark-speculation
@@ -132,6 +136,8 @@ executable benchmark-speculation
       ghc-options: -O2 -fspec-constr -funbox-strict-fields -fdicts-cheap
     build-depends:
       base >= 4 && < 6, 
+      ghc-prim >= 0.2 && < 0.3,
+      parallel >= 2.2 && < 2.3,
       stm >= 2.1 && < 2.2,
       containers >= 0.3.0 && < 0.4,
       criterion >= 0.5 && < 0.6
@@ -139,4 +145,5 @@ executable benchmark-speculation
       Control.Concurrent.Speculation.Internal
       Control.Concurrent.Speculation
       Data.Foldable.Speculation
+      Data.Traversable.Speculation
       Data.List.Speculation