Merge pull request #854 from an0/master

Add O(1) `contains()` implementation for Range with Comparable Element

Author: gribozavr

stdlib/private/StdlibUnittest/StdlibUnittest.swift.gyb

@@ -2607,6 +2607,48 @@ public func < (
   return MinimalComparableValue.lessImpl.value(lhs.value, rhs.value)
 }
 
+/// A type that conforms only to `Comparable` and `ForwardIndexType`.
+///
+/// This type can be used to check that generic functions don't rely on any
+/// other conformances.
+public struct MinimalComparableIndexValue : Comparable, ForwardIndexType {
+  public static var timesEqualEqualWasCalled = 0
+  public static var timesLessWasCalled = 0
+
+  public func successor() -> MinimalComparableIndexValue {
+    return MinimalComparableIndexValue(value.successor())
+  }
+
+  public var value: Int
+  public var identity: Int
+
+  public init(_ value: Int) {
+    self.value = value
+    self.identity = 0
+  }
+
+  public init(_ value: Int, identity: Int) {
+    self.value = value
+    self.identity = identity
+  }
+}
+
+public func == (
+  lhs: MinimalComparableIndexValue,
+  rhs: MinimalComparableIndexValue
+) -> Bool {
+  MinimalComparableIndexValue.timesEqualEqualWasCalled += 1
+  return lhs.value == rhs.value
+}
+
+public func < (
+  lhs: MinimalComparableIndexValue,
+  rhs: MinimalComparableIndexValue
+) -> Bool {
+  MinimalComparableIndexValue.timesLessWasCalled += 1
+  return lhs.value < rhs.value
+}
+
 // ${'Local Variables'}:
 // eval: (read-only-mode 1)
 // End:

stdlib/public/core/Range.swift

@@ -153,6 +153,29 @@ extension Range : CustomReflectable {
   }
 }
 
+/// O(1) implementation of `contains()` for ranges of comparable elements.
+extension Range where Element: Comparable {
+  @warn_unused_result
+  func _customContainsEquatableElement(element: Generator.Element) -> Bool? {
+    return element >= self.startIndex && element < self.endIndex
+  }
+    
+  // FIXME: copied from SequenceAlgorithms as a workaround for https://bugs.swift.org/browse/SR-435
+  @warn_unused_result
+  public func contains(element: Generator.Element) -> Bool {
+    if let result = _customContainsEquatableElement(element) {
+      return result
+    }
+
+    for e in self {
+      if e == element {
+        return true
+      }
+    }
+    return false
+  }
+}
+
 @warn_unused_result
 public func == <Element>(lhs: Range<Element>, rhs: Range<Element>) -> Bool {
   return lhs.startIndex == rhs.startIndex &&
@@ -207,8 +230,6 @@ public func ... <Pos : ForwardIndexType where Pos : Comparable> (
 public func ~= <I : ForwardIndexType where I : Comparable> (
   pattern: Range<I>, value: I
 ) -> Bool {
-  // Intervals can check for containment in O(1).
-  return 
-    HalfOpenInterval(pattern.startIndex, pattern.endIndex).contains(value)
+  return pattern.contains(value)
 }
 

validation-test/stdlib/SequenceType.swift.gyb

@@ -608,6 +608,21 @@ SequenceTypeTests.test("Set<T>.contains/CustomImplementation/${dispatch}") {
 
 % end
 
+SequenceTypeTests.test("Range<Element>.contains/WhereElementIsComparable/dispatch") {
+  MinimalComparableIndexValue.timesLessWasCalled = 0
+  let start = 0
+  let end = 10
+  let range = Range(start: MinimalComparableIndexValue(start), end: MinimalComparableIndexValue(end))
+  let count = 20
+  for test in 0..<count {
+    expectEqual(
+      test >= start && test < end,
+      range.contains(MinimalComparableIndexValue(test)))
+  }
+  expectEqual(
+    count * 2, MinimalComparableIndexValue.timesLessWasCalled)
+}
+
 SequenceTypeTests.test("contains/Predicate") {
   for test in findTests {
     let s = MinimalSequence<OpaqueValue<Int>>(