SE-0045 non-lazy implementation and tests

Author: Russ Bishop

stdlib/private/StdlibCollectionUnittest/CheckSequenceType.swift

@@ -1614,6 +1614,24 @@ self.test("\(testNamePrefix).dropLast/semantics/negative") {
   _ = s.dropLast(-1)
 }
 
+//===----------------------------------------------------------------------===//
+// drop(while:)
+//===----------------------------------------------------------------------===//
+
+self.test("\(testNamePrefix).drop(while:)/semantics") {
+  for test in findTests {
+    let s = makeWrappedSequenceWithEquatableElement(test.sequence)
+    let closureLifetimeTracker = LifetimeTracked(0)
+    let remainingSequence = s.drop {
+      _blackHole(closureLifetimeTracker)
+      return $0 == wrapValueIntoEquatable(test.element)
+    }
+    let remaining = Array(remainingSequence)
+    expectEqual(test.expectedLeftoverSequence.count, remaining.count)
+    expectEqualSequence(test.expectedLeftoverSequence, remaining)
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // prefix()
 //===----------------------------------------------------------------------===//
@@ -1653,6 +1671,25 @@ self.test("\(testNamePrefix).prefix/semantics/negative") {
   _ = s.prefix(-1)
 }
 
+//===----------------------------------------------------------------------===//
+// prefix(while:)
+//===----------------------------------------------------------------------===//
+
+self.test("\(testNamePrefix).prefix(while:)/semantics") {
+  for test in findTests {
+    let s = makeWrappedSequenceWithEquatableElement(test.sequence)
+    let closureLifetimeTracker = LifetimeTracked(0)
+    let remainingSequence = s.prefix {
+      _blackHole(closureLifetimeTracker)
+      return $0 == wrapValueIntoEquatable(test.element)
+    }
+    let expectedPrefix = test.sequence.dropLast(expectedLeftoverSequence.count)
+    let remaining = Array(remainingSequence)
+    expectEqual(expectedPrefix.count, remaining.count)
+    expectEqualSequence(expectedPrefix, remaining)
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // suffix()
 //===----------------------------------------------------------------------===//

stdlib/private/StdlibCollectionUnittest/LoggingWrappers.swift.gyb

@@ -84,6 +84,8 @@ public class SequenceLog {
   public static var first = TypeIndexed(0)
   public static var dropFirst = TypeIndexed(0)
   public static var dropLast = TypeIndexed(0)
+  public static var dropWhile = TypeIndexed(0)
+  public static var prefixWhile = TypeIndexed(0)
   public static var prefixMaxLength = TypeIndexed(0)
   public static var suffixMaxLength = TypeIndexed(0)
   public static var split = TypeIndexed(0)
@@ -260,11 +262,25 @@ public struct ${Self}<
     Log.dropLast[selfType] += 1
     return base.dropLast(n)
   }
+  
+  public func drop(
+    while predicate: @noescape (Base.Iterator.Element) throws -> Bool
+  ) rethrows -> SubSequence {
+    Log.dropWhile[selfType] += 1
+    return base.drop(while: predicate)
+  }
 
   public func prefix(_ maxLength: Int) -> SubSequence {
     Log.prefixMaxLength[selfType] += 1
     return base.prefix(maxLength)
   }
+  
+  public func prefix(
+    while predicate: @noescape (Base.Iterator.Element) throws -> Bool
+  ) rethrows -> SubSequence {
+    Log.prefixWhile[selfType] += 1
+    return base.prefix(while: predicate)
+  }
 
   public func suffix(_ maxLength: Int) -> SubSequence {
     Log.suffixMaxLength[selfType] += 1

stdlib/public/core/Collection.swift

@@ -1338,6 +1338,23 @@ extension Collection {
       offsetBy: numericCast(amount), limitedBy: endIndex) ?? endIndex
     return self[startIndex..<end]
   }
+  
+  /// Returns a subsequence by skipping elements while `predicate` returns
+  /// `true` and returning the remaining elements.
+  ///
+  /// - Parameter while: A closure that takes an element of the
+  ///   sequence as its argument and returns `true` if the element should
+  ///   be skipped or `false` if it should be included. Once the predicate
+  ///   returns `false` it will not be called again.
+  public func drop(
+    while predicate: @noescape (Iterator.Element) throws -> Bool
+  ) rethrows -> SubSequence {
+    var start = startIndex
+    while try start < endIndex && predicate(self[start]) {
+      formIndex(after: &start)
+    } 
+    return self[start..<endIndex]
+  }
 
   /// Returns a subsequence, up to the specified maximum length, containing
   /// the initial elements of the collection.
@@ -1363,6 +1380,23 @@ extension Collection {
       offsetBy: numericCast(maxLength), limitedBy: endIndex) ?? endIndex
     return self[startIndex..<end]
   }
+  
+  /// Returns a subsequence containing the initial elements until `predicate`
+  /// returns `false` and skipping the remaining elements.
+  ///
+  /// - Parameter while: A closure that takes an element of the
+  ///   sequence as its argument and returns `true` if the element should
+  ///   be included or `false` if it should be excluded. Once the predicate
+  ///   returns `false` it will not be called again.
+  public func prefix(
+    while predicate: @noescape (Iterator.Element) throws -> Bool
+  ) rethrows -> SubSequence {
+    var end = startIndex
+    while try end < endIndex && predicate(self[end]) {
+      formIndex(after: &end)
+    }
+    return self[startIndex..<end]
+  }
 
   /// Returns a subsequence, up to the given maximum length, containing the
   /// final elements of the collection.

stdlib/public/core/Sequence.swift

@@ -470,6 +470,17 @@ public protocol Sequence {
   ///
   /// - Complexity: O(*n*), where *n* is the length of the sequence.
   func dropLast(_ n: Int) -> SubSequence
+  
+
+  /// Returns a subsequence by skipping elements while `predicate` returns
+  /// `true` and returning the remaining elements.
+  ///
+  /// - Parameter while: A closure that takes an element of the
+  ///   sequence as its argument and returns a Boolean value indicating
+  ///   whether the element is a match.
+  func drop(
+    while predicate: @noescape (Iterator.Element) throws -> Bool
+  ) rethrows -> SubSequence
 
   /// Returns a subsequence, up to the specified maximum length, containing
   /// the initial elements of the sequence.
@@ -488,6 +499,16 @@ public protocol Sequence {
   /// - Returns: A subsequence starting at the beginning of this sequence
   ///   with at most `maxLength` elements.
   func prefix(_ maxLength: Int) -> SubSequence
+  
+  /// Returns a subsequence containing the initial elements until `predicate`
+  /// returns `false` and skipping the remaining elements.
+  ///
+  /// - Parameter while: A closure that takes an element of the
+  ///   sequence as its argument and returns a Boolean value indicating
+  ///   whether the element is a match.
+  func prefix(
+    while predicate: @noescape (Iterator.Element) throws -> Bool
+  ) rethrows -> SubSequence
 
   /// Returns a subsequence, up to the given maximum length, containing the
   /// final elements of the sequence.
@@ -1126,6 +1147,31 @@ extension Sequence where
     }
     return AnySequence(result)
   }
+  
+  /// Returns a subsequence by skipping elements while `predicate` returns
+  /// `true` and returning the remaining elements.
+  ///
+  /// - Parameter while: A closure that takes an element of the
+  ///   sequence as its argument and returns `true` if the element should
+  ///		be skipped or `false` if it should be included. Once the predicate
+  ///		returns `false` it will not be called again.
+  public func drop(
+    while predicate: @noescape (Iterator.Element) throws -> Bool
+  ) rethrows -> AnySequence<Iterator.Element> {
+    var iterator = makeIterator()
+    var nextElement: Iterator.Element? = iterator.next()
+    repeat {
+      nextElement = iterator.next()
+    } while try nextElement.flatMap(predicate) == true
+
+    return AnySequence {
+      AnyIterator {
+        guard let next = nextElement else { return nil }
+        nextElement = iterator.next()
+        return next
+      }
+    }
+  }
 
   /// Returns a subsequence, up to the specified maximum length, containing the
   /// initial elements of the sequence.
@@ -1153,6 +1199,27 @@ extension Sequence where
     return AnySequence(
       _PrefixSequence(_iterator: makeIterator(), maxLength: maxLength))
   }
+  
+  /// Returns a subsequence containing the initial elements until `predicate`
+  /// returns `false` and skipping the remaining elements.
+  ///
+  /// - Parameter while: A closure that takes an element of the
+  ///   sequence as its argument and returns `true` if the element should
+  ///		be included or `false` if it should be excluded. Once the predicate
+  ///		returns `false` it will not be called again.
+  public func prefix(
+    while predicate: @noescape (Iterator.Element) throws -> Bool
+  ) rethrows -> AnySequence<Iterator.Element> {
+    var result: [Iterator.Element] = []
+
+    for element in self {
+      guard try predicate(element) else {
+        break
+      }
+      result.append(element)
+    }
+    return AnySequence(result)
+  }
 }
 
 extension Sequence {

validation-test/stdlib/SequenceType.swift.gyb

@@ -971,6 +971,16 @@ SequenceTypeTests.test("dropLast/dispatch") {
   expectCustomizable(tester, tester.log.dropLast)
 }
 
+//===----------------------------------------------------------------------===//
+// drop(while:)
+//===----------------------------------------------------------------------===//
+
+SequenceTypeTests.test("drop(while:)/dispatch") {
+  var tester = SequenceLog.dispatchTester([OpaqueValue(1)])
+  tester.drop { _ in return false }
+  expectCustomizable(tester, tester.log.dropWhile)
+}
+
 //===----------------------------------------------------------------------===//
 // prefix()
 //===----------------------------------------------------------------------===//
@@ -981,6 +991,16 @@ SequenceTypeTests.test("prefix/dispatch") {
   expectCustomizable(tester, tester.log.prefixMaxLength)
 }
 
+//===----------------------------------------------------------------------===//
+// prefix(while:)
+//===----------------------------------------------------------------------===//
+
+SequenceTypeTests.test("prefix(while:)/dispatch") {
+  var tester = SequenceLog.dispatchTester([OpaqueValue(1)])
+  tester.prefix { _ in return false }
+  expectCustomizable(tester, tester.log.prefixWhile)
+}
+
 //===----------------------------------------------------------------------===//
 // suffix()
 //===----------------------------------------------------------------------===//