Eager tree building for primaries

include/swift/AST/ASTScope.h

@@ -521,6 +521,7 @@ public:
   NullablePtr<DeclContext> getDeclContext() const override;
 
   void addNewDeclsToScopeTree();
+  void buildScopeTreeEagerly();
 
   const SourceFile *getSourceFile() const override;
   NullablePtr<const void> addressForPrinting() const override { return SF; }

include/swift/AST/Module.h

@@ -1241,6 +1241,8 @@ public:
 
   bool canBeParsedInFull() const;
 
+  bool isSuitableForASTScopes() const { return canBeParsedInFull(); }
+
   syntax::SourceFileSyntax getSyntaxRoot() const;
   void setSyntaxRoot(syntax::SourceFileSyntax &&Root);
   bool hasSyntaxRoot() const;

include/swift/AST/NameLookup.h

@@ -592,6 +592,10 @@ class ASTScope {
 public:
   ASTScope(SourceFile *);
 
+  /// Cannot be lazy during type-checking because it mutates the AST.
+  /// So build eagerly before type-checking
+  void buildScopeTreeEagerly();
+
   /// \return the scopes traversed
   static llvm::SmallVector<const ast_scope::ASTScopeImpl *, 0>
   unqualifiedLookup(SourceFile *, DeclName, SourceLoc,

lib/AST/ASTScopeCreation.cpp

@@ -185,9 +185,27 @@ class ScopeCreator final {
   /// For allocating scopes.
   ASTContext &ctx;
 
+public:
+  /// Because type checking can mutate the AST, eagerly build the tree, then
+  /// freeze it
+  enum class Temperature {
+    Warm,     // Can be lazy
+    Freezing, // Should expand everything eagerly
+    Frozen    // No more changes, except when Decls are added to the source file
+  };
+
+private:
+  /// Because type checking can mutate the AST, eagerly build the tree, then
+  /// freeze it
+  Temperature temperature = Temperature::Warm;
+
 public:
   ASTSourceFileScope *const sourceFileScope;
   ASTContext &getASTContext() const { return ctx; }
+  bool getIsFrozen() const { return temperature == Temperature::Frozen; }
+  bool getIsFreezing() const { return temperature == Temperature::Freezing; }
+  void beFreezing() { temperature = Temperature::Freezing; }
+  void beFrozen() { temperature = Temperature::Frozen; }
 
   /// The AST can have duplicate nodes, and we don't want to create scopes for
   /// those.
@@ -627,7 +645,9 @@ public:
     return !n.isDecl(DeclKind::Var);
   }
 
-  bool shouldBeLazy() const { return ctx.LangOpts.LazyASTScopes; }
+  bool shouldBeLazy() const {
+    return !getIsFreezing() && ctx.LangOpts.LazyASTScopes;
+  }
 
 public:
   /// For debugging. Return true if scope tree contains all the decl contexts in
@@ -637,9 +657,11 @@ public:
     auto allDeclContexts = findLocalizableDeclContextsInAST();
     llvm::DenseMap<const DeclContext *, const ASTScopeImpl *> bogusDCs;
     bool rebuilt = false;
-    sourceFileScope->preOrderDo([&](ASTScopeImpl *scope) {
-      rebuilt |= scope->reexpandIfObsolete(*this);
-    });
+    if (!getIsFrozen()) {
+      sourceFileScope->preOrderDo([&](ASTScopeImpl *scope) {
+        rebuilt |= scope->reexpandIfObsolete(*this);
+      });
+    }
     sourceFileScope->postOrderDo([&](ASTScopeImpl *scope) {
       if (auto *dc = scope->getDeclContext().getPtrOrNull()) {
         auto iter = allDeclContexts.find(dc);
@@ -719,12 +741,24 @@ public:
 
 ASTScope::ASTScope(SourceFile *SF) : impl(createScopeTree(SF)) {}
 
+void ASTScope::buildScopeTreeEagerly() {
+  impl->buildScopeTreeEagerly();
+}
+
 ASTSourceFileScope *ASTScope::createScopeTree(SourceFile *SF) {
   ScopeCreator *scopeCreator = new (SF->getASTContext()) ScopeCreator(SF);
   scopeCreator->sourceFileScope->addNewDeclsToScopeTree();
   return scopeCreator->sourceFileScope;
 }
 
+void ASTSourceFileScope::buildScopeTreeEagerly() {
+  scopeCreator->beFreezing();
+  // Eagerly expand any decls already in the tree.
+  preOrderDo([&](ASTScopeImpl *s) { s->reexpandIfObsolete(*scopeCreator); });
+  addNewDeclsToScopeTree();
+  scopeCreator->beFrozen();
+}
+
 void ASTSourceFileScope::addNewDeclsToScopeTree() {
   assert(SF && scopeCreator);
   ArrayRef<Decl *> decls = SF->Decls;
@@ -1697,7 +1731,7 @@ void IterableTypeScope::expandBody(ScopeCreator &scopeCreator) {
 #pragma mark - reexpandIfObsolete
 
 bool ASTScopeImpl::reexpandIfObsolete(ScopeCreator &scopeCreator) {
-  if (isCurrent())
+  if (scopeCreator.getIsFrozen() || isCurrent())
     return false;
   reexpand(scopeCreator);
   return true;

lib/AST/Module.cpp

@@ -1745,6 +1745,8 @@ StringRef SourceFile::getFilename() const {
 }
 
 ASTScope &SourceFile::getScope() {
+  assert(getASTContext().LangOpts.EnableASTScopeLookup &&
+         isSuitableForASTScopes() && "Should not be creating scope tree");
   if (!Scope)
     Scope = std::unique_ptr<ASTScope>(new (getASTContext()) ASTScope(this));
   return *Scope.get();

lib/AST/UnqualifiedLookup.cpp

@@ -536,9 +536,10 @@ bool UnqualifiedLookupFactory::useASTScopesForExperimentalLookup() const {
 
 bool UnqualifiedLookupFactory::useASTScopesForExperimentalLookupIfEnabled()
     const {
-  return Loc.isValid() && DC->getParentSourceFile() &&
-         DC->getParentSourceFile()->Kind != SourceFileKind::REPL &&
-         DC->getParentSourceFile()->Kind != SourceFileKind::SIL;
+  if (!Loc.isValid())
+    return false;
+  const auto *const SF = DC->getParentSourceFile();
+  return SF && SF->isSuitableForASTScopes();
 }
 
 #pragma mark context-based lookup definitions

lib/Sema/TypeChecker.cpp

@@ -351,6 +351,13 @@ void swift::performTypeChecking(SourceFile &SF, TopLevelContext &TLC,
   if (SF.ASTStage == SourceFile::TypeChecked)
     return;
 
+  // Eagerly build a scope tree before type checking
+  // because type-checking mutates the AST and that throws off the scope-based
+  // lookups.
+  if (SF.getASTContext().LangOpts.EnableASTScopeLookup &&
+      SF.isSuitableForASTScopes())
+    SF.getScope().buildScopeTreeEagerly();
+
   auto &Ctx = SF.getASTContext();
   BufferIndirectlyCausingDiagnosticRAII cpr(SF);
 

test/NameBinding/scope_map_top_level.swift

@@ -37,7 +37,7 @@ var i: Int = b.my_identity()
 // CHECK-EXPANDED-NEXT:       `-BraceStmtScope {{.*}}, [8:1 - 21:28]
 // CHECK-EXPANDED-NEXT:         `-GuardStmtScope {{.*}}, [8:1 - 21:28]
 // CHECK-EXPANDED-NEXT:           |-ConditionalClauseScope, [8:7 - 8:22] index 0
-// CHECK-EXPANDED-NEXT:             `-ConditionalClausePatternUseScope, [8:22 - 8:22] let b?
+// CHECK-EXPANDED-NEXT:             `-ConditionalClausePatternUseScope, [8:22 - 8:22] let b{{.*}}
 // CHECK-EXPANDED-NEXT:           |-BraceStmtScope {{.*}}, [8:22 - 9:1]
 // CHECK-EXPANDED-NEXT:           `-LookupParentDiversionScope, [9:1 - 21:28]
 // CHECK-EXPANDED-NEXT:             |-AbstractFunctionDeclScope {{.*}}, [11:1 - 11:13] 'foo()'
@@ -51,11 +51,14 @@ var i: Int = b.my_identity()
 // CHECK-EXPANDED-NEXT:                 |-TypeAliasDeclScope {{.*}}, [15:1 - 15:15]
 // CHECK-EXPANDED-NEXT:                 |-ExtensionDeclScope {{.*}}, [17:1 - 19:1]
 // CHECK-EXPANDED-NEXT:                   `-ExtensionBodyScope {{.*}}, [17:15 - 19:1]
+// CHECK-EXPANDED-NEXT:                     `-AbstractFunctionDeclScope {{.*}}, [18:3 - 18:43] 'my_identity()'
+// CHECK-EXPANDED-NEXT:                       `-ParameterListScope {{.*}}, [18:19 - 18:43]
+// CHECK-EXPANDED-NEXT:                         `-MethodBodyScope {{.*}}, [18:29 - 18:43]
+// CHECK-EXPANDED-NEXT:                           `-BraceStmtScope {{.*}}, [18:29 - 18:43]
 // CHECK-EXPANDED-NEXT:                 `-TopLevelCodeScope {{.*}}, [21:1 - 21:28]
 // CHECK-EXPANDED-NEXT:                   `-BraceStmtScope {{.*}}, [21:1 - 21:28]
 // CHECK-EXPANDED-NEXT:                     `-PatternEntryDeclScope {{.*}}, [21:5 - 21:28] entry 0 'i'
 // CHECK-EXPANDED-NEXT:                       `-PatternEntryInitializerScope {{.*}}, [21:14 - 21:28] entry 0 'i'
 
-
 // REQUIRES: asserts
 // absence of assertions can change the "uncached" bit and cause failures