[AST] Use flat array storage for SubstitutionMap's replacement types.

The replacement types in a SubstitutionMap correspond with the generic
parameters of its generic signature, so replace the DenseMap storage
with a flat array of Types, one element for each generic parameter.

Author: DougGregor

include/swift/AST/SubstitutionMap.h

@@ -55,30 +55,49 @@ class SubstitutionMap {
   /// The generic signature for which we are performing substitutions.
   GenericSignature *genericSig;
 
-  // FIXME: Switch to a more efficient representation.
-  llvm::DenseMap<GenericTypeParamType *, Type> subMap;
+  /// The replacement types for the generic type parameters.
+  std::unique_ptr<Type[]> replacementTypes;
+
+  // FIXME: Switch to a more efficient representation that corresponds to
+  // the conformance requirements in the GenericSignature.
   llvm::DenseMap<TypeBase *, SmallVector<ProtocolConformanceRef, 1>>
     conformanceMap;
 
+  /// Retrieve the array of replacement types, which line up with the
+  /// generic parameters.
+  ///
+  /// Note that the types may be null, for cases where the generic parameter
+  /// is concrete but hasn't been queried yet.
+  ArrayRef<Type> getReplacementTypes() const;
+
+  MutableArrayRef<Type> getReplacementTypes();
+
 public:
   SubstitutionMap()
     : SubstitutionMap(static_cast<GenericSignature *>(nullptr)) { }
 
-  SubstitutionMap(GenericSignature *genericSig)
-    : genericSig(genericSig) { }
+  SubstitutionMap(GenericSignature *genericSig);
 
   SubstitutionMap(GenericEnvironment *genericEnv);
 
+  SubstitutionMap(SubstitutionMap &&other) = default;
+  SubstitutionMap &operator=(SubstitutionMap &&other) = default;
+
+  SubstitutionMap(const SubstitutionMap &other);
+
+  SubstitutionMap &operator=(const SubstitutionMap &other);
+
+  ~SubstitutionMap();
+
   /// Retrieve the generic signature describing the environment in which
   /// substitutions occur.
   GenericSignature *getGenericSignature() const { return genericSig; }
 
   Optional<ProtocolConformanceRef>
   lookupConformance(CanType type, ProtocolDecl *proto) const;
 
-  bool empty() const {
-    return subMap.empty();
-  }
+  /// Whether the substitution map is empty.
+  bool empty() const { return getGenericSignature() == nullptr; }
 
   /// Query whether any replacement types in the map contain archetypes.
   bool hasArchetypes() const;

lib/AST/SubstitutionMap.cpp

@@ -33,27 +33,68 @@
 
 using namespace swift;
 
+ArrayRef<Type> SubstitutionMap::getReplacementTypes() const {
+  if (empty()) return { };
+
+  return llvm::makeArrayRef(replacementTypes.get(),
+                            genericSig->getGenericParams().size());
+}
+
+MutableArrayRef<Type> SubstitutionMap::getReplacementTypes() {
+  if (empty()) return { };
+
+  return MutableArrayRef<Type>(replacementTypes.get(),
+                               genericSig->getGenericParams().size());
+
+}
+
+SubstitutionMap::SubstitutionMap(GenericSignature *genericSig) : genericSig(genericSig) {
+  if (genericSig) {
+    replacementTypes.reset(new Type [genericSig->getGenericParams().size()]);
+  }
+}
+
 SubstitutionMap::SubstitutionMap(GenericEnvironment *genericEnv)
   : SubstitutionMap(genericEnv->getGenericSignature()) { }
 
+SubstitutionMap::SubstitutionMap(const SubstitutionMap &other)
+  : SubstitutionMap(other.getGenericSignature())
+{
+  std::copy(other.getReplacementTypes().begin(),
+            other.getReplacementTypes().end(),
+            getReplacementTypes().begin());
+
+  conformanceMap = other.conformanceMap;
+}
+
+SubstitutionMap &SubstitutionMap::operator=(const SubstitutionMap &other) {
+  *this = SubstitutionMap(other);
+  return *this;
+}
+
+SubstitutionMap::~SubstitutionMap() { }
+
 bool SubstitutionMap::hasArchetypes() const {
-  for (auto &entry : subMap)
-    if (entry.second->hasArchetype())
+  for (Type replacementTy : getReplacementTypes()) {
+    if (replacementTy && replacementTy->hasArchetype())
       return true;
+  }
   return false;
 }
 
 bool SubstitutionMap::hasOpenedExistential() const {
-  for (auto &entry : subMap)
-    if (entry.second->hasOpenedExistential())
+  for (Type replacementTy : getReplacementTypes()) {
+    if (replacementTy && replacementTy->hasOpenedExistential())
       return true;
+  }
   return false;
 }
 
 bool SubstitutionMap::hasDynamicSelf() const {
-  for (auto &entry : subMap)
-    if (entry.second->hasDynamicSelfType())
+  for (Type replacementTy : getReplacementTypes()) {
+    if (replacementTy && replacementTy->hasDynamicSelfType())
       return true;
+  }
   return false;
 }
 
@@ -70,21 +111,36 @@ Type SubstitutionMap::lookupSubstitution(CanSubstitutableType type) const {
       genericEnv->mapTypeOutOfContext(archetype)->getCanonicalType());
   }
 
