ASTScope: Fix TopLevelCodeScope source range

Top-level code can contain guard statements which introduce
bindings until the end of the parent scope, so plumb that
through.

Author: slavapestov

include/swift/AST/ASTScope.h

@@ -1189,8 +1189,10 @@ class ClosureParametersScope final : public ASTScopeImpl {
 class TopLevelCodeScope final : public ASTScopeImpl {
 public:
   TopLevelCodeDecl *const decl;
+  SourceLoc endLoc;
 
-  TopLevelCodeScope(TopLevelCodeDecl *e) : decl(e) {}
+  TopLevelCodeScope(TopLevelCodeDecl *e, SourceLoc endLoc)
+      : decl(e), endLoc(endLoc) {}
   virtual ~TopLevelCodeScope() {}
 
 protected:
@@ -1643,13 +1645,21 @@ class BraceStmtScope final : public AbstractStmtScope {
   /// definition.
   SmallVector<VarDecl *, 2> localVars;
 
+  /// The end location for bindings introduced in this scope. This can
+  /// extend past the actual end of the BraceStmt in top-level code,
+  /// where every TopLevelCodeDecl introduces a new scope through the
+  /// end of the buffer.
+  SourceLoc endLoc;
+
 public:
   BraceStmtScope(BraceStmt *e,
                  SmallVector<ValueDecl *, 2> localFuncsAndTypes,
-                 SmallVector<VarDecl *, 2> localVars)
+                 SmallVector<VarDecl *, 2> localVars,
+                 SourceLoc endLoc)
       : stmt(e),
         localFuncsAndTypes(localFuncsAndTypes),
-        localVars(localVars) {}
+        localVars(localVars),
+        endLoc(endLoc) {}
   virtual ~BraceStmtScope() {}
 
 protected:

lib/AST/ASTScopeCreation.cpp

@@ -662,8 +662,9 @@ class NodeAdder
   NullablePtr<ASTScopeImpl> visitTopLevelCodeDecl(TopLevelCodeDecl *d,
                                                   ASTScopeImpl *p,
                                                   ScopeCreator &scopeCreator) {
-    return scopeCreator.ifUniqueConstructExpandAndInsert<TopLevelCodeScope>(p,
-                                                                            d);
+    ASTScopeAssert(endLoc.hasValue(), "TopLevelCodeDecl in wrong place?");
+    return scopeCreator.ifUniqueConstructExpandAndInsert<TopLevelCodeScope>(
+        p, d, *endLoc);
   }
 
 #pragma mark special-case creation
@@ -708,9 +709,14 @@ class NodeAdder
       }
     }
 
+    SourceLoc endLocForBraceStmt = bs->getEndLoc();
+    if (endLoc.hasValue())
+      endLocForBraceStmt = *endLoc;
+
     auto maybeBraceScope =
         scopeCreator.ifUniqueConstructExpandAndInsert<BraceStmtScope>(
-            p, bs, std::move(localFuncsAndTypes), std::move(localVars));
+            p, bs, std::move(localFuncsAndTypes), std::move(localVars),
+            endLocForBraceStmt);
     if (auto *s = scopeCreator.getASTContext().Stats)
       ++s->getFrontendCounters().NumBraceStmtASTScopes;
 
@@ -989,12 +995,17 @@ ASTSourceFileScope::expandAScopeThatCreatesANewInsertionPoint(
   ASTScopeAssert(SF, "Must already have a SourceFile.");
   ArrayRef<Decl *> decls = SF->getTopLevelDecls();
   // Assume that decls are only added at the end, in source order
+  Optional<SourceLoc> endLoc = None;
+  if (!decls.empty())
+    endLoc = decls.back()->getEndLoc();
+
   std::vector<ASTNode> newNodes(decls.begin(), decls.end());
   insertionPoint =
       scopeCreator.addSiblingsToScopeTree(insertionPoint,
                                           scopeCreator.sortBySourceRange(
                                             scopeCreator.cull(newNodes)),
-                                          None);
+                                          endLoc);
+
   // Too slow to perform all the time:
   //    ASTScopeAssert(scopeCreator->containsAllDeclContextsFromAST(),
   //           "ASTScope tree missed some DeclContexts or made some up");
@@ -1123,7 +1134,7 @@ BraceStmtScope::expandAScopeThatCreatesANewInsertionPoint(
                                           scopeCreator.sortBySourceRange(
                                             scopeCreator.cull(
                                               stmt->getElements())),
-                                          stmt->getEndLoc());
+                                          endLoc);
   if (auto *s = scopeCreator.getASTContext().Stats)
     ++s->getFrontendCounters().NumBraceStmtASTScopeExpansions;
   return {
@@ -1137,7 +1148,7 @@ TopLevelCodeScope::expandAScopeThatCreatesANewInsertionPoint(ScopeCreator &
 
   if (auto *body =
           scopeCreator
-              .addToScopeTreeAndReturnInsertionPoint(decl->getBody(), this, None)
+              .addToScopeTreeAndReturnInsertionPoint(decl->getBody(), this, endLoc)
               .getPtrOrNull())
     return {body, "So next top level code scope and put its decls in its body "
                   "under a guard statement scope (etc) from the last top level "

lib/AST/ASTScopeSourceRange.cpp

@@ -221,7 +221,7 @@ SourceRange FunctionBodyScope::getSourceRangeOfThisASTNode(
 
 SourceRange TopLevelCodeScope::getSourceRangeOfThisASTNode(
     const bool omitAssertions) const {
-  return decl->getSourceRange();
+  return SourceRange(decl->getStartLoc(), endLoc);
 }
 
 SourceRange SubscriptDeclScope::getSourceRangeOfThisASTNode(
@@ -401,9 +401,9 @@ BraceStmtScope::getSourceRangeOfThisASTNode(const bool omitAssertions) const {
   // 'in' keyword, when present.
   if (auto closure = parentClosureIfAny()) {
     if (closure.get()->getInLoc().isValid())
-      return SourceRange(closure.get()->getInLoc(), stmt->getEndLoc());
+      return SourceRange(closure.get()->getInLoc(), endLoc);
   }
-  return stmt->getSourceRange();
+  return SourceRange(stmt->getStartLoc(), endLoc);
 }
 
 SourceRange ConditionalClauseInitializerScope::getSourceRangeOfThisASTNode(