[AST] Requestify local type declarations

Instead of recording local type declarations while parsing, make a
request to collect them from the AST. This reduces the side effect of
parsing.

Author: rintaro

include/swift/AST/SourceFile.h

@@ -313,9 +313,6 @@ class SourceFile final : public FileUnit {
   /// includes the SourceFileSyntax node corresponding to this source file.
   void *getExportedSourceFile() const;
 
-  /// The list of local type declarations in the source file.
-  llvm::SetVector<TypeDecl *> LocalTypeDecls;
-
   /// Defer type checking of `AFD` to the end of `Sema`
   void addDelayedFunction(AbstractFunctionDecl *AFD);
 
@@ -762,6 +759,8 @@ class SourceFile final : public FileUnit {
   /// Returns true if the source file contains concurrency in the top-level
   bool isAsyncTopLevelSourceFile() const;
 
+  ArrayRef<TypeDecl *> getLocalTypeDecls() const;
+
 private:
 
   /// If not \c None, the underlying vector contains the parsed tokens of this

include/swift/AST/TypeCheckRequests.h

@@ -4732,6 +4732,23 @@ class UniqueUnderlyingTypeSubstitutionsRequest
   void cacheResult(llvm::Optional<SubstitutionMap>) const;
 };
 
+/// Collect all local type declarations in the SourceFile.
+class LocalTypeDeclsRequest
+    : public SimpleRequest<LocalTypeDeclsRequest,
+                           ArrayRef<TypeDecl *>(SourceFile *),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  ArrayRef<TypeDecl *> evaluate(Evaluator &evaluator, SourceFile *sf) const;
+
+public:
+  bool isCached() const { return true; }
+};
+
 #define SWIFT_TYPEID_ZONE TypeChecker
 #define SWIFT_TYPEID_HEADER "swift/AST/TypeCheckerTypeIDZone.def"
 #include "swift/Basic/DefineTypeIDZone.h"

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -542,3 +542,6 @@ SWIFT_REQUEST(TypeChecker, SemanticDeclAttrsRequest,
 SWIFT_REQUEST(TypeChecker, UniqueUnderlyingTypeSubstitutionsRequest,
               llvm::Optional<SubstitutionMap>(const Decl *),
               SeparatelyCached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, LocalTypeDeclsRequest,
+              ArrayRef<TypeDecl *>(SourceFile *),
+              Cached, NoLocationInfo)

include/swift/Parse/Parser.h

@@ -1015,8 +1015,6 @@ class Parser {
   /// 'isLine = true' indicates parsing #line instead of #sourcelocation
   ParserStatus parseLineDirective(bool isLine = false);
 
-  void recordLocalType(TypeDecl *TD);
-
   /// Skip an `#if` configuration block containing only attributes.
   ///
   /// \returns true if the skipping was successful, false otherwise.

lib/AST/Decl.cpp

@@ -9398,12 +9398,8 @@ void AbstractFunctionDecl::setBodyToBeReparsed(SourceRange bodyRange) {
   BodyRange = bodyRange;
   setBodyKind(BodyKind::Unparsed);
 
-  // Remove local type decls from the source file.
   if (auto SF = getParentSourceFile()) {
-    SF->LocalTypeDecls.remove_if([&](TypeDecl *TD) {
-      auto dc = TD->getDeclContext();
-      return dc == this || dc->isChildContextOf(this);
-    });
+    SF->getASTContext().evaluator.clearCachedOutput(LocalTypeDeclsRequest{SF});
   }
 }
 

lib/AST/Module.cpp

@@ -1444,7 +1444,8 @@ void SourceFile::getPrecedenceGroups(
 }
 
 void SourceFile::getLocalTypeDecls(SmallVectorImpl<TypeDecl*> &Results) const {
-  Results.append(LocalTypeDecls.begin(), LocalTypeDecls.end());
+  auto decls = getLocalTypeDecls();
+  Results.append(decls.begin(), decls.end());
 }
 
 void
@@ -1456,7 +1457,7 @@ const {
 
 TypeDecl *SourceFile::lookupLocalType(llvm::StringRef mangledName) const {
   ASTContext &ctx = getASTContext();
-  for (auto typeDecl : LocalTypeDecls) {
+  for (auto typeDecl : getLocalTypeDecls()) {
     auto typeMangledName = evaluateOrDefault(ctx.evaluator,
                                              MangleLocalTypeDeclRequest { typeDecl },
                                              std::string());
@@ -4245,6 +4246,50 @@ ASTNode GetSourceFileAsyncNode::evaluate(Evaluator &eval,
   return ASTNode();
 }
 
+ArrayRef<TypeDecl *> SourceFile::getLocalTypeDecls() const {
+  auto *mutableThis = const_cast<SourceFile *>(this);
+  return evaluateOrDefault(getASTContext().evaluator,
+                           LocalTypeDeclsRequest{mutableThis}, {});
+}
+
+namespace {
+class LocalTypeDeclCollector : public ASTWalker {
+  SmallVectorImpl<TypeDecl *> &results;
+
+public:
+  LocalTypeDeclCollector(SmallVectorImpl<TypeDecl *> &results)
+      : results(results) {}
+
+  MacroWalking getMacroWalkingBehavior() const override {
+    return MacroWalking::Expansion;
+  }
+
+  PreWalkAction walkToDeclPre(Decl *D) override {
+    switch (D->getKind()) {
+    case DeclKind::Enum:
+    case DeclKind::Struct:
+    case DeclKind::Class:
+    case DeclKind::Protocol:
+    case DeclKind::TypeAlias:
+      if (D->getDeclContext()->isLocalContext())
+        results.push_back(cast<TypeDecl>(D));
+      break;
+    default:
+      break;
+    }
+    return PreWalkAction::Continue;
+  }
+};
+} // namespace
+
+ArrayRef<TypeDecl *> LocalTypeDeclsRequest::evaluate(Evaluator &evaluator,
+                                                     SourceFile *sf) const {
+  SmallVector<TypeDecl *> results;
+  LocalTypeDeclCollector collector(results);
+  sf->walk(collector);
+  return sf->getASTContext().AllocateCopy(results);
+}
+
 //===----------------------------------------------------------------------===//
 // SynthesizedFileUnit Implementation
 //===----------------------------------------------------------------------===//

lib/FrontendTool/FrontendTool.cpp

@@ -280,7 +280,7 @@ static void countStatsOfSourceFile(UnifiedStatsReporter &Stats,
   auto &C = Stats.getFrontendCounters();
   auto &SM = Instance.getSourceMgr();
   C.NumDecls += SF->getTopLevelDecls().size();
-  C.NumLocalTypeDecls += SF->LocalTypeDecls.size();
+  C.NumLocalTypeDecls += SF->getLocalTypeDecls().size();
   C.NumObjCMethods += SF->ObjCMethods.size();
 
   SmallVector<OperatorDecl *, 2> operators;

lib/IRGen/GenDecl.cpp

@@ -465,7 +465,7 @@ void IRGenModule::emitSourceFile(SourceFile &SF) {
     emitGlobalDecl(decl);
   for (auto *decl : SF.getHoistedDecls())
     emitGlobalDecl(decl);
-  for (auto *localDecl : SF.LocalTypeDecls)
+  for (auto *localDecl : SF.getLocalTypeDecls())
     emitGlobalDecl(localDecl);
   for (auto *opaqueDecl : SF.getOpaqueReturnTypeDecls())
     maybeEmitOpaqueTypeDecl(opaqueDecl);

lib/Parse/ParseDecl.cpp

@@ -5728,15 +5728,6 @@ void Parser::consumeDecl(ParserPosition BeginParserPosition, bool IsTopLevel) {
   }
 }
 
-void Parser::recordLocalType(TypeDecl *TD) {
-  // If we're not in a local context, this is unnecessary.
-  if (!TD->getDeclContext()->isLocalContext())
-    return;
-
-  if (!InInactiveClauseEnvironment && !InFreestandingMacroArgument)
-    SF.getOutermostParentSourceFile()->LocalTypeDecls.insert(TD);
-}
-
 /// Set the original declaration in `@differentiable` attributes.
 ///
 /// Necessary because `Parser::parseNewDeclAttribute` (which calls
@@ -7200,7 +7191,6 @@ parseDeclTypeAlias(Parser::ParseDeclOptions Flags, DeclAttributes &Attributes) {
 
   auto *TAD = new (Context) TypeAliasDecl(TypeAliasLoc, EqualLoc, Id, IdLoc,
                                           genericParams, CurDeclContext);
-  recordLocalType(TAD);
   ParserResult<TypeRepr> UnderlyingTy;
 
   if (Tok.is(tok::colon) || Tok.is(tok::equal)) {
@@ -8885,7 +8875,6 @@ ParserResult<EnumDecl> Parser::parseDeclEnum(ParseDeclOptions Flags,
 
   EnumDecl *ED = new (Context) EnumDecl(EnumLoc, EnumName, EnumNameLoc,
                                         { }, GenericParams, CurDeclContext);
-  recordLocalType(ED);
   ED->getAttrs() = Attributes;
 
   ContextChange CC(*this, ED);
@@ -9152,7 +9141,6 @@ ParserResult<StructDecl> Parser::parseDeclStruct(ParseDeclOptions Flags,
                                             { },
                                             GenericParams,
                                             CurDeclContext);
-  recordLocalType(SD);
   SD->getAttrs() = Attributes;
 
   ContextChange CC(*this, SD);
@@ -9246,7 +9234,6 @@ ParserResult<ClassDecl> Parser::parseDeclClass(ParseDeclOptions Flags,
   ClassDecl *CD = new (Context) ClassDecl(ClassLoc, ClassName, ClassNameLoc,
                                           { }, GenericParams, CurDeclContext,
                                           isExplicitActorDecl);
-  recordLocalType(CD);
   CD->getAttrs() = Attributes;
 
   // Parsed classes never have missing vtable entries.
@@ -9425,7 +9412,6 @@ parseDeclProtocol(ParseDeclOptions Flags, DeclAttributes &Attributes) {
       ProtocolDecl(CurDeclContext, ProtocolLoc, NameLoc, ProtocolName,
                    Context.AllocateCopy(PrimaryAssociatedTypeNames),
                    Context.AllocateCopy(InheritedProtocols), TrailingWhere);
-  recordLocalType(Proto);
 
   Proto->getAttrs() = Attributes;
   if (whereClauseHadCodeCompletion && CodeCompletionCallbacks)

lib/SILGen/SILGen.cpp

@@ -2057,7 +2057,7 @@ class SILGenModuleRAII {
       SGM.visit(D);
     }
 
-    for (TypeDecl *TD : sf->LocalTypeDecls) {
+    for (TypeDecl *TD : sf->getLocalTypeDecls()) {
       FrontendStatsTracer StatsTracer(SGM.getASTContext().Stats,
                                       "SILgen-tydecl", TD);
       // FIXME: Delayed parsing would prevent these types from being added to