Serialization: Lazily deserialize OpaqueTypeDecl's underlying substitutions

We need to serialize the underlying type substitution map for an
inlinable function. However, there is no reason to deserialize it
eagerly, since doing so can lead to cycles. It is better for
correctness and performance to only deserialize it when needed.

Technically this fixes a regression from #84299, but the actual
problem was there all along, it was just exposed by my change
on a specific project.

Fixes rdar://163301203.

Author: slavapestov

include/swift/AST/Decl.h

@@ -761,6 +761,11 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl>, public Swi
     HasAnyUnavailableDuringLoweringValues : 1
   );
 
+  SWIFT_INLINE_BITFIELD(OpaqueTypeDecl, GenericTypeDecl, 1,
+    /// Whether we have lazily-loaded underlying substitutions.
+    HasLazyUnderlyingSubstitutions : 1
+  );
+
   SWIFT_INLINE_BITFIELD(ModuleDecl, TypeDecl, 1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+1+8,
     /// If the module is compiled as static library.
     StaticLibrary : 1,
@@ -3646,7 +3651,8 @@ class OpaqueTypeDecl final :
   OpaqueTypeDecl(ValueDecl *NamingDecl, GenericParamList *GenericParams,
                  DeclContext *DC,
                  GenericSignature OpaqueInterfaceGenericSignature,
-                 ArrayRef<TypeRepr *> OpaqueReturnTypeReprs);
+                 ArrayRef<TypeRepr *> OpaqueReturnTypeReprs,
+                 bool hasLazyUnderlyingSubstitutions);
 
   unsigned getNumOpaqueReturnTypeReprs() const {
     return NamingDeclAndHasOpaqueReturnTypeRepr.getInt()
@@ -3658,13 +3664,21 @@ class OpaqueTypeDecl final :
     return getNumOpaqueReturnTypeReprs();
   }
 
+  void loadLazyUnderlyingSubstitutions();
+
 public:
-  static OpaqueTypeDecl *get(
+  static OpaqueTypeDecl *create(
       ValueDecl *NamingDecl, GenericParamList *GenericParams,
       DeclContext *DC,
       GenericSignature OpaqueInterfaceGenericSignature,
       ArrayRef<TypeRepr *> OpaqueReturnTypeReprs);
 
+  static OpaqueTypeDecl *createDeserialized(
+      GenericParamList *GenericParams,
+      DeclContext *DC,
+      GenericSignature OpaqueInterfaceGenericSignature,
+      LazyMemberLoader *lazyLoader, uint64_t underlyingSubsData);
+
   ValueDecl *getNamingDecl() const {
     return NamingDeclAndHasOpaqueReturnTypeRepr.getPointer();
   }
@@ -3720,19 +3734,19 @@ class OpaqueTypeDecl final :
       bool typeCheckFunctionBodies=true) const;
 
   void setUniqueUnderlyingTypeSubstitutions(SubstitutionMap subs) {
-    assert(!UniqueUnderlyingType.has_value() && "resetting underlying type?!");
+    ASSERT(!Bits.OpaqueTypeDecl.HasLazyUnderlyingSubstitutions);
+    ASSERT(!UniqueUnderlyingType.has_value() && "resetting underlying type?!");
     UniqueUnderlyingType = subs;
   }
 
   bool hasConditionallyAvailableSubstitutions() const {
+    const_cast<OpaqueTypeDecl *>(this)->loadLazyUnderlyingSubstitutions();
+
     return ConditionallyAvailableTypes.has_value();
   }
 
   ArrayRef<ConditionallyAvailableSubstitutions *>
-  getConditionallyAvailableSubstitutions() const {
-    assert(ConditionallyAvailableTypes);
-    return ConditionallyAvailableTypes.value();
-  }
+  getConditionallyAvailableSubstitutions() const;
 
   void setConditionallyAvailableSubstitutions(
       ArrayRef<ConditionallyAvailableSubstitutions *> substitutions);

include/swift/AST/LazyResolver.h

@@ -63,6 +63,12 @@ class LazyAssociatedTypeData : public LazyContextData {
   uint64_t defaultDefinitionTypeData = 0;
 };
 
