diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,7 @@ +# 0.3.1.0 + +* Added `ireplicateM` and `ireplicateM_`. + # 0.3.0.0 * `ifind` now returns the index alongside with the value (same as in `lens`). diff --git a/bench/Functions.hs b/bench/Functions.hs --- a/bench/Functions.hs +++ b/bench/Functions.hs @@ -1,6 +1,7 @@ {-# LANGUAGE MagicHash, -BangPatterns +BangPatterns, +CPP #-} @@ -16,6 +17,7 @@ import Data.List import Data.List.Index import Control.Monad +import qualified Control.Loop as Loop indexed_zip :: [a] -> [(Int, a)] @@ -93,6 +95,20 @@ f (I# i) x go (i +# 1#) xs {-# INLINE imapM__rec #-} + +#if __GLASGOW_HASKELL__ < 710 +ireplicateM__loop + :: (Monad m, Functor m) => Int -> (Int -> m a) -> m () +#else +ireplicateM__loop + :: Monad m => Int -> (Int -> m a) -> m () +#endif +ireplicateM__loop n f = Loop.numLoop 0 (n-1) (void . f) +{-# INLINE ireplicateM__loop #-} + +ireplicateM__for :: Monad m => Int -> (Int -> m a) -> m () +ireplicateM__for n f = forM_ [0..n-1] f +{-# INLINE ireplicateM__for #-} iall_zip :: (Int -> a -> Bool) -> [a] -> Bool iall_zip p xs = and (zipWith p [0..] xs) diff --git a/bench/Main.hs b/bench/Main.hs --- a/bench/Main.hs +++ b/bench/Main.hs @@ -106,6 +106,16 @@ bench "lens" $ nf (\n -> flip runState [] $ L.imapM_ (\i x -> modify ((i+x):) >> return (i-x)) [0..n]) 100000, bench "our" $ nf (\n -> flip runState [] $ imapM_ (\i x -> modify ((i+x):) >> return (i-x)) [0..n]) 100000 ] ], + bgroup "ireplicateM_" [ + bgroup "Just" [ + bench "loop" $ nf (\n -> ireplicateM__loop n (\i -> if i<50000 then Just i else Nothing)) 100000, + bench "for" $ nf (\n -> ireplicateM__for n (\i -> if i<50000 then Just i else Nothing)) 100000, + bench "our" $ nf (\n -> ireplicateM_ n (\i -> if i<50000 then Just i else Nothing)) 100000 ], + bgroup "State" [ + bench "loop" $ nf (\n -> flip runState 0 $ ireplicateM__loop n (\i -> modify' (i+))) 100000, + bench "for" $ nf (\n -> flip runState 0 $ ireplicateM__for n (\i -> modify' (i+))) 100000, + bench "our" $ nf (\n -> flip runState 0 $ ireplicateM_ n (\i -> modify' (i+))) 100000 ] ], + bgroup "ifilter" [ bgroup "consume" [ bench "rec" $ nf (\n -> sum $ ifilter_rec (\i x -> rem (i+x) 5000 == 0) [0..n]) 100000, diff --git a/ilist.cabal b/ilist.cabal --- a/ilist.cabal +++ b/ilist.cabal @@ -1,8 +1,8 @@ name: ilist -version: 0.3.0.0 +version: 0.3.1.0 synopsis: Optimised list functions for doing index-related things description: - Optimised list functions for doing index-related things. They're faster than common idioms in all cases, and sometimes they fuse better as well. + Optimised list functions for doing index-related things. They're faster than common idioms in all cases, they avoid , and sometimes they fuse better as well. homepage: http://github.com/aelve/ilist bug-reports: http://github.com/aelve/ilist/issues license: BSD3 @@ -49,6 +49,7 @@ , ilist -- imapM_ is broken in 4.13.2 , lens >= 4.13.2.1 + , loop , transformers , vector ghc-options: -O2 -Wall -fno-warn-unused-do-bind diff --git a/lib/Data/List/Index.hs b/lib/Data/List/Index.hs --- a/lib/Data/List/Index.hs +++ b/lib/Data/List/Index.hs @@ -58,6 +58,7 @@ -- ** Monadic functions iforM, iforM_, itraverse, itraverse_, + ireplicateM, ireplicateM_, ifoldrM, ifoldlM, @@ -341,6 +342,31 @@ ifor_ :: Applicative m => [a] -> (Int -> a -> m b) -> m () ifor_ = flip itraverse_ {-# INLINE ifor_ #-} + +{- | +Perform a given action @n@ times. Behaves like @for_ [0..n-1]@, but avoids . + +If you want more complicated loops (e.g. counting downwards), consider the package. +-} +ireplicateM :: Applicative m => Int -> (Int -> m a) -> m [a] +ireplicateM cnt f = go 0 + where + go !i | i >= cnt = pure [] + | otherwise = (:) <$> f i <*> go (i + 1) +{-# INLINE ireplicateM #-} + +{- | +NB. This function intentionally uses 'Monad' even though 'Applicative' is enough. That's because the @transformers@ package didn't have an optimized definition of ('*>') for 'StateT' prior to 0.5.3.0, so for a common case of 'StateT' this function would be 40 times slower with the 'Applicative' constraint. +-} +ireplicateM_ :: Monad m => Int -> (Int -> m a) -> m () +ireplicateM_ cnt f = if cnt > 0 then go 0 else return () + where + -- this is 30% faster for Maybe than the simpler + -- go i | i == cnt = return () + -- | otherwise = f i >> go (i + 1) + cnt_ = cnt-1 + go !i = if i == cnt_ then f i >> return () else f i >> go (i + 1) +{-# INLINE ireplicateM_ #-} -- Using unboxed ints here doesn't seem to result in any benefit ifoldr :: (Int -> a -> b -> b) -> b -> [a] -> b diff --git a/tests/Main.hs b/tests/Main.hs --- a/tests/Main.hs +++ b/tests/Main.hs @@ -136,7 +136,7 @@ specify "basic" $ do let f i x = modify ((i,x):) >> return (i-x) let (resA, stA) = runState (imapM f [1,3..9]) [] - let (resB, stB) = runState (itraverse f [1,3..9]) [] + (resB, stB) = runState (itraverse f [1,3..9]) [] resA `shouldBe` [0-1,1-3,2-5,3-7,4-9] resB `shouldBe` [0-1,1-3,2-5,3-7,4-9] stA `shouldBe` reverse (zip [0..4] [1,3..9]) @@ -154,9 +154,24 @@ specify "basic" $ do let f i x = modify ((i,x):) >> return (i-x) let stA = execState (imapM_ f [1,3..9]) [] - let stB = execState (itraverse_ f [1,3..9]) [] + stB = execState (itraverse_ f [1,3..9]) [] stA `shouldBe` reverse (zip [0..4] [1,3..9]) stB `shouldBe` reverse (zip [0..4] [1,3..9]) + + describe "ireplicateM" $ do + describe "State" $ do + specify "basic" $ do + let f i = modify (i:) >> return ((i+1)*2) + let (res, st) = runState (ireplicateM 5 f) [] + res `shouldBe` [2,4..10] + st `shouldBe` reverse [0..4] + + describe "ireplicateM_" $ do + describe "State" $ do + specify "basic" $ do + let f i = modify (i:) >> return ((i+1)*2) + let st = execState (ireplicateM_ 5 f) [] + st `shouldBe` reverse [0..4] specialFolds :: Spec specialFolds = describe "special folds" $ do