Merge pull request #82171 from j-hui/j-hui/6.2/jump-to-def-for-macro-expanded-clang-imports

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

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 @@ public:
   /// \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 @@ public:
   /// 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"
@@ -8648,17 +8649,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;
     }
   }
@@ -8668,20 +8668,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);
@@ -8704,8 +8701,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

@@ -1057,9 +1057,9 @@ public:
   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 @@ private:
         !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]]

tools/SourceKit/lib/SwiftLang/SwiftSourceDocInfo.cpp

@@ -26,6 +26,7 @@
 #include "swift/AST/ModuleNameLookup.h"
 #include "swift/AST/NameLookup.h"
 #include "swift/AST/SwiftNameTranslation.h"
+#include "swift/Basic/Assertions.h"
 #include "swift/Basic/SourceManager.h"
 #include "swift/Frontend/Frontend.h"
 #include "swift/Frontend/PrintingDiagnosticConsumer.h"
@@ -48,6 +49,7 @@
 #include "clang/Index/USRGeneration.h"
 #include "clang/Lex/Lexer.h"
 
+#include "llvm/Support/Compiler.h"
 #include "llvm/Support/MemoryBuffer.h"
 
 #include <numeric>
@@ -877,56 +879,110 @@ static void setLocationInfoForClangNode(ClangNode ClangNode,
   }
 }
 
