From a67b66e663d159c219750a5044ccf553c4b21bdb Mon Sep 17 00:00:00 2001
From: Ben Gamari <bgamari.foss@gmail.com>
Date: Thu, 24 Aug 2017 11:47:40 -0400
Subject: Add strict variant of iterate

Summary: This closes the nearly-eight-year-old #3474.

Test Plan: Validate

Reviewers: RyanGlScott, austin, hvr

Subscribers: rwbarton, thomie

GHC Trac Issues: #3474

Differential Revision: https://phabricator.haskell.org/D3870
---
 libraries/base/Data/List.hs       |  1 +
 libraries/base/Data/OldList.hs    |  1 +
 libraries/base/GHC/List.hs        | 30 ++++++++++++++++++++++++++++--
 libraries/base/changelog.md       |  3 +++
 libraries/base/tests/T3474.hs     |  5 +++++
 libraries/base/tests/T3474.stdout |  1 +
 libraries/base/tests/all.T        |  4 ++++
 7 files changed, 43 insertions(+), 2 deletions(-)
 create mode 100644 libraries/base/tests/T3474.hs
 create mode 100644 libraries/base/tests/T3474.stdout

diff --git a/libraries/base/Data/List.hs b/libraries/base/Data/List.hs
index 693c0dd151..2ac04a9165 100644
--- a/libraries/base/Data/List.hs
+++ b/libraries/base/Data/List.hs
@@ -76,6 +76,7 @@ module Data.List
 
    -- ** Infinite lists
    , iterate
+   , iterate'
    , repeat
    , replicate
    , cycle
diff --git a/libraries/base/Data/OldList.hs b/libraries/base/Data/OldList.hs
index d03c0bcc96..c4c38d4029 100644
--- a/libraries/base/Data/OldList.hs
+++ b/libraries/base/Data/OldList.hs
@@ -77,6 +77,7 @@ module Data.OldList
 
    -- ** Infinite lists
    , iterate
+   , iterate'
    , repeat
    , replicate
    , cycle
diff --git a/libraries/base/GHC/List.hs b/libraries/base/GHC/List.hs
index 70bfbe4de0..37bba9ac3e 100644
--- a/libraries/base/GHC/List.hs
+++ b/libraries/base/GHC/List.hs
@@ -23,7 +23,7 @@ module GHC.List (
    map, (++), filter, concat,
    head, last, tail, init, uncons, null, length, (!!),
    foldl, foldl', foldl1, foldl1', scanl, scanl1, scanl', foldr, foldr1,
-   scanr, scanr1, iterate, repeat, replicate, cycle,
+   scanr, scanr1, iterate, iterate', repeat, replicate, cycle,
    take, drop, sum, product, maximum, minimum, splitAt, takeWhile, dropWhile,
    span, break, reverse, and, or,
    any, all, elem, notElem, lookup,
@@ -442,7 +442,10 @@ minimum xs              =  foldl1 min xs
 -- of @f@ to @x@:
 --
 -- > iterate f x == [x, f x, f (f x), ...]
-
+--
+-- Note that 'iterate' is lazy, potentially leading to thunk build-up if
+-- the consumer doesn't force each iterate. See 'iterate\'' for a strict
+-- variant of this function.
 {-# NOINLINE [1] iterate #-}
 iterate :: (a -> a) -> a -> [a]
 iterate f x =  x : iterate f (f x)
@@ -458,6 +461,29 @@ iterateFB c f x0 = go x0
  #-}
 
 
+-- | 'iterate\'' is the strict version of 'iterate'.
+--
+-- It ensures that the result of each application of force to weak head normal
+-- form before proceeding.
+{-# NOINLINE [1] iterate' #-}
+iterate' :: (a -> a) -> a -> [a]
+iterate' f x =
+    let x' = f x
+    in x' `seq` (x : iterate' f x')
+
+{-# INLINE [0] iterate'FB #-} -- See Note [Inline FB functions]
+iterate'FB :: (a -> b -> b) -> (a -> a) -> a -> b
+iterate'FB c f x0 = go x0
+  where go x =
+            let x' = f x
+            in x' `seq` (x `c` go x')
+
+{-# RULES
+"iterate'"    [~1] forall f x.   iterate' f x = build (\c _n -> iterate'FB c f x)
+"iterate'FB"  [1]                iterate'FB (:) = iterate'
+ #-}
+
+
 -- | 'repeat' @x@ is an infinite list, with @x@ the value of every element.
 repeat :: a -> [a]
 {-# INLINE [0] repeat #-}
diff --git a/libraries/base/changelog.md b/libraries/base/changelog.md
index cce9fba5fc..a8915cbbeb 100644
--- a/libraries/base/changelog.md
+++ b/libraries/base/changelog.md
@@ -21,6 +21,9 @@
     be able to successfully parse more strings containing `"Proxy"` _et al._
     without surrounding parentheses (e.g., `"Thing Proxy"`) (#12874).
 
+  * Add `iterate'`, a strict version of `iterate`, to `Data.List`
+    and `Data.OldList` (#3474)
+
 ## 4.10.0.0 *April 2017*
   * Bundled with GHC *TBA*
 
diff --git a/libraries/base/tests/T3474.hs b/libraries/base/tests/T3474.hs
new file mode 100644
index 0000000000..dbd59011b4
--- /dev/null
+++ b/libraries/base/tests/T3474.hs
@@ -0,0 +1,5 @@
+import Data.List
+
+-- this should evaluate in constant space
+main :: IO ()
+main = print $ iterate' (+1) 1 !! 100000000
diff --git a/libraries/base/tests/T3474.stdout b/libraries/base/tests/T3474.stdout
new file mode 100644
index 0000000000..2e8da1af1b
--- /dev/null
+++ b/libraries/base/tests/T3474.stdout
@@ -0,0 +1 @@
+100000001
\ No newline at end of file
diff --git a/libraries/base/tests/all.T b/libraries/base/tests/all.T
index 970fb7e454..9055bd5b45 100644
--- a/libraries/base/tests/all.T
+++ b/libraries/base/tests/all.T
@@ -217,3 +217,7 @@ test('T13191',
 test('T13525', when(opsys('mingw32'), skip), compile_and_run, [''])
 test('T13097', normal, compile_and_run, [''])
 test('functorOperators', normal, compile_and_run, [''])
+test('T3474',
+     [stats_num_field('max_bytes_used', [ (wordsize(64), 44504, 5) ]),
+      only_ways(['normal'])],
+     compile_and_run, ['-O'])
-- 
cgit v1.2.1