[SourceKit] Support location info for macro-expanded Clang imports

Currently, when we jump-to-definition for decls that are macro-expanded
from Clang imported decls (e.g., safe overloads generated by
@_SwiftifyImport), setLocationInfo() emits a bongus location pointing to
a generated buffer, leading the IDE to try to jump to a file that does
not exist.

The root cause here is that setLocationInfo() calls getOriginalRange()
(earlier, getOriginalLocation()), which was not written to account for
such cases where a macro is generated from another generated buffer
whose kind is 'AttributeFromClang'.

This patch fixes setLocationInfo() with some refactoring:

-   getOriginalRange() is inlined into setLocationInfo(), so that the
    generated buffer-handling logic is localized to that function. This
    includes how it handles buffers generated for ReplacedFunctionBody.

-   getOriginalLocation() is used in a couple of other places that only
    care about macros expanded from the same buffer (so other generated
    buffers not not relevant). This "macro-chasing" logic is simplified
    and moved from ModuleDecl::getOriginalRange() to a free-standing
    function, getMacroUnexpandedRange() (there is no reason for it to be
    a method of ModuleDecl).

-   GeneratedSourceInfo now carries an extra ClangNode field, which is
    populated by getClangSwiftAttrSourceFile() when constructing
    a generated buffer for an 'AttributeFromClang'. This could probably
    be union'ed with one or more of the other fields in the future.

rdar://151020332

Author: j-hui

include/swift/AST/ASTNode.h

@@ -21,6 +21,7 @@
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/PointerUnion.h"
 #include "swift/Basic/Debug.h"
+#include "swift/Basic/SourceManager.h"
 #include "swift/AST/TypeAlignments.h"
 
 namespace llvm {
@@ -98,6 +99,12 @@ namespace swift {
       return llvm::hash_value(N.getOpaqueValue());
     }
   };
