[ClangImporter] Import underlying Clang locations

include/swift/AST/Decl.h

@@ -3432,7 +3432,9 @@ class EnumDecl final : public NominalTypeDecl {
 
   SourceLoc getStartLoc() const { return EnumLoc; }
   SourceRange getSourceRange() const {
-    return SourceRange(EnumLoc, getBraces().End);
+    if (getBraces().isValid())
+      return SourceRange(EnumLoc, getBraces().End);
+    return SourceRange();
   }
 
 public:
@@ -3600,7 +3602,9 @@ class StructDecl final : public NominalTypeDecl {
 
   SourceLoc getStartLoc() const { return StructLoc; }
   SourceRange getSourceRange() const {
-    return SourceRange(StructLoc, getBraces().End);
+    if (getBraces().isValid())
+      return SourceRange(StructLoc, getBraces().End);
+    return SourceRange();
   }
 
   // Implement isa/cast/dyncast/etc.
@@ -3747,7 +3751,9 @@ class ClassDecl final : public NominalTypeDecl {
 
   SourceLoc getStartLoc() const { return ClassLoc; }
   SourceRange getSourceRange() const {
-    return SourceRange(ClassLoc, getBraces().End);
+    if (getBraces().isValid())
+      return SourceRange(ClassLoc, getBraces().End);
+    return SourceRange();
   }
 
   /// Determine whether the member area of this class's metadata (which consists
@@ -4216,7 +4222,9 @@ class ProtocolDecl final : public NominalTypeDecl {
 
   SourceLoc getStartLoc() const { return ProtocolLoc; }
   SourceRange getSourceRange() const {
-    return SourceRange(ProtocolLoc, getBraces().End);
+    if (getBraces().isValid())
+      return SourceRange(ProtocolLoc, getBraces().End);
+    return SourceRange();
   }
 
   /// True if this protocol can only be conformed to by class types.

lib/AST/Decl.cpp

@@ -690,6 +690,7 @@ static_assert(sizeof(checkSourceLocType(&ID##Decl::getLoc)) == 2, \
     return getLocFromSource();
   switch(File->getKind()) {
   case FileUnitKind::Source:
+  case FileUnitKind::ClangModule:
     return getLocFromSource();
   case FileUnitKind::SerializedAST: {
     if (!SerializedOK)
@@ -698,7 +699,6 @@ static_assert(sizeof(checkSourceLocType(&ID##Decl::getLoc)) == 2, \
   }
   case FileUnitKind::Builtin:
   case FileUnitKind::Synthesized:
-  case FileUnitKind::ClangModule:
   case FileUnitKind::DWARFModule:
     return SourceLoc();
   }

lib/Basic/SourceLoc.cpp

@@ -165,9 +165,14 @@ SourceManager::getVirtualFile(SourceLoc Loc) const {
   if (CachedVFile.first == p)
     return CachedVFile.second;
 
-  // Returns the first element that is >p.
-  auto VFileIt = VirtualFiles.upper_bound(p);
-  if (VFileIt != VirtualFiles.end() && VFileIt->second.Range.contains(Loc)) {
+  // Returns the first element that is >=p.
+  auto VFileIt = VirtualFiles.lower_bound(p);
+  if (VFileIt == VirtualFiles.end())
+    return nullptr;
+
+  // Consider the end as within the range in order to handle EOF locations
+  const CharSourceRange &Range = VFileIt->second.Range;
+  if (Loc == Range.getEnd() || VFileIt->second.Range.contains(Loc)) {
     CachedVFile = { p, &VFileIt->second };
     return CachedVFile.second;
   }

lib/ClangImporter/ClangDiagnosticConsumer.cpp

@@ -123,14 +123,15 @@ void ClangDiagnosticConsumer::HandleDiagnostic(
 
   const ASTContext &ctx = ImporterImpl.SwiftContext;
   ClangSourceBufferImporter &bufferImporter =
-      ImporterImpl.getBufferImporterForDiagnostics();
+      ImporterImpl.getBufferImporter();
 
   if (clangDiag.getID() == clang::diag::err_module_not_built &&
       CurrentImport && clangDiag.getArgStdStr(0) == CurrentImport->getName()) {
     SourceLoc loc = DiagLoc;
     if (clangDiag.getLocation().isValid()) {
-      loc = bufferImporter.resolveSourceLocation(clangDiag.getSourceManager(),
-                                                 clangDiag.getLocation());
+      loc = bufferImporter.importSourceLoc(clangDiag.getSourceManager(),
+                                           clangDiag.getLocation(),
+                                           /*forDiagnostics=*/true);
     }
 
     ctx.Diags.diagnose(loc, diag::clang_cannot_build_module,
@@ -172,8 +173,9 @@ void ClangDiagnosticConsumer::HandleDiagnostic(
 
     SourceLoc noteLoc;
     if (clangNoteLoc.isValid())
-      noteLoc = bufferImporter.resolveSourceLocation(clangNoteLoc.getManager(),
-                                                     clangNoteLoc);
+      noteLoc = bufferImporter.importSourceLoc(clangNoteLoc.getManager(),
+                                               clangNoteLoc,
+                                               /*forDiagnostics=*/true);
     ctx.Diags.diagnose(noteLoc, diagKind, message);
   };
 

lib/ClangImporter/ClangImporter.cpp

@@ -1106,7 +1106,7 @@ ClangImporter::create(ASTContext &ctx,
       std::make_shared<SwiftNameLookupExtension>(
           importer->Impl.BridgingHeaderLookupTable,
           importer->Impl.LookupTables, importer->Impl.SwiftContext,
-          importer->Impl.getBufferImporterForDiagnostics(),
+          importer->Impl.getBufferImporter(),
           importer->Impl.platformAvailability));
 
   // Create a compiler instance.
@@ -1446,7 +1446,7 @@ bool ClangImporter::Implementation::importHeader(
 
   // Finalize the lookup table, which may fail.
   finalizeLookupTable(*BridgingHeaderLookupTable, getNameImporter(),
-                      getBufferImporterForDiagnostics());
+                      getBufferImporter());
 
   // FIXME: What do we do if there was already an error?
   if (!hadError && clangDiags.hasErrorOccurred()) {
@@ -2205,7 +2205,7 @@ ClangImporter::Implementation::Implementation(
       IsReadingBridgingPCH(false),
       CurrentVersion(ImportNameVersion::fromOptions(ctx.LangOpts)),
       BridgingHeaderLookupTable(new SwiftLookupTable(nullptr)),
-      BuffersForDiagnostics(ctx.SourceMgr),
+      BufferImporter(ctx.SourceMgr),
       platformAvailability(ctx.LangOpts), nameImporter(),
       DisableSourceImport(ctx.ClangImporterOpts.DisableSourceImport),
       DWARFImporter(dwarfImporterDelegate) {}
@@ -2263,14 +2263,14 @@ ClangImporter::Implementation::exportSourceLoc(SourceLoc loc) {
 
 SourceLoc
 ClangImporter::Implementation::importSourceLoc(clang::SourceLocation loc) {
-  // FIXME: Implement!
-  return SourceLoc();
+  return BufferImporter.importSourceLoc(
+      getClangASTContext().getSourceManager(), loc);
 }
 
 SourceRange
-ClangImporter::Implementation::importSourceRange(clang::SourceRange loc) {
-  // FIXME: Implement!
-  return SourceRange();
+ClangImporter::Implementation::importSourceRange(clang::SourceRange range) {
+  return SourceRange(importSourceLoc(range.getBegin()),
+                     importSourceLoc(range.getEnd()));
 }
 
 #pragma mark Importing names

lib/ClangImporter/ClangSourceBufferImporter.cpp

@@ -13,6 +13,7 @@
 #include "ClangSourceBufferImporter.h"
 #include "swift/Basic/SourceManager.h"
 #include "clang/Basic/SourceManager.h"
+#include "clang/Basic/SourceManagerInternals.h"
 #include "llvm/Support/MemoryBuffer.h"
 
 using namespace swift;
@@ -29,69 +30,105 @@ static SourceLoc findEndOfLine(SourceManager &SM, SourceLoc loc,
   return loc.getAdvancedLoc(newlineOffset);
 }
 
-SourceLoc ClangSourceBufferImporter::resolveSourceLocation(
-    const clang::SourceManager &clangSrcMgr,
-    clang::SourceLocation clangLoc) {
-  SourceLoc loc;
+SourceLoc ClangSourceBufferImporter::importSourceLoc(
+    const clang::SourceManager &clangSourceManager,
+    clang::SourceLocation clangLoc, bool forDiagnostics) {
+  // Map the location to one in an actual file (as opposed to within a macro)
+  clangLoc = clangSourceManager.getFileLoc(clangLoc);
 
-  clangLoc = clangSrcMgr.getFileLoc(clangLoc);
-  auto decomposedLoc = clangSrcMgr.getDecomposedLoc(clangLoc);
-  if (decomposedLoc.first.isInvalid())
-    return loc;
+  clang::FileID clangFileID;
+  unsigned offset;
+  std::tie(clangFileID, offset) =
+      clangSourceManager.getDecomposedLoc(clangLoc);
+  if (clangFileID.isInvalid())
+    return SourceLoc();
 
-  auto clangFileID = decomposedLoc.first;
-  auto buffer = clangSrcMgr.getBufferOrFake(clangFileID);
-  unsigned mirrorID;
+  llvm::Optional<llvm::MemoryBufferRef> buffer =
+      clangSourceManager.getBufferOrNone(clangFileID);
+  if (!buffer)
+    return SourceLoc();
 
-  auto mirrorIter = mirroredBuffers.find(buffer.getBufferStart());
+  // Grab the already mapped buffer or add a new one (without copying) if it
+  // hasn't been added before
+  unsigned mirrorID;
+  auto mirrorIter = mirroredBuffers.find(buffer->getBufferStart());
+  bool allDirectives = false;
   if (mirrorIter != mirroredBuffers.end()) {
-    mirrorID = mirrorIter->second;
+    mirrorID = mirrorIter->second.FileID;
+    allDirectives = mirrorIter->second.Complete;
   } else {
     std::unique_ptr<llvm::MemoryBuffer> mirrorBuffer{
-      llvm::MemoryBuffer::getMemBuffer(buffer.getBuffer(),
-                                       buffer.getBufferIdentifier(),
+      llvm::MemoryBuffer::getMemBuffer(buffer->getBuffer(),
+                                       buffer->getBufferIdentifier(),
                                        /*RequiresNullTerminator=*/true)
     };
     mirrorID = swiftSourceManager.addNewSourceBuffer(std::move(mirrorBuffer));
-    mirroredBuffers[buffer.getBufferStart()] = mirrorID;
+    mirroredBuffers[buffer->getBufferStart()] = { mirrorID, !forDiagnostics };
+
+    // Make sure to keep a reference to the underlying manager
+    if (clangSourceManagers.insert(&clangSourceManager).second) {
+      clangSourceManagerRefs.emplace_back(&clangSourceManager);
+    }
   }
-  loc = swiftSourceManager.getLocForOffset(mirrorID, decomposedLoc.second);
 
-  auto presumedLoc = clangSrcMgr.getPresumedLoc(clangLoc);
-  if (!presumedLoc.getFilename())
+  SourceLoc loc = swiftSourceManager.getLocForOffset(mirrorID, offset);
+  if (allDirectives)
+    return swiftSourceManager.getLocForOffset(mirrorID, offset);
+
+  // Now add any line directives, if any. Note that the line table will be
+  // incomplete if a module is broken and diagnostics are output, but otherwise
+  // we'll be reading from the serialized module and hence have the full line
+  // table.
+  //
+  // Always add the enclosing line directive if importing a source location for
+  // a diagnostic, but otherwise assume we have the full line table and add
+  // all directives at once to avoid adding a virtual file for every location.
+
+  bool Invalid = false;
+  const clang::SrcMgr::SLocEntry &Entry =
+      clangSourceManager.getSLocEntry(clangFileID, &Invalid);
+  if (Invalid || !Entry.isFile())
     return loc;
-  if (presumedLoc.getLine() == 0)
-    return SourceLoc();
 
-  unsigned bufferLineNumber =
-    clangSrcMgr.getLineNumber(decomposedLoc.first, decomposedLoc.second);
-
-  StringRef presumedFile = presumedLoc.getFilename();
-  SourceLoc startOfLine = loc.getAdvancedLoc(-presumedLoc.getColumn() + 1);
-
-  // FIXME: Virtual files can't actually model the EOF position correctly, so
-  // if this virtual file would start at EOF, just hope the physical location
-  // will do.
-  if (startOfLine != swiftSourceManager.getRangeForBuffer(mirrorID).getEnd()) {
-    bool isNewVirtualFile = swiftSourceManager.openVirtualFile(
-        startOfLine, presumedFile, presumedLoc.getLine() - bufferLineNumber);
-    if (isNewVirtualFile) {
-      SourceLoc endOfLine = findEndOfLine(swiftSourceManager, loc, mirrorID);
-      swiftSourceManager.closeVirtualFile(endOfLine);
-    }
+  // TODO: The LineTable is obstensibly part of the "internals", should we be
+  //       relying on it here?
+  if (!clangSourceManager.hasLineTable() ||
+      !Entry.getFile().hasLineDirectives())
+    return loc;
+
+  clang::LineTableInfo &lineTable =
+      *clangSourceManager.getCurrentLineTable();
+
+  ArrayRef<clang::LineEntry> lineEntries;
+  if (!forDiagnostics) {
+    mirroredBuffers[buffer->getBufferStart()].Complete = true;
+
+    lineEntries = lineTable.entries(clangFileID);
+    if (lineEntries.empty())
+      return loc;
+  } else {
+    const clang::LineEntry *entry =
+        lineTable.FindNearestLineEntry(clangFileID, offset);
+    if (!entry)
+      return loc;
+    lineEntries = llvm::makeArrayRef(*entry);
   }
 
-  using SourceManagerRef = llvm::IntrusiveRefCntPtr<const clang::SourceManager>;
-  auto iter = std::lower_bound(sourceManagersWithDiagnostics.begin(),
-                               sourceManagersWithDiagnostics.end(),
-                               &clangSrcMgr,
-                               [](const SourceManagerRef &inArray,
-                                  const clang::SourceManager *toInsert) {
-    return std::less<const clang::SourceManager *>()(inArray.get(), toInsert);
-  });
-  if (iter == sourceManagersWithDiagnostics.end() ||
-      iter->get() != &clangSrcMgr) {
-    sourceManagersWithDiagnostics.insert(iter, &clangSrcMgr);
+  for (const clang::LineEntry &entry : lineEntries) {
+    SourceLoc fileLoc =
+        swiftSourceManager.getLocForOffset(mirrorID, entry.FileOffset);
+    const SourceManager::VirtualFile *file = swiftSourceManager.getVirtualFile(fileLoc);
+    if (file) {
+      if (file->Range.getStart() == fileLoc)
+        continue;
+      swiftSourceManager.closeVirtualFile(fileLoc);
+    }
+
+    StringRef fileName = lineTable.getFilename(entry.FilenameID);
+    unsigned line =
+        swiftSourceManager.getLineAndColumnInBuffer(fileLoc).first;
+    swiftSourceManager.openVirtualFile(fileLoc, fileName,
+                                       entry.LineNo - line - 1);
   }
 
   return loc;

lib/ClangImporter/ClangSourceBufferImporter.h

@@ -18,6 +18,7 @@
 #include "clang/Basic/SourceLocation.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 
 namespace llvm {
@@ -33,29 +34,34 @@ class SourceManager;
 
 namespace importer {
 
-/// A helper class used to keep alive the Clang source managers where
-/// diagnostics have been reported.
+/// A helper class used to import clang::SourceLocation to swift::SourceLoc,
+/// for eg. attaching to declarations or for diagnostics.
 ///
-/// This is a bit of a hack, but LLVM's source manager (and by extension
-/// Swift's) does not support buffers going away, so if we want to report
-/// diagnostics in them we have to do it this way.
+/// Rather than copying the underlying MemoryBuffer, instead keep any
+/// clang::SourceManager with imported locations alive and reference its
+/// underlying buffer directly.
 class ClangSourceBufferImporter {
-  // This is not using SmallPtrSet or similar because we need the
-  // IntrusiveRefCntPtr to stay a ref-counting pointer.
+  struct MirrorEntry {
+    unsigned FileID;
+    bool Complete;
+  };
+
+  llvm::SmallPtrSet<const clang::SourceManager *, 4> clangSourceManagers;
   SmallVector<llvm::IntrusiveRefCntPtr<const clang::SourceManager>, 4>
-    sourceManagersWithDiagnostics;
-  llvm::DenseMap<const char *, unsigned> mirroredBuffers;
+    clangSourceManagerRefs;
+  llvm::DenseMap<const char *, MirrorEntry> mirroredBuffers;
   SourceManager &swiftSourceManager;
 
 public:
   explicit ClangSourceBufferImporter(SourceManager &sourceMgr)
-    : swiftSourceManager(sourceMgr) {}
+      : swiftSourceManager(sourceMgr) {}
 
   /// Returns a Swift source location that points into a Clang buffer.
   ///
   /// This will keep the Clang buffer alive as long as this object.
-  SourceLoc resolveSourceLocation(const clang::SourceManager &clangSrcMgr,
-                                  clang::SourceLocation clangLoc);
+  SourceLoc importSourceLoc(
+      const clang::SourceManager &clangSourceManager,
+      clang::SourceLocation clangLoc, bool forDiagnostics=false);
 };
 
 } // end namespace importer