ASTGen: Translate associated type declarations

AnthonyLatsis · AnthonyLatsis · commit b6be6da27759 · 2023-09-12T20:37:50.000+03:00
Plus tweak `DefaultDefinitionTypeRequest` caching to support querying the
cached type when dumping. This fixes a crash where type computation is
triggered in the dumper before import resolution in `-dump-parse` mode.
diff --git a/include/swift/AST/ASTContext.h b/include/swift/AST/ASTContext.h
@@ -1347,11 +1347,14 @@ class ASTContext final {
   getInheritedConformance(Type type, ProtocolConformance *inherited);
 
   /// Get the lazy data for the given declaration.
+  LazyContextData *getLazyContextData(const Decl *decl) const;
+
+  /// Get or otherwise allocate the lazy data for the given declaration.
   ///
   /// \param lazyLoader If non-null, the lazy loader to use when creating the
   /// lazy data. The pointer must either be null or be consistent
-  /// across all calls for the same \p func.
-  LazyContextData *getOrCreateLazyContextData(const DeclContext *decl,
+  /// across all calls for the same \p decl.
+  LazyContextData *getOrCreateLazyContextData(const Decl *decl,
                                               LazyMemberLoader *lazyLoader);
 
   /// Get the lazy iterable context for the given iterable declaration context.
diff --git a/include/swift/AST/CASTBridging.h b/include/swift/AST/CASTBridging.h
@@ -372,6 +372,14 @@ BridgedDeclContextAndDecl ProtocolDecl_create(
     BridgedSourceLoc cNameLoc, BridgedArrayRef cPrimaryAssociatedTypeNames,
     BridgedArrayRef cInheritedTypes, BridgedSourceRange cBraceRange);
 
+void *AssociatedTypeDecl_create(BridgedASTContext cContext,
+                                BridgedDeclContext cDeclContext,
+                                BridgedSourceLoc cAssociatedtypeKeywordLoc,
+                                BridgedIdentifier cName,
+                                BridgedSourceLoc cNameLoc,
+                                BridgedArrayRef cInheritedTypes,
+                                void *_Nullable opaqueDefaultType);
+
 void *GenericParamList_create(BridgedASTContext cContext,
                               BridgedSourceLoc cLAngleLoc,
                               BridgedArrayRef params,
diff --git a/include/swift/AST/Decl.h b/include/swift/AST/Decl.h
@@ -548,6 +548,11 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl> {
     IsOpaqueType : 1
   );
 
+  SWIFT_INLINE_BITFIELD_FULL(AssociatedTypeDecl, TypeDecl, 1,
+    /// Whether we have computed the default type.
+    IsDefaultDefinitionTypeComputed : 1
+  );
+
   SWIFT_INLINE_BITFIELD_EMPTY(GenericTypeDecl, TypeDecl);
 
   SWIFT_INLINE_BITFIELD(TypeAliasDecl, GenericTypeDecl, 1+1,
@@ -3677,24 +3682,28 @@ class AssociatedTypeDecl : public TypeDecl {
   SourceLoc KeywordLoc;
 
   /// The default definition.
-  TypeRepr *DefaultDefinition;
+  TypeLoc DefaultDefinition;
 
   /// The where clause attached to the associated type.
   TrailingWhereClause *TrailingWhere;
 
-  LazyMemberLoader *Resolver = nullptr;
-  uint64_t ResolverContextData;
-
   friend class DefaultDefinitionTypeRequest;
 
-public:
   AssociatedTypeDecl(DeclContext *dc, SourceLoc keywordLoc, Identifier name,
                      SourceLoc nameLoc, TypeRepr *defaultDefinition,
                      TrailingWhereClause *trailingWhere);
-  AssociatedTypeDecl(DeclContext *dc, SourceLoc keywordLoc, Identifier name,
-                     SourceLoc nameLoc, TrailingWhereClause *trailingWhere,
-                     LazyMemberLoader *definitionResolver,
-                     uint64_t resolverData);
+
+public:
+  static AssociatedTypeDecl *createParsed(ASTContext &ctx, DeclContext *dc,
+                                          SourceLoc keywordLoc, Identifier name,
+                                          SourceLoc nameLoc,
+                                          TypeRepr *defaultDefinition,
+                                          TrailingWhereClause *trailingWhere);
+
+  static AssociatedTypeDecl *createDeserialized(
+      ASTContext &ctx, DeclContext *dc, SourceLoc keywordLoc, Identifier name,
+      SourceLoc nameLoc, TrailingWhereClause *trailingWhere,
+      LazyMemberLoader *lazyLoader, uint64_t defaultDefinitionTypeData);
 
   /// Get the protocol in which this associated type is declared.
   ProtocolDecl *getProtocol() const {
@@ -3705,15 +3714,25 @@ class AssociatedTypeDecl : public TypeDecl {
   bool hasDefaultDefinitionType() const {
     // If we have a TypeRepr, return true immediately without kicking off
     // a request.
-    return DefaultDefinition || getDefaultDefinitionType();
+    return DefaultDefinition.getTypeRepr() || getDefaultDefinitionType();
   }
 
   /// Retrieve the default definition type.
   Type getDefaultDefinitionType() const;
 
+  /// Retrieve the default definition type if computed, `None` otherwise.
+  ///
+  /// \Note Should only be used for dumping.
+  llvm::Optional<Type> getCachedDefaultDefinitionType() const;
+
+private:
+  /// Set the computed default definition type.
+  void setDefaultDefinitionType(Type ty);
+
+public:
   /// Retrieve the default definition as written in the source.
   TypeRepr *getDefaultDefinitionTypeRepr() const {
-    return DefaultDefinition;
+    return DefaultDefinition.getTypeRepr();
   }
 
   /// Retrieve the trailing where clause for this associated type, if any.
diff --git a/include/swift/AST/LazyResolver.h b/include/swift/AST/LazyResolver.h
@@ -57,6 +57,12 @@ class LazyIterableDeclContextData : public LazyContextData {
   uint64_t allConformancesData = 0;
 };
 
+class LazyAssociatedTypeData : public LazyContextData {
+public:
+  /// The context data used for loading the default type.
+  uint64_t defaultDefinitionTypeData = 0;
+};
+
 /// Context data for protocols.
 class LazyProtocolData : public LazyIterableDeclContextData {
 public:
diff --git a/include/swift/AST/TypeCheckRequests.h b/include/swift/AST/TypeCheckRequests.h
@@ -528,10 +528,10 @@ class RequirementSignatureRequest :
 };
 
 /// Compute the default definition type of an associated type.
-class DefaultDefinitionTypeRequest :
-    public SimpleRequest<DefaultDefinitionTypeRequest,
-                         Type(AssociatedTypeDecl *),
-                         RequestFlags::Cached> {
+class DefaultDefinitionTypeRequest
+    : public SimpleRequest<DefaultDefinitionTypeRequest,
+                           Type(AssociatedTypeDecl *),
+                           RequestFlags::SeparatelyCached> {
 public:
   using SimpleRequest::SimpleRequest;
 
@@ -544,6 +544,8 @@ class DefaultDefinitionTypeRequest :
 public:
   // Caching.
   bool isCached() const { return true; }
+  llvm::Optional<Type> getCachedResult() const;
+  void cacheResult(Type value) const;
 };
 
 /// Describes the owner of a where clause, from which we can extract
diff --git a/lib/AST/ASTContext.cpp b/lib/AST/ASTContext.cpp
@@ -341,7 +341,7 @@ struct ASTContext::Implementation {
     DelayedMissingWitnesses;
 
   /// Stores information about lazy deserialization of various declarations.
-  llvm::DenseMap<const DeclContext *, LazyContextData *> LazyContexts;
+  llvm::DenseMap<const Decl *, LazyContextData *> LazyContexts;
 
   /// A fake generic parameter list <Self> for parsing @opened archetypes
   /// in textual SIL.
@@ -2829,23 +2829,30 @@ PackConformance *PackConformance::get(PackType *conformingType,
   return result;
 }
 
-LazyContextData *ASTContext::getOrCreateLazyContextData(
-                                                const DeclContext *dc,
-                                                LazyMemberLoader *lazyLoader) {
-  LazyContextData *&entry = getImpl().LazyContexts[dc];
-  if (entry) {
+LazyContextData *ASTContext::getLazyContextData(const Decl *decl) const {
+  return getImpl().LazyContexts.lookup(decl);
+}
+
+LazyContextData *
+ASTContext::getOrCreateLazyContextData(const Decl *decl,
+                                       LazyMemberLoader *lazyLoader) {
+  if (auto *data = getLazyContextData(decl)) {
     // Make sure we didn't provide an incompatible lazy loader.
-    assert(!lazyLoader || lazyLoader == entry->loader);
-    return entry;
+    assert(!lazyLoader || lazyLoader == data->loader);
+    return data;
   }
 
+  LazyContextData *&entry = getImpl().LazyContexts[decl];
+
   // Create new lazy context data with the given loader.
   assert(lazyLoader && "Queried lazy data for non-lazy iterable context");
-  if (isa<ProtocolDecl>(dc))
+  if (isa<ProtocolDecl>(decl)) {
     entry = Allocate<LazyProtocolData>();
-  else {
-    assert(isa<NominalTypeDecl>(dc) || isa<ExtensionDecl>(dc));
+  } else if (isa<NominalTypeDecl>(decl) || isa<ExtensionDecl>(decl)) {
     entry = Allocate<LazyIterableDeclContextData>();
+  } else {
+    assert(isa<AssociatedTypeDecl>(decl));
+    entry = Allocate<LazyAssociatedTypeData>();
   }
 
   entry->loader = lazyLoader;
diff --git a/lib/AST/ASTDumper.cpp b/lib/AST/ASTDumper.cpp
@@ -1110,9 +1110,14 @@ namespace {
 
     void visitAssociatedTypeDecl(AssociatedTypeDecl *decl, StringRef label) {
       printCommon(decl, "associated_type_decl", label);
-      if (auto defaultDef = decl->getDefaultDefinitionType()) {
-        printFieldQuoted(defaultDef, "default");
+
+      StringRef fieldName("default");
+      if (auto defaultDef = decl->getCachedDefaultDefinitionType()) {
+        printFieldQuoted(*defaultDef, fieldName);
+      } else {
+        printField("<not computed>", fieldName);
       }
+
       printWhereRequirements(decl);
       if (decl->overriddenDeclsComputed()) {
         printFieldQuotedRaw([&](raw_ostream &OS) {
diff --git a/lib/AST/CASTBridging.cpp b/lib/AST/CASTBridging.cpp
@@ -625,6 +625,26 @@ BridgedDeclContextAndDecl ProtocolDecl_create(
   return {bridgeDeclContext(decl), static_cast<Decl *>(decl)};
 }
 
+void *AssociatedTypeDecl_create(BridgedASTContext cContext,
+                                BridgedDeclContext cDeclContext,
+                                BridgedSourceLoc cAssociatedtypeKeywordLoc,
+                                BridgedIdentifier cName,
+                                BridgedSourceLoc cNameLoc,
+                                BridgedArrayRef cInheritedTypes,
+                                void *_Nullable opaqueDefaultType) {
+  ASTContext &context = convertASTContext(cContext);
+
+  auto *decl = AssociatedTypeDecl::createParsed(
+      context, convertDeclContext(cDeclContext),
+      convertSourceLoc(cAssociatedtypeKeywordLoc), convertIdentifier(cName),
+      convertSourceLoc(cNameLoc), static_cast<TypeRepr *>(opaqueDefaultType),
+      nullptr);
+  decl->setInherited(
+      context.AllocateCopy(convertToInheritedEntries(cInheritedTypes)));
+
+  return static_cast<Decl *>(decl);
+}
+
 void *OptionalTypeRepr_create(BridgedASTContext cContext, void *base,
                               BridgedSourceLoc cQuestionLoc) {
   ASTContext &context = convertASTContext(cContext);
diff --git a/lib/AST/Decl.cpp b/lib/AST/Decl.cpp
@@ -5234,18 +5234,48 @@ AssociatedTypeDecl::AssociatedTypeDecl(DeclContext *dc, SourceLoc keywordLoc,
                                        TrailingWhereClause *trailingWhere)
     : TypeDecl(DeclKind::AssociatedType, dc, name, nameLoc, { }),
       KeywordLoc(keywordLoc), DefaultDefinition(defaultDefinition),
-      TrailingWhere(trailingWhere) {}
+      TrailingWhere(trailingWhere) {
+  Bits.AssociatedTypeDecl.IsDefaultDefinitionTypeComputed = false;
+}
 
-AssociatedTypeDecl::AssociatedTypeDecl(DeclContext *dc, SourceLoc keywordLoc,
-                                       Identifier name, SourceLoc nameLoc,
-                                       TrailingWhereClause *trailingWhere,
-                                       LazyMemberLoader *definitionResolver,
-                                       uint64_t resolverData)
-    : TypeDecl(DeclKind::AssociatedType, dc, name, nameLoc, { }),
-      KeywordLoc(keywordLoc), DefaultDefinition(nullptr),
-      TrailingWhere(trailingWhere), Resolver(definitionResolver),
-      ResolverContextData(resolverData) {
-  assert(Resolver && "missing resolver");
+AssociatedTypeDecl *
+AssociatedTypeDecl::createParsed(ASTContext &ctx, DeclContext *dc,
+                                 SourceLoc keywordLoc, Identifier name,
+                                 SourceLoc nameLoc, TypeRepr *defaultDefinition,
+                                 TrailingWhereClause *trailingWhere) {
+  auto *decl = new (ctx) AssociatedTypeDecl(dc, keywordLoc, name, nameLoc,
+                                            defaultDefinition, trailingWhere);
+
+  // Sort out this trivial case now to enable the AST dumper to differentiate
+  // between a nonexistent and null default type without having to trigger a
+  // request.
+  if (!defaultDefinition)
+    decl->setDefaultDefinitionType(nullptr);
+
+  return decl;
+}
+
+AssociatedTypeDecl *AssociatedTypeDecl::createDeserialized(
+    ASTContext &ctx, DeclContext *dc, SourceLoc keywordLoc, Identifier name,
+    SourceLoc nameLoc, TrailingWhereClause *trailingWhere,
+    LazyMemberLoader *lazyLoader, uint64_t defaultDefinitionTypeData) {
+  assert(lazyLoader && "missing lazy member loader");
+  auto *decl = new (ctx)
+      AssociatedTypeDecl(dc, keywordLoc, name, nameLoc,
+                         /*defaultDefinition*/ nullptr, trailingWhere);
+
+  // Sort out this trivial case now to enable the AST dumper to differentiate
+  // between a nonexistent and null default type without having to trigger a
+  // request. '0' is the sentinel ID for no data.
+  if (defaultDefinitionTypeData == 0) {
+    decl->setDefaultDefinitionType(nullptr);
+  } else {
+    auto *data = static_cast<LazyAssociatedTypeData *>(
+        ctx.getOrCreateLazyContextData(decl, lazyLoader));
+    data->defaultDefinitionTypeData = defaultDefinitionTypeData;
+  }
+
+  return decl;
 }
 
 Type AssociatedTypeDecl::getDefaultDefinitionType() const {
@@ -5254,6 +5284,19 @@ Type AssociatedTypeDecl::getDefaultDefinitionType() const {
            Type());
 }
 
+llvm::Optional<Type>
+AssociatedTypeDecl::getCachedDefaultDefinitionType() const {
+  if (Bits.AssociatedTypeDecl.IsDefaultDefinitionTypeComputed)
+    return DefaultDefinition.getType();
+
+  return llvm::None;
+}
+
+void AssociatedTypeDecl::setDefaultDefinitionType(Type ty) {
+  DefaultDefinition.setType(ty);
+  Bits.AssociatedTypeDecl.IsDefaultDefinitionTypeComputed = true;
+}
+
 SourceRange AssociatedTypeDecl::getSourceRange() const {
   SourceLoc endLoc;
   if (auto TWC = getTrailingWhereClause()) {
diff --git a/lib/AST/TypeCheckRequests.cpp b/lib/AST/TypeCheckRequests.cpp
@@ -375,6 +375,20 @@ void RequirementSignatureRequest::cacheResult(RequirementSignature value) const
   proto->setRequirementSignature(value);
 }
 
+//----------------------------------------------------------------------------//
+// DefaultDefinitionTypeRequest computation.
+//----------------------------------------------------------------------------//
+
+llvm::Optional<Type> DefaultDefinitionTypeRequest::getCachedResult() const {
+  auto *decl = std::get<0>(getStorage());
+  return decl->getCachedDefaultDefinitionType();
+}
+
+void DefaultDefinitionTypeRequest::cacheResult(Type value) const {
+  auto *decl = std::get<0>(getStorage());
+  decl->setDefaultDefinitionType(value);
+}
+
 //----------------------------------------------------------------------------//
 // Requirement computation.
 //----------------------------------------------------------------------------//
diff --git a/lib/ASTGen/Sources/ASTGen/Decls.swift b/lib/ASTGen/Sources/ASTGen/Decls.swift
@@ -89,6 +89,22 @@ extension ASTGenVisitor {
 
     return .decl(decl.asDecl)
   }
+
+  func visit(_ node: AssociatedTypeDeclSyntax) -> ASTNode {
+    let (name, nameLoc) = node.name.bridgedIdentifierAndSourceLoc(in: self)
+
+    return .decl(
+      AssociatedTypeDecl_create(
+        self.ctx,
+        self.declContext,
+        self.bridgedSourceLoc(for: node.associatedtypeKeyword),
+        name,
+        nameLoc,
+        self.visit(node.inheritanceClause?.inheritedTypes),
+        self.visit(node.initializer?.value)?.rawValue
+      )
+    )
+  }
 }
 
 extension ASTGenVisitor {
diff --git a/lib/Parse/ParseDecl.cpp b/lib/Parse/ParseDecl.cpp
@@ -7051,9 +7051,9 @@ ParserResult<TypeDecl> Parser::parseDeclAssociatedType(Parser::ParseDeclOptions
     return Status;
   }
 
-  auto assocType = new (Context)
-      AssociatedTypeDecl(CurDeclContext, AssociatedTypeLoc, Id, IdLoc,
-                         UnderlyingTy.getPtrOrNull(), TrailingWhere);
+  auto assocType = AssociatedTypeDecl::createParsed(
+      Context, CurDeclContext, AssociatedTypeLoc, Id, IdLoc,
+      UnderlyingTy.getPtrOrNull(), TrailingWhere);
   assocType->getAttrs() = Attributes;
   if (!Inherited.empty())
     assocType->setInherited(Context.AllocateCopy(Inherited));
diff --git a/lib/Sema/TypeCheckDecl.cpp b/lib/Sema/TypeCheckDecl.cpp
@@ -1024,11 +1024,11 @@ IsDynamicRequest::evaluate(Evaluator &evaluator, ValueDecl *decl) const {
 Type
 DefaultDefinitionTypeRequest::evaluate(Evaluator &evaluator,
                                        AssociatedTypeDecl *assocType) const {
-  if (assocType->Resolver) {
-    auto defaultType = assocType->Resolver->loadAssociatedTypeDefault(
-                                    assocType, assocType->ResolverContextData);
-    assocType->Resolver = nullptr;
-    return defaultType;
+  auto &ctx = assocType->getASTContext();
+  if (auto *data = static_cast<LazyAssociatedTypeData *>(
+          ctx.getLazyContextData(assocType))) {
+    return data->loader->loadAssociatedTypeDefault(
+        assocType, data->defaultDefinitionTypeData);
   }
 
   TypeRepr *defaultDefinition = assocType->getDefaultDefinitionTypeRepr();
diff --git a/lib/Serialization/Deserialization.cpp b/lib/Serialization/Deserialization.cpp
diff --git a/test/ASTGen/verify-parse.swift b/test/ASTGen/verify-parse.swift
