[ASTScope] Re-enable checkSourceRangeBeforeAddingChild

Update the logic to correctly handle replacement ranges, and re-enable
the assertion. For now, carve out an exception for attributes on
extensions, I'll try and tackle those in a follow-up.

Author: hamishknight

include/swift/AST/ASTScope.h

@@ -239,6 +239,9 @@ class ASTScopeImpl : public ASTAllocated<ASTScopeImpl> {
   NullablePtr<Stmt> getStmtIfAny() const;
   NullablePtr<Expr> getExprIfAny() const;
 
+  /// Whether this scope is for a decl attribute.
+  bool isDeclAttribute() const;
+
 #pragma mark - debugging and printing
 
 public:

lib/AST/ASTScope.cpp

@@ -297,6 +297,20 @@ NullablePtr<Expr> ASTScopeImpl::getExprIfAny() const {
   }
 }
 
+bool ASTScopeImpl::isDeclAttribute() const {
+  switch (getKind()) {
+#define DECL_ATTRIBUTE_SCOPE_NODE(Name) \
+    case ScopeKind::Name: return true;
+#define SCOPE_NODE(Name)
+#include "swift/AST/ASTScopeNodes.def"
+
+#define DECL_ATTRIBUTE_SCOPE_NODE(Name)
+#define SCOPE_NODE(Name) case ScopeKind::Name:
+#include "swift/AST/ASTScopeNodes.def"
+      return false;
+  }
+}
+
 SourceManager &ASTScopeImpl::getSourceManager() const {
   return getASTContext().SourceMgr;
 }

lib/AST/ASTScopeCreation.cpp

@@ -714,7 +714,7 @@ void ASTScopeImpl::addChild(ASTScopeImpl *child, ASTContext &ctx) {
   child->parentAndWasExpanded.setPointer(this);
 
 #ifndef NDEBUG
-  // checkSourceRangeBeforeAddingChild(child, ctx);
+  checkSourceRangeBeforeAddingChild(child, ctx);
 #endif
 
   // If this is the first time we've added children, notify the ASTContext

lib/AST/ASTScopeSourceRange.cpp

@@ -39,16 +39,15 @@ 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 attributes on extensions, currently they exist outside of the
+  // extension's source range due to the way we've setup the scope for
+  // extension binding.
+  // FIXME: We ought to fix the source range for extension scopes.
+  if (isa<ExtensionScope>(this) && child->isDeclAttribute())
+    return;
+
   // Ignore debugger bindings - they're a special mix of user code and implicit
   // wrapper code that is too difficult to check for consistency.
   if (auto d = getDeclIfAny().getPtrOrNull())
@@ -60,14 +59,14 @@ void ASTScopeImpl::checkSourceRangeBeforeAddingChild(ASTScopeImpl *child,
 
   auto range = getCharSourceRangeOfScope(sourceMgr);
 
-  std::function<bool(SourceRange)> containedInParent;
-  containedInParent = [&](SourceRange childCharRange) {
+  auto containedInParent = [&](SourceRange childCharRange) {
     // HACK: For code completion. Handle replaced range.
+    // Note that the replaced SourceRanges here are already disguised
+    // CharSourceRanges, we don't need to adjust them. We use `rangeContains`
+    // since we're only interested in comparing within a single buffer.
     for (const auto &pair : sourceMgr.getReplacedRanges()) {
-      auto originalRange = getCharSourceRange(sourceMgr, pair.first);
-      auto newRange = getCharSourceRange(sourceMgr, pair.second);
-      if (sourceMgr.encloses(range, originalRange) &&
-          sourceMgr.encloses(newRange, childCharRange))
+      if (sourceMgr.rangeContains(range, pair.first) &&
+          sourceMgr.rangeContains(pair.second, childCharRange))
         return true;
     }