ASTScope: Directly calculate end location of GuardStmtScope and PatternEntryDeclScope

Today, the reported source range of a GuardStmtScope is just that of
the statement itself, ending after the 'else' block. Similarly, a
PatternEntryDeclScope ends immediately after the initializer
expression.

Since these scopes introduce names until the end of the parent
BraceStmt (meaning they change the insertion point), we were
relying on the addition of child scopes to extend the source
range.

While we still need the hack for extending source ranges to deal
with string interpolation, at least in the other cases we can
get rid of it.

Note that this fixes a bug where a GuardStmtScope would end
before an inactive '#if' block if the inactive '#if' block was
the last element of a BraceStmt.

Author: slavapestov

include/swift/AST/ASTScope.h

@@ -1040,6 +1040,7 @@ class AbstractPatternEntryScope : public ASTScopeImpl {
 
 class PatternEntryDeclScope final : public AbstractPatternEntryScope {
   const bool isLocalBinding;
+  Optional<SourceLoc> endLoc;
 
 public:
   PatternEntryDeclScope(PatternBindingDecl *pbDecl, unsigned entryIndex,
@@ -1400,7 +1401,8 @@ class WhileStmtScope final : public LabeledConditionalStmtScope {
 class GuardStmtScope final : public LabeledConditionalStmtScope {
 public:
   GuardStmt *const stmt;
-  GuardStmtScope(GuardStmt *e) : stmt(e) {}
+  SourceLoc endLoc;
+  GuardStmtScope(GuardStmt *e, SourceLoc endLoc) : stmt(e), endLoc(endLoc) {}
   virtual ~GuardStmtScope() {}
 
 protected:
@@ -1413,6 +1415,8 @@ class GuardStmtScope final : public LabeledConditionalStmtScope {
 public:
   std::string getClassName() const override;
   LabeledConditionalStmt *getLabeledConditionalStmt() const override;
+  SourceRange
+  getSourceRangeOfThisASTNode(bool omitAssertions = false) const override;
 };
 
 /// A scope after a guard statement that follows lookups into the conditions
@@ -1427,9 +1431,11 @@ class LookupParentDiversionScope final : public ASTScopeImpl {
 public:
   ASTScopeImpl *const lookupParent;
   const SourceLoc startLoc;
+  const SourceLoc endLoc;
 
-  LookupParentDiversionScope(ASTScopeImpl *lookupParent, SourceLoc startLoc)
-      : lookupParent(lookupParent), startLoc(startLoc) {}
+  LookupParentDiversionScope(ASTScopeImpl *lookupParent,
+                             SourceLoc startLoc, SourceLoc endLoc)
+      : lookupParent(lookupParent), startLoc(startLoc), endLoc(endLoc) {}
 
   SourceRange
   getSourceRangeOfThisASTNode(bool omitAssertions = false) const override;

lib/AST/ASTScopeCreation.cpp

@@ -213,14 +213,17 @@ class ScopeCreator final {
 
   /// Given an array of ASTNodes or Decl pointers, add them
   /// Return the resultant insertionPoint
+  ///
+  /// \param endLoc The end location for any "scopes until the end" that
+  /// we introduce here, such as PatternEntryDeclScope and GuardStmtScope
   ASTScopeImpl *
   addSiblingsToScopeTree(ASTScopeImpl *const insertionPoint,
-                         ASTScopeImpl *const organicInsertionPoint,
-                         ArrayRef<ASTNode> nodesOrDeclsToAdd) {
+                         ArrayRef<ASTNode> nodesOrDeclsToAdd,
+                         Optional<SourceLoc> endLoc) {
     auto *ip = insertionPoint;
     for (auto nd : nodesOrDeclsToAdd) {
       auto *const newIP =
-          addToScopeTreeAndReturnInsertionPoint(nd, ip).getPtrOr(ip);
+          addToScopeTreeAndReturnInsertionPoint(nd, ip, endLoc).getPtrOr(ip);
       ip = newIP;
     }
     return ip;
@@ -229,12 +232,16 @@ class ScopeCreator final {
 public:
   /// For each of searching, call this unless the insertion point is needed
   void addToScopeTree(ASTNode n, ASTScopeImpl *parent) {
-    (void)addToScopeTreeAndReturnInsertionPoint(n, parent);
+    (void)addToScopeTreeAndReturnInsertionPoint(n, parent, None);
   }
   /// Return new insertion point if the scope was not a duplicate
   /// For ease of searching, don't call unless insertion point is needed
+  ///
+  /// \param endLoc The end location for any "scopes until the end" that
+  /// we introduce here, such as PatternEntryDeclScope and GuardStmtScope
   NullablePtr<ASTScopeImpl>
-  addToScopeTreeAndReturnInsertionPoint(ASTNode, ASTScopeImpl *parent);
+  addToScopeTreeAndReturnInsertionPoint(ASTNode, ASTScopeImpl *parent,
+                                        Optional<SourceLoc> endLoc);
 
   bool isWorthTryingToCreateScopeFor(ASTNode n) const {
     if (!n)
@@ -419,6 +426,18 @@ class ScopeCreator final {
   void addChildrenForKnownAttributes(ValueDecl *decl,
                                      ASTScopeImpl *parent);
 
+  /// Add PatternEntryDeclScopes for each pattern binding entry.
+  ///
+  /// Returns the new insertion point.
+  ///
+  /// \param endLoc Must be valid iff the pattern binding is in a local
+  /// scope, in which case this is the last source location where the
+  /// pattern bindings are going to be visible.
+  NullablePtr<ASTScopeImpl>
+  addPatternBindingToScopeTree(PatternBindingDecl *patternBinding,
+                               ASTScopeImpl *parent,
+                               Optional<SourceLoc> endLoc);
+
   /// Remove VarDecls because we'll find them when we expand the
   /// PatternBindingDecls. Remove EnunCases
   /// because they overlap EnumElements and AST includes the elements in the
@@ -541,7 +560,10 @@ class NodeAdder
     : public ASTVisitor<NodeAdder, NullablePtr<ASTScopeImpl>,
                         NullablePtr<ASTScopeImpl>, NullablePtr<ASTScopeImpl>,
                         void, void, void, ASTScopeImpl *, ScopeCreator &> {
+  Optional<SourceLoc> endLoc;
+
 public:
+  explicit NodeAdder(Optional<SourceLoc> endLoc) : endLoc(endLoc) {}
 
 #pragma mark ASTNodes that do not create scopes
 
@@ -633,7 +655,9 @@ class NodeAdder
   // the deferred nodes.
   NullablePtr<ASTScopeImpl> visitGuardStmt(GuardStmt *e, ASTScopeImpl *p,
                                            ScopeCreator &scopeCreator) {
-    return scopeCreator.ifUniqueConstructExpandAndInsert<GuardStmtScope>(p, e);
+    ASTScopeAssert(endLoc.hasValue(), "GuardStmt outside of a BraceStmt?");
+    return scopeCreator.ifUniqueConstructExpandAndInsert<GuardStmtScope>(
+      p, e, *endLoc);
   }
   NullablePtr<ASTScopeImpl> visitTopLevelCodeDecl(TopLevelCodeDecl *d,
                                                   ASTScopeImpl *p,
@@ -694,28 +718,8 @@ class NodeAdder
   visitPatternBindingDecl(PatternBindingDecl *patternBinding,
                           ASTScopeImpl *parentScope,
                           ScopeCreator &scopeCreator) {
-    if (auto *var = patternBinding->getSingleVar())
-      scopeCreator.addChildrenForKnownAttributes(var, parentScope);
-
-    auto *insertionPoint = parentScope;
-    for (auto i : range(patternBinding->getNumPatternEntries())) {
-      bool isLocalBinding = false;
-      if (auto *varDecl = patternBinding->getAnchoringVarDecl(i)) {
-        isLocalBinding = varDecl->getDeclContext()->isLocalContext();
-      }
-
-      insertionPoint =
-          scopeCreator
-              .ifUniqueConstructExpandAndInsert<PatternEntryDeclScope>(
-                  insertionPoint, patternBinding, i, isLocalBinding)
-              .getPtrOr(insertionPoint);
-
-      ASTScopeAssert(isLocalBinding || insertionPoint == parentScope,
-                     "Bindings at the top-level or members of types should "
-                     "not change the insertion point");
-    }
-
-    return insertionPoint;
+    return scopeCreator.addPatternBindingToScopeTree(
+        patternBinding, parentScope, endLoc);
   }
 
   NullablePtr<ASTScopeImpl> visitEnumElementDecl(EnumElementDecl *eed,
@@ -769,15 +773,18 @@ class NodeAdder
 // NodeAdder
 NullablePtr<ASTScopeImpl>
 ScopeCreator::addToScopeTreeAndReturnInsertionPoint(ASTNode n,
-                                                    ASTScopeImpl *parent) {
+                                                    ASTScopeImpl *parent,
+                                                    Optional<SourceLoc> endLoc) {
   if (!isWorthTryingToCreateScopeFor(n))
     return parent;
+
+  NodeAdder adder(endLoc);
   if (auto *p = n.dyn_cast<Decl *>())
-    return NodeAdder().visit(p, parent, *this);
+    return adder.visit(p, parent, *this);
   if (auto *p = n.dyn_cast<Expr *>())
-    return NodeAdder().visit(p, parent, *this);
+    return adder.visit(p, parent, *this);
   auto *p = n.get<Stmt *>();
-  return NodeAdder().visit(p, parent, *this);
+  return adder.visit(p, parent, *this);
 }
 
 void ScopeCreator::addChildrenForAllLocalizableAccessorsInSourceOrder(
@@ -827,6 +834,41 @@ void ScopeCreator::addChildrenForKnownAttributes(ValueDecl *decl,
   }
 }
 
+NullablePtr<ASTScopeImpl>
+ScopeCreator::addPatternBindingToScopeTree(PatternBindingDecl *patternBinding,
+                                           ASTScopeImpl *parentScope,
+                                           Optional<SourceLoc> endLoc) {
+  if (auto *var = patternBinding->getSingleVar())
+    addChildrenForKnownAttributes(var, parentScope);
+
+  auto *insertionPoint = parentScope;
+  for (auto i : range(patternBinding->getNumPatternEntries())) {
+    bool isLocalBinding = false;
+    if (auto *varDecl = patternBinding->getAnchoringVarDecl(i)) {
+      isLocalBinding = varDecl->getDeclContext()->isLocalContext();
+    }
+
+    Optional<SourceLoc> endLocForBinding = None;
+    if (isLocalBinding) {
+      endLocForBinding = endLoc;
+      ASTScopeAssert(endLoc.hasValue() && endLoc->isValid(),
+                     "PatternBindingDecl in local context outside of BraceStmt?");
+    }
+
+    insertionPoint =
+        ifUniqueConstructExpandAndInsert<PatternEntryDeclScope>(
+            insertionPoint, patternBinding, i,
+            isLocalBinding, endLocForBinding)
+        .getPtrOr(insertionPoint);
+
+    ASTScopeAssert(isLocalBinding || insertionPoint == parentScope,
+                   "Bindings at the top-level or members of types should "
+                   "not change the insertion point");
+  }
+
+  return insertionPoint;
+}
+
 #pragma mark creation helpers
 
 void ASTScopeImpl::addChild(ASTScopeImpl *child, ASTContext &ctx) {
@@ -946,9 +988,10 @@ ASTSourceFileScope::expandAScopeThatCreatesANewInsertionPoint(
   // Assume that decls are only added at the end, in source order
   std::vector<ASTNode> newNodes(decls.begin(), decls.end());
   insertionPoint =
-      scopeCreator.addSiblingsToScopeTree(insertionPoint, this,
+      scopeCreator.addSiblingsToScopeTree(insertionPoint,
                                           scopeCreator.sortBySourceRange(
-                                            scopeCreator.cull(newNodes)));
+                                            scopeCreator.cull(newNodes)),
+                                          None);
   // Too slow to perform all the time:
   //    ASTScopeAssert(scopeCreator->containsAllDeclContextsFromAST(),
   //           "ASTScope tree missed some DeclContexts or made some up");
@@ -1048,7 +1091,7 @@ GuardStmtScope::expandAScopeThatCreatesANewInsertionPoint(ScopeCreator &
   auto *const lookupParentDiversionScope =
       scopeCreator
           .constructExpandAndInsertUncheckable<LookupParentDiversionScope>(
-              this, conditionLookupParent, stmt->getEndLoc());
+              this, conditionLookupParent, stmt->getEndLoc(), endLoc);
   return {lookupParentDiversionScope,
           "Succeeding code must be in scope of guard variables"};
 }
@@ -1066,10 +1109,11 @@ BraceStmtScope::expandAScopeThatCreatesANewInsertionPoint(
   // TODO: remove the sort after fixing parser to create brace statement
   // elements in source order
   auto *insertionPoint =
-      scopeCreator.addSiblingsToScopeTree(this, this,
+      scopeCreator.addSiblingsToScopeTree(this,
                                           scopeCreator.sortBySourceRange(
                                             scopeCreator.cull(
-                                              stmt->getElements())));
+                                              stmt->getElements())),
+                                          stmt->getEndLoc());
   if (auto *s = scopeCreator.getASTContext().Stats)
     ++s->getFrontendCounters().NumBraceStmtASTScopeExpansions;
   return {
@@ -1083,7 +1127,7 @@ TopLevelCodeScope::expandAScopeThatCreatesANewInsertionPoint(ScopeCreator &
 
   if (auto *body =
           scopeCreator
-              .addToScopeTreeAndReturnInsertionPoint(decl->getBody(), this)
+              .addToScopeTreeAndReturnInsertionPoint(decl->getBody(), this, None)
               .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 "
@@ -1414,7 +1458,7 @@ void GenericTypeOrExtensionScope::expandBody(ScopeCreator &) {}
 void IterableTypeScope::expandBody(ScopeCreator &scopeCreator) {
   auto nodes = asNodeVector(getIterableDeclContext().get()->getMembers());
   nodes = scopeCreator.sortBySourceRange(scopeCreator.cull(nodes));
-  scopeCreator.addSiblingsToScopeTree(this, this, nodes);
+  scopeCreator.addSiblingsToScopeTree(this, nodes, None);
   if (auto *s = scopeCreator.getASTContext().Stats)
     ++s->getFrontendCounters().NumIterableTypeBodyASTScopeExpansions;
 }

lib/AST/ASTScopeSourceRange.cpp

@@ -248,7 +248,13 @@ SourceRange DefaultArgumentInitializerScope::getSourceRangeOfThisASTNode(
 
 SourceRange PatternEntryDeclScope::getSourceRangeOfThisASTNode(
     const bool omitAssertions) const {
-  return getPatternEntry().getSourceRange();
+  SourceRange range = getPatternEntry().getSourceRange();
+  if (endLoc.hasValue()) {
+    ASTScopeAssert(endLoc->isValid(),
+                   "BraceStmt ends before pattern binding entry?");
+    range.End = *endLoc;
+  }
+  return range;
 }
 
 SourceRange PatternEntryInitializerScope::getSourceRangeOfThisASTNode(
@@ -445,9 +451,14 @@ SourceRange AttachedPropertyWrapperScope::getSourceRangeOfThisASTNode(
   return sourceRangeWhenCreated;
 }
 
+SourceRange GuardStmtScope::getSourceRangeOfThisASTNode(
+    const bool omitAssertions) const {
+  return SourceRange(getStmt()->getStartLoc(), endLoc);
+}
+
 SourceRange LookupParentDiversionScope::getSourceRangeOfThisASTNode(
     const bool omitAssertions) const {
-  return SourceRange(startLoc);
+  return SourceRange(startLoc, endLoc);
 }
 
 #pragma mark source range caching

test/decl/var/usage.swift

@@ -242,6 +242,15 @@ func testBuildConfigs() {
 #endif
 }
 
+// same as above, but with a guard statement
+func testGuardWithPoundIf(x: Int?) {
+  guard let x = x else { return }
+
+#if false
+  _ = x
+#endif
+}
+
 // <rdar://problem/21091625> Bogus 'never mutated' warning when protocol variable is mutated only by mutating method
 protocol Fooable {
   mutating func mutFoo()