AST: UnqualifiedLookup only finds forward references to outer local bindings when IncludeOuterResults is set

The old behavior was that ASTScope would introduce all VarDecls
defined in a BraceStmt at the beginning of the BraceStmt.

I recently enabled the use of PatternEntryDeclScopes, which
introduce the binding at its actual source location instead of
at the beginning of the parent statement.

This patch now makes use of the new information by having
UnqualifiedLookupFlags::IncludeOuterResults toggle between
the two behaviors. When searching for outer results, we also
consider all VarDecls in a BraceStmt, not just those in scope.

This is implemented by giving AbstractASTScopeDeclConsumer a
new entry point, consumePossiblyNotInScope(). When looking up
into a BraceStmt, all VarDecls are passed in to this entry
point.

The default implementation does nothing, which means that
ASTScope::lookupSingleLocalDecl() now respects source locations
when searching for bindings, just like parse-time lookup.

However, Sema's preCheckExpression() pass, which sets
Flags::IgnoreOuterResults, will continue to find
forward-referenced VarDecls, just as it did with the old
context-based DeclContext lookup.

Author: slavapestov

include/swift/AST/ASTScope.h

@@ -1658,10 +1658,22 @@ class CaseStmtBodyScope final : public ASTScopeImpl {
 };
 
 class BraceStmtScope final : public AbstractStmtScope {
+  BraceStmt *const stmt;
+
+  /// Declarations which are in scope from the beginning of the statement.
+  SmallVector<ValueDecl *, 2> localFuncsAndTypes;
+
+  /// Declarations that are normally in scope only after their
+  /// definition.
+  SmallVector<VarDecl *, 2> localVars;
 
 public:
-  BraceStmt *const stmt;
-  BraceStmtScope(BraceStmt *e) : stmt(e) {}
+  BraceStmtScope(BraceStmt *e,
+                 SmallVector<ValueDecl *, 2> localFuncsAndTypes,
+                 SmallVector<VarDecl *, 2> localVars)
+      : stmt(e),
+        localFuncsAndTypes(localFuncsAndTypes),
+        localVars(localVars) {}
   virtual ~BraceStmtScope() {}
 
 protected:

include/swift/AST/NameLookup.h

@@ -613,6 +613,14 @@ class AbstractASTScopeDeclConsumer {
   lookInMembers(DeclContext *const scopeDC,
                 NominalTypeDecl *const nominal) = 0;
 
+  /// Called for local VarDecls that might not yet be in scope.
+  ///
+  /// Note that the set of VarDecls visited here are going to be a
+  /// superset of those visited in consume().
+  virtual bool consumePossiblyNotInScope(ArrayRef<VarDecl *> values) {
+    return false;
+  }
+
   /// Called right before looking at the parent scope of a BraceStmt.
   ///
   /// \return true if the lookup should be stopped at this point.

lib/AST/ASTScopeCreation.cpp

@@ -667,10 +667,30 @@ class NodeAdder
 
   NullablePtr<ASTScopeImpl> visitBraceStmt(BraceStmt *bs, ASTScopeImpl *p,
                                            ScopeCreator &scopeCreator) {
+    SmallVector<ValueDecl *, 2> localFuncsAndTypes;
+    SmallVector<VarDecl *, 2> localVars;
+
+    // All types and functions are visible anywhere within a brace statement
+    // scope. When ordering matters (i.e. var decl) we will have split the brace
+    // statement into nested scopes.
+    for (auto braceElement : bs->getElements()) {
+      if (auto localBinding = braceElement.dyn_cast<Decl *>()) {
+        if (auto *vd = dyn_cast<ValueDecl>(localBinding)) {
+          if (isa<FuncDecl>(vd)  || isa<TypeDecl>(vd)) {
+            localFuncsAndTypes.push_back(vd);
+          } else if (auto *var = dyn_cast<VarDecl>(localBinding)) {
+            localVars.push_back(var);
+          }
+        }
+      }
+    }
+
     auto maybeBraceScope =
-        scopeCreator.ifUniqueConstructExpandAndInsert<BraceStmtScope>(p, bs);
+        scopeCreator.ifUniqueConstructExpandAndInsert<BraceStmtScope>(
+            p, bs, std::move(localFuncsAndTypes), std::move(localVars));
     if (auto *s = scopeCreator.getASTContext().Stats)
       ++s->getFrontendCounters().NumBraceStmtASTScopes;
+
     return maybeBraceScope.getPtrOr(p);
   }
 

lib/AST/ASTScopeLookup.cpp

@@ -371,20 +371,10 @@ bool DifferentiableAttributeScope::lookupLocalsOrMembers(
 }
 
 bool BraceStmtScope::lookupLocalsOrMembers(DeclConsumer consumer) const {
-  // All types and functions are visible anywhere within a brace statement
-  // scope. When ordering matters (i.e. var decl) we will have split the brace
-  // statement into nested scopes.
-  //
-  // Don't stop at the first one, there may be local funcs with same base name
-  // and want them all.
-  SmallVector<ValueDecl *, 32> localBindings;
-  for (auto braceElement : stmt->getElements()) {
-    if (auto localBinding = braceElement.dyn_cast<Decl *>()) {
-      if (auto *vd = dyn_cast<ValueDecl>(localBinding))
-        localBindings.push_back(vd);
-    }
-  }
-  if (consumer.consume(localBindings, DeclVisibilityKind::LocalVariable))
+  if (consumer.consume(localFuncsAndTypes, DeclVisibilityKind::LocalVariable))
+    return true;
+
+  if (consumer.consumePossiblyNotInScope(localVars))
     return true;
 
   if (consumer.finishLookupInBraceStmt(stmt))

lib/AST/UnqualifiedLookup.cpp

@@ -133,7 +133,12 @@ namespace {
     /// Can lookup stop searching for results, assuming hasn't looked for outer
     /// results yet?
     bool isFirstResultEnough() const;
-    
+
+    /// Do we want precise scoping of VarDecls? If IncludeOuterResults is on,
+    /// this is true, which allows us to resolve forward references to
+    /// local VarDecls from inside local function and closure bodies.
+    bool hasPreciseScopingOfVarDecls() const;
+
     /// Every time lookup finishes searching a scope, call me
     /// to record the dividing line between results from first fruitful scope and
     /// the result.
@@ -209,6 +214,8 @@ class ASTScopeDeclConsumerForUnqualifiedLookup
   bool consume(ArrayRef<ValueDecl *> values, DeclVisibilityKind vis,
                NullablePtr<DeclContext> baseDC = nullptr) override;
 
+  bool consumePossiblyNotInScope(ArrayRef<VarDecl *> vars) override;
+
   /// returns true if finished
   bool lookInMembers(DeclContext *const scopeDC,
                      NominalTypeDecl *const nominal) override;
@@ -462,6 +469,10 @@ bool UnqualifiedLookupFactory::isFirstResultEnough() const {
   return !Results.empty() && !options.contains(Flags::IncludeOuterResults);
 }
 
+bool UnqualifiedLookupFactory::hasPreciseScopingOfVarDecls() const {
+  return !options.contains(Flags::IncludeOuterResults);
+}
+
 void UnqualifiedLookupFactory::recordCompletionOfAScope() {
   // OK to call (NOOP) if there are more inner results and Results is empty
   if (IndexOfFirstOuterResult == 0)
@@ -556,12 +567,18 @@ bool ASTScopeDeclConsumerForUnqualifiedLookup::consume(
     if (factory.isOriginallyTypeLookup && !isa<TypeDecl>(value))
       continue;
 
-    // Try to resolve the base for unqualified instance member
-    // references. This is used by lookInMembers().
     if (auto *var = dyn_cast<VarDecl>(value)) {
+      // Try to resolve the base for unqualified instance member
+      // references. This is used by lookInMembers().
       if (var->getName() == factory.Ctx.Id_self) {
         maybeUpdateSelfDC(var);
       }
+
+      // Local VarDecls with a pattern binding are visited as part of their
+      // BraceStmt when hasPreciseScopingOfVarDecls() is off.
+      if (var->getParentPatternBinding() &&
+          !factory.hasPreciseScopingOfVarDecls())
+        continue;
     }
 
     if (!value->getName().matchesRef(factory.Name.getFullName()))
@@ -587,6 +604,21 @@ bool ASTScopeDeclConsumerForUnqualifiedLookup::consume(
   return factory.isFirstResultEnough();
 }
 
+bool ASTScopeDeclConsumerForUnqualifiedLookup::consumePossiblyNotInScope(
+    ArrayRef<VarDecl *> vars) {
+  if (factory.hasPreciseScopingOfVarDecls())
+    return false;
+
+  for (auto *var : vars) {
+    if (!factory.Name.getFullName().isSimpleName(var->getName()))
+      continue;
+
+    factory.Results.push_back(LookupResultEntry(var));
+  }
+
+  return false;
+}
+
 bool ASTScopeDeclGatherer::consume(ArrayRef<ValueDecl *> valuesArg,
                                    DeclVisibilityKind,
                                    NullablePtr<DeclContext>) {