Add takeIterateM

Tests/test-monad-loops.hs

@@ -32,6 +32,20 @@ testTakeWhileMEdgeCase3 = do
     let expected = takeWhile  (const undefined) emptyList
     actual @?= expected
 
+testTakeIterateM :: Assertion
+testTakeIterateM = do
+  let n = 3 :: Int
+  actual <- takeIterateM n (return . (+ 1)) 0
+  expected <- takeIterateM 1 (return . id) n
+  actual @?= expected
+
+testTakeIterateMEdgeCase1 :: Assertion
+testTakeIterateMEdgeCase1 = do
+  let n = 0 :: Int
+  actual <- takeIterateM n (return . (+ 1)) 0
+  expected <- return n
+  actual @?= expected
+
 tests :: TestTree
 tests = testGroup "unit tests"
     [ testCase
@@ -46,7 +60,14 @@ tests = testGroup "unit tests"
     , testCase
         "Testing `takeWhileM (edge case 3)`"
         testTakeWhileMEdgeCase3
+    , testCase
+        "Testing `takeIterateM`"
+        testTakeIterateM
+    , testCase
+        "Testing `takeIterateM` (edge case 1)"
+        testTakeIterateMEdgeCase1
     ]
 
+
 main :: IO ()
-main = defaultMain tests
+main = defaultMain tests

src/Control/Monad/Loops.hs

@@ -132,6 +132,12 @@ iterateM_ :: Monad m => (a -> m a) -> a -> m b
 iterateM_ f = g
     where g x = f x >>= g
 
+{-# SPECIALIZE takeIterateM :: Int -> (a -> IO a) -> a -> IO a #-}
+-- |Like 'iterateM_', but only executes the action n times.
+takeIterateM :: Monad m => Int -> (a -> m a) -> a -> m a
+takeIterateM 0 _ a = return a
+takeIterateM n f a = f a >>= takeIterateM (n - 1) f
+
 {-# SPECIALIZE untilM  :: IO a -> IO Bool -> IO [a] #-}
 {-# SPECIALIZE untilM' :: Monad m => m a -> m Bool -> m [a] #-}
 {-# SPECIALIZE untilM' :: IO a -> IO Bool -> IO [a] #-}