ASTGen: Cache parsed IterableDeclContext members using the request evaluator

...instead of adding them to the iterable context lest they be re-parsed using
the legacy parser and later re-added through `IterableDeclContext::loadAllMembers`.

Author: AnthonyLatsis

include/swift/AST/CASTBridging.h

@@ -202,9 +202,8 @@ typedef struct BridgedFuncDecl {
 } BridgedFuncDecl;
 
 typedef struct BridgedDeclContextAndDecl {
-  BridgedDeclContext declContext;
-  void *nominalDecl;
-  void *decl;
+  BridgedDeclContext asDeclContext;
+  void *asDecl;
 } BridgedDeclContextAndDecl;
 
 typedef struct BridgedTypeAttributes {
@@ -338,7 +337,8 @@ void *TypeAliasDecl_create(
     BridgedSourceLoc cNameLoc, void *_Nullable opaqueGenericParamList,
     BridgedSourceLoc cEqualLoc, void *opaqueUnderlyingType);
 
-void NominalTypeDecl_setMembers(void *decl, BridgedArrayRef members);
+void IterableDeclContext_setParsedMembers(BridgedArrayRef members,
+                                          void *opaqueDecl);
 
 BridgedDeclContextAndDecl
 StructDecl_create(BridgedASTContext cContext, BridgedDeclContext cDeclContext,

lib/AST/CASTBridging.cpp

@@ -8,6 +8,7 @@
 #include "swift/AST/GenericParamList.h"
 #include "swift/AST/Identifier.h"
 #include "swift/AST/ParameterList.h"
+#include "swift/AST/ParseRequests.h"
 #include "swift/AST/Pattern.h"
 #include "swift/AST/PluginRegistry.h"
 #include "swift/AST/Stmt.h"
@@ -517,10 +518,15 @@ void *TypeAliasDecl_create(
   return static_cast<Decl *>(decl);
 }
 
-void NominalTypeDecl_setMembers(void *decl, BridgedArrayRef members) {
-  auto declMembers = convertArrayRef<Decl *>(members);
-  for (auto m : declMembers)
-    ((NominalTypeDecl *)decl)->addMember(m);
+void IterableDeclContext_setParsedMembers(BridgedArrayRef bridgedMembers,
+                                          void *opaqueDecl) {
+  auto *decl = static_cast<Decl *>(opaqueDecl);
+  auto &ctx = decl->getASTContext();
+  auto members = convertArrayRef<Decl *>(bridgedMembers);
+
+  ctx.evaluator.cacheOutput(
+      ParseMembersRequest{cast<IterableDeclContext>(decl)},
+      FingerprintAndMembers{llvm::None, ctx.AllocateCopy(members)});
 }
 
 BridgedDeclContextAndDecl
@@ -530,14 +536,13 @@ StructDecl_create(BridgedASTContext cContext, BridgedDeclContext cDeclContext,
                   void *_Nullable opaqueGenericParamList) {
   ASTContext &context = convertASTContext(cContext);
 
-  auto *out = new (context)
+  auto *decl = new (context)
       StructDecl(convertSourceLoc(cStructKeywordLoc), convertIdentifier(cName),
                  convertSourceLoc(cNameLoc), {},
                  static_cast<GenericParamList *>(opaqueGenericParamList),
                  convertDeclContext(cDeclContext));
 
-  return {bridgeDeclContext(out), static_cast<NominalTypeDecl *>(out),
-          static_cast<Decl *>(out)};
+  return {bridgeDeclContext(decl), static_cast<Decl *>(decl)};
 }
 
 BridgedDeclContextAndDecl ClassDecl_create(BridgedASTContext cContext,
@@ -547,13 +552,12 @@ BridgedDeclContextAndDecl ClassDecl_create(BridgedASTContext cContext,
                                            BridgedSourceLoc cNameLoc) {
   ASTContext &context = convertASTContext(cContext);
 
-  auto *out = new (context)
+  auto *decl = new (context)
       ClassDecl(convertSourceLoc(cClassKeywordLoc), convertIdentifier(cName),
                 convertSourceLoc(cNameLoc), {}, nullptr,
                 convertDeclContext(cDeclContext), false);
 
-  return {bridgeDeclContext(out), static_cast<NominalTypeDecl *>(out),
-          static_cast<Decl *>(out)};
+  return {bridgeDeclContext(decl), static_cast<Decl *>(decl)};
 }
 
 void *OptionalTypeRepr_create(BridgedASTContext cContext, void *base,

lib/ASTGen/Sources/ASTGen/ASTGen.swift

@@ -68,7 +68,7 @@ struct ASTGenVisitor: SyntaxTransformVisitor {
 
   let base: UnsafeBufferPointer<UInt8>
 
-  @Boxed var declContext: BridgedDeclContext
+  @Boxed private(set) var declContext: BridgedDeclContext
 
   let ctx: BridgedASTContext
 
@@ -123,6 +123,17 @@ struct ASTGenVisitor: SyntaxTransformVisitor {
   }
 }
 
+extension ASTGenVisitor {
+  /// Replaces the current declaration context with `declContext` for the duration of its execution, and calls `body`.
+  @inline(__always)
+  func withDeclContext(_ declContext: BridgedDeclContext, _ body: () -> Void) {
+    let oldDeclContext = self.declContext
+    self.declContext = declContext
+    body()
+    self.declContext = oldDeclContext
+  }
+}
+
 extension ASTGenVisitor {
   /// Emits the given diagnostic via the C++ diagnostic engine.
   @inline(__always)

lib/ASTGen/Sources/ASTGen/Decls.swift

@@ -4,6 +4,8 @@ import CASTBridging
 @_spi(SyntaxTransformVisitor)
 import SwiftSyntax
 
+// MARK: - TypeDecl
+
 extension ASTGenVisitor {
   public func visit(_ node: TypeAliasDeclSyntax) -> ASTNode {
     let (name, nameLoc) = node.name.bridgedIdentifierAndSourceLoc(in: self)
@@ -25,7 +27,7 @@ extension ASTGenVisitor {
   public func visit(_ node: StructDeclSyntax) -> ASTNode {
     let (name, nameLoc) = node.name.bridgedIdentifierAndSourceLoc(in: self)
 
-    let out = StructDecl_create(
+    let decl = StructDecl_create(
       self.ctx,
       self.declContext,
       self.bridgedSourceLoc(for: node.structKeyword),
@@ -34,34 +36,33 @@ extension ASTGenVisitor {
       self.visit(node.genericParameterClause)?.rawValue
     )
 
-    let oldDeclContext = declContext
-    declContext = out.declContext
-    defer { declContext = oldDeclContext }
-
-    NominalTypeDecl_setMembers(out.nominalDecl, self.visit(node.memberBlock.members))
+    self.withDeclContext(decl.asDeclContext) {
+      IterableDeclContext_setParsedMembers(self.visit(node.memberBlock.members), decl.asDecl)
+    }
 
-    return .decl(out.decl)
+    return .decl(decl.asDecl)
   }
 
   public func visit(_ node: ClassDeclSyntax) -> ASTNode {
     let (name, nameLoc) = node.name.bridgedIdentifierAndSourceLoc(in: self)
 
-    let out = ClassDecl_create(
+    let decl = ClassDecl_create(
       self.ctx,
       self.declContext,
       self.bridgedSourceLoc(for: node.classKeyword),
       name,
       nameLoc
     )
-    let oldDeclContext = declContext
-    declContext = out.declContext
-    defer { declContext = oldDeclContext }
 
-    NominalTypeDecl_setMembers(out.nominalDecl, self.visit(node.memberBlock.members))
+    self.withDeclContext(decl.asDeclContext) {
+      IterableDeclContext_setParsedMembers(self.visit(node.memberBlock.members), decl.asDecl)
+    }
 
-    return .decl(out.decl)
+    return .decl(decl.asDecl)
   }
+}
 
+extension ASTGenVisitor {
   public func visit(_ node: VariableDeclSyntax) -> ASTNode {
     let pattern = visit(node.bindings.first!.pattern).rawValue
     let initializer = visit(node.bindings.first!.initializer!).rawValue
@@ -131,19 +132,10 @@ extension ASTGenVisitor {
       self.visit(node.signature.returnClause?.type)?.rawValue
     )
 
-    let oldDeclContext = declContext
-    declContext = out.declContext
-    defer { declContext = oldDeclContext }
-
-    let body: ASTNode?
-    if let nodeBody = node.body {
-      body = visit(nodeBody)
-    } else {
-      body = nil
-    }
-
-    if let body = body {
-      FuncDecl_setBody(out.funcDecl, body.rawValue)
+    if let body = node.body {
+      self.withDeclContext(out.declContext) {
+        FuncDecl_setBody(out.funcDecl, self.visit(body).rawValue)
+      }
     }
 
     return .decl(out.decl)