Rework the relationship between generic environments and opaque archetypes.

Teach `GenericEnvironment` how to lazily create opaque type archetypes,
performing the contextual substitutions as required but without
building the "bound" generic signature until required. To get here,
augment `GenericEnvironment` with knowledge of the purpose of the
environment, whether it is for normal cases (any signature), an opened
existential type, or an opaque type. For opaque types, store the
opaque type declaration and substitution map, which we also use to
uniquely find the generic environment. Among other things, this
ensures that different opaque type archetypes within the same opaque
type declaration are properly sharing a generic environment, which
wasn't happening before.

Author: DougGregor

include/swift/AST/GenericEnvironment.h

@@ -33,6 +33,8 @@ namespace swift {
 class ArchetypeType;
 class ASTContext;
 class GenericTypeParamType;
+class OpaqueTypeDecl;
+class OpenedArchetypeType;
 class SILModule;
 class SILType;
 
@@ -57,11 +59,28 @@ class QueryInterfaceTypeSubstitutions {
 /// TypeBase::mapTypeOutOfContext().
 ///
 class alignas(1 << DeclAlignInBits) GenericEnvironment final
-        : private llvm::TrailingObjects<GenericEnvironment, Type> {
-  GenericSignature Signature = GenericSignature();
+    : private llvm::TrailingObjects<
+        GenericEnvironment, OpaqueTypeDecl *, SubstitutionMap, Type> {
+public:
+  enum class Kind {
+    /// A normal generic environment, determined only by its generic
+    /// signature.
+    Normal,
+    /// A generic environment describing an opened existential archetype.
+    OpenedExistential,
+    /// A generic environment describing an opaque type archetype.
+    Opaque,
+  };
+
+private:
+  mutable llvm::PointerIntPair<GenericSignature, 2, Kind> SignatureAndKind{
+      GenericSignature(), Kind::Normal};
 
   friend TrailingObjects;
+  friend OpaqueTypeArchetypeType;
 
+  size_t numTrailingObjects(OverloadToken<OpaqueTypeDecl *>) const;
+  size_t numTrailingObjects(OverloadToken<SubstitutionMap>) const;
   size_t numTrailingObjects(OverloadToken<Type>) const;
 
   /// Retrieve the array containing the context types associated with the
@@ -74,7 +93,9 @@ class alignas(1 << DeclAlignInBits) GenericEnvironment final
   /// generic signature.
   ArrayRef<Type> getContextTypes() const;
 
-  explicit GenericEnvironment(GenericSignature signature);
+  explicit GenericEnvironment(GenericSignature signature, Kind kind);
+  explicit GenericEnvironment(
+      GenericSignature signature, OpaqueTypeDecl *opaque, SubstitutionMap subs);
 
   friend ArchetypeType;
   friend QueryInterfaceTypeSubstitutions;
@@ -86,18 +107,39 @@ class alignas(1 << DeclAlignInBits) GenericEnvironment final
   /// This is only useful when lazily populating a generic environment.
   Optional<Type> getMappingIfPresent(GenericParamKey key) const;
 
+  /// Get the "raw" generic signature, without substituting into opaque
+  /// type's signature.
+  GenericSignature getRawGenericSignature() const;
+
 public:
-  GenericSignature getGenericSignature() const {
-    return Signature;
-  }
+  GenericSignature getGenericSignature() const;
+
+  Kind getKind() const { return SignatureAndKind.getInt(); }
 
   TypeArrayView<GenericTypeParamType> getGenericParams() const;
 
+  /// Retrieve the opaque type declaration for a generic environment describing
+  /// opaque types.
+  OpaqueTypeDecl *getOpaqueTypeDecl() const;
+
+  /// Retrieve the substitutions applied to an opaque type declaration to
+  /// create a generic environment.
+  SubstitutionMap getOpaqueSubstitutions() const;
+
   /// Create a new, "incomplete" generic environment that will be populated
   /// by calls to \c addMapping().
   static
   GenericEnvironment *getIncomplete(GenericSignature signature);
 
+  /// Create a new generic environment for an opened existential.
+  static GenericEnvironment *forOpenedExistential(
+      GenericSignature signature, const OpenedArchetypeType *type);
+
+  /// Create a new generic environment for an opaque type with the given set of
+  /// outer substitutions.
+  static GenericEnvironment *forOpaqueType(
+      OpaqueTypeDecl *opaque, SubstitutionMap subs, AllocationArena arena);
+
   /// Add a mapping of a generic parameter to a specific type (which may be
   /// an archetype)
   void addMapping(GenericParamKey key, Type contextType);
@@ -112,6 +154,12 @@ class alignas(1 << DeclAlignInBits) GenericEnvironment final
     return Mem; 
   }
 
+  /// For an opaque archetype environment, apply the substitutions.
+  Type maybeApplyOpaqueTypeSubstitutions(Type type) const;
+
+  /// Compute the canonical interface type within this environment.
+  Type getCanonicalInterfaceType(Type interfaceType);
+
   /// Map an interface type to a contextual type.
   static Type mapTypeIntoContext(GenericEnvironment *genericEnv,
                                  Type type);

include/swift/AST/Types.h

@@ -5548,37 +5548,36 @@ END_CAN_TYPE_WRAPPER(PrimaryArchetypeType, ArchetypeType)
 
 /// An archetype that represents an opaque type.
 class OpaqueTypeArchetypeType final : public ArchetypeType,
-    public llvm::FoldingSetNode,
     private ArchetypeTrailingObjects<OpaqueTypeArchetypeType>
 {
   friend TrailingObjects;
   friend ArchetypeType;
   friend GenericSignatureBuilder;
 
-  /// The declaration that defines the opaque type.
-  OpaqueTypeDecl *OpaqueDecl;
-  /// The substitutions into the interface signature of the opaque type.
-  SubstitutionMap Substitutions;
-  
   /// A GenericEnvironment with this opaque archetype bound to the interface
   /// type of the output type from the OpaqueDecl.
-  mutable GenericEnvironment *Environment = nullptr;
-  
+  GenericEnvironment *Environment = nullptr;
+
+  friend class GenericEnvironment;
+
+  static OpaqueTypeArchetypeType *getNew(
+      GenericEnvironment *environment, Type interfaceType,
+      ArrayRef<ProtocolDecl*> conformsTo, Type superclass,
+      LayoutConstraint layout);
+
 public:
   /// Get an opaque archetype representing the underlying type of the given
   /// opaque type decl's opaque param with ordinal `ordinal`. For example, in
   /// `(some P, some Q)`, `some P`'s type param would have ordinal 0 and `some
   /// Q`'s type param would have ordinal 1.
-  static OpaqueTypeArchetypeType *get(OpaqueTypeDecl *Decl, unsigned ordinal,
-                                      SubstitutionMap Substitutions);
+  static Type get(OpaqueTypeDecl *decl, unsigned ordinal, SubstitutionMap subs);
+
+  /// Retrieve the opaque type declaration.
+  OpaqueTypeDecl *getDecl() const;
+
+  /// Retrieve the set of substitutions applied to the opaque type.
+  SubstitutionMap getSubstitutions() const;
 
-  OpaqueTypeDecl *getDecl() const {
-    return OpaqueDecl;
-  }
-  SubstitutionMap getSubstitutions() const {
-    return Substitutions;
-  }
-  
   /// Get the generic signature used to build out this archetype. This is
   /// equivalent to the OpaqueTypeDecl's interface generic signature, with
   /// all of the generic parameters aside from the opaque type's interface
@@ -5605,22 +5604,12 @@ class OpaqueTypeArchetypeType final : public ArchetypeType,
   /// then the underlying type of `some P` would be ordinal 0, and `some Q` would be ordinal 1.
   unsigned getOrdinal() const;
 
-  static void Profile(llvm::FoldingSetNodeID &ID,
-                      OpaqueTypeDecl *OpaqueDecl,
-                      unsigned ordinal,
-                      SubstitutionMap Substitutions);
-  
-  void Profile(llvm::FoldingSetNodeID &ID) {
-    Profile(ID, getDecl(), getOrdinal(), getSubstitutions());
-  };
-  
 private:
-  OpaqueTypeArchetypeType(OpaqueTypeDecl *OpaqueDecl,
-                          SubstitutionMap Substitutions,
-                          RecursiveTypeProperties Props,
-                          Type InterfaceType,
-                          ArrayRef<ProtocolDecl*> ConformsTo,
-                          Type Superclass, LayoutConstraint Layout);
+  OpaqueTypeArchetypeType(GenericEnvironment *environment,
+                          RecursiveTypeProperties properties,
+                          Type interfaceType,
+                          ArrayRef<ProtocolDecl*> conformsTo,
+                          Type superclass, LayoutConstraint layout);
 };
 BEGIN_CAN_TYPE_WRAPPER(OpaqueTypeArchetypeType, ArchetypeType)
 END_CAN_TYPE_WRAPPER(OpaqueTypeArchetypeType, ArchetypeType)

lib/AST/ASTContext.cpp

@@ -414,7 +414,8 @@ struct ASTContext::Implementation {
     llvm::FoldingSet<BoundGenericType> BoundGenericTypes;
     llvm::FoldingSet<ProtocolCompositionType> ProtocolCompositionTypes;
     llvm::FoldingSet<LayoutConstraintInfo> LayoutConstraints;
-    llvm::FoldingSet<OpaqueTypeArchetypeType> OpaqueArchetypes;
+    llvm::DenseMap<std::pair<OpaqueTypeDecl *, SubstitutionMap>,
+                   GenericEnvironment *> OpaqueArchetypeEnvironments;
 
     /// The set of function types.
     llvm::FoldingSet<FunctionType> FunctionTypes;
@@ -4280,77 +4281,62 @@ DependentMemberType *DependentMemberType::get(Type base,
   return known;
 }
 
-OpaqueTypeArchetypeType *
-OpaqueTypeArchetypeType::get(OpaqueTypeDecl *Decl, unsigned ordinal,
-                             SubstitutionMap Substitutions) {
-  auto opaqueParamType = Decl->getOpaqueGenericParams()[ordinal];
-
-  // TODO: We could attempt to preserve type sugar in the substitution map.
-  // Currently archetypes are assumed to be always canonical in many places,
-  // though, so doing so would require fixing those places.
-  Substitutions = Substitutions.getCanonical();
-
-  llvm::FoldingSetNodeID id;
-  Profile(id, Decl, ordinal, Substitutions);
-  
-  auto &ctx = Decl->getASTContext();
-
+/// Compute the recursive type properties of an opaque type archetype.
+static RecursiveTypeProperties getOpaqueTypeArchetypeProperties(
+    SubstitutionMap subs) {
   // An opaque type isn't contextually dependent like other archetypes, so
   // by itself, it doesn't impose the "Has Archetype" recursive property,
   // but the substituted types might. A disjoint "Has Opaque Archetype" tracks
   // the presence of opaque archetypes.
   RecursiveTypeProperties properties =
     RecursiveTypeProperties::HasOpaqueArchetype;
-  for (auto type : Substitutions.getReplacementTypes()) {
+  for (auto type : subs.getReplacementTypes()) {
     properties |= type->getRecursiveProperties();
   }
-  
+  return properties;
+}
+
+OpaqueTypeArchetypeType *OpaqueTypeArchetypeType::getNew(
+    GenericEnvironment *environment, Type interfaceType,
+    ArrayRef<ProtocolDecl *> conformsTo, Type superclass,
+    LayoutConstraint layout) {
+  auto properties = getOpaqueTypeArchetypeProperties(
+      environment->getOpaqueSubstitutions());
   auto arena = getArena(properties);
-  
-  llvm::FoldingSet<OpaqueTypeArchetypeType> &set
-    = ctx.getImpl().getArena(arena).OpaqueArchetypes;
-  
-  {
-    void *insertPos; // Discarded because the work below may invalidate the
-                     // insertion point inside the folding set
-    if (auto existing = set.FindNodeOrInsertPos(id, insertPos)) {
-      return existing;
-    }
-  }
-  
-  auto reqs = Decl->getOpaqueInterfaceGenericSignature()
-      ->getLocalRequirements(opaqueParamType);
-  auto superclass = reqs.superclass;
-  #if !DO_IT_CORRECTLY
-    // Ad-hoc substitute the generic parameters of the superclass.
-    // If we correctly applied the substitutions to the generic signature
-    // constraints above, this would be unnecessary.
-    if (superclass && superclass->hasTypeParameter()) {
-      superclass = superclass.subst(Substitutions);
-    }
-  #endif
+  auto size = OpaqueTypeArchetypeType::totalSizeToAlloc<
+      ProtocolDecl *, Type, LayoutConstraint>(
+         conformsTo.size(), superclass ? 1 : 0, layout ? 1 : 0);
+  ASTContext &ctx = interfaceType->getASTContext();
+  auto mem = ctx.Allocate(size, alignof(OpaqueTypeArchetypeType), arena);
+  return ::new (mem)
+      OpaqueTypeArchetypeType(environment, properties, interfaceType,
+                              conformsTo, superclass, layout);
+}
+
+Type OpaqueTypeArchetypeType::get(
+    OpaqueTypeDecl *Decl, unsigned ordinal, SubstitutionMap Substitutions) {
+  // TODO: We could attempt to preserve type sugar in the substitution map.
+  // Currently archetypes are assumed to be always canonical in many places,
+  // though, so doing so would require fixing those places.
+  Substitutions = Substitutions.getCanonical();
 
-  auto mem = ctx.Allocate(
-    OpaqueTypeArchetypeType::totalSizeToAlloc<ProtocolDecl *, Type, LayoutConstraint>(
-      reqs.protos.size(), superclass ? 1 : 0, reqs.layout ? 1 : 0),
-      alignof(OpaqueTypeArchetypeType),
-      arena);
+  auto &ctx = Decl->getASTContext();
 
-  auto newOpaque = ::new (mem)
-      OpaqueTypeArchetypeType(Decl, Substitutions, properties, opaqueParamType,
-                              reqs.protos, superclass, reqs.layout);
+  // Look for an opaque archetype environment in the appropriate arena.
+  auto properties = getOpaqueTypeArchetypeProperties(Substitutions);
+  auto arena = getArena(properties);
+  auto &environments
+    = ctx.getImpl().getArena(arena).OpaqueArchetypeEnvironments;
+  GenericEnvironment *env = environments[{Decl, Substitutions}];
 
-  // Look up the insertion point in the folding set again in case something
-  // invalidated it above.
-  {
-    void *insertPos;
-    auto existing = set.FindNodeOrInsertPos(id, insertPos);
-    (void)existing;
-    assert(!existing && "race to create opaque archetype?!");
-    set.InsertNode(newOpaque, insertPos);
+  // Create the environment if it's missing.
+  if (!env) {
+    env = GenericEnvironment::forOpaqueType(Decl, Substitutions, arena);
+    environments[{Decl, Substitutions}] = env;
   }
 
-  return newOpaque;
+  auto opaqueParamType = Decl->getOpaqueGenericParams()[ordinal];
+  return env->getOrCreateArchetypeFromInterfaceType(opaqueParamType);
 }
 
 CanOpenedArchetypeType OpenedArchetypeType::get(Type existential,
@@ -4411,15 +4397,11 @@ GenericEnvironment *OpenedArchetypeType::getGenericEnvironment() const {
   if (Environment)
     return Environment;
 
-  auto thisType = Type(const_cast<OpenedArchetypeType*>(this));
+  auto thisType = const_cast<OpenedArchetypeType*>(this);
   auto &ctx = thisType->getASTContext();
   // Create a generic environment to represent the opened type.
   auto signature = ctx.getOpenedArchetypeSignature(Opened);
-  auto *env = GenericEnvironment::getIncomplete(signature);
-  env->addMapping(signature.getGenericParams().front().getPointer(), thisType);
-  Environment = env;
-  
-  return env;
+  return GenericEnvironment::forOpenedExistential(signature, thisType);
 }
 
 CanType OpenedArchetypeType::getAny(Type existential) {
@@ -4575,9 +4557,43 @@ GenericEnvironment *GenericEnvironment::getIncomplete(
 
   // Allocate and construct the new environment.
   unsigned numGenericParams = signature.getGenericParams().size();
-  size_t bytes = totalSizeToAlloc<Type>(numGenericParams);
+  size_t bytes = totalSizeToAlloc<OpaqueTypeDecl *, SubstitutionMap, Type>(
+      0, 0, numGenericParams);
+  void *mem = ctx.Allocate(bytes, alignof(GenericEnvironment));
+  return new (mem) GenericEnvironment(signature, Kind::Normal);
+}
+
+/// Create a new generic environment for an opened archetype.
+GenericEnvironment *GenericEnvironment::forOpenedExistential(
+    GenericSignature signature, const OpenedArchetypeType *type) {
+  auto &ctx = signature->getASTContext();
+
+  // Allocate and construct the new environment.
+  unsigned numGenericParams = signature.getGenericParams().size();
+  size_t bytes = totalSizeToAlloc<OpaqueTypeDecl *, SubstitutionMap, Type>(
+      0, 0, numGenericParams);
   void *mem = ctx.Allocate(bytes, alignof(GenericEnvironment));
-  return new (mem) GenericEnvironment(signature);
+  auto env = new (mem) GenericEnvironment(signature, Kind::OpenedExistential);
+  env->addMapping(
+      signature.getGenericParams().front(),
+      Type(const_cast<OpenedArchetypeType *>(type)));
+  return env;
+}
+
+/// Create a new generic environment for an opaque type with the given set of
+/// outer substitutions.
+GenericEnvironment *GenericEnvironment::forOpaqueType(
+    OpaqueTypeDecl *opaque, SubstitutionMap subs, AllocationArena arena) {
+  auto &ctx = opaque->getASTContext();
+
+  // Allocate and construct the new environment.
+  auto signature = opaque->getOpaqueInterfaceGenericSignature();
+  unsigned numGenericParams = signature.getGenericParams().size();
+  size_t bytes = totalSizeToAlloc<OpaqueTypeDecl *, SubstitutionMap, Type>(
+      1, 1, numGenericParams);
+  void *mem = ctx.Allocate(bytes, alignof(GenericEnvironment), arena);
+  auto env = new (mem) GenericEnvironment(GenericSignature(), opaque, subs);
+  return env;
 }
 
 void DeclName::CompoundDeclName::Profile(llvm::FoldingSetNodeID &id,