AST: Add SubstitutionMap::subst()

This will replace Substitution::subst().

Author: slavapestov

include/swift/AST/SubstitutionMap.h

@@ -93,6 +93,15 @@ class SubstitutionMap {
   /// Query whether any replacement type sin the map contain dynamic Self.
   bool hasDynamicSelf() const;
 
+  /// Apply a substitution to all replacement types in the map. Does not
+  /// change keys.
+  SubstitutionMap subst(const SubstitutionMap &subMap) const;
+
+  /// Apply a substitution to all replacement types in the map. Does not
+  /// change keys.
+  SubstitutionMap subst(TypeSubstitutionFn subs,
+                        LookupConformanceFn conformances) const;
+
   /// Create a substitution map for a protocol conformance.
   static SubstitutionMap
   getProtocolSubstitutions(ProtocolDecl *protocol,

lib/AST/ProtocolConformance.cpp

@@ -638,7 +638,7 @@ ProtocolConformance::subst(Type substType,
         subMap = genericSig->getSubstitutionMap(
           [&](SubstitutableType *t) -> Type {
             return genericEnv->mapTypeIntoContext(
-              t).subst(subs, conformances);
+              t).subst(subs, conformances, SubstFlags::UseErrorType);
           },
           [&](CanType origType, Type substType, ProtocolType *protoType)
             -> Optional<ProtocolConformanceRef> {
@@ -675,14 +675,15 @@ ProtocolConformance::subst(Type substType,
   case ProtocolConformanceKind::Specialized: {
     // Substitute the substitutions in the specialized conformance.
     auto spec = cast<SpecializedProtocolConformance>(this);
-    SmallVector<Substitution, 8> newSubs;
-    newSubs.reserve(spec->getGenericSubstitutions().size());
-    for (auto &sub : spec->getGenericSubstitutions())
-      newSubs.push_back(sub.subst(subs, conformances));
-    
+    auto genericConformance
+      = cast<SpecializedProtocolConformance>(this)->getGenericConformance();
+    auto subMap =
+      genericConformance->getGenericSignature()
+        ->getSubstitutionMap(spec->getGenericSubstitutions());
+
     return substType->getASTContext()
-      .getSpecializedConformance(substType, spec->getGenericConformance(),
-                                 newSubs);
+      .getSpecializedConformance(substType, genericConformance,
+                                 subMap.subst(subs, conformances));
   }
   }
   llvm_unreachable("bad ProtocolConformanceKind");
@@ -713,7 +714,7 @@ ProtocolConformance::getInheritedConformance(ProtocolDecl *protocol) const {
            && "substitution didn't produce conformance for same type?!");
     return r;
   }
-    
+
   case ProtocolConformanceKind::Inherited: {
     auto classInherited = cast<InheritedProtocolConformance>(this);
     auto protoInherited = classInherited->getInheritedConformance()

lib/AST/SubstitutionMap.cpp

@@ -222,6 +222,38 @@ addParent(CanType type, CanType parent, AssociatedTypeDecl *assocType) {
   parentMap[type.getPointer()].push_back(std::make_pair(parent, assocType));
 }
 
+
+SubstitutionMap SubstitutionMap::subst(const SubstitutionMap &subMap) const {
+  return subst(QuerySubstitutionMap{subMap},
+               LookUpConformanceInSubstitutionMap(subMap));
+}
+
+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 iter = result.conformanceMap.begin(),
+            end = result.conformanceMap.end();
+       iter != end; ++iter) {
+    auto origType = Type(iter->first).subst(
+        *this, SubstFlags::UseErrorType);
+    for (auto citer = iter->second.begin(),
+              cend = iter->second.end();
+         citer != cend; ++citer) {
+      *citer = citer->subst(origType, subs, conformances);
+    }
+  }
+
+  return result;
+}
+
 SubstitutionMap
 SubstitutionMap::getProtocolSubstitutions(ProtocolDecl *protocol,
                                           Type selfType,