Re-implement ASTScope's child-insertion assertions

ASTScope assertions that a new child node inserted into its parent list
is after all other child nodes in the source, which included some bespoke
code for dealing with macro expansions. Replace this bespoke code with
uses of the newer SourceManager operations that check ordering of
source locations in a manner that respects macro expansions.

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

@@ -255,6 +255,20 @@ class SourceManager {
   /// 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.
   ///

lib/AST/ASTScopeLookup.cpp

@@ -120,20 +120,16 @@ ASTScopeImpl::findChildContaining(ModuleDecl *parentModule,
   auto *const *child = llvm::lower_bound(
       getChildren(), loc,
       [&](const ASTScopeImpl *scope, SourceLoc loc) {
-        auto rangeOfScope = scope->getSourceRangeOfThisASTNode();
-        auto endOfScope =
-            Lexer::getLocForEndOfToken(sourceMgr, rangeOfScope.End);
-
+        auto rangeOfScope = scope->getCharSourceRangeOfScope(sourceMgr);
         loc = translateLocForReplacedRange(sourceMgr, rangeOfScope.Start, loc);
-        return sourceMgr.isAtOrBefore(endOfScope, loc);
+        return sourceMgr.isAtOrBefore(rangeOfScope.End, loc);
       });
 
   if (child != getChildren().end()) {
-    auto rangeOfScope = (*child)->getSourceRangeOfThisASTNode();
-    auto endOfScope = Lexer::getLocForEndOfToken(sourceMgr, rangeOfScope.End);
+    auto rangeOfScope = (*child)->getCharSourceRangeOfScope(sourceMgr);
     loc = translateLocForReplacedRange(sourceMgr, rangeOfScope.Start, loc);
     if (sourceMgr.isAtOrBefore(rangeOfScope.Start, loc) &&
-        sourceMgr.isBefore(loc, endOfScope))
+        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/Basic/SourceLoc.cpp

@@ -817,3 +817,12 @@ bool SourceManager::isAtOrBefore(SourceLoc first, SourceLoc second) const {
 bool SourceManager::containsTokenLoc(SourceRange range, SourceLoc loc) const {
   return isAtOrBefore(range.Start, loc) && isAtOrBefore(loc, range.End);
 }
+
+bool SourceManager::containsLoc(SourceRange range, SourceLoc loc) const {
+  return isAtOrBefore(range.Start, loc) && isBefore(loc, range.End);
+}
+
+bool SourceManager::encloses(SourceRange enclosing, SourceRange inner) const {
+  return containsLoc(enclosing, inner.Start) &&
+      isAtOrBefore(inner.End, enclosing.End);
+}