[Stdlib] Implement sequence(first:next:) and sequence(state:next:)

stdlib/public/core/CMakeLists.txt

@@ -132,6 +132,7 @@ set(SWIFTLIB_SOURCES
   Process.swift
   SliceBuffer.swift
   Tuple.swift.gyb
+  UnfoldSequence.swift
   VarArgs.swift
   Zip.swift
   )

stdlib/public/core/GroupInfo.json

@@ -49,6 +49,7 @@
     "Indices.swift",
     "Existential.swift",
     "WriteBackMutableSlice.swift",
+    "UnfoldSequence.swift",
     {
       "Type-erased": [
         "ExistentialCollection.swift"

stdlib/public/core/UnfoldSequence.swift

@@ -0,0 +1,126 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2016 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See http://swift.org/LICENSE.txt for license information
+// See http://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+/// Returns a sequence formed from `first` and repeated lazy applications of
+/// `next`.
+///
+/// The first element in the sequence is always `first`, and each successive
+/// element is the result of invoking `next` with the previous element. The
+/// sequence ends when `next` returns `nil`. If `next` never returns `nil`, the
+/// sequence is infinite.
+///
+/// This function can be used to replace many cases that were previously handled
+/// using C-style `for` loops.
+///
+/// Example:
+///
+///     // Walk the elements of a tree from a node up to the root
+///     for node in sequence(first: leaf, next: { $0.parent }) {
+///       // node is leaf, then leaf.parent, then leaf.parent.parent, etc.
+///     }
+///
+///     // Iterate over all powers of two (ignoring overflow)
+///     for value in sequence(first: 1, next: { $0 * 2 }) {
+///       // value is 1, then 2, then 4, then 8, etc.
+///     }
+///
+/// - Parameter first: The first element to be returned from the sequence.
+/// - Parameter next: A closure that accepts the previous sequence element and
+///   returns the next element.
+/// - Returns: A sequence that starts with `first` and continues with every
+///   value returned by passing the previous element to `next`.
+///
+/// - SeeAlso: `sequence(state:next:)`
+public func sequence<T>(first: T, next: (T) -> T?) -> UnfoldFirstSequence<T> {
+  // The trivial implementation where the state is the next value to return
+  // has the downside of being unnecessarily eager (it evaluates `next` one
+  // step in advance). We solve this by using a boolean value to disambiguate
+  // between the first value (that's computed in advance) and the rest.
+  return sequence(state: (first, true), next: { (state: inout (T?, Bool)) -> T? in
+    switch state {
+    case (let value, true):
+      state.1 = false
+      return value
+    case (let value?, _):
+      let nextValue = next(value)
+      state.0 = nextValue
+      return nextValue
+    case (nil, _):
+      return nil
+    }
+  })
+}
+
+/// Returns a sequence formed from repeated lazy applications of `next` to a
+/// mutable `state`.
+///
+/// The elements of the sequence are obtaned by invoking `next` with a mutable
+/// state. The same state is passed to all invocations of `next`, so subsequent
+/// calls will see any mutations made by previous calls. The sequence ends when
+/// `next` returns `nil`. If `next` never returns `nil`, the sequence is
+/// infinite.
+///
+/// This function can be used to replace many instances of `AnyIterator` that
+/// wrap a closure.
+///
+/// Example:
+///
+///     // Interleave two sequences that yield the same element type
+///     sequence(state: (false, seq1.makeIterator(), seq2.makeIterator()), next: { iters in
+///       iters.0 = !iters.0
+///       return iters.0 ? iters.1.next() : iters.2.next()
+///     })
+///
+/// - Parameter state: The initial state that will be passed to the closure.
+/// - Parameter next: A closure that accepts an `inout` state and returns the
+///   next element of the sequence.
+/// - Returns: A sequence that yields each successive value from `next`.
+///
+/// - SeeAlso: `sequence(first:next:)`
+public func sequence<T, State>(state: State, next: (inout State) -> T?)
+  -> UnfoldSequence<T, State> {
+  return UnfoldSequence(_state: state, _next: next)
+}
+
+/// The return type of `sequence(first:next:)`.
+public typealias UnfoldFirstSequence<T> = UnfoldSequence<T, (T?, Bool)>
+
+/// A sequence whose elements are produced via repeated applications of a
+/// closure to some mutable state.
+///
+/// The elements of the sequence are computed lazily and the sequence may
+/// potentially be infinite in length.
+///
+/// Instances of `UnfoldSequence` are created with the functions
+/// `sequence(first:next:)` and `sequence(state:next:)`.
+///
+/// - SeeAlso: `sequence(first:next:)`, `sequence(state:next:)`
+public struct UnfoldSequence<Element, State> : Sequence, IteratorProtocol {
+  public mutating func next() -> Element? {
+    guard !_done else { return nil }
+    if let elt = _next(&_state) {
+        return elt
+    } else {
+        _done = true
+        return nil
+    }
+  }
+
+  internal init(_state: State, _next: (inout State) -> Element?) {
+    self._state = _state
+    self._next = _next
+  }
+
+  internal var _state: State
+  internal let _next: (inout State) -> Element?
+  internal var _done = false
+}

validation-test/stdlib/UnfoldSequence.swift

@@ -0,0 +1,80 @@
+// RUN: %target-run-simple-swift
+// REQUIRES: executable_test
+
+import StdlibUnittest
+import StdlibCollectionUnittest
+
+var UnfoldTests = TestSuite("UnfoldSequence")
+
+UnfoldTests.test("sequence(state:next:)") {
+  // FIXME: The full type signatures on these closures should not be
+  // necessary, but at the moment the compiler gives very confusing errors if
+  // we don't have them.
+
+  let s0 = sequence(state: 1, next: { (val: inout Int) -> Int? in
+    defer { val *= 2 }; return val > 16 ? nil : val
+  })
+  checkSequence([1,2,4,8,16], s0)
+
+  let s1 = sequence(state: (), next: { (_: inout ()) -> Int? in 1 })
+  checkSequence([1, 1, 1], s1.prefix(3))
+
+  let s2 = sequence(state: (1..<6).makeIterator(), next: {
+    (iter: inout CountableRange<Int>.Iterator) in
+    iter.next()
+  })
+  checkSequence(1..<6, s2)
+
+  // Make sure we don't evaluate any step in advance
+  var calls = 0
+  var s3 = sequence(state: 0, next: { (val: inout Int) -> Int? in
+    calls += 1; val += 1; return val
+  })
+  for i in 1..<6 {
+    expectEqual(i, s3.next())
+    expectEqual(i, calls)
+  }
+
+  // Make sure we don't invoke next() after it returns nil
+  calls = 0
+  var s4 = sequence(state: 1, next : { (val: inout Int) -> Int? in
+    calls += 1; defer { val *= 2 }; return val > 16 ? nil : val
+  })
+  checkSequence([1,2,4,8,16], s4)
+  expectEqual(6, calls)
+
+  let s5 = sequence(state: (), next: { (_: inout ()) -> Int? in nil })
+  checkSequence([], s5)
+
+  // This is similar to s0 except calling the `next` closure after it returns
+  // nil starts returning values again.
+  let s6 = sequence(state: 1, next: { (val: inout Int) -> Int? in
+    defer { val *= 2 }
+    return val == 32 ? nil : val
+  })
+  checkSequence([1,2,4,8,16], s6)
+}
+
+UnfoldTests.test("sequence(first:next:)") {
+  let s0 = sequence(first: 1, next: { $0 < 50 ? $0 * 2 : nil })
+  expectEqualSequence([1, 2, 4, 8, 16, 32, 64], s0)
+
+  // Make sure we don't evaluate any step in advance
+  var calls = 0
+  var s1 = sequence(first: 0, next: { calls += 1; return $0 + 1 })
+  for i in 0..<5 {
+    expectEqual(i, s1.next())
+    expectEqual(i, calls)
+  }
+
+  // Make sure we don't invoke next() after it returns nil
+  calls = 0
+  let s2 = sequence(first: 1, next: {
+    calls += 1
+    return $0 >= 16 ? nil : $0 * 2
+  })
+  checkSequence([1,2,4,8,16], s2)
+  expectEqual(5, calls)
+}
+
+runAllTests()