[SourceKit] Properly handle cursor info range for macro expansions

Make `getOriginalLocation` work with source ranges, and adjust the
cursor info logic to map the range into the original buffer. This
fixes the case where we were using bogus range lengths for macro
expansion decls.

rdar://151411756

Author: hamishknight

include/swift/AST/Module.h

@@ -434,10 +434,18 @@ class ModuleDecl
   /// \c nullptr if the source location isn't in this module.
   SourceFile *getSourceFileContainingLocation(SourceLoc loc);
 
+  // Retrieve the buffer ID and source range of the outermost node that
+  // caused the generation of the buffer containing \p range. \p range and its
+  // buffer if it isn't in a generated buffer or has no original range.
+  std::pair<unsigned, SourceRange> getOriginalRange(SourceRange range) const;
+
   // Retrieve the buffer ID and source location of the outermost location that
   // caused the generation of the buffer containing \p loc. \p loc and its
   // buffer if it isn't in a generated buffer or has no original location.
-  std::pair<unsigned, SourceLoc> getOriginalLocation(SourceLoc loc) const;
+  std::pair<unsigned, SourceLoc> getOriginalLocation(SourceLoc loc) const {
+    auto [buffer, range] = getOriginalRange(loc);
+    return std::make_pair(buffer, range.Start);
+  }
 
   /// Creates a map from \c #filePath strings to corresponding \c #fileID
   /// strings, diagnosing any conflicts.

lib/AST/Module.cpp

@@ -847,14 +847,14 @@ SourceFile *ModuleDecl::getSourceFileContainingLocation(SourceLoc loc) {
   return nullptr;
 }
 
-std::pair<unsigned, SourceLoc>
-ModuleDecl::getOriginalLocation(SourceLoc loc) const {
-  assert(loc.isValid());
+std::pair<unsigned, SourceRange>
+ModuleDecl::getOriginalRange(SourceRange range) const {
+  assert(range.isValid());
 
   SourceManager &SM = getASTContext().SourceMgr;
-  unsigned bufferID = SM.findBufferContainingLoc(loc);
+  unsigned bufferID = SM.findBufferContainingLoc(range.Start);
 
-  SourceLoc startLoc = loc;
+  auto startRange = range;
   unsigned startBufferID = bufferID;
   while (const GeneratedSourceInfo *info =
              SM.getGeneratedSourceInfo(bufferID)) {
@@ -866,12 +866,12 @@ ModuleDecl::getOriginalLocation(SourceLoc loc) const {
       // Location was within a macro expansion, return the expansion site, not
       // the insertion location.
       if (info->attachedMacroCustomAttr) {
-        loc = info->attachedMacroCustomAttr->getLocation();
+        range = info->attachedMacroCustomAttr->getRange();
       } else {
         ASTNode expansionNode = ASTNode::getFromOpaqueValue(info->astNode);
-        loc = expansionNode.getStartLoc();
+        range = expansionNode.getSourceRange();
       }
-      bufferID = SM.findBufferContainingLoc(loc);
+      bufferID = SM.findBufferContainingLoc(range.Start);
       break;
     }
     case GeneratedSourceInfo::DefaultArgument:
@@ -883,11 +883,11 @@ ModuleDecl::getOriginalLocation(SourceLoc loc) const {
     case GeneratedSourceInfo::PrettyPrinted:
     case GeneratedSourceInfo::AttributeFromClang:
       // No original location, return the original buffer/location
-      return {startBufferID, startLoc};
+      return {startBufferID, startRange};
     }
   }
 
-  return {bufferID, loc};
+  return {bufferID, range};
 }
 
 ArrayRef<SourceFile *>

test/SourceKit/Macros/expansion_decl_cursor_info.swift

@@ -0,0 +1,46 @@
+// RUN: %empty-directory(%t)
+// REQUIRES: swift_swift_parser
+
+// RUN: split-file --leading-lines %s %t
+
+// RUN: %host-build-swift -swift-version 5 -emit-library -o %t/%target-library-name(MacroDefinition) -module-name=MacroDefinition %t/MacroDefinition.swift -g -no-toolchain-stdlib-rpath
+
+//--- MacroDefinition.swift
+
+import SwiftSyntax
+import SwiftSyntaxMacros
+
+public struct AddFuncMacro: DeclarationMacro, PeerMacro {
+  public static func expansion(
+    of macro: some FreestandingMacroExpansionSyntax,
+    in context: some MacroExpansionContext
+  ) -> [DeclSyntax] {
+    ["func foo(_ x: String) {}"]
+  }
+  public static func expansion(
+    of node: AttributeSyntax,
+    providingPeersOf declaration: some DeclSyntaxProtocol,
+    in context: some MacroExpansionContext
+  ) throws -> [DeclSyntax] {
+    ["func foo(_ x: Int) {}"]
+  }
+}
+
+//--- main.swift
+
+@attached(peer, names: named(foo))
+macro AddFuncPeer(x: Int) = #externalMacro(module: "MacroDefinition", type: "AddFuncMacro")
+
+@freestanding(declaration, names: named(foo))
+macro AddFunc(x: Int) = #externalMacro(module: "MacroDefinition", type: "AddFuncMacro")
+
+@AddFuncPeer(x: 0)
+#AddFunc(x: 0)
+
+foo(0)
+// RUN: %sourcekitd-test -req=cursor %t/main.swift -pos=%(line-1):1 -- %t/main.swift -load-plugin-library %t/%target-library-name(MacroDefinition) -module-name MacroUser | %FileCheck --check-prefix PEER %s
+// PEER: source.lang.swift.ref.function.free ([[@LINE-5]]:1-[[@LINE-5]]:19)
+
+foo("")
+// RUN: %sourcekitd-test -req=cursor %t/main.swift -pos=%(line-1):1 -- %t/main.swift -load-plugin-library %t/%target-library-name(MacroDefinition) -module-name MacroUser | %FileCheck --check-prefix FREESTANDING %s
+// FREESTANDING: source.lang.swift.ref.function.free ([[@LINE-8]]:1-[[@LINE-8]]:15)

tools/SourceKit/lib/SwiftLang/SwiftSourceDocInfo.cpp

@@ -877,11 +877,6 @@ static void setLocationInfoForClangNode(ClangNode ClangNode,
   }
 }
 
