Add mapAccumL' (#22)

Author: Daniel-Diaz

src/Data/List/Infinite.hs

@@ -45,6 +45,7 @@ module Data.List.Infinite (
   scanl',
   scanl1,
   mapAccumL,
+  mapAccumL',
   traverse_,
   for_,
 
@@ -621,6 +622,26 @@ mapAccumLFB f = \x r -> oneShot (\s -> let (s', y) = f s x in y :< r s')
     mapAccumL f s xs
   #-}
 
+-- | Same as 'mapAccumL', but strict in accumulator.
+mapAccumL' :: (acc -> x -> (acc, y)) -> acc -> Infinite x -> Infinite y
+mapAccumL' f = flip (foldr (\x acc !s -> let (s', y) = f s x in y :< acc s'))
+
+mapAccumL'FB :: (acc -> x -> (acc, y)) -> x -> (acc -> Infinite y) -> acc -> Infinite y
+mapAccumL'FB f = \x r -> oneShot (\(!s) -> let (s', y) = f s x in y :< r s')
+
+{-# NOINLINE [1] mapAccumL' #-}
+
+{-# INLINE [0] mapAccumL'FB #-}
+
+{-# RULES
+"mapAccumL'" [~1] forall f s xs.
+  mapAccumL' f s xs = foldr (mapAccumL'FB f) xs s
+
+"mapAccumL'List" [1] forall f s xs.
+  foldr (mapAccumL'FB f) xs s = mapAccumL' f s xs
+  #-}
+
+
 -- | Generate an infinite list of repeated applications.
 iterate :: (a -> a) -> a -> Infinite a
 iterate f = go

test/Fusion.hs

@@ -84,11 +84,13 @@ mapAccumLRepeat :: Int -> Infinite Int
 mapAccumLRepeat n =
   I.mapAccumL (\acc x -> (acc, fromIntegral x)) 'q' (I.repeat (fromIntegral n :: Word))
 
+mapAccumLRepeat' :: Int -> Infinite Int
+mapAccumLRepeat' n =
+  I.mapAccumL' (\acc x -> (acc, fromIntegral x)) 'q' (I.repeat (fromIntegral n :: Word))
 
 takeFilterIterate :: [Int]
 takeFilterIterate = I.take 100 $ I.filter odd $ I.iterate (+ 1) 0
 
-
 sumTakeFilterIterate :: Int
 sumTakeFilterIterate = sum $ I.take 100 $ I.filter odd $ I.iterate (+ 1) 0
 
@@ -259,6 +261,7 @@ main = defaultMain $ testGroup "All"
   , $(inspectTest $ 'foldrCycle `hasNoType` ''Infinite)
   , $(inspectTest $ 'foldrWordsCycle `hasNoType` ''NonEmpty)
   , $(inspectTest $ 'mapAccumLRepeat `hasNoType` ''Word)
+  , $(inspectTest $ 'mapAccumLRepeat' `hasNoType` ''Word)
 
   , $(inspectTest $ 'takeFilterIterate `hasNoType` ''Infinite)
   , $(inspectTest $ 'sumTakeFilterIterate `hasNoTypes` [''Infinite, ''[]])

test/Properties.hs

@@ -233,6 +233,10 @@ main = defaultMain $ testGroup "All"
   , testProperty "mapAccumL laziness" $ once $
     I.head (I.mapAccumL (\_ x -> (undefined, x)) undefined ('q' :< undefined)) === 'q'
 
+  , testProperty "mapAccumL'" $
+    \(curry . applyFun -> (f :: Bool -> Int -> (Bool, Word))) (Blind xs) ->
+      trim (I.mapAccumL' f False xs) === snd (L.mapAccumL f False (trim xs))
+
   , testProperty "iterate" $
     \(applyFun -> (f :: Int -> Int)) s ->
       trim (I.iterate f s) === L.take 10 (L.iterate f s)