+class LazyOpaqueTypeData : public LazyContextData {
+public:
+  /// The context data used for loading the underlying type substitution map.
+  uint64_t underlyingSubsData = 0;
+};
+
 /// Context data for protocols.
 class LazyProtocolData : public LazyIterableDeclContextData {
 public:
@@ -138,6 +144,10 @@ class alignas(void*) LazyMemberLoader {
   // Returns the target parameter of the `@_specialize` attribute or null.
   virtual ValueDecl *loadTargetFunctionDecl(const AbstractSpecializeAttr *attr,
                                             uint64_t contextData) = 0;
+
+  /// Loads the underlying type substitution map of an opaque result declaration.
+  virtual void
+  finishOpaqueTypeDecl(OpaqueTypeDecl *opaqueDecl, uint64_t contextData) = 0;
 };
 
 /// A class that can lazily load conformances from a serialized format.

lib/AST/ASTContext.cpp

@@ -3174,20 +3174,22 @@ 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 == data->loader);
+    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");
+  ASSERT(lazyLoader && "Queried lazy data for non-lazy iterable context");
   if (isa<ProtocolDecl>(decl)) {
     entry = Allocate<LazyProtocolData>();
   } else if (isa<NominalTypeDecl>(decl) || isa<ExtensionDecl>(decl)) {
     entry = Allocate<LazyIterableDeclContextData>();
+  } else if (isa<OpaqueTypeDecl>(decl)) {
+    entry = Allocate<LazyOpaqueTypeData>();
   } else {
-    assert(isa<AssociatedTypeDecl>(decl));
+    ASSERT(isa<AssociatedTypeDecl>(decl));
     entry = Allocate<LazyAssociatedTypeData>();
   }
 

lib/AST/Decl.cpp

@@ -10810,13 +10810,16 @@ bool AbstractFunctionDecl::isResilient(ModuleDecl *M,
 OpaqueTypeDecl::OpaqueTypeDecl(ValueDecl *NamingDecl,
                                GenericParamList *GenericParams, DeclContext *DC,
                                GenericSignature OpaqueInterfaceGenericSignature,
-                               ArrayRef<TypeRepr *>
-                                   OpaqueReturnTypeReprs)
+                               ArrayRef<TypeRepr *> OpaqueReturnTypeReprs,
+                               bool hasLazyUnderlyingSubstitutions)
     : GenericTypeDecl(DeclKind::OpaqueType, DC, Identifier(), SourceLoc(), {},
                       GenericParams),
       NamingDeclAndHasOpaqueReturnTypeRepr(
         NamingDecl, !OpaqueReturnTypeReprs.empty()),
       OpaqueInterfaceGenericSignature(OpaqueInterfaceGenericSignature) {
+  Bits.OpaqueTypeDecl.HasLazyUnderlyingSubstitutions
+      = hasLazyUnderlyingSubstitutions;
+
   // Always implicit.
   setImplicit();
 
@@ -10829,7 +10832,7 @@ OpaqueTypeDecl::OpaqueTypeDecl(ValueDecl *NamingDecl,
                           OpaqueReturnTypeReprs.end(), getTrailingObjects());
 }
 