+  // Find the index of the replacement type based on the generic parameter we
+  // have.
   auto genericParam = cast<GenericTypeParamType>(type);
-  auto known = subMap.find(genericParam);
-  if (known != subMap.end() && known->second)
-    return known->second;
+  auto mutableThis = const_cast<SubstitutionMap *>(this);
+  auto replacementTypes = mutableThis->getReplacementTypes();
+  auto genericParams = getGenericSignature()->getGenericParams();
+  auto replacementIndex =
+    GenericParamKey(genericParam).findIndexIn(genericParams);
+
+  // If this generic parameter isn't represented, we don't have a replacement
+  // type for it.
+  if (replacementIndex == genericParams.size())
+    return Type();
+
+  // If we already have a replacement type, return it.
+  Type &replacementType = replacementTypes[replacementIndex];
+  if (replacementType)
+    return replacementType;
 
   // The generic parameter may have been made concrete by the generic signature,
   // substitute into the concrete type.
   ModuleDecl &anyModule = *genericParam->getASTContext().getStdlibModule();
   auto genericSig = getGenericSignature();
   if (auto concreteType = genericSig->getConcreteType(genericParam, anyModule)){
-    auto mutableThis = const_cast<SubstitutionMap *>(this);
-    mutableThis->subMap[genericParam] = ErrorType::get(concreteType);
-    Type result = concreteType.subst(*this);
-    mutableThis->subMap[genericParam] = result;
-    return result;
+    // Set the replacement type to an error, to block infinite recursion.
+    replacementType = ErrorType::get(concreteType);
+
+    // Substitute into the replacement type.
+    replacementType = concreteType.subst(*this);
+    return replacementType;
   }
 
   // Not known.
@@ -95,9 +151,15 @@ void SubstitutionMap::
 addSubstitution(CanGenericTypeParamType type, Type replacement) {
   assert(getGenericSignature() &&
          "cannot add entries to empty substitution map");
-  auto result = subMap.insert(std::make_pair(type, replacement));
-  assert(result.second || result.first->second->isEqual(replacement));
-  (void) result;
+
+  auto replacementTypes = getReplacementTypes();
+  auto genericParams = getGenericSignature()->getGenericParams();
+  auto replacementIndex = GenericParamKey(type).findIndexIn(genericParams);
+
+  assert((!replacementTypes[replacementIndex] ||
+          replacementTypes[replacementIndex]->isEqual(replacement)));
+
+  replacementTypes[replacementIndex] = replacement;
 }
 
 Optional<ProtocolConformanceRef>
@@ -215,11 +277,11 @@ SubstitutionMap SubstitutionMap::subst(TypeSubstitutionFn subs,
                                        LookupConformanceFn conformances) const {
   SubstitutionMap result(*this);
 
-  for (auto iter = result.subMap.begin(),
-            end = result.subMap.end();
-       iter != end; ++iter) {
-    iter->second = iter->second.subst(subs, conformances,
-                                      SubstFlags::UseErrorType);
+  for (auto &replacementType : result.getReplacementTypes()) {
+    if (replacementType) {
+      replacementType = replacementType.subst(subs, conformances,
+                                              SubstFlags::UseErrorType);
+    }
   }
 
   for (auto iter = result.conformanceMap.begin(),
@@ -425,11 +487,16 @@ void SubstitutionMap::dump(llvm::raw_ostream &out) const {
   genericSig->print(out);
   out << "\n";
   out << "Substitutions:\n";
-  for (const auto &sub : subMap) {
+  auto genericParams = genericSig->getGenericParams();
+  auto replacementTypes = getReplacementTypes();
+  for (unsigned i : indices(genericParams)) {
     out.indent(2);
-    sub.first->print(out);
+    genericParams[i]->print(out);
     out << " -> ";
-    sub.second->print(out);
+    if (replacementTypes[i])
+      replacementTypes[i]->print(out);
+    else
+      out << "<<unresolved concrete type>>";
     out << "\n";
   }