[NFC NameLookup ASTScope] Fixes for large app, eager primary tree creation, memberCount fix.

include/swift/AST/ASTScope.h

@@ -215,6 +215,10 @@ class ASTScopeImpl {
 
   bool checkSourceRangeOfThisASTNode() const;
 
+  /// For debugging
+  bool doesRangeMatch(unsigned start, unsigned end, StringRef file = "",
+                      StringRef className = "");
+
 private:
   SourceRange computeSourceRangeOfScope(bool omitAssertions = false) const;
   SourceRange
@@ -264,7 +268,7 @@ class ASTScopeImpl {
   getSourceRangeOfEnclosedParamsOfASTNode(bool omitAssertions) const;
 
 private:
-  bool checkSourceRangeAfterExpansion() const;
+  bool checkSourceRangeAfterExpansion(const ASTContext &) const;
 
 #pragma mark common queries
 public:
@@ -318,9 +322,9 @@ class ASTScopeImpl {
 private:
   /// Compare the pre-expasion range with the post-expansion range and return
   /// false if lazyiness couild miss lookups.
-  bool checkLazySourceRange() const;
+  bool checkLazySourceRange(const ASTContext &) const;
 
-protected:
+public:
   /// Some scopes can be expanded lazily.
   /// Such scopes must: not change their source ranges after expansion, and
   /// their expansion must return an insertion point outside themselves.
@@ -517,6 +521,7 @@ class ASTSourceFileScope final : public ASTScopeImpl {
   NullablePtr<DeclContext> getDeclContext() const override;
 
   void addNewDeclsToScopeTree();
+  void buildScopeTreeEagerly();
 
   const SourceFile *getSourceFile() const override;
   NullablePtr<const void> addressForPrinting() const override { return SF; }
@@ -1002,6 +1007,10 @@ class AbstractFunctionBodyScope : public ASTScopeImpl {
                              DeclConsumer) const override;
   Optional<bool>
   resolveIsCascadingUseForThisScope(Optional<bool>) const override;
+
+public:
+  NullablePtr<ASTScopeImpl> insertionPointForDeferredExpansion() override;
+  SourceRange sourceRangeForDeferredExpansion() const override;
 };
 
 /// Body of methods, functions in types.

include/swift/AST/DeclContext.h

@@ -760,7 +760,7 @@ class IterableDeclContext {
   void addMember(Decl *member, Decl *hint = nullptr);
 
   /// See \c MemberCount
-  unsigned getMemberCount() const { return MemberCount; }
+  unsigned getMemberCount() const;
 
   /// Check whether there are lazily-loaded members.
   bool hasLazyMembers() const {

include/swift/AST/ExperimentalDependencies.h

@@ -222,6 +222,7 @@ class BiIndexedTwoStageMap {
   bool insert(const Key1 &k1, const Key2 &k2, Value &v) {
     const bool r1 = map1.insert(k1, k2, v);
     const bool r2 = map2.insert(k2, k1, v);
+    (void)r2;
     assertConsistent(r1, r2);
     return r1;
   }

include/swift/AST/Module.h

@@ -1241,6 +1241,8 @@ class SourceFile final : public FileUnit {
 
   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,

include/swift/Basic/LangOptions.h

@@ -254,6 +254,9 @@ namespace swift {
     /// Should  we compare to ASTScope-based resolution for debugging?
     bool CompareToASTScopeLookup = false;
 
+    /// Should  we stress ASTScope-based resolution for debugging?
+    bool StressASTScopeLookup = false;
+
     /// Since some tests fail if the warning is output, use a flag to decide
     /// whether it is. The warning is useful for testing.
     bool WarnIfASTScopeLookup = false;

include/swift/Driver/ExperimentalDependencyDriverGraph.h

@@ -223,6 +223,7 @@ class ModuleDepGraph {
     ModuleDepGraphNode *nodeActuallyErased = nodeMap.findAndErase(
         nodeToErase->getSwiftDeps().getValueOr(std::string()),
         nodeToErase->getKey());
+    (void)nodeActuallyErased;
     assert(
         nodeToErase == nodeActuallyErased ||
         mapCorruption("Node found from key must be same as node holding key."));

include/swift/Option/FrontendOptions.td

@@ -124,6 +124,9 @@ def disable_target_os_checking :
 
 def compare_to_astscope_lookup : Flag<["-"], "compare-to-astscope-lookup">,
   HelpText<"Compare legacy to ASTScope-based unqualified name lookup (for debugging)">;
+
+def stress_astscope_lookup : Flag<["-"], "stress-astscope-lookup">,
+  HelpText<"Stress ASTScope-based unqualified name lookup (for testing)">;
 
 def warn_if_astscope_lookup : Flag<["-"], "warn-if-astscope-lookup">,
   HelpText<"Print a warning if ASTScope lookup is used">;

lib/AST/ASTScopeCreation.cpp

@@ -93,6 +93,26 @@ static void dumpRangeable(SpecializeAttr *r, llvm::raw_ostream &f) {
   llvm::errs() << "SpecializeAttr\n";
 }
 
+/// For Debugging
+template <typename T>
+bool doesRangeableRangeMatch(const T *x, const SourceManager &SM,
+                             unsigned start, unsigned end,
+                             StringRef file = "") {
+  auto const r = getRangeableSourceRange(x);
+  if (r.isInvalid())
+    return false;
+  if (start && SM.getLineNumber(r.Start) != start)
+    return false;
+  if (end && SM.getLineNumber(r.End) != end)
+    return false;
+  if (file.empty())
+    return true;
+  const auto buf = SM.findBufferContainingLoc(r.Start);
+  return SM.getIdentifierForBuffer(buf).endswith(file);
+}
+
+#pragma mark end of rangeable
+
 static std::vector<ASTNode> asNodeVector(DeclRange dr) {
   std::vector<ASTNode> nodes;
   llvm::transform(dr, std::back_inserter(nodes),
@@ -165,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.
@@ -366,7 +404,7 @@ class ScopeCreator final {
             parent, vd);
     });
   }
-  // HERE
+
 public:
   /// Create the matryoshka nested generic param scopes (if any)
   /// that are subscopes of the receiver. Return
@@ -607,7 +645,9 @@ class ScopeCreator final {
     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
@@ -617,9 +657,11 @@ class ScopeCreator final {
     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);
@@ -699,12 +741,24 @@ class ScopeCreator final {
 
 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;
@@ -805,6 +859,7 @@ class NodeAdder
   VISIT_AND_CREATE_WHOLE_PORTION(ExtensionDecl, ExtensionScope)
   VISIT_AND_CREATE_WHOLE_PORTION(StructDecl, NominalTypeScope)
   VISIT_AND_CREATE_WHOLE_PORTION(ClassDecl, NominalTypeScope)
+  VISIT_AND_CREATE_WHOLE_PORTION(ProtocolDecl, NominalTypeScope)
   VISIT_AND_CREATE_WHOLE_PORTION(EnumDecl, NominalTypeScope)
   VISIT_AND_CREATE_WHOLE_PORTION(TypeAliasDecl, TypeAliasScope)
   VISIT_AND_CREATE_WHOLE_PORTION(OpaqueTypeDecl, OpaqueTypeScope)
@@ -817,12 +872,6 @@ class NodeAdder
     return visitAbstractFunctionDecl(ad, p, scopeCreator);
   }
 
-  NullablePtr<ASTScopeImpl> visitProtocolDecl(ProtocolDecl *e, ASTScopeImpl *p,
-                                              ScopeCreator &scopeCreator) {
-    return scopeCreator.ifUniqueConstructWithPortionExpandAndInsert<
-        NominalTypeScope, GenericTypeOrExtensionWholePortion>(p, e);
-  }
-
 #pragma mark simple creation with deferred nodes
 
   // Each of the following creates a new scope, so that nodes which were parsed
@@ -1011,6 +1060,9 @@ void ASTScopeImpl::addChild(ASTScopeImpl *child, ASTContext &ctx) {
   assert(!child->getParent() && "child should not already have parent");
   child->parent = this;
   clearCachedSourceRangesOfMeAndAncestors();
+  // It's possible that some callees do lookups back into the tree.
+  // So make sure childrenCountWhenLastExpanded is up to date.
+  setChildrenCountWhenLastExpanded();
 }
 
 void ASTScopeImpl::removeChildren() {
@@ -1036,8 +1088,7 @@ ASTScopeImpl *ASTScopeImpl::expandAndBeCurrent(ScopeCreator &scopeCreator) {
            insertionPointForDeferredExpansion().get() == insertionPoint);
   }
   beCurrent();
-  setChildrenCountWhenLastExpanded();
-  assert(checkSourceRangeAfterExpansion());
+  assert(checkSourceRangeAfterExpansion(scopeCreator.getASTContext()));
   return insertionPoint;
 }
 
@@ -1448,25 +1499,30 @@ void AttachedPropertyWrapperScope::
 
 ASTScopeImpl *GenericTypeOrExtensionWholePortion::expandScope(
     GenericTypeOrExtensionScope *scope, ScopeCreator &scopeCreator) const {
+  // Get now in case recursion emancipates scope
+  auto *const ip = scope->getParent().get();
+
   // Prevent circular request bugs caused by illegal input and
   // doing lookups that getExtendedNominal in the midst of getExtendedNominal.
   // rdar://53972776
   if (scope->shouldHaveABody() && !scope->doesDeclHaveABody())
-    return scope->getParent().get();
+    return ip;
 
   auto *deepestScope = scopeCreator.addNestedGenericParamScopesToTree(
       scope->getDecl(), scope->getGenericContext()->getGenericParams(), scope);
   if (scope->getGenericContext()->getTrailingWhereClause())
     scope->createTrailingWhereClauseScope(deepestScope, scopeCreator);
   scope->createBodyScope(deepestScope, scopeCreator);
-  return scope->getParent().get();
+  return ip;
 }
 
 ASTScopeImpl *
 IterableTypeBodyPortion::expandScope(GenericTypeOrExtensionScope *scope,
                                      ScopeCreator &scopeCreator) const {
+  // Get it now in case of recursion and this one gets emancipated
+  auto *const ip = scope->getParent().get();
   scope->expandBody(scopeCreator);
-  return scope->getParent().get();
+  return ip;
 }
 
 ASTScopeImpl *GenericTypeOrExtensionWherePortion::expandScope(
@@ -1670,7 +1726,9 @@ void IterableTypeScope::expandBody(ScopeCreator &scopeCreator) {
 #pragma mark - reexpandIfObsolete
 
 bool ASTScopeImpl::reexpandIfObsolete(ScopeCreator &scopeCreator) {
-  if (isCurrent())
+  if (scopeCreator.getIsFrozen() ||
+      (isCurrent() &&
+          !scopeCreator.getASTContext().LangOpts.StressASTScopeLookup))
     return false;
   reexpand(scopeCreator);
   return true;
@@ -1731,6 +1789,11 @@ NullablePtr<ASTScopeImpl> ASTScopeImpl::insertionPointForDeferredExpansion() {
   return nullptr;
 }
 
+NullablePtr<ASTScopeImpl>
+AbstractFunctionBodyScope::insertionPointForDeferredExpansion() {
+  return getParent().get();
+}
+
 NullablePtr<ASTScopeImpl>
 IterableTypeScope::insertionPointForDeferredExpansion() {
   return portion->insertionPointForDeferredExpansion(this);
@@ -1782,20 +1845,29 @@ bool TopLevelCodeScope::isCurrent() const {
   return bodyWhenLastExpanded == decl->getBody();
 }
 
+// Try to avoid the work of counting
+static const bool assumeVarsDoNotGetAdded = true;
+
+static unsigned countVars(const PatternBindingEntry &entry) {
+  unsigned varCount = 0;
+  entry.getPattern()->forEachVariable([&](VarDecl *) { ++varCount; });
+  return varCount;
+}
+
 void PatternEntryDeclScope::beCurrent() {
   initWhenLastExpanded = getPatternEntry().getOriginalInit();
-  unsigned varCount = 0;
-  getPatternEntry().getPattern()->forEachVariable(
-      [&](VarDecl *) { ++varCount; });
-  varCountWhenLastExpanded = varCount;
+  if (assumeVarsDoNotGetAdded && varCountWhenLastExpanded)
+    return;
+  varCountWhenLastExpanded = countVars(getPatternEntry());
 }
 bool PatternEntryDeclScope::isCurrent() const {
   if (initWhenLastExpanded != getPatternEntry().getOriginalInit())
     return false;
-  unsigned varCount = 0;
-  getPatternEntry().getPattern()->forEachVariable(
-      [&](VarDecl *) { ++varCount; });
-  return varCount == varCountWhenLastExpanded;
+  if (assumeVarsDoNotGetAdded && varCountWhenLastExpanded) {
+    assert(varCountWhenLastExpanded == countVars(getPatternEntry()));
+    return true;
+  }
+  return countVars(getPatternEntry()) == varCountWhenLastExpanded;
 }
 
 void WholeClosureScope::beCurrent() {

lib/AST/ASTScopeLookup.cpp

@@ -58,7 +58,7 @@ const ASTScopeImpl *ASTScopeImpl::findStartingScopeForLookup(
 
   auto *const fileScope = sourceFile->getScope().impl;
   // Parser may have added decls to source file, since previous lookup
-  sourceFile->getScope().impl->addNewDeclsToScopeTree();
+  fileScope->addNewDeclsToScopeTree();
   if (name.isOperator())
     return fileScope; // operators always at file scope