Stop using SourceManager::getBufferIdentifierForLoc to find buffer IDs

The right way is findBufferContainingLoc. getBufferIdentifierForLoc is
both slower and wrong in the presence of #sourceLocation.

I couldn't come up with a test for the change in IDE/Utils.cpp because
refactoring still seems to be broken around #sourceLocation. I'll file
bugs for that.

Author: jrose-apple

lib/IDE/Utils.cpp

@@ -904,8 +904,7 @@ class ClangFileRewriterHelper {
   }
 
   void replaceText(SourceManager &SM, CharSourceRange Range, StringRef Text) {
-    auto BufferId = SM.getIDForBufferIdentifier(SM.
-      getBufferIdentifierForLoc(Range.getStart())).getValue();
+    auto BufferId = SM.findBufferContainingLoc(Range.getStart());
     if (BufferId == InterestedId) {
       HasChange = true;
       auto StartLoc = SM.getLocOffsetInBuffer(Range.getStart(), BufferId);

lib/Parse/Lexer.cpp

@@ -2579,12 +2579,11 @@ static const char *findStartOfLine(const char *bufStart, const char *current) {
 }
 
 SourceLoc Lexer::getLocForStartOfToken(SourceManager &SM, SourceLoc Loc) {
-  Optional<unsigned> BufferIdOp = SM.getIDForBufferIdentifier(SM.
-    getBufferIdentifierForLoc(Loc));
-  if (!BufferIdOp.hasValue())
+  if (!Loc.isValid())
     return SourceLoc();
-  return getLocForStartOfToken(SM, BufferIdOp.getValue(),
-    SM.getLocOffsetInBuffer(Loc, BufferIdOp.getValue()));
+  unsigned BufferId = SM.findBufferContainingLoc(Loc);
+  return getLocForStartOfToken(SM, BufferId,
+                               SM.getLocOffsetInBuffer(Loc, BufferId));
 }
 
 SourceLoc Lexer::getLocForStartOfToken(SourceManager &SM, unsigned BufferID,

test/SourceKit/CursorInfo/cursor_info_param.swift

@@ -0,0 +1,45 @@
+func simpleParams(a: Int, b theB: Int) {}
+
+struct AccessorTest {
+  subscript(a: Int, b theB: Int) -> Int { return a }
+  var prop: Int {
+    get { return 0 }
+    set(v) {}
+  }
+}
+
+#sourceLocation(file: "custom.swuft", line: 2000)
+func customSourceLocation(a: Int) {}
+#sourceLocation()
+
+// RUN: %sourcekitd-test -req=cursor -pos=1:19 %s -- %s | %FileCheck -check-prefix=CHECK-FUNC-A %s
+// CHECK-FUNC-A: s:17cursor_info_param12simpleParams1a1bySi_SitFACL_Sivp
+// CHECK-FUNC-A: PARENT OFFSET: 5
+
+// RUN: %sourcekitd-test -req=cursor -pos=1:27 %s -- %s | %FileCheck -check-prefix=CHECK-FUNC-B %s
+// CHECK-FUNC-B: s:17cursor_info_param12simpleParams1a1bySi_SitF{{$}}
+
+// RUN: %sourcekitd-test -req=cursor -pos=1:29 %s -- %s | %FileCheck -check-prefix=CHECK-FUNC-THEB %s
+// CHECK-FUNC-THEB: s:17cursor_info_param12simpleParams1a1bySi_SitF4theBL_Sivp
+// CHECK-FUNC-THEB-NOT: PARENT OFFSET
+
+// RUN: %sourcekitd-test -req=cursor -pos=4:13 %s -- %s | %FileCheck -check-prefix=CHECK-SUBSCRIPT-A %s
+// FIXME: This USR is wrong; see https://bugs.swift.org/browse/SR-8660.
+// CHECK-SUBSCRIPT-A: s:17cursor_info_param12AccessorTestV1aL_Sivp
+// CHECK-SUBSCRIPT-A: PARENT OFFSET: 67
+
+// RUN: %sourcekitd-test -req=cursor -pos=4:21 %s -- %s | %FileCheck -check-prefix=CHECK-SUBSCRIPT-B %s
+// CHECK-SUBSCRIPT-B: s:17cursor_info_param12AccessorTestV
+
+// RUN: %sourcekitd-test -req=cursor -pos=4:23 %s -- %s | %FileCheck -check-prefix=CHECK-SUBSCRIPT-THEB %s
+// FIXME: This USR is wrong; see https://bugs.swift.org/browse/SR-8660.
+// CHECK-SUBSCRIPT-THEB: s:17cursor_info_param12AccessorTestV4theBL_Sivp
+// CHECK-SUBSCRIPT-THEB-NOT: PARENT OFFSET
+
+// RUN: %sourcekitd-test -req=cursor -pos=7:9 %s -- %s | %FileCheck -check-prefix=CHECK-SETTER-V %s
+// CHECK-SETTER-V: s:17cursor_info_param12AccessorTestV4propSivs1vL_Sivp
+// CHECK-SETTER-V: PARENT OFFSET: 161
+
+// RUN: %sourcekitd-test -req=cursor -pos=12:27 %s -- %s | %FileCheck -check-prefix=CHECK-CUSTOM-SOURCELOCATION %s
+// CHECK-CUSTOM-SOURCELOCATION: s:17cursor_info_param20customSourceLocation1aySi_tFACL_Sivp
+// CHECK-CUSTOM-SOURCELOCATION: PARENT OFFSET: 233

tools/SourceKit/lib/SwiftLang/SwiftSourceDocInfo.cpp

@@ -684,8 +684,7 @@ getParamParentNameOffset(const ValueDecl *VD, SourceLoc Cursor) {
   if (Loc.isInvalid())
     return None;
   auto &SM = VD->getASTContext().SourceMgr;
-  return SM.getLocOffsetInBuffer(Loc, SM.getIDForBufferIdentifier(SM.
-    getBufferIdentifierForLoc(Loc)).getValue());
+  return SM.getLocOffsetInBuffer(Loc, SM.findBufferContainingLoc(Loc));
 }
 
 /// Returns true for failure to resolve.

unittests/Parse/LexerTests.cpp

@@ -635,6 +635,29 @@ TEST_F(LexerTest, getLocForStartOfToken) {
   }
 }
 
+TEST_F(LexerTest, getLocForStartOfTokenWithCustomSourceLocation) {
+  const char *Source =
+      "aaa \n"
+      // This next line is exactly 50 bytes to make it easy to compare with the
+      // previous test.
+      "#sourceLocation(file: \"custom-50.swuft\", line: 9)\n"
+      " \tbbb \"hello\" \"-\\(val)-\"";
+
+  unsigned BufferID = SourceMgr.addMemBufferCopy(Source);
+
+  // First is character offset, second is its token offset.
+  unsigned Offs[][2] =
+    { {1, 0}, {2, 0}, {3, 3}, {4, 4},
+      {56, 56}, {59, 57}, {64, 61},
+      // interpolated string
+      {70, 69}, {73, 73}, {74, 73}, {75, 73}, {76, 76}, {77, 69} };
+
+  for (auto Pair : Offs) {
+    ASSERT_EQ(Lexer::getLocForStartOfToken(SourceMgr, BufferID, Pair[0]),
+              SourceMgr.getLocForOffset(BufferID, Pair[1]));
+  }
+}
+
 TEST_F(LexerTest, NestedSubLexers) {
   const char *Source = "aaa0 bbb1 ccc2 ddd3 eee4 fff5 ggg6";