[AST] Tail-allocate GenericEnvironment's archetypes-to-interface types mapping

include/swift/AST/GenericEnvironment.h

@@ -22,6 +22,7 @@
 #include "swift/AST/GenericSignature.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/Support/TrailingObjects.h"
+#include <utility>
 
 namespace swift {
 
@@ -34,17 +35,36 @@ class SILType;
 /// Describes the mapping between archetypes and interface types for the
 /// generic parameters of a DeclContext.
 class alignas(1 << DeclAlignInBits) GenericEnvironment final
-                  : private llvm::TrailingObjects<GenericEnvironment, Type> {
+        : private llvm::TrailingObjects<GenericEnvironment, Type,
+                                        std::pair<ArchetypeType *,
+                                                  GenericTypeParamType *>> {
   GenericSignature *Signature;
   ArchetypeBuilder *Builder;
-  TypeSubstitutionMap ArchetypeToInterfaceMap;
+
+  // The number of generic type parameter -> context type mappings we have
+  // recorded so far. This saturates at the number of generic type parameters,
+  // at which point the archetype-to-interface trailing array is sorted.
+  unsigned NumMappingsRecorded : 16;
+
+  // The number of archetype-to-interface type mappings. This is always <=
+  // \c NumMappingsRecorded.
+  unsigned NumArchetypeToInterfaceMappings : 16;
 
   friend TrailingObjects;
 
+  /// An entry in the array mapping from archetypes to their corresponding
+  /// generic type parameters.
+  typedef std::pair<ArchetypeType *, GenericTypeParamType *>
+                                                    ArchetypeToInterfaceMapping;
+
   size_t numTrailingObjects(OverloadToken<Type>) const {
     return Signature->getGenericParams().size();
   }
 
+  size_t numTrailingObjects(OverloadToken<ArchetypeToInterfaceMapping>) const {
+    return Signature->getGenericParams().size();
+  }
+
   /// Retrieve the array containing the context types associated with the
   /// generic parameters, stored in parallel with the generic parameters of the
   /// generic signature.
@@ -61,6 +81,30 @@ class alignas(1 << DeclAlignInBits) GenericEnvironment final
                           Signature->getGenericParams().size());
   }
 
+  /// Retrieve the active set of archetype-to-interface mappings.
+  ArrayRef<ArchetypeToInterfaceMapping>
+                                getActiveArchetypeToInterfaceMappings() const {
+    return { getTrailingObjects<ArchetypeToInterfaceMapping>(),
+             NumArchetypeToInterfaceMappings };
+  }
+
+  /// Retrieve the active set of archetype-to-interface mappings.
+  MutableArrayRef<ArchetypeToInterfaceMapping>
+                                      getActiveArchetypeToInterfaceMappings() {
+    return { getTrailingObjects<ArchetypeToInterfaceMapping>(),
+             NumArchetypeToInterfaceMappings };
+  }
+
+  /// Retrieve the buffer for the archetype-to-interface mappings.
+  ///
+  /// Only the first \c NumArchetypeToInterfaceMappings elements in the buffer
+  /// are valid.
+  MutableArrayRef<ArchetypeToInterfaceMapping>
+                                      getArchetypeToInterfaceMappingsBuffer() {
+    return { getTrailingObjects<ArchetypeToInterfaceMapping>(),
+             Signature->getGenericParams().size() };
+  }
+
   GenericEnvironment(GenericSignature *signature,
                      ArchetypeBuilder *builder,
                      TypeSubstitutionMap interfaceToArchetypeMap);
@@ -83,7 +127,20 @@ class alignas(1 << DeclAlignInBits) GenericEnvironment final
     Type operator()(SubstitutableType *type) const;
   };
   friend class QueryInterfaceTypeSubstitutions;