+
+  /// Find the outermost range that \p range was originally generated from.
+  /// Returns an invalid source range if \p range wasn't generated from a macro.
+  SourceRange getUnexpandedMacroRange(const SourceManager &SM,
+                                      SourceRange range);
+
 } // namespace swift
 
 namespace llvm {

include/swift/AST/Module.h

@@ -434,19 +434,6 @@ 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 {
-    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.
   ///

include/swift/Basic/SourceManager.h

@@ -13,6 +13,7 @@
 #ifndef SWIFT_BASIC_SOURCEMANAGER_H
 #define SWIFT_BASIC_SOURCEMANAGER_H
 
+#include "swift/AST/ClangNode.h"
 #include "swift/Basic/FileSystem.h"
 #include "swift/Basic/SourceLoc.h"
 #include "clang/Basic/FileManager.h"
@@ -22,6 +23,7 @@
 #include "llvm/Support/SourceMgr.h"
 #include <map>
 #include <optional>
+#include <utility>
 #include <vector>
 
 namespace swift {
@@ -122,6 +124,10 @@ class GeneratedSourceInfo {
   /// Contains the ancestors of this source buffer, starting with the root source
   /// buffer and ending at this source buffer.
   mutable llvm::ArrayRef<unsigned> ancestors = llvm::ArrayRef<unsigned>();
+
+  /// Clang node where this buffer comes from. This should be set when this is
+  /// an 'AttributeFromClang'.
+  ClangNode clangNode = ClangNode();
 };
 
 /// This class manages and owns source buffers.

lib/AST/ASTNode.cpp

@@ -157,3 +157,29 @@ FUNC(Expr)
 FUNC(Decl)
 FUNC(Pattern)
 #undef FUNC
+
+SourceRange swift::getUnexpandedMacroRange(const SourceManager &SM,
+                                           SourceRange range) {
+  unsigned bufferID = SM.findBufferContainingLoc(range.Start);
+  SourceRange outerRange;
+  while (const auto *info = SM.getGeneratedSourceInfo(bufferID)) {
+    switch (info->kind) {
+#define MACRO_ROLE(Name, Description)                                          \
+  case GeneratedSourceInfo::Name##MacroExpansion:
+#include "swift/Basic/MacroRoles.def"
+      if (auto *customAttr = info->attachedMacroCustomAttr)
+        outerRange = customAttr->getRange();
+      else
+        outerRange =
+            ASTNode::getFromOpaqueValue(info->astNode).getSourceRange();
+      bufferID = SM.findBufferContainingLoc(outerRange.Start);
+      continue;
+    case GeneratedSourceInfo::ReplacedFunctionBody:
+    case GeneratedSourceInfo::PrettyPrinted:
+    case GeneratedSourceInfo::DefaultArgument:
+    case GeneratedSourceInfo::AttributeFromClang:
+      return SourceRange();
+    }
+  }
+  return outerRange;
+}

lib/AST/Module.cpp

@@ -68,6 +68,7 @@
 #include "llvm/Support/SaveAndRestore.h"
 #include "llvm/Support/YAMLTraits.h"
 #include "llvm/Support/raw_ostream.h"
+#include <optional>
 
 using namespace swift;
 
@@ -847,49 +848,6 @@ SourceFile *ModuleDecl::getSourceFileContainingLocation(SourceLoc loc) {
   return nullptr;
 }
 
-std::pair<unsigned, SourceRange>
-ModuleDecl::getOriginalRange(SourceRange range) const {
-  assert(range.isValid());
-
-  SourceManager &SM = getASTContext().SourceMgr;
-  unsigned bufferID = SM.findBufferContainingLoc(range.Start);
-
-  auto startRange = range;
-  unsigned startBufferID = bufferID;
-  while (const GeneratedSourceInfo *info =
-             SM.getGeneratedSourceInfo(bufferID)) {
-    switch (info->kind) {
-#define MACRO_ROLE(Name, Description)  \
-    case GeneratedSourceInfo::Name##MacroExpansion:
-#include "swift/Basic/MacroRoles.def"
-    {
-      // Location was within a macro expansion, return the expansion site, not
-      // the insertion location.
-      if (info->attachedMacroCustomAttr) {
-        range = info->attachedMacroCustomAttr->getRange();
-      } else {
-        ASTNode expansionNode = ASTNode::getFromOpaqueValue(info->astNode);
-        range = expansionNode.getSourceRange();
-      }
-      bufferID = SM.findBufferContainingLoc(range.Start);
-      break;
-    }
-    case GeneratedSourceInfo::DefaultArgument:
-      // No original location as it's not actually in any source file
-    case GeneratedSourceInfo::ReplacedFunctionBody:
-      // There's not really any "original" location for locations within
-      // replaced function bodies. The body is actually different code to the
-      // original file.
-    case GeneratedSourceInfo::PrettyPrinted:
-    case GeneratedSourceInfo::AttributeFromClang:
-      // No original location, return the original buffer/location
-      return {startBufferID, startRange};
-    }
-  }
-
-  return {bufferID, range};
-}
-
 ArrayRef<SourceFile *>
 PrimarySourceFilesRequest::evaluate(Evaluator &evaluator,
                                     ModuleDecl *mod) const {
@@ -1431,11 +1389,11 @@ SourceFile::getExternalRawLocsForDecl(const Decl *D) const {
   bool InGeneratedBuffer =
       !SM.rangeContainsTokenLoc(SM.getRangeForBuffer(BufferID), MainLoc);
   if (InGeneratedBuffer) {
-    unsigned UnderlyingBufferID;
-    std::tie(UnderlyingBufferID, MainLoc) =
-        D->getModuleContext()->getOriginalLocation(MainLoc);
-    if (BufferID != UnderlyingBufferID)
-      return std::nullopt;
+    if (auto R = getUnexpandedMacroRange(SM, MainLoc)) {
+      if (BufferID != SM.findBufferContainingLoc(R.Start))
+        return std::nullopt;
+      MainLoc = R.Start;
+    }
   }
 
   auto setLoc = [&](ExternalSourceLocs::RawLoc &RawLoc, SourceLoc Loc) {

lib/ClangImporter/ImportDecl.cpp

@@ -46,6 +46,7 @@
 #include "swift/Basic/Defer.h"
 #include "swift/Basic/PrettyStackTrace.h"
 #include "swift/Basic/SourceLoc.h"
+#include "swift/Basic/SourceManager.h"
 #include "swift/Basic/Statistic.h"
 #include "swift/Basic/StringExtras.h"
 #include "swift/Basic/Version.h"
@@ -8669,17 +8670,16 @@ bool importer::hasSameUnderlyingType(const clang::Type *a,
 }
 
 SourceFile &ClangImporter::Implementation::getClangSwiftAttrSourceFile(
-    ModuleDecl &module,
-    StringRef attributeText,
-    bool cached
-) {
+    Decl *MappedDecl, StringRef attributeText, bool cached) {
+  auto *module = MappedDecl->getDeclContext()->getParentModule();
+
   ::TinyPtrVector<SourceFile *> *sourceFiles = nullptr;
   if (cached) {
     sourceFiles = &ClangSwiftAttrSourceFiles[attributeText];
 
     // Check whether we've already created a source file.
     for (auto sourceFile : *sourceFiles) {
-      if (sourceFile->getParentModule() == &module)
+      if (sourceFile->getParentModule() == module)
         return *sourceFile;
     }
   }
@@ -8689,20 +8689,17 @@ SourceFile &ClangImporter::Implementation::getClangSwiftAttrSourceFile(
   auto &sourceMgr = SwiftContext.SourceMgr;
   auto bufferID = sourceMgr.addMemBufferCopy(attributeText);
 
-  // Note that this is for an attribute.
-  sourceMgr.setGeneratedSourceInfo(
-      bufferID,
-      {
-        GeneratedSourceInfo::AttributeFromClang,
-        CharSourceRange(),
-        sourceMgr.getRangeForBuffer(bufferID),
-        &module
-      }
-  );
+  auto info = GeneratedSourceInfo{GeneratedSourceInfo::AttributeFromClang,
+                                  CharSourceRange(),
+                                  sourceMgr.getRangeForBuffer(bufferID)};
+  info.astNode = static_cast<void *>(module);
+  info.clangNode = MappedDecl->getClangNode();
+
+  sourceMgr.setGeneratedSourceInfo(bufferID, info);
 
   // Create the source file.
-  auto sourceFile = new (SwiftContext)
-      SourceFile(module, SourceFileKind::Library, bufferID);
+  auto sourceFile =
+      new (SwiftContext) SourceFile(*module, SourceFileKind::Library, bufferID);
 
   if (cached)
     sourceFiles->push_back(sourceFile);
@@ -8725,8 +8722,8 @@ void ClangImporter::Implementation::importNontrivialAttribute(
   bool cached = true;
   while (true) {
     // Dig out a source file we can use for parsing.
-    auto &sourceFile = getClangSwiftAttrSourceFile(
-        *MappedDecl->getDeclContext()->getParentModule(), AttrString, cached);
+    auto &sourceFile =
+        getClangSwiftAttrSourceFile(MappedDecl, AttrString, cached);
 
     auto topLevelDecls = sourceFile.getTopLevelDecls();
 

lib/ClangImporter/ImporterImpl.h

@@ -1065,9 +1065,9 @@ class LLVM_LIBRARY_VISIBILITY ClangImporter::Implementation
   StringRef getSwiftNameFromClangName(StringRef name);
 
   /// Retrieve the placeholder source file for use in parsing Swift attributes
-  /// in the given module.
-  SourceFile &getClangSwiftAttrSourceFile(
-      ModuleDecl &module, StringRef attributeText, bool cached);
+  /// of the given Decl.
+  SourceFile &getClangSwiftAttrSourceFile(Decl *MappedDecl,
+                                          StringRef attributeText, bool cached);
 
   /// Create attribute with given text and attach it to decl, creating or
   /// retrieving a chached source file as needed.

lib/Index/Index.cpp

@@ -38,6 +38,7 @@
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/FileSystem.h"
+#include <optional>
 #include <tuple>
 
 using namespace swift;
@@ -1093,10 +1094,12 @@ class IndexSwiftASTWalker : public SourceEntityWalker {
         !SrcMgr.rangeContainsTokenLoc(SrcMgr.getRangeForBuffer(bufferID), loc);
 
     if (inGeneratedBuffer) {
-      std::tie(bufferID, loc) = CurrentModule->getOriginalLocation(loc);
-      if (BufferID.value() != bufferID) {
-        assert(false && "Location is not within file being indexed");
-        return std::nullopt;
+      if (auto unexpandedRange = getUnexpandedMacroRange(SrcMgr, loc)) {
+        loc = unexpandedRange.Start;
+        if (bufferID != SrcMgr.findBufferContainingLoc(loc)) {
+          assert(false && "Location should be within file being indexed");
+          return std::nullopt;
+        }
       }
     }
 

test/SourceKit/Macros/clang-overload-cursor-info.swift

@@ -0,0 +1,61 @@
+// RUN: %empty-directory(%t)
+// RUN: split-file %s %t
+//--- Main.swift
+import FromClang // NOTE: line offset = -4
+
+// REQUIRES: swift_feature_SafeInteropWrappers
+// REQUIRES: swift_feature_LifetimeDependence
+
+// The macro-generated interface we're looking up source info for
+// (this is more so for documentation than checking correctness)
+//
+// INTERFACE:      @_alwaysEmitIntoClient @_disfavoredOverload public func hasBufferOverload(_ p: UnsafeMutableBufferPointer<Int32>)
+// INTERFACE:      @{{_?}}lifetime(p: copy p)
+// INTERFACE-NEXT: @_alwaysEmitIntoClient @_disfavoredOverload public func hasSpanOverload(_ p: inout MutableSpan<Int32>)
+// RUN: %target-swift-ide-test \
+// RUN:   -print-module -module-to-print=FromClang -source-filename=x \
+// RUN:   -plugin-path %swift-plugin-dir -I %t/Inputs \
+// RUN:   -enable-experimental-feature SafeInteropWrappers \
+// RUN:   -enable-experimental-feature LifetimeDependence \
+// RUN: | %FileCheck %t/Main.swift --check-prefix INTERFACE
+
+@inlinable
+public func callWithBufferPtr(_ p: UnsafeMutableBufferPointer<CInt>) {
+  hasBufferOverload(p)
+// RUN: %sourcekitd-test -req=cursor -pos=%(line-4):3 %t/Main.swift -- -I %t/Inputs %t/Main.swift \
+// RUN:   -enable-experimental-feature SafeInteropWrappers \
+// RUN:   -enable-experimental-feature LifetimeDependence \
+// RUN: | %FileCheck %t/Inputs/from-clang.h --check-prefix BUFFER-OVERLOAD
+}
+
+@available(visionOS 1.1, tvOS 12.2, watchOS 5.2, iOS 12.2, macOS 10.14.4, *)
+@lifetime(p: copy p)
+@inlinable
+public func callReturnLifetimeBound(_ p: inout MutableSpan<CInt>) {
+  hasSpanOverload(p)
+// RUN: %sourcekitd-test -req=cursor -pos=%(line-4):3 %t/Main.swift -- -I %t/Inputs %t/Main.swift \
+// RUN:   -enable-experimental-feature SafeInteropWrappers \
+// RUN:   -enable-experimental-feature LifetimeDependence \
+// RUN: | %FileCheck %t/Inputs/from-clang.h --check-prefix SPAN-OVERLOAD
+}
+
+//--- Inputs/module.modulemap
+module FromClang {
+    header "from-clang.h"
+    export *
+}
+
+//--- Inputs/from-clang.h
+#pragma once
+
+#define __counted_by(x) __attribute__((__counted_by__(x)))
+#define __noescape __attribute__((noescape))
+#define __lifetimebound __attribute__((lifetimebound))
+
+void hasBufferOverload(int len, int * __counted_by(len) p);
+// BUFFER-OVERLOAD:      source.lang.swift.ref.function.free
+// BUFFER-OVERLOAD-SAME: from-clang.h:[[@LINE-2]]
+
+void hasSpanOverload(int len, int * __counted_by(len) __noescape p);
+// SPAN-OVERLOAD:      source.lang.swift.ref.function.free
+// SPAN-OVERLOAD-SAME: from-clang.h:[[@LINE-2]]