From 88ef487ca58411b3691f23973ed0b75eb030f271 Mon Sep 17 00:00:00 2001
From: Bodigrim <andrew.lelechenko@gmail.com>
Date: Sun, 30 Jul 2023 22:05:12 +0100
Subject: [PATCH] #33 Implement iterateWithIndex

---
 changelog.md        |  1 +
 src/Data/Chimera.hs | 40 ++++++++++++++++++++++++++++++++++++++++
 test/Test.hs        | 16 ++++++++++++++++
 3 files changed, 57 insertions(+)

diff --git a/changelog.md b/changelog.md
index af9736c..7e4153b 100644
--- a/changelog.md
+++ b/changelog.md
@@ -1,6 +1,7 @@
 # 0.4.0.0
 
 * Add `foldr` catamorphism and `fromInfinite` / `toInfinite` conversions.
+* Add `iterateWithIndex` and `iterateWithIndexM`.
 * Remove instances `Foldable` and `Traversable`, they are too dangerous to diverge.
 
 # 0.3.3.0
diff --git a/src/Data/Chimera.hs b/src/Data/Chimera.hs
index 2c9508c..3ebc31f 100644
--- a/src/Data/Chimera.hs
+++ b/src/Data/Chimera.hs
@@ -32,6 +32,7 @@ module Data.Chimera (
   tabulateFix,
   tabulateFix',
   iterate,
+  iterateWithIndex,
   unfoldr,
   cycle,
   fromListWithDef,
@@ -53,6 +54,7 @@ module Data.Chimera (
   tabulateFixM,
   tabulateFixM',
   iterateM,
+  iterateWithIndexM,
   unfoldrM,
 
   -- * Subvectors
@@ -359,6 +361,8 @@ tabulateFixM_ strat f = result
 -- | 'iterate' @f@ @x@ returns an infinite stream, generated by
 -- repeated applications of @f@ to @x@.
 --
+-- It holds that 'index' ('iterate' @f@ @x@) 0 is equal to @x@.
+--
 -- >>> ch = iterate (+ 1) 0 :: UChimera Int
 -- >>> take 10 (toList ch)
 -- [0,1,2,3,4,5,6,7,8,9]
@@ -432,6 +436,42 @@ unfoldrM f seed = do
   (z, seed') <- unfoldrExactVecNM 1 f seed
   zs <- go 0 seed'
   pure $ Chimera $ A.fromListN (bits + 1) (z : zs)
+{-# SPECIALIZE unfoldrM :: G.Vector v b => (a -> Identity (b, a)) -> a -> Identity (Chimera v b) #-}
+
+-- | 'iterateWithIndex' @f@ @x@ returns an infinite stream, generated by
+-- applications of @f@ to a current index and previous value, starting from @x@.
+--
+-- It holds that 'index' ('iterateWithIndex' @f@ @x@) 0 is equal to @x@.
+--
+-- >>> ch = iterateWithIndex (+) 100 :: UChimera Word
+-- >>> take 10 (toList ch)
+-- [100,101,103,106,110,115,121,128,136,145]
+--
+-- @since 0.4.0.0
+iterateWithIndex :: G.Vector v a => (Word -> a -> a) -> a -> Chimera v a
+iterateWithIndex f = runIdentity . iterateWithIndexM ((pure .) . f)
+
+iterateWithIndexExactVecNM :: forall m a v. (Monad m, G.Vector v a) => Int -> (Word -> a -> m a) -> a -> m (v a)
+iterateWithIndexExactVecNM n f s = G.unfoldrExactNM n go (int2word n, s)
+  where
+    go :: (Word, a) -> m (a, (Word, a))
+    go (i, x) = do
+      x' <- f i x
+      pure (x', (i + 1, x'))
+
+-- | Monadic version of 'iterateWithIndex'.
+--
+-- @since 0.4.0.0
+iterateWithIndexM :: (Monad m, G.Vector v a) => (Word -> a -> m a) -> a -> m (Chimera v a)
+iterateWithIndexM f seed = do
+  nextSeed <- f 1 seed
+  let z = G.singleton seed
+  zs <- iterateListNM bits go (G.singleton nextSeed)
+  pure $ Chimera $ A.fromListN (bits + 1) (z : zs)
+  where
+    go vec =
+      iterateWithIndexExactVecNM (G.length vec `shiftL` 1) f (G.unsafeLast vec)
+{-# SPECIALIZE iterateWithIndexM :: G.Vector v a => (Word -> a -> Identity a) -> a -> Identity (Chimera v a) #-}
 
 interleaveVec :: G.Vector v a => v a -> v a -> v a
 interleaveVec as bs =
diff --git a/test/Test.hs b/test/Test.hs
index ec449db..cd5ae25 100644
--- a/test/Test.hs
+++ b/test/Test.hs
@@ -103,6 +103,19 @@ chimeraTests = testGroup "Chimera"
       let jx = ix `mod` 65536 in
         iterate f seed !! fromIntegral jx === Ch.index (Ch.iterate f seed :: UChimera Word) jx
 
+  , QC.testProperty "head . iterate" $
+    \(Fun _ (f :: Word -> Word)) seed ->
+        seed === Ch.index (Ch.iterate f seed :: UChimera Word) 0
+
+  , QC.testProperty "iterateWithIndex" $
+    \(Fun _ (f :: (Word, Int) -> Int)) seed ix ->
+      let jx = ix `mod` 65536 in
+        iterateWithIndex (curry f) seed !! fromIntegral jx === Ch.index (Ch.iterateWithIndex (curry f) seed :: UChimera Int) jx
+
+  , QC.testProperty "head . iterateWithIndex" $
+    \(Fun _ (f :: (Word, Int) -> Int)) seed ->
+        seed === Ch.index (Ch.iterateWithIndex (curry f) seed :: UChimera Int) 0
+
   , QC.testProperty "unfoldr" $
     \(Fun _ (f :: Word -> (Int, Word))) seed ix ->
       let jx = ix `mod` 65536 in
@@ -171,3 +184,6 @@ mkUnfix splt f x
   $ map f
   $ takeWhile (\y -> 0 <= y && y < x)
   $ splt x
+
+iterateWithIndex :: (Word -> a -> a) -> a -> [a]
+iterateWithIndex f seed = L.unfoldr (\(ix, a) -> let a' = f (ix + 1) a in Just (a, (ix + 1, a'))) (0, seed)