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

DougGregor · DougGregor · commit 5eddbc7e8d36 · 2016-12-01T16:36:03.000-08:00
Store the archetype-to-interface-type mapping (which is used to map
*out* of a generic environment) is a tail-allocated array of
(archetype, generic type parameter) pairs. This array is built up
lazily, as we compute the context types for generic parameters.

Searching in this array is linear while it is being constructed. Once
it is complete, it is sorted so that future searches are logarithmic.

Aside from the space savings of not having a DenseMap lying around,
this means we no longer need to register a destructor of a
GenericEnvironment with the ASTContext, which saves us tear-down
time.
diff --git a/include/swift/AST/GenericEnvironment.h b/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;
diff --git a/lib/AST/ASTContext.cpp b/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);
 }
diff --git a/lib/AST/GenericEnvironment.cpp b/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);
