Implement ordering of source locations that accounts for macro expansion

SourceManager's `isBeforeInBuffer` in buffer takes two source locations
that are assumed to be within the same source buffer and determines
whether the first precedes the second. However, due to macro
expansion, there can be source locations in different source buffers
that are nonetheless conceptually part of the same source file. For
such cases, `isBeforeInBuffer` will silently produce wrong results.

Extend `SourceManager` with a new `isBefore` operation that determines
whether one source location precedes another in the same conceptual
source file, even if they are in different source buffers that (e.g.)
represent macro expansions within that source file. Also introduce
`isAtOrBefore` to cover the case where the source locations might be
equivalent, which was previously handled via a separate equality check
on the source locations that won't work with macro expansions.

Add caching for the list of accessors of a buffer, so we aren't doing
a full walk up the source-buffer chain ever time we do a comparison of
source locations. LCA is still linear-time, but this eliminates
extraneous hash table lookups along the way.

Over time, we should both move more clients over to these new variants
and introduce more assertions into the old (in-buffer) versions to
catch improper uses of them.

Author: DougGregor

include/swift/Basic/SourceManager.h

@@ -16,6 +16,7 @@
 #include "swift/Basic/FileSystem.h"
 #include "swift/Basic/SourceLoc.h"
 #include "clang/Basic/FileManager.h"
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/None.h"
 #include "llvm/ADT/Optional.h"
@@ -69,6 +70,10 @@ class GeneratedSourceInfo {
   /// The name of the source file on disk that was created to hold the
   /// contents of this file for external clients.
   StringRef onDiskBufferCopyFileName = StringRef();
+
+  /// 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>();
 };
 
 /// This class manages and owns source buffers.
