Serialization: Preserve identity of opened generic environments

We used to create a new environment for each opened archetype, which is
incorrect when deserializing a nested type of another opened archetype.

include/swift/AST/GenericEnvironment.h

@@ -54,6 +54,7 @@ public:
 /// Extra data in a generic environment for an opened existential.
 struct OpenedGenericEnvironmentData {
   Type existential;
+  GenericSignature parentSig;
   UUID uuid;
 };
 
@@ -110,7 +111,8 @@ private:
 
   explicit GenericEnvironment(GenericSignature signature);
   explicit GenericEnvironment(
-      GenericSignature signature, Type existential, UUID uuid);
+      GenericSignature signature,
+      Type existential, GenericSignature parentSig, UUID uuid);
   explicit GenericEnvironment(
       GenericSignature signature, OpaqueTypeDecl *opaque, SubstitutionMap subs);
 
@@ -143,6 +145,9 @@ public:
   /// Retrieve the UUID for an opened existential environment.
   UUID getOpenedExistentialUUID() const;
 
+  /// Retrieve the parent signature for an opened existential environment.
+  GenericSignature getOpenedExistentialParentSignature() const;
+
   /// Retrieve the opaque type declaration for a generic environment describing
   /// opaque types.
   OpaqueTypeDecl *getOpaqueTypeDecl() const;
@@ -156,29 +161,12 @@ public:
 
   /// Create a new generic environment for an opened existential.
   ///
-  /// This function uses the provided parent signature to construct a new
-  /// signature suitable for use with an opened archetype. If you have an
-  /// existing generic signature from e.g. deserialization use
-  /// \c GenericEnvironment::forOpenedArchetypeSignature instead.
-  ///
   /// \param existential The subject existential type
   /// \param parentSig The signature of the context where this existential type is being opened
   /// \param uuid The unique identifier for this opened existential
   static GenericEnvironment *
   forOpenedExistential(Type existential, GenericSignature parentSig, UUID uuid);
 