-static unsigned getCharLength(SourceManager &SM, SourceRange TokenRange) {
-  SourceLoc CharEndLoc = Lexer::getLocForEndOfToken(SM, TokenRange.End);
-  return SM.getByteDistance(TokenRange.Start, CharEndLoc);
-}
-
 static void setLocationInfo(const ValueDecl *VD,
                             LocationInfo &Location) {
   ASTContext &Ctx = VD->getASTContext();
@@ -891,29 +886,24 @@ static void setLocationInfo(const ValueDecl *VD,
 
   auto Loc = VD->getLoc(/*SerializedOK=*/true);
   if (Loc.isValid()) {
-    auto getParameterListRange =
-        [&](const ValueDecl *VD) -> std::optional<unsigned> {
-      if (auto FD = dyn_cast<AbstractFunctionDecl>(VD)) {
-        SourceRange R = FD->getParameterListSourceRange();
-        if (R.isValid())
-          return getCharLength(SM, R);
-      }
-      return std::nullopt;
-    };
-    unsigned NameLen;
-    if (auto SigLen = getParameterListRange(VD)) {
-      NameLen = SigLen.value();
-    } else if (VD->hasName()) {
-      NameLen = VD->getBaseName().userFacingName().size();
-    } else {
-      NameLen = getCharLength(SM, Loc);
+    // For most cases we just want the range of the name itself. One exception
+    // is for functions, where we also want to include the parameter list.
+    SourceRange Range = Loc;
+    if (auto *FD = dyn_cast<AbstractFunctionDecl>(VD)) {
+      if (auto R = FD->getParameterListSourceRange())
+        Range = R;
     }
 
-    auto [DeclBufID, DeclLoc] =
-        VD->getModuleContext()->getOriginalLocation(Loc);
+    auto [DeclBufID, DeclRange] =
+        VD->getModuleContext()->getOriginalRange(Range);
+
+    auto DeclCharRange =
+        Lexer::getCharSourceRangeFromSourceRange(SM, DeclRange);
+    auto DeclLoc = DeclCharRange.getStart();
+
     Location.Filename = SM.getIdentifierForBuffer(DeclBufID);
     Location.Offset = SM.getLocOffsetInBuffer(DeclLoc, DeclBufID);
-    Location.Length = NameLen;
+    Location.Length = DeclCharRange.getByteLength();
     std::tie(Location.Line, Location.Column) =
         SM.getLineAndColumnInBuffer(DeclLoc, DeclBufID);
     if (auto GeneratedSourceInfo = SM.getGeneratedSourceInfo(DeclBufID)) {