-
+
+  /// Query function suitable for use as a \c TypeSubstitutionFn that queries
+  /// the mapping of archetypes back to interface types.
+  class QueryArchetypeToInterfaceSubstitutions {
+    const GenericEnvironment *self;
+
+  public:
+    QueryArchetypeToInterfaceSubstitutions(const GenericEnvironment *self)
+      : self(self) { }
+
+    Type operator()(SubstitutableType *type) const;
+  };
+  friend class QueryArchetypeToInterfaceSubstitutions;
+
 public:
   GenericSignature *getGenericSignature() const {
     return Signature;

lib/AST/ASTContext.cpp

@@ -3561,14 +3561,16 @@ GenericSignature *GenericSignature::get(ArrayRef<GenericTypeParamType *> params,
 GenericEnvironment *
 GenericEnvironment::get(GenericSignature *signature,
                         TypeSubstitutionMap interfaceToArchetypeMap) {
+  unsigned numGenericParams = signature->getGenericParams().size();
   assert(!interfaceToArchetypeMap.empty());
-  assert(interfaceToArchetypeMap.size() == signature->getGenericParams().size()
+  assert(interfaceToArchetypeMap.size() == numGenericParams
          && "incorrect number of parameters");
 
   ASTContext &ctx = signature->getASTContext();
 
   // Allocate and construct the new environment.
-  size_t bytes = totalSizeToAlloc<Type>(signature->getGenericParams().size());
+  size_t bytes = totalSizeToAlloc<Type, ArchetypeToInterfaceMapping>(
+                                           numGenericParams, numGenericParams);
   void *mem = ctx.Allocate(bytes, alignof(GenericEnvironment));
   return new (mem) GenericEnvironment(signature, nullptr,
                                       interfaceToArchetypeMap);
@@ -3579,8 +3581,11 @@ GenericEnvironment *GenericEnvironment::getIncomplete(
                                                   ArchetypeBuilder *builder) {
   TypeSubstitutionMap empty;
   auto &ctx = signature->getASTContext();
+
   // Allocate and construct the new environment.
-  size_t bytes = totalSizeToAlloc<Type>(signature->getGenericParams().size());
+  unsigned numGenericParams = signature->getGenericParams().size();
+  size_t bytes = totalSizeToAlloc<Type, ArchetypeToInterfaceMapping>(
+                                           numGenericParams, numGenericParams);
   void *mem = ctx.Allocate(bytes, alignof(GenericEnvironment));
   return new (mem) GenericEnvironment(signature, builder, empty);
 }

lib/AST/GenericEnvironment.cpp

@@ -26,6 +26,9 @@ GenericEnvironment::GenericEnvironment(
     TypeSubstitutionMap interfaceToArchetypeMap)
   : Signature(signature), Builder(builder)
 {
+  NumMappingsRecorded = 0;
+  NumArchetypeToInterfaceMappings = 0;
+
   // Clear out the memory that holds the context types.
   std::uninitialized_fill(getContextTypes().begin(), getContextTypes().end(), Type());
 
@@ -35,9 +38,6 @@ GenericEnvironment::GenericEnvironment(
   // mapTypeOutOfContext() produces a human-readable type.
   for (auto entry : interfaceToArchetypeMap)
     addMapping(entry.first->castTo<GenericTypeParamType>(), entry.second);
-
-  // Make sure this generic environment gets destroyed.
-  signature->getASTContext().addDestructorCleanup(*this);
 }
 
 void GenericEnvironment::addMapping(GenericParamKey key,
@@ -54,22 +54,52 @@ void GenericEnvironment::addMapping(GenericParamKey key,
 
   // If we mapped the generic parameter to an archetype, add it to the
   // reverse mapping.
-  auto *archetype = contextType->getAs<ArchetypeType>();
-  if (!archetype) return;
-
-  // Check whether we've already recorded an interface type for this archetype.
-  // If not, record one and we're done.
-  auto genericParam = genericParams[index];
-  auto result = ArchetypeToInterfaceMap.insert({archetype, genericParam});
-  if (result.second) return;
-
-  // Multiple generic parameters map to the same archetype. If the
-  // existing entry comes from a later generic parameter, replace it with
-  // the earlier generic parameter. This gives us a deterministic reverse
-  // mapping.
-  auto otherGP = result.first->second->castTo<GenericTypeParamType>();
-  if (GenericParamKey(genericParam) < GenericParamKey(otherGP))
-    result.first->second = genericParam;
+  if (auto *archetype = contextType->getAs<ArchetypeType>()) {
+    auto genericParam = genericParams[index];
+
+    // Check whether we've already recorded a generic parameter for this
+    // archetype. Note that we always perform a linear search, because we
+    // won't have sorted the list yet.
+    bool found = false;
+    for (auto &mapping : getActiveArchetypeToInterfaceMappings()) {
+      if (mapping.first != archetype) continue;
+
+      // Multiple generic parameters map to the same archetype. If the
+      // existing entry comes from a later generic parameter, replace it with
+      // the earlier generic parameter. This gives us a deterministic reverse
+      // mapping.
+      auto otherGP = mapping.second->castTo<GenericTypeParamType>();
+      if (GenericParamKey(genericParam) < GenericParamKey(otherGP))
+        mapping.second = genericParam;
+      found = true;
+      break;
+    }
+
+    // If we haven't recorded a generic parameter for this archetype, do so now.
+    if (!found) {
+      void *ptr = getArchetypeToInterfaceMappingsBuffer().data()
+                + NumArchetypeToInterfaceMappings;
+      new (ptr) ArchetypeToInterfaceMapping(archetype, genericParam);
+      ++NumArchetypeToInterfaceMappings;
+    }
+  }
+
+  // Note that we've recorded this mapping.
+  ++NumMappingsRecorded;
+
+  // If we've recorded all of the mappings, go ahead and sort the array of
+  // archetype-to-interface-type mappings.
+  if (NumMappingsRecorded == genericParams.size()) {
+    llvm::array_pod_sort(getActiveArchetypeToInterfaceMappings().begin(),
+                         getActiveArchetypeToInterfaceMappings().end(),
+                         [](const ArchetypeToInterfaceMapping *lhs,
+                            const ArchetypeToInterfaceMapping *rhs) -> int {
+                           std::less<ArchetypeType *> compare;
+                           if (compare(lhs->first, rhs->first)) return -1;
+                           if (compare(rhs->first, lhs->first)) return 1;
+                           return 0;
+                         });
+  }
 }
 
 Optional<Type> GenericEnvironment::getMappingIfPresent(
@@ -88,11 +118,13 @@ Optional<Type> GenericEnvironment::getMappingIfPresent(
 
 bool GenericEnvironment::containsPrimaryArchetype(
                                               ArchetypeType *archetype) const {
-  return ArchetypeToInterfaceMap.count(archetype) > 0;
+  return static_cast<bool>(
+                       QueryArchetypeToInterfaceSubstitutions(this)(archetype));
 }
 
 Type GenericEnvironment::mapTypeOutOfContext(ModuleDecl *M, Type type) const {
-  type = type.subst(M, ArchetypeToInterfaceMap, SubstFlags::AllowLoweredTypes);
+  type = type.subst(M, QueryArchetypeToInterfaceSubstitutions(this),
+                    SubstFlags::AllowLoweredTypes);
   assert(!type->hasArchetype() && "not fully substituted");
   return type;
 }
@@ -132,6 +164,70 @@ Type GenericEnvironment::QueryInterfaceTypeSubstitutions::operator()(
   return Type();
 }
 
+Type GenericEnvironment::QueryArchetypeToInterfaceSubstitutions::operator()(
+                                                SubstitutableType *type) const {
+  auto archetype = type->getAs<ArchetypeType>();
+  if (!archetype) return Type();
+
+  // If not all generic parameters have had their context types recorded,
+  // perform a linear search.
+  auto genericParams = self->Signature->getGenericParams();
+  unsigned numGenericParams = genericParams.size();
+  if (self->NumMappingsRecorded < numGenericParams) {
+    // Search through all of the active archetype-to-interface mappings.
+    for (auto &mapping : self->getActiveArchetypeToInterfaceMappings())
+      if (mapping.first == archetype) return mapping.second;
+
+    // We don't know if the archetype is from a different context or if we
+    // simply haven't recorded it yet. Spin through all of the generic
+    // parameters looking for one that provides this mapping.
+    for (auto gp : genericParams) {
+      // Map the generic parameter into our context. If we get back an
+      // archetype that matches, we're done.
+      auto gpArchetype = self->mapTypeIntoContext(gp)->getAs<ArchetypeType>();
+      if (gpArchetype == archetype) return gp;
+    }
+
+    // We have checked all of the generic parameters and not found anything;
+    // there is no substitution.
+    return Type();
+  }
+
+  // All generic parameters have ad their context types recorded, which means
+  // that the archetypes-to-interface-types array is sorted by address. Use a
+  // binary search.
+  struct MappingComparison {
+    bool operator()(const ArchetypeToInterfaceMapping &lhs,
+                    const ArchetypeType *rhs) const {
+      std::less<const ArchetypeType *> compare;
+
+      return compare(lhs.first, rhs);
+    }
+
+    bool operator()(const ArchetypeType *lhs,
+                    const ArchetypeToInterfaceMapping &rhs) const {
+      std::less<const ArchetypeType *> compare;
+
+      return compare(lhs, rhs.first);
+    }
+
+    bool operator()(const ArchetypeToInterfaceMapping &lhs,
+                    const ArchetypeToInterfaceMapping &rhs) const {
+      std::less<const ArchetypeType *> compare;
+
+      return compare(lhs.first, rhs.first);
+    }
+  } mappingComparison;
+
+  auto mappings = self->getActiveArchetypeToInterfaceMappings();
+  auto known = std::lower_bound(mappings.begin(), mappings.end(), archetype,
+                                mappingComparison);
+  if (known != mappings.end() && known->first == archetype)
+    return known->second;
+
+  return Type();
+}
+
 Type GenericEnvironment::mapTypeIntoContext(ModuleDecl *M, Type type) const {
   type = type.subst(M, QueryInterfaceTypeSubstitutions(this),
                     SubstFlags::AllowLoweredTypes);