-OpaqueTypeDecl *OpaqueTypeDecl::get(
+OpaqueTypeDecl *OpaqueTypeDecl::create(
       ValueDecl *NamingDecl, GenericParamList *GenericParams,
       DeclContext *DC,
       GenericSignature OpaqueInterfaceGenericSignature,
@@ -10840,7 +10843,33 @@ OpaqueTypeDecl *OpaqueTypeDecl::get(
   auto mem = ctx.Allocate(size, alignof(OpaqueTypeDecl));
   return new (mem) OpaqueTypeDecl(
       NamingDecl, GenericParams, DC, OpaqueInterfaceGenericSignature,
-      OpaqueReturnTypeReprs);
+      OpaqueReturnTypeReprs, /*hasLazyUnderlyingSubstitutions=*/false);
+}
+
+OpaqueTypeDecl *OpaqueTypeDecl::createDeserialized(
+    GenericParamList *GenericParams, DeclContext *DC,
+    GenericSignature OpaqueInterfaceGenericSignature,
+    LazyMemberLoader *lazyLoader, uint64_t underlyingSubsData) {
+  bool hasLazyUnderlyingSubstitutions = (underlyingSubsData != 0);
+
+  ASTContext &ctx = DC->getASTContext();
+  auto size = totalSizeToAlloc<TypeRepr *>(0);
+  auto mem = ctx.Allocate(size, alignof(OpaqueTypeDecl));
+
+  // NamingDecl is set later by deserialization
+  auto *decl = new (mem) OpaqueTypeDecl(
+      /*namingDecl=*/nullptr, GenericParams, DC,
+      OpaqueInterfaceGenericSignature, { },
+      hasLazyUnderlyingSubstitutions);
+
+  if (hasLazyUnderlyingSubstitutions) {
+    auto &ctx = DC->getASTContext();
+    auto *data = static_cast<LazyOpaqueTypeData *>(
+          ctx.getOrCreateLazyContextData(decl, lazyLoader));
+    data->underlyingSubsData = underlyingSubsData;
+  }
+
+  return decl;
 }
 
 bool OpaqueTypeDecl::isOpaqueReturnTypeOf(const Decl *ownerDecl) const {
@@ -10896,16 +10925,41 @@ bool OpaqueTypeDecl::exportUnderlyingType() const {
   llvm_unreachable("The naming decl is expected to be either an AFD or ASD");
 }
 
+void OpaqueTypeDecl::loadLazyUnderlyingSubstitutions() {
+  if (!Bits.OpaqueTypeDecl.HasLazyUnderlyingSubstitutions)
+    return;
+
+  Bits.OpaqueTypeDecl.HasLazyUnderlyingSubstitutions = 0;
+
+  auto &ctx = getASTContext();
+  auto *data = static_cast<LazyOpaqueTypeData *>(
+          ctx.getLazyContextData(this));
+  ASSERT(data != nullptr);
+
+  data->loader->finishOpaqueTypeDecl(
+      this, data->underlyingSubsData);
+}
+
 std::optional<SubstitutionMap>
 OpaqueTypeDecl::getUniqueUnderlyingTypeSubstitutions(
     bool typeCheckFunctionBodies) const {
+  const_cast<OpaqueTypeDecl *>(this)->loadLazyUnderlyingSubstitutions();
+
   if (!typeCheckFunctionBodies)
     return UniqueUnderlyingType;
 
   return evaluateOrDefault(getASTContext().evaluator,
                            UniqueUnderlyingTypeSubstitutionsRequest{this}, {});
 }
 
+ArrayRef<OpaqueTypeDecl::ConditionallyAvailableSubstitutions *>
+OpaqueTypeDecl::getConditionallyAvailableSubstitutions() const {
+  const_cast<OpaqueTypeDecl *>(this)->loadLazyUnderlyingSubstitutions();
+
+  assert(ConditionallyAvailableTypes);
+  return ConditionallyAvailableTypes.value();
+}
+
 std::optional<unsigned>
 OpaqueTypeDecl::getAnonymousOpaqueParamOrdinal(TypeRepr *repr) const {
   assert(NamingDeclAndHasOpaqueReturnTypeRepr.getInt() &&
@@ -10935,7 +10989,8 @@ Identifier OpaqueTypeDecl::getOpaqueReturnTypeIdentifier() const {
 
 void OpaqueTypeDecl::setConditionallyAvailableSubstitutions(
     ArrayRef<ConditionallyAvailableSubstitutions *> substitutions) {
-  assert(!ConditionallyAvailableTypes &&
+  ASSERT(!Bits.OpaqueTypeDecl.HasLazyUnderlyingSubstitutions);
+  ASSERT(!ConditionallyAvailableTypes &&
          "resetting conditionally available substitutions?!");
   ConditionallyAvailableTypes = getASTContext().AllocateCopy(substitutions);
 }

lib/ClangImporter/ImporterImpl.h

@@ -1765,6 +1765,11 @@ class LLVM_LIBRARY_VISIBILITY ClangImporter::Implementation
     llvm_unreachable("unimplemented for ClangImporter");
   }
 
+  void finishOpaqueTypeDecl(OpaqueTypeDecl *decl,
+                            uint64_t contextData) override {
+    llvm_unreachable("unimplemented for ClangImporter");
+  }
+
   template <typename DeclTy, typename ...Targs>
   DeclTy *createDeclWithClangNode(ClangNode ClangN, AccessLevel access,
                                   Targs &&... Args) {

lib/Sema/TypeCheckGeneric.cpp

@@ -215,7 +215,7 @@ OpaqueResultTypeRequest::evaluate(Evaluator &evaluator,
   }
 
   // Create the OpaqueTypeDecl for the result type.
-  auto opaqueDecl = OpaqueTypeDecl::get(
+  auto opaqueDecl = OpaqueTypeDecl::create(
       originatingDecl, genericParams, parentDC, interfaceSignature,
       opaqueReprs);
   if (auto originatingSig = originatingDC->getGenericSignatureOfContext()) {

lib/Serialization/DeclTypeRecordNodes.def

@@ -215,10 +215,11 @@ OTHER(ERROR_FLAG, 155)
 OTHER(ABI_ONLY_COUNTERPART, 156)
 OTHER(DECL_NAME_REF, 157)
 
+TRAILING_INFO(UNDERLYING_SUBSTITUTION)
 TRAILING_INFO(CONDITIONAL_SUBSTITUTION)
 TRAILING_INFO(CONDITIONAL_SUBSTITUTION_COND)
 
-OTHER(LIFETIME_DEPENDENCE, 160)
+OTHER(LIFETIME_DEPENDENCE, 161)
 TRAILING_INFO(INHERITED_PROTOCOLS)
 
 #ifndef DECL_ATTR