@@ -205,6 +210,15 @@ class SourceManager {
   llvm::Optional<GeneratedSourceInfo>
   getGeneratedSourceInfo(unsigned bufferID) const;
 
+  /// Retrieve the list of ancestors of the given source buffer, starting with
+  /// the root buffer and proceding to the given buffer ID at the end.
+  ///
+  /// The scratch parameter will be used to avoid allocation in the case where
+  /// the given buffer ID is the top-level buffer, in which case bufferID will
+  /// be written into scratch at the returned array will contain that one
+  /// element.
+  ArrayRef<unsigned> getAncestors(unsigned bufferID, unsigned &scratch) const;
+
   /// Record the starting source location of a regex literal.
   void recordRegexLiteralStartLoc(SourceLoc loc) {
     RegexLiteralStartLocs.insert(loc);
@@ -216,13 +230,35 @@ class SourceManager {
     return RegexLiteralStartLocs.contains(loc);
   }
 
-  /// Returns true if \c LHS is before \c RHS in the source buffer.
+  /// Returns true if \c first is before \c second in the same source file,
+  /// accounting for the possibility that first and second are in different
+  /// source buffers that are conceptually within the same source file,
+  /// due to macro expansion.
+  bool isBefore(SourceLoc first, SourceLoc second) const;
+
+  /// Returns true if \c first is at or before \c second in the same source
+  /// file, accounting for the possibility that first and second are in
+  /// different source buffers that are conceptually within the same source
+  /// file, due to macro expansion.
+  bool isAtOrBefore(SourceLoc first, SourceLoc second) const;
+
+  /// Returns true if \c LHS is before \c RHS in the same source buffer.
   bool isBeforeInBuffer(SourceLoc LHS, SourceLoc RHS) const {
     return LHS.Value.getPointer() < RHS.Value.getPointer();
   }
 
-  /// Returns true if range \c R contains the location \c Loc.  The location
-  /// \c Loc should point at the beginning of the token.
+  /// Returns true if \c range contains the location \c loc.  The location
+  /// \c loc should point at the beginning of the token.
+  ///
+  /// This function accounts for the possibility that the source locations
+  /// provided might come from different source buffers that are conceptually
+  /// part of the same source file, for example due to macro expansion.
+  bool containsTokenLoc(SourceRange range, SourceLoc loc) const;
+
+  /// Returns true if range \c R contains the location \c Loc, where all
+  /// locations are known to be in the same source buffer.
+  ///
+  /// The location \c Loc should point at the beginning of the token.
   bool rangeContainsTokenLoc(SourceRange R, SourceLoc Loc) const {
     return Loc == R.Start || Loc == R.End ||
            (isBeforeInBuffer(R.Start, Loc) && isBeforeInBuffer(Loc, R.End));

lib/AST/ASTScopeLookup.cpp

@@ -109,63 +109,6 @@ static SourceLoc translateLocForReplacedRange(SourceManager &sourceMgr,
   return loc;
 }
 
-/// Populate the ancestors list for this buffer, with the root source buffer
-/// at the beginning and the given source buffer at the end.
-static void populateAncestors(
-    SourceManager &sourceMgr, unsigned bufferID,
-    SmallVectorImpl<unsigned> &ancestors) {
-  if (auto info = sourceMgr.getGeneratedSourceInfo(bufferID)) {
-    auto ancestorLoc = info->originalSourceRange.getStart();
-    if (ancestorLoc.isValid()) {
-      auto ancestorBufferID = sourceMgr.findBufferContainingLoc(ancestorLoc);
-      populateAncestors(sourceMgr, ancestorBufferID, ancestors);
-    }
-  }
-
-  ancestors.push_back(bufferID);
-}
-
-/// Determine whether the first source location precedes the second, accounting
-/// for macro expansions.
-static bool isBeforeInSource(
-    SourceManager &sourceMgr, SourceLoc firstLoc, SourceLoc secondLoc,
-    bool allowEqual) {
-  // If the two locations are in the same source buffer, compare their pointers.
-  unsigned firstBufferID = sourceMgr.findBufferContainingLoc(firstLoc);
-  unsigned secondBufferID = sourceMgr.findBufferContainingLoc(secondLoc);
-  if (firstBufferID == secondBufferID) {
-    return sourceMgr.isBeforeInBuffer(firstLoc, secondLoc) ||
-        (allowEqual && firstLoc == secondLoc);
-  }
-
-  // If the two locations are in different source buffers, we need to compute
-  // the least common ancestor.
-  SmallVector<unsigned, 4> firstAncestors;
-  populateAncestors(sourceMgr, firstBufferID, firstAncestors);
-  SmallVector<unsigned, 4> secondAncestors;
-  populateAncestors(sourceMgr, secondBufferID, secondAncestors);
-
-  // Find the first mismatch between the two ancestor lists; this is the
-  // point of divergence.
-  auto [firstMismatch, secondMismatch] = std::mismatch(
-      firstAncestors.begin(), firstAncestors.end(),
-      secondAncestors.begin(), secondAncestors.end());
-  assert(firstMismatch != firstAncestors.begin() &&
-         secondMismatch != secondAncestors.begin() &&
-         "Ancestors don't have the same root source file");
-
-  SourceLoc firstLocInLCA = firstMismatch == firstAncestors.end()
-      ? firstLoc
-      : sourceMgr.getGeneratedSourceInfo(*firstMismatch)
-          ->originalSourceRange.getStart();
-  SourceLoc secondLocInLCA = secondMismatch == secondAncestors.end()
-      ? secondLoc
-      : sourceMgr.getGeneratedSourceInfo(*secondMismatch)
-          ->originalSourceRange.getStart();
-  return sourceMgr.isBeforeInBuffer(firstLocInLCA, secondLocInLCA) ||
-    (allowEqual && firstLocInLCA == secondLocInLCA);
-}
-
 NullablePtr<ASTScopeImpl>
 ASTScopeImpl::findChildContaining(ModuleDecl *parentModule,
                                   SourceLoc loc,
@@ -182,17 +125,15 @@ ASTScopeImpl::findChildContaining(ModuleDecl *parentModule,
             Lexer::getLocForEndOfToken(sourceMgr, rangeOfScope.End);
 
         loc = translateLocForReplacedRange(sourceMgr, rangeOfScope.Start, loc);
-        return isBeforeInSource(
-            sourceMgr, endOfScope, loc, /*allowEqual=*/true);
+        return sourceMgr.isAtOrBefore(endOfScope, loc);
       });
 
   if (child != getChildren().end()) {
     auto rangeOfScope = (*child)->getSourceRangeOfThisASTNode();
     auto endOfScope = Lexer::getLocForEndOfToken(sourceMgr, rangeOfScope.End);
     loc = translateLocForReplacedRange(sourceMgr, rangeOfScope.Start, loc);
-    if (isBeforeInSource(sourceMgr, rangeOfScope.Start, loc,
-                         /*allowEqual=*/true) &&
-        isBeforeInSource(sourceMgr, loc, endOfScope, /*allowEqual=*/false))
+    if (sourceMgr.isAtOrBefore(rangeOfScope.Start, loc) &&
+        sourceMgr.isBefore(loc, endOfScope))
       return *child;
   }
 

lib/Basic/SourceLoc.cpp

@@ -216,6 +216,10 @@ SourceManager::getIDForBufferIdentifier(StringRef BufIdentifier) const {
 SourceManager::~SourceManager() {
   for (auto &generated : GeneratedSourceInfos) {
     free((void*)generated.second.onDiskBufferCopyFileName.data());
+
+    if (generated.second.ancestors.size() > 0) {
+      delete [] generated.second.ancestors.data();
+    }
   }
 }
 
@@ -717,3 +721,99 @@ SourceManager::getLocForForeignLoc(SourceLoc otherLoc,
 
   return SourceLoc();
 }
+
+/// Populate the ancestors list for this buffer, with the root source buffer
+/// at the beginning and the given source buffer at the end.
+static void populateAncestors(
+    const SourceManager &sourceMgr, unsigned bufferID,
+    SmallVectorImpl<unsigned> &ancestors) {
+  if (auto info = sourceMgr.getGeneratedSourceInfo(bufferID)) {
+    auto ancestorLoc = info->originalSourceRange.getStart();
+    if (ancestorLoc.isValid()) {
+      auto ancestorBufferID = sourceMgr.findBufferContainingLoc(ancestorLoc);
+      populateAncestors(sourceMgr, ancestorBufferID, ancestors);
+    }
+  }
+
+  ancestors.push_back(bufferID);
+}
+
+ArrayRef<unsigned> SourceManager::getAncestors(
+    unsigned bufferID, unsigned &scratch
+) const {
+  // If there is no generated source information for this buffer, then this is
+  // the only buffer here. Avoid memory allocation by using the scratch space
+  // we were given.
+  auto knownInfo = GeneratedSourceInfos.find(bufferID);
+  if (knownInfo == GeneratedSourceInfos.end()) {
+    scratch = bufferID;
+    return ArrayRef<unsigned>(&scratch, 1);
+  }
+
+  // If we already have the ancestors cached, use them.
+  if (!knownInfo->second.ancestors.empty())
+    return knownInfo->second.ancestors;
+
+  // Compute all of the ancestors. We only do this once for a given buffer.
+  SmallVector<unsigned, 4> ancestors;
+  populateAncestors(*this, bufferID, ancestors);
+
+  // Cache the ancestors in the generated source info record.
+  unsigned *ancestorsPtr = new unsigned [ancestors.size()];
+  std::copy(ancestors.begin(), ancestors.end(), ancestorsPtr);
+  knownInfo->second.ancestors = llvm::makeArrayRef(ancestorsPtr, ancestors.size());
+  return knownInfo->second.ancestors;
+}
+
+
+/// Determine whether the first source location precedes the second, accounting
+/// for macro expansions.
+static bool isBeforeInSource(
+    const SourceManager &sourceMgr, SourceLoc firstLoc, SourceLoc secondLoc,
+    bool allowEqual) {
+  // If the two locations are in the same source buffer, compare their pointers.
+  unsigned firstBufferID = sourceMgr.findBufferContainingLoc(firstLoc);
+  unsigned secondBufferID = sourceMgr.findBufferContainingLoc(secondLoc);
+  if (firstBufferID == secondBufferID) {
+    return sourceMgr.isBeforeInBuffer(firstLoc, secondLoc) ||
+        (allowEqual && firstLoc == secondLoc);
+  }
+
+  // If the two locations are in different source buffers, we need to compute
+  // the least common ancestor.
+  unsigned firstScratch, secondScratch;
+  auto firstAncestors = sourceMgr.getAncestors(firstBufferID, firstScratch);
+  auto secondAncestors = sourceMgr.getAncestors(secondBufferID, secondScratch);
+
+  // Find the first mismatch between the two ancestor lists; this is the
+  // point of divergence.
+  auto [firstMismatch, secondMismatch] = std::mismatch(
+      firstAncestors.begin(), firstAncestors.end(),
+      secondAncestors.begin(), secondAncestors.end());
+  assert(firstMismatch != firstAncestors.begin() &&
+         secondMismatch != secondAncestors.begin() &&
+         "Ancestors don't have the same root source file");
+
+  SourceLoc firstLocInLCA = firstMismatch == firstAncestors.end()
+      ? firstLoc
+      : sourceMgr.getGeneratedSourceInfo(*firstMismatch)
+          ->originalSourceRange.getStart();
+  SourceLoc secondLocInLCA = secondMismatch == secondAncestors.end()
+      ? secondLoc
+      : sourceMgr.getGeneratedSourceInfo(*secondMismatch)
+          ->originalSourceRange.getStart();
+  return sourceMgr.isBeforeInBuffer(firstLocInLCA, secondLocInLCA) ||
+    (allowEqual && firstLocInLCA == secondLocInLCA);
+}
+
+bool SourceManager::isBefore(SourceLoc first, SourceLoc second) const {
+  return isBeforeInSource(*this, first, second, /*allowEqual=*/false);
+}
+
+bool SourceManager::isAtOrBefore(SourceLoc first, SourceLoc second) const {
+  return isBeforeInSource(*this, first, second, /*allowEqual=*/true);
+}
+
+bool SourceManager::containsTokenLoc(SourceRange range, SourceLoc loc) const {
+  return isAtOrBefore(range.Start, loc) && isAtOrBefore(loc, range.End);
+}