Merge pull request swiftlang#28075 from davidungar/rdar-54906120-health-bug

[NameLookup, ASTScope] Eagerly expand function bodies before type-checking them

Author: David Ungar

include/swift/AST/ASTScope.h

@@ -431,12 +431,12 @@ class ASTScopeImpl {
 
 protected:
   /// Not const because may reexpand some scopes.
-  const ASTScopeImpl *findInnermostEnclosingScope(SourceLoc,
-                                                  NullablePtr<raw_ostream>);
-  const ASTScopeImpl *findInnermostEnclosingScopeImpl(SourceLoc,
-                                                      NullablePtr<raw_ostream>,
-                                                      SourceManager &,
-                                                      ScopeCreator &);
+  ASTScopeImpl *findInnermostEnclosingScope(SourceLoc,
+                                            NullablePtr<raw_ostream>);
+  ASTScopeImpl *findInnermostEnclosingScopeImpl(SourceLoc,
+                                                NullablePtr<raw_ostream>,
+                                                SourceManager &,
+                                                ScopeCreator &);
 
 private:
   NullablePtr<ASTScopeImpl> findChildContaining(SourceLoc loc,
@@ -567,6 +567,8 @@ class ASTSourceFileScope final : public ASTScopeImpl {
   void
   buildEnoughOfTreeForTopLevelExpressionsButDontRequestGenericsOrExtendedNominals();
 
+  void expandFunctionBody(AbstractFunctionDecl *AFD);
+
   const SourceFile *getSourceFile() const override;
   NullablePtr<const void> addressForPrinting() const override { return SF; }
   bool crossCheckWithAST();

include/swift/AST/NameLookup.h

@@ -596,6 +596,8 @@ class ASTScope {
   void
   buildEnoughOfTreeForTopLevelExpressionsButDontRequestGenericsOrExtendedNominals();
 
+  static void expandFunctionBody(AbstractFunctionDecl *);
+
   /// Flesh out the tree for dumping
   void buildFullyExpandedTree();
 
@@ -631,6 +633,8 @@ class ASTScope {
 
 private:
   static ast_scope::ASTSourceFileScope *createScopeTree(SourceFile *);
+
+  void expandFunctionBodyImpl(AbstractFunctionDecl *);
 };
 
 } // end namespace swift

lib/AST/ASTScopeCreation.cpp

@@ -735,6 +735,16 @@ bool ASTScope::areInactiveIfConfigClausesSupported() {
   return ScopeCreator::includeInactiveIfConfigClauses;
 }
 
+void ASTScope::expandFunctionBody(AbstractFunctionDecl *AFD) {
+  auto *const SF = AFD->getParentSourceFile();
+  if (SF->isSuitableForASTScopes())
+    SF->getScope().expandFunctionBodyImpl(AFD);
+}
+
+void ASTScope::expandFunctionBodyImpl(AbstractFunctionDecl *AFD) {
+  impl->expandFunctionBody(AFD);
+}
+
 ASTSourceFileScope *ASTScope::createScopeTree(SourceFile *SF) {
   ScopeCreator *scopeCreator = new (SF->getASTContext()) ScopeCreator(SF);
   return scopeCreator->sourceFileScope;
@@ -752,6 +762,15 @@ void ASTSourceFileScope::
       expandAndBeCurrentDetectingRecursion(*scopeCreator);
 }
 
+void ASTSourceFileScope::expandFunctionBody(AbstractFunctionDecl *AFD) {
+  if (!AFD)
+    return;
+  auto sr = AFD->getBodySourceRange();
+  if (sr.isInvalid())
+    return;
+  ASTScopeImpl *bodyScope = findInnermostEnclosingScope(sr.Start, nullptr);
+  bodyScope->expandAndBeCurrentDetectingRecursion(*scopeCreator);
+}
 
 ASTSourceFileScope::ASTSourceFileScope(SourceFile *SF,
                                        ScopeCreator *scopeCreator)
@@ -1059,6 +1078,8 @@ void ASTScopeImpl::disownDescendants(ScopeCreator &scopeCreator) {
 
 ASTScopeImpl *
 ASTScopeImpl::expandAndBeCurrentDetectingRecursion(ScopeCreator &scopeCreator) {
+  assert(scopeCreator.getASTContext().LangOpts.EnableASTScopeLookup &&
+         "Should not be getting here if ASTScopes are disabled");
   return evaluateOrDefault(scopeCreator.getASTContext().evaluator,
                            ExpandASTScopeRequest{this, &scopeCreator}, nullptr);
 }

lib/AST/ASTScopeLookup.cpp

@@ -63,7 +63,7 @@ const ASTScopeImpl *ASTScopeImpl::findStartingScopeForLookup(
   if (name.isOperator())
     return fileScope; // operators always at file scope
 
-  const auto innermost = fileScope->findInnermostEnclosingScope(loc, nullptr);
+  const auto *innermost = fileScope->findInnermostEnclosingScope(loc, nullptr);
   ASTScopeAssert(innermost->getWasExpanded(),
                  "If looking in a scope, it must have been expanded.");
 
@@ -117,14 +117,14 @@ const ASTScopeImpl *ASTScopeImpl::findStartingScopeForLookup(
   return startingScope;
 }
 
-const ASTScopeImpl *
+ASTScopeImpl *
 ASTScopeImpl::findInnermostEnclosingScope(SourceLoc loc,
                                           NullablePtr<raw_ostream> os) {
   return findInnermostEnclosingScopeImpl(loc, os, getSourceManager(),
                                          getScopeCreator());
 }
 
-const ASTScopeImpl *ASTScopeImpl::findInnermostEnclosingScopeImpl(
+ASTScopeImpl *ASTScopeImpl::findInnermostEnclosingScopeImpl(
     SourceLoc loc, NullablePtr<raw_ostream> os, SourceManager &sourceMgr,
     ScopeCreator &scopeCreator) {
   expandAndBeCurrentDetectingRecursion(scopeCreator);

lib/AST/UnqualifiedLookup.cpp

@@ -232,7 +232,7 @@ namespace {
     bool useASTScopesForLookup() const;
 
     /// For testing, assume this lookup is enabled:
-    bool useASTScopesForLookupIfEnabled() const;
+    bool wouldUseASTScopesForLookupIfItWereEnabled() const;
 
     void lookUpTopLevelNamesInModuleScopeContext(DeclContext *);
 
@@ -506,7 +506,8 @@ void UnqualifiedLookupFactory::performUnqualifiedLookup() {
       lookupNamesIntroducedBy(contextAndIsCascadingUse);
   }
 
-  if (crosscheckUnqualifiedLookup && useASTScopesForLookupIfEnabled()) {
+  if (crosscheckUnqualifiedLookup &&
+      wouldUseASTScopesForLookupIfItWereEnabled()) {
     ResultsVector results;
     size_t indexOfFirstOuterResult = 0;
     UnqualifiedLookupFactory altLookup(Name, DC, Loc, options, results,
@@ -550,10 +551,12 @@ void UnqualifiedLookupFactory::lookUpTopLevelNamesInModuleScopeContext(
 }
 
 bool UnqualifiedLookupFactory::useASTScopesForLookup() const {
-  return Ctx.LangOpts.EnableASTScopeLookup && useASTScopesForLookupIfEnabled();
+  return Ctx.LangOpts.EnableASTScopeLookup &&
+         wouldUseASTScopesForLookupIfItWereEnabled();
 }
 
-bool UnqualifiedLookupFactory::useASTScopesForLookupIfEnabled() const {
+bool UnqualifiedLookupFactory::wouldUseASTScopesForLookupIfItWereEnabled()
+    const {
   if (!Loc.isValid())
     return false;
   const auto *const SF = DC->getParentSourceFile();

lib/Sema/TypeCheckStmt.cpp

@@ -2142,6 +2142,12 @@ TypeCheckFunctionBodyUntilRequest::evaluate(Evaluator &evaluator,
                              body->getRBraceLoc(), body->isImplicit());
   }
 
+  // Typechecking, in particular ApplySolution is going to replace closures
+  // with OpaqueValueExprs and then try to do lookups into the closures.
+  // So, build out the body now.
+  if (ctx.LangOpts.EnableASTScopeLookup)
+    ASTScope::expandFunctionBody(AFD);
+
   StmtChecker SC(tc, AFD);
   SC.EndTypeCheckLoc = endTypeCheckLoc;
   bool hadError = SC.typeCheckBody(body);