+static void setLocationInfoForRange(SourceManager &SM, SourceRange R,
+                                    unsigned BufID,
+                                    LocationInfo &Location,
+                                    bool Presumed = false) {
+  CONDITIONAL_ASSERT(BufID == SM.findBufferContainingLoc(R.Start) &&
+                     "SourceRange R should be in BufID");
+
+  auto CR = Lexer::getCharSourceRangeFromSourceRange(SM, R);
+
+  Location.Filename = SM.getIdentifierForBuffer(BufID);
+  Location.Length = CR.getByteLength();
+  Location.Offset = SM.getLocOffsetInBuffer(CR.getStart(), BufID);
+
+  if (Presumed)
+    std::tie(Location.Line, Location.Column) =
+        SM.getPresumedLineAndColumnForLoc(CR.getStart(), BufID);
+  else
+    std::tie(Location.Line, Location.Column) =
+        SM.getLineAndColumnInBuffer(CR.getStart(), BufID);
+}
+
 static void setLocationInfo(const ValueDecl *VD,
                             LocationInfo &Location) {
   ASTContext &Ctx = VD->getASTContext();
   SourceManager &SM = Ctx.SourceMgr;
+  auto *Importer = static_cast<ClangImporter *>(Ctx.getClangModuleLoader());
 
-  auto ClangNode = VD->getClangNode();
+  if (auto loc = VD->getLoc(/*SerializedOK=*/true)) {
+    // For most cases we just want the range of the name itself; it suffices to
+    // make Range from just Loc because Lexer::getCharSourceRangeFromSourceRange
+    // will grow the range to encompass the end of the token at Loc.
+    SourceRange range = loc;
 
-  auto Loc = VD->getLoc(/*SerializedOK=*/true);
-  if (Loc.isValid()) {
-    // 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;
+    // One exception is for functions, where we also want to include the range
+    // of the parameter list.
     if (auto *FD = dyn_cast<AbstractFunctionDecl>(VD)) {
-      if (auto R = FD->getParameterListSourceRange())
-        Range = R;
+      if (auto FDR = FD->getParameterListSourceRange())
+        range = FDR;
     }
 
-    auto [DeclBufID, DeclRange] =
-        VD->getModuleContext()->getOriginalRange(Range);
+    unsigned bufID = SM.findBufferContainingLoc(range.Start);
 
-    auto DeclCharRange =
-        Lexer::getCharSourceRangeFromSourceRange(SM, DeclRange);
-    auto DeclLoc = DeclCharRange.getStart();
+    // If this range is from a generated buffer, recursively "unexpand" macros
+    // to chase after where the macro was originally expanded from.
+    //
+    // However, we don't care about certain kinds of generated buffers, so save
+    // the original range and buffer ID so we can set location according to VD.
+    auto VDRange = range;
+    auto VDBufID = bufID;
+    while (auto *info = SM.getGeneratedSourceInfo(bufID)) {
+      switch (info->kind) {
+#define MACRO_ROLE(Name, Description)                                          \
+  case GeneratedSourceInfo::Name##MacroExpansion:
+#include "swift/Basic/MacroRoles.def"
+        if (auto *customAttr = info->attachedMacroCustomAttr)
+          range = customAttr->getRange();
+        else
+          range = ASTNode::getFromOpaqueValue(info->astNode).getSourceRange();
+        bufID = SM.findBufferContainingLoc(range.Start);
+        continue; // Continue while-loop to recursively un-expand macros
 
-    Location.Filename = SM.getIdentifierForBuffer(DeclBufID);
-    Location.Offset = SM.getLocOffsetInBuffer(DeclLoc, DeclBufID);
-    Location.Length = DeclCharRange.getByteLength();
-    std::tie(Location.Line, Location.Column) =
-        SM.getLineAndColumnInBuffer(DeclLoc, DeclBufID);
-    if (auto GeneratedSourceInfo = SM.getGeneratedSourceInfo(DeclBufID)) {
-      if (GeneratedSourceInfo->kind ==
-          GeneratedSourceInfo::ReplacedFunctionBody) {
-        // The location was in a temporary source buffer that just contains the
-        // function body and which we created while reusing the ASTContext for
-        // the rest of the file. Map the location back to the original file.
-        unsigned OriginalBufID = SM.findBufferContainingLoc(
-            GeneratedSourceInfo->originalSourceRange.getStart());
-        auto OriginalStartOffset = SM.getLocOffsetInBuffer(
-            GeneratedSourceInfo->originalSourceRange.getStart(), OriginalBufID);
-        auto GeneratedStartOffset = SM.getLocOffsetInBuffer(
-            GeneratedSourceInfo->generatedSourceRange.getStart(), DeclBufID);
-        Location.Offset += OriginalStartOffset - GeneratedStartOffset;
-        std::tie(Location.Line, Location.Column) =
-            SM.getPresumedLineAndColumnForLoc(DeclLoc, DeclBufID);
+      case GeneratedSourceInfo::ReplacedFunctionBody:
+        if (bufID == VDBufID) {
+          // The location was in a temporary source buffer that just contains
+          // the function body, which we created while reusing the ASTContext
+          // for the rest of the file. Set the location so that it maps back to
+          // the original file.
+          setLocationInfoForRange(SM, VDRange, VDBufID, Location,
+                                  /*Presumed=*/true);
+          // Adjust offset according to generated buffer.
+          auto originalLoc = info->originalSourceRange.getStart();
+          auto originalBufID = SM.findBufferContainingLoc(originalLoc);
+          auto generatedLoc = info->generatedSourceRange.getStart();
+          auto generatedBufID = SM.findBufferContainingLoc(generatedLoc);
+          Location.Offset +=
+              SM.getLocOffsetInBuffer(originalLoc, originalBufID) -
+              SM.getLocOffsetInBuffer(generatedLoc, generatedBufID);
+        } else {
+          // We somehow encountered a replaced function body while looking
+          // through other macro expansions. Fall back to setting location based
+          // on VD's original source range because mapping location back to the
+          // original file might be tricky.
+          setLocationInfoForRange(SM, VDRange, VDBufID, Location);
+        }
+        return;
+      case GeneratedSourceInfo::AttributeFromClang:
+        // This buffer was generated for an imported ClangNode, so set location
+        // info according to that.
+        if (auto node = info->clangNode)
+          setLocationInfoForClangNode(node, Importer, Location);
+        else
+          setLocationInfoForRange(SM, VDRange, VDBufID, Location);
+        return;
+      case GeneratedSourceInfo::DefaultArgument:
+      case GeneratedSourceInfo::PrettyPrinted:
+        setLocationInfoForRange(SM, VDRange, VDBufID, Location);
+        return;
       }
+      llvm_unreachable("All switch cases either explicitly continue or return");
     }
-  } else if (ClangNode) {
-    ClangImporter *Importer =
-        static_cast<ClangImporter*>(Ctx.getClangModuleLoader());
-    setLocationInfoForClangNode(ClangNode, Importer, Location);
+
+    setLocationInfoForRange(SM, range, bufID, Location);
+  } else if (auto CNode = VD->getClangNode()) {
+    setLocationInfoForClangNode(CNode, Importer, Location);
   }
 }