[6.2] [SourceKit] Properly handle cursor info range for macro expansions

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)) {