[overlay] Fix Foundation extensions to Substring

When a substring gets bridged to NSString, it loses the initial offset,
therefore APIs that accept or return StringIndex ranges should handle
this case explicitly by adding/subtracting the substring start offset.

Fixes <rdar://problem/33873277>

(cherry picked from commit 7322b63)

Author: Max Moiseev

stdlib/public/SDK/Foundation/NSStringAPI.swift

@@ -30,12 +30,6 @@ func _toNSArray<T, U : AnyObject>(_ a: [T], f: (T) -> U) -> NSArray {
   return result
 }
 
-func _toNSRange(_ r: Range<String.Index>) -> NSRange {
-  return NSRange(
-    location: r.lowerBound.encodedOffset,
-    length: r.upperBound.encodedOffset - r.lowerBound.encodedOffset)
-}
-
 // We only need this for UnsafeMutablePointer, but there's not currently a way
 // to write that constraint.
 extension Optional {
@@ -417,10 +411,26 @@ extension StringProtocol where Index == String.Index {
     return self._ephemeralString as NSString
   }
 
+  // self can be a Substring so we need to subtract/add this offset when
+  // passing _ns to the Foundation APIs. Will be 0 if self is String.
+  @_inlineable
+  @_versioned
+  internal var _substringOffset: Int {
+    return self.startIndex.encodedOffset
+  }
+
   /// Return an `Index` corresponding to the given offset in our UTF-16
   /// representation.
   func _index(_ utf16Index: Int) -> Index {
-    return Index(encodedOffset: utf16Index)
+    return Index(encodedOffset: utf16Index + _substringOffset)
+  }
+
+  @_inlineable
+  @_versioned
+  internal func _toRelativeNSRange(_ r: Range<String.Index>) -> NSRange {
+    return NSRange(
+      location: r.lowerBound.encodedOffset - _substringOffset,
+      length: r.upperBound.encodedOffset - r.lowerBound.encodedOffset)
   }
 
   /// Return a `Range<Index>` corresponding to the given `NSRange` of
@@ -581,7 +591,7 @@ extension StringProtocol where Index == String.Index {
     return locale != nil ? _ns.compare(
       aString,
       options: mask,
-      range: _toNSRange(
+      range: _toRelativeNSRange(
         range ?? startIndex..<endIndex
       ),
       locale: locale
@@ -590,7 +600,7 @@ extension StringProtocol where Index == String.Index {
     : range != nil ? _ns.compare(
       aString,
       options: mask,
-      range: _toNSRange(range!)
+      range: _toRelativeNSRange(range!)
     )
 
     : !mask.isEmpty ? _ns.compare(aString, options: mask)
@@ -1008,7 +1018,7 @@ extension StringProtocol where Index == String.Index {
     T : StringProtocol, R : RangeExpression
   >(in range: R, with replacement: T) -> String where R.Bound == Index {
     return _ns.replacingCharacters(
-      in: _toNSRange(range.relative(to: self)),
+      in: _toRelativeNSRange(range.relative(to: self)),
       with: replacement._ephemeralString)
   }
 
@@ -1041,7 +1051,7 @@ extension StringProtocol where Index == String.Index {
       of: target,
       with: replacement,
       options: options,
-      range: _toNSRange(
+      range: _toRelativeNSRange(
         searchRange ?? startIndex..<endIndex
       )
     )
@@ -1163,7 +1173,7 @@ extension StringProtocol where Index == String.Index {
   ) where R.Bound == Index {
     let range = range.relative(to: self)
     _ns.enumerateLinguisticTags(
-      in: _toNSRange(range),
+      in: _toRelativeNSRange(range),
       scheme: tagScheme._ephemeralString,
       options: opts,
       orthography: orthography != nil ? orthography! : nil
@@ -1227,7 +1237,7 @@ extension StringProtocol where Index == String.Index {
     ) -> Void
   ) where R.Bound == Index {
     _ns.enumerateSubstrings(
-      in: _toNSRange(range.relative(to: self)), options: opts) {
+      in: _toRelativeNSRange(range.relative(to: self)), options: opts) {
       var stop_ = false
 
       body($0,
@@ -1300,7 +1310,7 @@ extension StringProtocol where Index == String.Index {
         usedLength: usedBufferCount,
         encoding: encoding.rawValue,
         options: options,
-        range: _toNSRange(range.relative(to: self)),
+        range: _toRelativeNSRange(range.relative(to: self)),
         remaining: $0)
     }
   }
@@ -1327,7 +1337,7 @@ extension StringProtocol where Index == String.Index {
           contentsEnd in self._ns.getLineStart(
             start, end: end,
             contentsEnd: contentsEnd,
-            for: _toNSRange(range.relative(to: self)))
+            for: _toRelativeNSRange(range.relative(to: self)))
         }
       }
     }
@@ -1355,7 +1365,7 @@ extension StringProtocol where Index == String.Index {
           contentsEnd in self._ns.getParagraphStart(
             start, end: end,
             contentsEnd: contentsEnd,
-            for: _toNSRange(range.relative(to: self)))
+            for: _toRelativeNSRange(range.relative(to: self)))
         }
       }
     }
@@ -1382,7 +1392,8 @@ extension StringProtocol where Index == String.Index {
   public func lineRange<
     R : RangeExpression
   >(for aRange: R) -> Range<Index> where R.Bound == Index {
-    return _range(_ns.lineRange(for: _toNSRange(aRange.relative(to: self))))
+    return _range(_ns.lineRange(
+      for: _toRelativeNSRange(aRange.relative(to: self))))
   }
 
   // - (NSArray *)
@@ -1406,7 +1417,7 @@ extension StringProtocol where Index == String.Index {
     var nsTokenRanges: NSArray?
     let result = tokenRanges._withNilOrAddress(of: &nsTokenRanges) {
       self._ns.linguisticTags(
-        in: _toNSRange(range.relative(to: self)),
+        in: _toRelativeNSRange(range.relative(to: self)),
         scheme: tagScheme._ephemeralString,
         options: opts,
         orthography: orthography,
@@ -1430,7 +1441,7 @@ extension StringProtocol where Index == String.Index {
     R : RangeExpression
   >(for aRange: R) -> Range<Index> where R.Bound == Index {
     return _range(
-      _ns.paragraphRange(for: _toNSRange(aRange.relative(to: self))))
+      _ns.paragraphRange(for: _toRelativeNSRange(aRange.relative(to: self))))
   }
 
   // - (NSRange)rangeOfCharacterFromSet:(NSCharacterSet *)aSet
@@ -1456,7 +1467,7 @@ extension StringProtocol where Index == String.Index {
       _ns.rangeOfCharacter(
         from: aSet,
         options: mask,
-        range: _toNSRange(
+        range: _toRelativeNSRange(
           aRange ?? startIndex..<endIndex
         )
       )
@@ -1487,7 +1498,7 @@ extension StringProtocol where Index == String.Index {
     // and output ranges due (if nothing else) to locale changes
     return _range(
       _ns.rangeOfComposedCharacterSequences(
-        for: _toNSRange(range.relative(to: self))))
+        for: _toRelativeNSRange(range.relative(to: self))))
   }
 
   // - (NSRange)rangeOfString:(NSString *)aString
@@ -1522,13 +1533,13 @@ extension StringProtocol where Index == String.Index {
       locale != nil ? _ns.range(
         of: aString,
         options: mask,
-        range: _toNSRange(
+        range: _toRelativeNSRange(
           searchRange ?? startIndex..<endIndex
         ),
         locale: locale
       )
       : searchRange != nil ? _ns.range(
-        of: aString, options: mask, range: _toNSRange(searchRange!)
+        of: aString, options: mask, range: _toRelativeNSRange(searchRange!)
       )
       : !mask.isEmpty ? _ns.range(of: aString, options: mask)
       : _ns.range(of: aString)
@@ -1637,7 +1648,7 @@ extension StringProtocol where Index == String.Index {
   @available(swift, deprecated: 4.0,
     message: "Please use String slicing subscript.")
   public func substring(with aRange: Range<Index>) -> String {
-    return _ns.substring(with: _toNSRange(aRange))
+    return _ns.substring(with: _toRelativeNSRange(aRange))
   }
 }
 

test/stdlib/NSStringAPI+Substring.swift

@@ -0,0 +1,146 @@
+// RUN: rm -rf %t ; mkdir -p %t
+// RUN: %target-build-swift %s -o %t/a.out4 -swift-version 4 && %target-run %t/a.out4
+// REQUIRES: executable_test
+
+// REQUIRES: objc_interop
+
+//
+// Tests for the NSString APIs on Substring
+//
+
+import StdlibUnittest
+
+import Foundation
+
+
+extension String {
+  func range(fromStart: Int, fromEnd: Int) -> Range<String.Index> {
+    return index(startIndex, offsetBy: fromStart) ..<
+           index(endIndex, offsetBy: fromEnd)
+  }
+  subscript(fromStart: Int, fromEnd: Int) -> SubSequence {
+    return self[range(fromStart: fromStart, fromEnd: fromEnd)]
+  }
+}
+
+var tests = TestSuite("NSStringAPIs/Substring")
+
+tests.test("range(of:)/NilRange") {
+  let ss = "aabcdd"[1, -1]
+  let range = ss.range(of: "bc")
+  expectOptionalEqual("bc", range.map { ss[$0] })
+}
+
+tests.test("range(of:)/NonNilRange") {
+  let s = "aabcdd"
+  let ss = s[1, -1]
+  let searchRange = s.range(fromStart: 2, fromEnd: -2)
+  let range = ss.range(of: "bc", range: searchRange)
+  expectOptionalEqual("bc", range.map { ss[$0] })
+}
+
+tests.test("rangeOfCharacter") {
+  let ss = "__hello__"[2, -2]
+  let range = ss.rangeOfCharacter(from: CharacterSet.alphanumerics)
+  expectOptionalEqual("h", range.map { ss[$0] })
+}
+
+tests.test("compare(_:options:range:locale:)/NilRange") {
+  let needle = "hello"
+  let haystack = "__hello__"[2, -2]
+  expectEqual(.orderedSame, haystack.compare(needle))
+}
+
+tests.test("compare(_:options:range:locale:)/NonNilRange") {
+  let needle = "hello"
+  let haystack = "__hello__"
+  let range = haystack.range(fromStart: 2, fromEnd: -2)
+  expectEqual(.orderedSame, haystack[range].compare(needle, range: range))
+}
+
+tests.test("replacingCharacters(in:with:)") {
+  let s = "__hello, world"
+  let range = s.range(fromStart: 2, fromEnd: -7)
+  let expected = "__goodbye, world"
+  let replacement = "goodbye"
+  expectEqual(expected,
+    s.replacingCharacters(in: range, with: replacement))
+  expectEqual(expected[2, 0],
+    s[2, 0].replacingCharacters(in: range, with: replacement))
+
+  expectEqual(replacement,
+    s.replacingCharacters(in: s.startIndex..., with: replacement))
+  expectEqual(replacement,
+    s.replacingCharacters(in: ..<s.endIndex, with: replacement))
+  expectEqual(expected[2, 0],
+    s[2, 0].replacingCharacters(in: range, with: replacement[...]))
+}
+
+tests.test("replacingOccurrences(of:with:options:range:)/NilRange") {
+  let s = "hello"
+
+  expectEqual("he11o", s.replacingOccurrences(of: "l", with: "1"))
+  expectEqual("he11o", s.replacingOccurrences(of: "l"[...], with: "1"))
+  expectEqual("he11o", s.replacingOccurrences(of: "l", with: "1"[...]))
+  expectEqual("he11o", s.replacingOccurrences(of: "l"[...], with: "1"[...]))
+
+  expectEqual("he11o",
+    s[...].replacingOccurrences(of: "l", with: "1"))
+  expectEqual("he11o",
+    s[...].replacingOccurrences(of: "l"[...], with: "1"))
+  expectEqual("he11o",
+    s[...].replacingOccurrences(of: "l", with: "1"[...]))
+  expectEqual("he11o",
+    s[...].replacingOccurrences(of: "l"[...], with: "1"[...]))
+}
+
+tests.test("replacingOccurrences(of:with:options:range:)/NonNilRange") {
+  let s = "hello"
+  let r = s.range(fromStart: 1, fromEnd: -2)
+
+  expectEqual("he1lo",
+    s.replacingOccurrences(of: "l", with: "1", range: r))
+  expectEqual("he1lo",
+    s.replacingOccurrences(of: "l"[...], with: "1", range: r))
+  expectEqual("he1lo",
+    s.replacingOccurrences(of: "l", with: "1"[...], range: r))
+  expectEqual("he1lo",
+    s.replacingOccurrences(of: "l"[...], with: "1"[...], range: r))
+
+  expectEqual("he1lo",
+    s[...].replacingOccurrences(of: "l", with: "1", range: r))
+  expectEqual("he1lo",
+    s[...].replacingOccurrences(of: "l"[...], with: "1", range: r))
+  expectEqual("he1lo",
+    s[...].replacingOccurrences(of: "l", with: "1"[...], range: r))
+  expectEqual("he1lo",
+    s[...].replacingOccurrences(of: "l"[...], with: "1"[...], range: r))
+
+  let ss = s[1, -1]
+  expectEqual("e1l",
+    ss.replacingOccurrences(of: "l", with: "1", range: r))
+  expectEqual("e1l",
+    ss.replacingOccurrences(of: "l"[...], with: "1", range: r))
+  expectEqual("e1l",
+    ss.replacingOccurrences(of: "l", with: "1"[...], range: r))
+  expectEqual("e1l",
+    ss.replacingOccurrences(of: "l"[...], with: "1"[...], range: r))
+}
+
+tests.test("substring(with:)") {
+  let s = "hello, world"
+  let r = s.range(fromStart: 7, fromEnd: 0)
+  expectEqual("world", s.substring(with: r))
+  expectEqual("world", s[...].substring(with: r))
+  expectEqual("world", s[1, 0].substring(with: r))
+}
+
+tests.test("substring(with:)/SubscriptEquivalence") {
+  let s = "hello, world"
+  let r = s.range(fromStart: 7, fromEnd: 0)
+  expectEqual(s[r], s.substring(with: r))
+  expectEqual(s[...][r], s[...].substring(with: r))
+  expectEqual(s[1, 0][r], s[1, 0].substring(with: r))
+}
+
+runAllTests()

test/stdlib/NSStringAPI.swift

@@ -1,4 +1,4 @@
-// RUN: %target-run-simple-swift -swift-version 3
+// RUN: %target-run-simple-swift
 // REQUIRES: executable_test
 
 // REQUIRES: objc_interop
@@ -1144,9 +1144,11 @@ NSStringAPIs.test("rangeOfComposedCharacterSequences(for:)") {
       for: s.index(s.startIndex, offsetBy: 8)..<s.index(s.startIndex, offsetBy: 10))])
 }
 
-func toIntRange(
-  _ string: String, _ maybeRange: Range<String.Index>?
-) -> Range<Int>? {
+func toIntRange<
+  S : StringProtocol
+>(
+  _ string: S, _ maybeRange: Range<String.Index>?
+) -> Range<Int>? where S.Index == String.Index, S.IndexDistance == Int {
   guard let range = maybeRange else { return nil }
 
   return