[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 that addresses
some issues with getOriginal{Range,Location}() that make it resistant to
a more localized patch:

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

-   getOriginalLocation() is used in a couple of other places that don't
    actually care about the original location, only the outermost buffer
    a macro is expanded from. The "macro-chasing" logic is now extracted
    to SourceManager::getMacroUnexpandedBufferID(); the method is moved
    to from ModuleDecl to SourceManager because there's no reason for it
    to be in 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/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.
@@ -554,6 +560,10 @@ class SourceManager {
   /// \p _SwiftifyImport macro.
   bool isImportMacroGeneratedLoc(SourceLoc loc);
 
+  /// Find the outermost buffer ID \p BufID was originally generated from.
+  /// Returns \p BufID if it wasn't generated from a macro.
+  unsigned getMacroUnexpandedBufferID(unsigned BufID) const;
+
 private:
   int getLineOffset(SourceLoc Loc) const {
     if (auto VFile = getVirtualFile(Loc))

lib/AST/Module.cpp

@@ -847,49 +847,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 {
@@ -1430,13 +1387,9 @@ SourceFile::getExternalRawLocsForDecl(const Decl *D) const {
   SourceManager &SM = getASTContext().SourceMgr;
   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 (InGeneratedBuffer && BufferID != SM.getMacroUnexpandedBufferID(
+                                           SM.findBufferContainingLoc(MainLoc)))
+    return std::nullopt;
 
   auto setLoc = [&](ExternalSourceLocs::RawLoc &RawLoc, SourceLoc Loc) {
     if (!Loc.isValid())

lib/Basic/SourceLoc.cpp

@@ -882,3 +882,26 @@ bool SourceManager::isImportMacroGeneratedLoc(SourceLoc loc) {
 
   return false;
 }
+
+unsigned SourceManager::getMacroUnexpandedBufferID(unsigned bufferID) const {
+  auto original = bufferID;
+  while (const auto *info = getGeneratedSourceInfo(bufferID)) {
+    switch (info->kind) {
+#define MACRO_ROLE(Name, Description)                                          \
+  case GeneratedSourceInfo::Name##MacroExpansion:
+#include "swift/Basic/MacroRoles.def"
+      if (auto *customAttr = info->attachedMacroCustomAttr)
+        bufferID = findBufferContainingLoc(customAttr->getStartLoc());
+      else
+        bufferID = findBufferContainingLoc(
+            ASTNode::getFromOpaqueValue(info->astNode).getStartLoc());
+      continue;
+    case GeneratedSourceInfo::ReplacedFunctionBody:
+    case GeneratedSourceInfo::PrettyPrinted:
+    case GeneratedSourceInfo::DefaultArgument:
+    case GeneratedSourceInfo::AttributeFromClang:
+      return original;
+    }
+  }
+  return bufferID;
+}

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

@@ -1092,12 +1092,11 @@ class IndexSwiftASTWalker : public SourceEntityWalker {
     bool inGeneratedBuffer =
         !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 (inGeneratedBuffer &&
+        bufferID != SrcMgr.getMacroUnexpandedBufferID(
+                        SrcMgr.findBufferContainingLoc(loc))) {
+      assert(false && "Location is not within file being indexed");
+      return std::nullopt;
     }
 
     auto [line, col] = SrcMgr.getLineAndColumnInBuffer(loc, bufferID);

test/SourceKit/Swiftify/Inputs/from-clang.h

@@ -0,0 +1,9 @@
+#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);
+
+void hasSpanOverload(int len, int * __counted_by(len) __noescape p);

test/SourceKit/Swiftify/Inputs/module.modulemap

@@ -0,0 +1,4 @@
+module FromClang {
+    header "from-clang.h"
+    export *
+}

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

@@ -0,0 +1,41 @@
+// REQUIRES: swift_feature_SafeInteropWrappers
+// REQUIRES: swift_feature_LifetimeDependence
+
+import FromClang
+
+// RUN: %target-swift-ide-test \
+// RUN:   -print-module -module-to-print=FromClang -source-filename=x \
+// RUN:   -plugin-path %swift-plugin-dir -I %S/Inputs \
+// RUN:   -enable-experimental-feature SafeInteropWrappers \
+// RUN:   -enable-experimental-feature LifetimeDependence \
+// RUN: | %FileCheck %s --check-prefix INTERFACE
+
+// 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: %sourcekitd-test -req=cursor -pos=26:3 %s -- -I %S/Inputs %s \
+// RUN:   -enable-experimental-feature SafeInteropWrappers \
+// RUN:   -enable-experimental-feature LifetimeDependence \
+// RUN: | %FileCheck %s --check-prefix CLANG-HEADER
+@inlinable
+public func callWithBufferPtr(_ p: UnsafeMutableBufferPointer<CInt>) {
+  hasBufferOverload(p)
+}
+
+// RUN: %sourcekitd-test -req=cursor -pos=37:3 %s -- -I %S/Inputs %s \
+// RUN:   -enable-experimental-feature SafeInteropWrappers \
+// RUN:   -enable-experimental-feature LifetimeDependence \
+// RUN: | %FileCheck %s --check-prefix CLANG-HEADER
+@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)
+}
+
+// Check that we get the file path of the header that Clang imported:
+// CLANG-HEADER: /Inputs/from-clang.h: