Merge pull request #70126 from DougGregor/source-loc-ordering-with-macros

Implement ordering of source locations that accounts for macro expansion

Author: DougGregor

include/swift/AST/ASTScope.h

@@ -155,7 +155,7 @@ class ASTScopeImpl : public ASTAllocated<ASTScopeImpl> {
   /// Child scopes, sorted by source range.
   Children storedChildren;
 
-  mutable llvm::Optional<CharSourceRange> cachedCharSourceRange;
+  mutable llvm::Optional<SourceRange> cachedCharSourceRange;
 
 #pragma mark - constructor / destructor
 public:
@@ -194,8 +194,17 @@ class ASTScopeImpl : public ASTAllocated<ASTScopeImpl> {
 #pragma mark - source ranges
 
 public:
-  CharSourceRange getCharSourceRangeOfScope(SourceManager &SM,
-                                            bool omitAssertions = false) const;
+  /// Retrieve the source range of the given scope, where the end location
+  /// is adjusted to refer to the end of the token.
+  ///
+  /// Since the adjustment to the end of the token requires lexing, this
+  /// routine also caches the result.
+  ///
+  /// Note that the start and end locations might be in different source
+  /// buffers, so we represent the result as SourceRange rather than
+  /// CharSourceRange.
+  SourceRange getCharSourceRangeOfScope(SourceManager &SM,
+                                        bool omitAssertions = false) const;
   bool isCharSourceRangeCached() const;
 
   /// Returns source range of this node alone, without factoring in any

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,49 @@ 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 \c range contains the location \c loc.
+  ///
+  /// 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 containsLoc(SourceRange range, SourceLoc loc) const;
+
+  /// Returns true if \c enclosing contains the whole range \c inner.
+  ///
+  /// 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 encloses(SourceRange enclosing, SourceRange inner) 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

@@ -47,7 +47,26 @@ void ASTScopeImpl::unqualifiedLookup(
 
 const ASTScopeImpl *ASTScopeImpl::findStartingScopeForLookup(
     SourceFile *sourceFile, const SourceLoc loc) {
-  auto *const fileScope = sourceFile->getScope().impl;
+  auto *fileScope = sourceFile->getScope().impl;
+
+  // Workaround for bad locations; just return the file scope.
+  if (loc.isInvalid())
+    return fileScope;
+
+  // Some callers get the actual source file wrong. Look for the actual
+  // source file containing this location.
+  auto actualSF =
+      sourceFile->getParentModule()->getSourceFileContainingLocation(loc);
+
+  // If there is no source file containing this location, just return the
+  // scope we have.
+  if (!actualSF)
+    return fileScope;
+
+  // Grab the new file scope.
+  if (actualSF != sourceFile)
+    fileScope = actualSF->getScope().impl;
+
   const auto *innermost = fileScope->findInnermostEnclosingScope(
       sourceFile->getParentModule(), loc, nullptr);
   ASTScopeAssert(innermost->getWasExpanded(),
@@ -79,11 +98,11 @@ ASTScopeImpl *ASTScopeImpl::findInnermostEnclosingScopeImpl(
 /// If the \p loc is in a new buffer but \p range is not, consider the location
 /// is at the start of replaced range. Otherwise, returns \p loc as is.
 static SourceLoc translateLocForReplacedRange(SourceManager &sourceMgr,
-                                              CharSourceRange range,
+                                              SourceLoc rangeStart,
                                               SourceLoc loc) {
   for (const auto &pair : sourceMgr.getReplacedRanges()) {
     if (sourceMgr.rangeContainsTokenLoc(pair.second, loc) &&
-        !sourceMgr.rangeContainsTokenLoc(pair.second, range.getStart())) {
+        !sourceMgr.rangeContainsTokenLoc(pair.second, rangeStart)) {
       return pair.first.Start;
     }
   }
@@ -94,89 +113,23 @@ NullablePtr<ASTScopeImpl>
 ASTScopeImpl::findChildContaining(ModuleDecl *parentModule,
                                   SourceLoc loc,
                                   SourceManager &sourceMgr) const {
-  auto *locSourceFile = parentModule->getSourceFileContainingLocation(loc);
+  if (loc.isInvalid())
+    return nullptr;
 
   // Use binary search to find the child that contains this location.
   auto *const *child = llvm::lower_bound(
       getChildren(), loc,
       [&](const ASTScopeImpl *scope, SourceLoc loc) {
         auto rangeOfScope = scope->getCharSourceRangeOfScope(sourceMgr);
-        ASTScopeAssert(!sourceMgr.isBeforeInBuffer(rangeOfScope.getEnd(),
-                                                   rangeOfScope.getStart()),
-                       "Source range is backwards");
-
-        // If the scope source range and the loc are in two different source
-        // files, one or both of them are in a macro expansion buffer.
-
-        // Note that `scope->getSourceFile()` returns the root of the source tree,
-        // not the source file containing the location of the ASTScope.
-        auto scopeStart = scope->getSourceRangeOfThisASTNode().Start;
-        auto *scopeSourceFile = parentModule->getSourceFileContainingLocation(scopeStart);
-
-        if (scopeSourceFile != locSourceFile) {
-          // To compare a source location that is possibly inside a macro expansion
-          // with a source range that is also possibly in a macro expansion (not
-          // necessarily the same one as before) we need to find the LCA in the
-          // source file tree of macro expansions, and compare the original source
-          // ranges within that common ancestor. We can't walk all the way up to the
-          // source file containing the parent scope we're searching the children of,
-          // because two independent (possibly nested) macro expansions can have the
-          // same original source range in that file; freestanding and peer macros
-          // mean that we can have arbitrarily nested macro expansions that all add
-          // declarations to the same scope, that all originate from a single macro
-          // invocation in the original source file.
-
-          // A map from enclosing source files to original source ranges of the macro
-          // expansions within that file, recording the chain of macro expansions for
-          // the given scope.
-          llvm::SmallDenseMap<const SourceFile *, CharSourceRange>
-              scopeExpansions;
-
-          // Walk up the chain of macro expansion buffers for the scope, recording the
-          // original source range of the macro expansion along the way using generated
-          // source info.
-          auto *scopeExpansion = scopeSourceFile;
-          scopeExpansions[scopeExpansion] =
-              Lexer::getCharSourceRangeFromSourceRange(
-                sourceMgr, scope->getSourceRangeOfThisASTNode());
-          while (auto *ancestor = scopeExpansion->getEnclosingSourceFile()) {
-            auto generatedInfo =
-                sourceMgr.getGeneratedSourceInfo(*scopeExpansion->getBufferID());
-            scopeExpansions[ancestor] = generatedInfo->originalSourceRange;
-            scopeExpansion = ancestor;
-          }
-
-          // Walk up the chain of macro expansion buffers for the source loc we're
-          // searching for to find the LCA using `scopeExpansions`.
-          auto *potentialLCA = locSourceFile;
-          auto expansionLoc = loc;
-          while (potentialLCA) {
-            auto scopeExpansion = scopeExpansions.find(potentialLCA);
-            if (scopeExpansion != scopeExpansions.end()) {
-              // Take the original expansion range within the LCA of the loc and
-              // the scope to compare.
-              rangeOfScope = scopeExpansion->second;
-              loc = expansionLoc;
-              break;
-            }
-
-            auto generatedInfo =
-                sourceMgr.getGeneratedSourceInfo(*potentialLCA->getBufferID());
-            if (generatedInfo)
-              expansionLoc = generatedInfo->originalSourceRange.getStart();
-            potentialLCA = potentialLCA->getEnclosingSourceFile();
-          }
-        }
-
-        loc = translateLocForReplacedRange(sourceMgr, rangeOfScope, loc);
-        return (rangeOfScope.getEnd() == loc ||
-                sourceMgr.isBeforeInBuffer(rangeOfScope.getEnd(), loc));
+        loc = translateLocForReplacedRange(sourceMgr, rangeOfScope.Start, loc);
+        return sourceMgr.isAtOrBefore(rangeOfScope.End, loc);
       });
 
   if (child != getChildren().end()) {
     auto rangeOfScope = (*child)->getCharSourceRangeOfScope(sourceMgr);
-    loc = translateLocForReplacedRange(sourceMgr, rangeOfScope, loc);
-    if (rangeOfScope.contains(loc))
+    loc = translateLocForReplacedRange(sourceMgr, rangeOfScope.Start, loc);
+    if (sourceMgr.isAtOrBefore(rangeOfScope.Start, loc) &&
+        sourceMgr.isBefore(loc, rangeOfScope.End))
       return *child;
   }
 

lib/AST/ASTScopeSourceRange.cpp

@@ -38,6 +38,14 @@ using namespace ast_scope;
 
 static SourceLoc getLocAfterExtendedNominal(const ExtensionDecl *);
 
+/// Retrieve the character-based source range for the given source range.
+static SourceRange getCharSourceRange(
+    SourceManager &sourceMgr, SourceRange range
+){
+  range.End = Lexer::getLocForEndOfToken(sourceMgr, range.End);
+  return range;
+}
+
 void ASTScopeImpl::checkSourceRangeBeforeAddingChild(ASTScopeImpl *child,
                                                      const ASTContext &ctx) const {
   // Ignore debugger bindings - they're a special mix of user code and implicit
@@ -51,35 +59,18 @@ void ASTScopeImpl::checkSourceRangeBeforeAddingChild(ASTScopeImpl *child,
 
   auto range = getCharSourceRangeOfScope(sourceMgr);
 
-  std::function<bool(CharSourceRange)> containedInParent;
-  containedInParent = [&](CharSourceRange childCharRange) {
+  std::function<bool(SourceRange)> containedInParent;
+  containedInParent = [&](SourceRange childCharRange) {
     // HACK: For code completion. Handle replaced range.
     for (const auto &pair : sourceMgr.getReplacedRanges()) {
-      auto originalRange =
-          Lexer::getCharSourceRangeFromSourceRange(sourceMgr, pair.first);
-      auto newRange =
-          Lexer::getCharSourceRangeFromSourceRange(sourceMgr, pair.second);
-      if (range.contains(originalRange) &&
-          newRange.contains(childCharRange))
+      auto originalRange = getCharSourceRange(sourceMgr, pair.first);
+      auto newRange = getCharSourceRange(sourceMgr, pair.second);
+      if (sourceMgr.encloses(range, originalRange) &&
+          sourceMgr.encloses(newRange, childCharRange))
         return true;
     }
 
-    if (range.contains(childCharRange))
-      return true;
-
-    // If this is from a macro expansion, look at the where the expansion
-    // occurred.
-    auto childBufferID =
-      sourceMgr.findBufferContainingLoc(childCharRange.getStart());
-    auto generatedInfo = sourceMgr.getGeneratedSourceInfo(childBufferID);
-    if (!generatedInfo)
-      return false;
-
-    CharSourceRange expansionRange = generatedInfo->originalSourceRange;
-    if (expansionRange.isInvalid())
-      return false;
-
-    return containedInParent(expansionRange);
+    return sourceMgr.encloses(range, childCharRange);
   };
 
   auto childCharRange = child->getCharSourceRangeOfScope(sourceMgr);
@@ -95,11 +86,9 @@ void ASTScopeImpl::checkSourceRangeBeforeAddingChild(ASTScopeImpl *child,
   if (!storedChildren.empty()) {
     auto previousChild = storedChildren.back();
     auto endOfPreviousChild = previousChild->getCharSourceRangeOfScope(
-        sourceMgr).getEnd();
+        sourceMgr).End;
 
-    if (childCharRange.getStart() != endOfPreviousChild &&
-        !sourceMgr.isBeforeInBuffer(endOfPreviousChild,
-                                    childCharRange.getStart())) {
+    if (!sourceMgr.isAtOrBefore(endOfPreviousChild, childCharRange.Start)) {
       auto &out = verificationError() << "child overlaps previous child:\n";
       child->print(out);
       out << "\n***Previous child\n";
@@ -384,18 +373,19 @@ SourceRange GuardStmtBodyScope::getSourceRangeOfThisASTNode(
 
 #pragma mark source range caching
 
-CharSourceRange
+SourceRange
 ASTScopeImpl::getCharSourceRangeOfScope(SourceManager &SM,
                                         bool omitAssertions) const {
   if (!isCharSourceRangeCached()) {
     auto range = getSourceRangeOfThisASTNode(omitAssertions);
     ASTScopeAssert(range.isValid(), "scope has invalid source range");
-    ASTScopeAssert(SM.isBeforeInBuffer(range.Start, range.End) ||
+    ASTScopeAssert(SM.isBefore(range.Start, range.End) ||
                    range.Start == range.End,
                    "scope source range ends before start");
 
-    cachedCharSourceRange =
-      Lexer::getCharSourceRangeFromSourceRange(SM, range);
+    range.End = Lexer::getLocForEndOfToken(SM, range.End);
+
+    cachedCharSourceRange = range;
   }
 
   return *cachedCharSourceRange;

lib/AST/DeclContext.cpp

@@ -1026,16 +1026,7 @@ void IterableDeclContext::addMemberSilently(Decl *member, Decl *hint,
     assert(nextStart.isValid() &&
            "Only implicit decls can have invalid source location");
 
-    if (getASTContext().SourceMgr.isBeforeInBuffer(prevEnd, nextStart))
-      return;
-
-    // Synthesized member macros can add new members in a macro expansion buffer.
-    SourceFile *memberSourceFile = member->getLoc()
-        ? member->getModuleContext()
-                ->getSourceFileContainingLocation(member->getLoc())
-        : member->getInnermostDeclContext()->getParentSourceFile();
-    if (memberSourceFile->getFulfilledMacroRole() == MacroRole::Member ||
-        memberSourceFile->getFulfilledMacroRole() == MacroRole::Peer)
+    if (getASTContext().SourceMgr.isAtOrBefore(prevEnd, nextStart))
       return;
 
     llvm::errs() << "Source ranges out of order in addMember():\n";

lib/AST/NameLookup.cpp

@@ -1665,6 +1665,15 @@ bool
 namelookup::isInMacroArgument(SourceFile *sourceFile, SourceLoc loc) {
   bool inMacroArgument = false;
 
+  // Make sure that the source location is actually within the given source
+  // file.
+  if (sourceFile && loc.isValid()) {
+    sourceFile =
+        sourceFile->getParentModule()->getSourceFileContainingLocation(loc);
+    if (!sourceFile)
+      return false;
+  }
+
   ASTScope::lookupEnclosingMacroScope(
       sourceFile, loc,
       [&](auto potentialMacro) -> bool {