-  /// Create a new generic environment for an opened existential.
-  ///
-  /// It is unlikely you want to use this function.
-  /// Call \c GenericEnvironment::forOpenedExistential instead.
-  ///
-  /// \param existential The subject existential type
-  /// \param signature The signature of the opened archetype
-  /// \param uuid The unique identifier for this opened existential
-  static GenericEnvironment *
-  forOpenedArchetypeSignature(Type existential,
-                              GenericSignature signature, UUID uuid);
-
   /// Create a new generic environment for an opaque type with the given set of
   /// outer substitutions.
   static GenericEnvironment *forOpaqueType(

lib/AST/ASTContext.cpp

@@ -4655,19 +4655,14 @@ GenericEnvironment::forOpenedExistential(
     Type existential, GenericSignature parentSig, UUID uuid) {
   auto &ctx = existential->getASTContext();
   auto signature = ctx.getOpenedExistentialSignature(existential, parentSig);
-  return GenericEnvironment::forOpenedArchetypeSignature(existential, signature, uuid);
-}
 
-GenericEnvironment *GenericEnvironment::forOpenedArchetypeSignature(
-    Type existential, GenericSignature signature, UUID uuid) {
   // Allocate and construct the new environment.
-  auto &ctx = existential->getASTContext();
   unsigned numGenericParams = signature.getGenericParams().size();
   size_t bytes = totalSizeToAlloc<OpaqueTypeDecl *, SubstitutionMap,
                                   OpenedGenericEnvironmentData, Type>(
       0, 0, 1, numGenericParams);
   void *mem = ctx.Allocate(bytes, alignof(GenericEnvironment));
-  return new (mem) GenericEnvironment(signature, existential, uuid);
+  return new (mem) GenericEnvironment(signature, existential, parentSig, uuid);
 }
 
 /// Create a new generic environment for an opaque type with the given set of

lib/AST/GenericEnvironment.cpp

@@ -104,6 +104,12 @@ UUID GenericEnvironment::getOpenedExistentialUUID() const {
   return getTrailingObjects<OpenedGenericEnvironmentData>()->uuid;
 }
 
+GenericSignature
+GenericEnvironment::getOpenedExistentialParentSignature() const {
+  assert(getKind() == Kind::OpenedExistential);
+  return getTrailingObjects<OpenedGenericEnvironmentData>()->parentSig;
+}
+
 GenericEnvironment::GenericEnvironment(GenericSignature signature)
   : SignatureAndKind(signature, Kind::Primary)
 {
@@ -113,11 +119,12 @@ GenericEnvironment::GenericEnvironment(GenericSignature signature)
 }
 
 GenericEnvironment::GenericEnvironment(
-    GenericSignature signature, Type existential, UUID uuid)
+    GenericSignature signature,
+    Type existential, GenericSignature parentSig, UUID uuid)
   : SignatureAndKind(signature, Kind::OpenedExistential)
 {
   new (getTrailingObjects<OpenedGenericEnvironmentData>())
-    OpenedGenericEnvironmentData{ existential, uuid };
+    OpenedGenericEnvironmentData{ existential, parentSig, uuid };
 
   // Clear out the memory that holds the context types.
   std::uninitialized_fill(getContextTypes().begin(), getContextTypes().end(),

lib/Serialization/DeclTypeRecordNodes.def

@@ -167,7 +167,7 @@ OTHER(TOP_LEVEL_CODE_DECL_CONTEXT, 112)
 OTHER(GENERIC_PARAM_LIST, 120)
 OTHER(GENERIC_SIGNATURE, 121)
 OTHER(REQUIREMENT_SIGNATURE, 122)
-// 123 is unused; was LAYOUT_REQUIREMENT
+OTHER(GENERIC_ENVIRONMENT, 123)
 OTHER(BUILTIN_PROTOCOL_CONFORMANCE, 124)
 OTHER(SIL_GENERIC_SIGNATURE, 125)
 OTHER(SUBSTITUTION_MAP, 126)

lib/Serialization/Deserialization.cpp

@@ -1188,6 +1188,53 @@ ModuleFile::getGenericSignatureChecked(serialization::GenericSignatureID ID) {
   return signature;
 }
 
+Expected<GenericEnvironment *>
+ModuleFile::getGenericEnvironmentChecked(serialization::GenericEnvironmentID ID) {
+  using namespace decls_block;
+
+  assert(ID <= GenericEnvironments.size() &&
+         "invalid GenericEnvironment ID");
+  auto &envOffset = GenericEnvironments[ID-1];
+
+  // If we've already deserialized this generic environment, return it.
+  if (envOffset.isComplete())
+    return envOffset.get();
+
+  // Read the generic environment.
+  BCOffsetRAII restoreOffset(DeclTypeCursor);
+  fatalIfNotSuccess(DeclTypeCursor.JumpToBit(envOffset));
+
+  llvm::BitstreamEntry entry =
+      fatalIfUnexpected(DeclTypeCursor.advance(AF_DontPopBlockAtEnd));
+  if (entry.Kind != llvm::BitstreamEntry::Record)
+    fatal();
+
+  StringRef blobData;
+  SmallVector<uint64_t, 8> scratch;
+  unsigned recordID = fatalIfUnexpected(
+      DeclTypeCursor.readRecord(entry.ID, scratch, &blobData));
+  if (recordID != GENERIC_ENVIRONMENT)
+    fatal();
+
+  GenericSignatureID parentSigID;
+  TypeID existentialID;
+  GenericEnvironmentLayout::readRecord(scratch, existentialID, parentSigID);
+
+  auto existentialTypeOrError = getTypeChecked(existentialID);
+  if (!existentialTypeOrError)
+    return existentialTypeOrError.takeError();
+
+  auto parentSigOrError = getGenericSignatureChecked(parentSigID);
+  if (!parentSigOrError)
+    return parentSigOrError.takeError();
+
+  auto *genericEnv = GenericEnvironment::forOpenedExistential(
+      existentialTypeOrError.get(), parentSigOrError.get(), UUID::fromTime());
+  envOffset = genericEnv;
+
+  return genericEnv;
+}
+
 SubstitutionMap ModuleFile::getSubstitutionMap(
                                         serialization::SubstitutionMapID id) {
   auto map = getSubstitutionMapChecked(id);
@@ -5747,29 +5794,22 @@ Expected<Type> DESERIALIZE_TYPE(PRIMARY_ARCHETYPE_TYPE)(
 
 Expected<Type> DESERIALIZE_TYPE(OPENED_ARCHETYPE_TYPE)(
     ModuleFile &MF, SmallVectorImpl<uint64_t> &scratch, StringRef blobData) {
-  TypeID existentialID;
   TypeID interfaceID;
-  GenericSignatureID sigID;
+  GenericEnvironmentID genericEnvID;
 
-  decls_block::OpenedArchetypeTypeLayout::readRecord(scratch, existentialID,
-                                                     interfaceID, sigID);
-
-  auto sigOrError = MF.getGenericSignatureChecked(sigID);
-  if (!sigOrError)
-    return sigOrError.takeError();
+  decls_block::OpenedArchetypeTypeLayout::readRecord(scratch,
+                                                     interfaceID,
+                                                     genericEnvID);
 
   auto interfaceTypeOrError = MF.getTypeChecked(interfaceID);
   if (!interfaceTypeOrError)
     return interfaceTypeOrError.takeError();
 
-  auto existentialTypeOrError = MF.getTypeChecked(existentialID);
-  if (!existentialTypeOrError)
-    return existentialTypeOrError.takeError();
+  auto envOrError = MF.getGenericEnvironmentChecked(genericEnvID);
+  if (!envOrError)
+    return envOrError.takeError();
 
-  auto env = GenericEnvironment::forOpenedArchetypeSignature(
-      existentialTypeOrError.get(), sigOrError.get(), UUID::fromTime());
-  return env->mapTypeIntoContext(interfaceTypeOrError.get())
-      ->castTo<OpenedArchetypeType>();
+  return envOrError.get()->mapTypeIntoContext(interfaceTypeOrError.get());
 }
 
 Expected<Type> DESERIALIZE_TYPE(OPAQUE_ARCHETYPE_TYPE)(

lib/Serialization/ModuleFile.cpp

@@ -113,6 +113,7 @@ ModuleFile::ModuleFile(std::shared_ptr<const ModuleFileSharedCore> core)
   allocateBuffer(Types, core->Types);
   allocateBuffer(ClangTypes, core->ClangTypes);
   allocateBuffer(GenericSignatures, core->GenericSignatures);
+  allocateBuffer(GenericEnvironments, core->GenericEnvironments);
   allocateBuffer(SubstitutionMaps, core->SubstitutionMaps);
   allocateBuffer(Identifiers, core->Identifiers);
 }

lib/Serialization/ModuleFile.h

@@ -259,6 +259,9 @@ private:
   /// Generic signatures referenced by this module.
   MutableArrayRef<Serialized<GenericSignature>> GenericSignatures;
 
+  /// Generic environments referenced by this module.
+  MutableArrayRef<Serialized<GenericEnvironment *>> GenericEnvironments;
+
   /// Substitution maps referenced by this module.
   MutableArrayRef<Serialized<SubstitutionMap>> SubstitutionMaps;
 
@@ -862,6 +865,10 @@ public:
   llvm::Expected<GenericSignature>
   getGenericSignatureChecked(serialization::GenericSignatureID ID);
 
+  /// Returns the generic environment for the given ID or the first error.
+  llvm::Expected<GenericEnvironment *>
+  getGenericEnvironmentChecked(serialization::GenericEnvironmentID ID);
+
   /// Returns the substitution map for the given ID, deserializing it if
   /// needed.
   SubstitutionMap getSubstitutionMap(serialization::SubstitutionMapID id);

lib/Serialization/ModuleFileSharedCore.cpp

@@ -839,6 +839,10 @@ bool ModuleFileSharedCore::readIndexBlock(llvm::BitstreamCursor &cursor) {
         assert(blobData.empty());
         allocateBuffer(GenericSignatures, scratch);
         break;
+      case index_block::GENERIC_ENVIRONMENT_OFFSETS:
+        assert(blobData.empty());
+        allocateBuffer(GenericEnvironments, scratch);
+        break;
       case index_block::SUBSTITUTION_MAP_OFFSETS:
         assert(blobData.empty());
         allocateBuffer(SubstitutionMaps, scratch);

lib/Serialization/ModuleFileSharedCore.h

@@ -218,6 +218,9 @@ private:
   /// Generic signatures referenced by this module.
   ArrayRef<RawBitOffset> GenericSignatures;
 
+  /// Generic environments referenced by this module.
+  ArrayRef<RawBitOffset> GenericEnvironments;
+
   /// Substitution maps referenced by this module.
   ArrayRef<RawBitOffset> SubstitutionMaps;
 

lib/Serialization/ModuleFormat.h

@@ -58,7 +58,7 @@ const uint16_t SWIFTMODULE_VERSION_MAJOR = 0;
 /// describe what change you made. The content of this comment isn't important;
 /// it just ensures a conflict if two people change the module format.
 /// Don't worry about adhering to the 80-column limit for this line.
-const uint16_t SWIFTMODULE_VERSION_MINOR = 700; // explicit_copy_addr
+const uint16_t SWIFTMODULE_VERSION_MINOR = 701; // opened archetype serialization
 
 /// A standard hash seed used for all string hashes in a serialized module.
 ///
@@ -145,6 +145,9 @@ using ProtocolConformanceIDField = DeclIDField;
 using GenericSignatureID = DeclID;
 using GenericSignatureIDField = DeclIDField;
 
+using GenericEnvironmentID = unsigned;
+using GenericEnvironmentIDField = BCFixed<32>;
+
 // SubstitutionMapID must be the same as DeclID because it is stored in the
 // same way.
 using SubstitutionMapID = DeclID;
@@ -1122,9 +1125,8 @@ namespace decls_block {
 
   TYPE_LAYOUT(OpenedArchetypeTypeLayout,
     OPENED_ARCHETYPE_TYPE,
-    TypeIDField,            // the existential type
-    TypeIDField,            // the interface type
-    GenericSignatureIDField // generic signature
+    TypeIDField,              // the interface type
+    GenericEnvironmentIDField // generic environment ID
   );
 
   TYPE_LAYOUT(OpaqueArchetypeTypeLayout,
@@ -1691,6 +1693,12 @@ namespace decls_block {
     BCArray<TypeIDField>         // generic parameter types
   >;
 
+  using GenericEnvironmentLayout = BCRecordLayout<
+    GENERIC_ENVIRONMENT,
+    TypeIDField,                 // existential type
+    GenericSignatureIDField      // parent signature
+  >;
+
   using SubstitutionMapLayout = BCRecordLayout<
     SUBSTITUTION_MAP,
     GenericSignatureIDField,     // generic signature
@@ -2178,6 +2186,7 @@ namespace index_block {
     LOCAL_TYPE_DECLS,
     OPAQUE_RETURN_TYPE_DECLS,
     GENERIC_SIGNATURE_OFFSETS,
+    GENERIC_ENVIRONMENT_OFFSETS,
     PROTOCOL_CONFORMANCE_OFFSETS,
     SIL_LAYOUT_OFFSETS,
 

lib/Serialization/Serialization.cpp

@@ -594,6 +594,13 @@ Serializer::addGenericSignatureRef(GenericSignature sig) {
   return GenericSignaturesToSerialize.addRef(sig);
 }
 
+GenericEnvironmentID
+Serializer::addGenericEnvironmentRef(GenericEnvironment *env) {
+  if (!env)
+    return 0;
+  return GenericEnvironmentsToSerialize.addRef(env);
+}
+
 SubstitutionMapID
 Serializer::addSubstitutionMapRef(SubstitutionMap substitutions) {
   return SubstitutionMapsToSerialize.addRef(substitutions);
@@ -878,6 +885,7 @@ void Serializer::writeBlockInfoBlock() {
   BLOCK_RECORD(index_block, ENTRY_POINT);
   BLOCK_RECORD(index_block, LOCAL_DECL_CONTEXT_OFFSETS);
   BLOCK_RECORD(index_block, GENERIC_SIGNATURE_OFFSETS);
+  BLOCK_RECORD(index_block, GENERIC_ENVIRONMENT_OFFSETS);
   BLOCK_RECORD(index_block, SUBSTITUTION_MAP_OFFSETS);
   BLOCK_RECORD(index_block, CLANG_TYPE_OFFSETS);
   BLOCK_RECORD(index_block, LOCAL_TYPE_DECLS);
@@ -1490,6 +1498,20 @@ void Serializer::writeASTBlockEntity(GenericSignature sig) {
   }
 }
 
+void Serializer::writeASTBlockEntity(const GenericEnvironment *genericEnv) {
+  using namespace decls_block;
+
+  assert(GenericEnvironmentsToSerialize.hasRef(genericEnv));
+
+  auto existentialTypeID = addTypeRef(genericEnv->getOpenedExistentialType());
+  auto parentSig = genericEnv->getOpenedExistentialParentSignature();
+  auto parentSigID = addGenericSignatureRef(parentSig);
+
+  auto genericEnvAbbrCode = DeclTypeAbbrCodes[GenericEnvironmentLayout::Code];
+  GenericEnvironmentLayout::emitRecord(Out, ScratchRecord, genericEnvAbbrCode,
+                                       existentialTypeID, parentSigID);
+}
+
 void Serializer::writeASTBlockEntity(const SubstitutionMap substitutions) {
   using namespace decls_block;
   assert(substitutions);
@@ -4603,14 +4625,12 @@ public:
 
   void visitOpenedArchetypeType(const OpenedArchetypeType *archetypeTy) {
     using namespace decls_block;
-    auto sig = archetypeTy->getGenericEnvironment()->getGenericSignature();
-    auto existentialTypeID = S.addTypeRef(archetypeTy->getExistentialType());
     auto interfaceTypeID = S.addTypeRef(archetypeTy->getInterfaceType());
-    auto sigID = S.addGenericSignatureRef(sig);
+    auto genericEnvID = S.addGenericEnvironmentRef(
+        archetypeTy->getGenericEnvironment());
     unsigned abbrCode = S.DeclTypeAbbrCodes[OpenedArchetypeTypeLayout::Code];
     OpenedArchetypeTypeLayout::emitRecord(S.Out, S.ScratchRecord, abbrCode,
-                                          existentialTypeID, interfaceTypeID,
-                                          sigID);
+                                          interfaceTypeID, genericEnvID);
   }
 
   void
@@ -5107,6 +5127,7 @@ void Serializer::writeAllDeclsAndTypes() {
   registerDeclTypeAbbr<InlinableBodyTextLayout>();
   registerDeclTypeAbbr<GenericParamListLayout>();
   registerDeclTypeAbbr<GenericSignatureLayout>();
+  registerDeclTypeAbbr<GenericEnvironmentLayout>();
   registerDeclTypeAbbr<RequirementSignatureLayout>();
   registerDeclTypeAbbr<SILGenericSignatureLayout>();
   registerDeclTypeAbbr<SubstitutionMapLayout>();
@@ -5161,6 +5182,8 @@ void Serializer::writeAllDeclsAndTypes() {
         writeASTBlockEntitiesIfNeeded(LocalDeclContextsToSerialize);
     wroteSomething |=
         writeASTBlockEntitiesIfNeeded(GenericSignaturesToSerialize);
+    wroteSomething |=
+        writeASTBlockEntitiesIfNeeded(GenericEnvironmentsToSerialize);
     wroteSomething |=
         writeASTBlockEntitiesIfNeeded(SubstitutionMapsToSerialize);
     wroteSomething |=
@@ -5764,6 +5787,7 @@ void Serializer::writeAST(ModuleOrSourceFile DC) {
     writeOffsets(Offsets, ClangTypesToSerialize);
     writeOffsets(Offsets, LocalDeclContextsToSerialize);
     writeOffsets(Offsets, GenericSignaturesToSerialize);
+    writeOffsets(Offsets, GenericEnvironmentsToSerialize);
     writeOffsets(Offsets, SubstitutionMapsToSerialize);
     writeOffsets(Offsets, ConformancesToSerialize);
     writeOffsets(Offsets, SILLayoutsToSerialize);

lib/Serialization/Serialization.h

@@ -219,6 +219,10 @@ class Serializer : public SerializerBase {
                        index_block::GENERIC_SIGNATURE_OFFSETS>
   GenericSignaturesToSerialize;
 
+  ASTBlockRecordKeeper<const GenericEnvironment *, GenericEnvironmentID,
+                       index_block::GENERIC_ENVIRONMENT_OFFSETS>
+  GenericEnvironmentsToSerialize;
+
   ASTBlockRecordKeeper<SubstitutionMap, SubstitutionMapID,
                        index_block::SUBSTITUTION_MAP_OFFSETS>
   SubstitutionMapsToSerialize;
@@ -353,6 +357,9 @@ private:
   /// Writes a generic signature.
   void writeASTBlockEntity(GenericSignature sig);
 
+  /// Writes a generic environment.
+  void writeASTBlockEntity(const GenericEnvironment *env);
+
   /// Writes a substitution map.
   void writeASTBlockEntity(const SubstitutionMap substitutions);
 
@@ -483,6 +490,9 @@ public:
   /// The GenericSignature will be scheduled for serialization if necessary.
   GenericSignatureID addGenericSignatureRef(GenericSignature sig);
 
+  /// Records the use of the given opened generic environment.
+  GenericEnvironmentID addGenericEnvironmentRef(GenericEnvironment *env);
+
   /// Records the use of the given substitution map.
   ///
   /// The SubstitutionMap will be scheduled for serialization if necessary.

test/Serialization/Inputs/nested_opened_archetype_other.swift

@@ -0,0 +1,11 @@
+public protocol P {
+  associatedtype A: Q
+
+  func foo() -> A
+}
+
+public protocol Q {}
+
+@_transparent public func f(p: any P) -> any Q {
+  return p.foo()
+}

test/Serialization/nested_opened_archetype.swift

@@ -0,0 +1,15 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend -emit-module -emit-module-path %t/nested_opened_archetype_other.swiftmodule %S/Inputs/nested_opened_archetype_other.swift
+// RUN: %target-swift-frontend -emit-sil %s -I %t
+
+import nested_opened_archetype_other
+
+struct SP : P {
+  func foo() -> SQ { return SQ() }
+}
+
+struct SQ : Q {}
+
+func g() {
+  f(p: SP())
+}