AST: Make GenericSignature and GenericEnvironment SubstitutionMaps interchangable

SubstitutionMap::lookupConformance() would map archetypes out
of context to compute a conformance path. Do the same thing
in SubstitutionMap::lookupSubstitution().

The DenseMap of replacement types in a SubstitutionMap now
always has GenericTypeParamTypes as keys.

This simplifies some code and brings us one step closer to
a more efficient representation of SubstitutionMaps.

Author: slavapestov

include/swift/AST/GenericEnvironment.h

@@ -226,12 +226,6 @@ class alignas(1 << DeclAlignInBits) GenericEnvironment final
   /// generic parameter types by their sugared form.
   Type getSugaredType(Type type) const;
 
-  /// Derive a contextual type substitution map from a substitution array.
-  /// This is just like GenericSignature::getSubstitutionMap(), except
-  /// with contextual types instead of interface types.
-  SubstitutionMap
-  getSubstitutionMap(SubstitutionList subs) const;
-
   /// Build a contextual type substitution map from a type substitution function
   /// and conformance lookup function.
   SubstitutionMap

include/swift/AST/SubstitutionMap.h

@@ -41,6 +41,7 @@ namespace swift {
 class GenericSignature;
 class GenericEnvironment;
 class SubstitutableType;
+typedef CanTypeWrapper<GenericTypeParamType> CanGenericTypeParamType;
 
 template<class Type> class CanTypeWrapper;
 typedef CanTypeWrapper<SubstitutableType> CanSubstitutableType;
@@ -55,7 +56,7 @@ class SubstitutionMap {
   GenericSignature *genericSig;
 
   // FIXME: Switch to a more efficient representation.
-  llvm::DenseMap<SubstitutableType *, Type> subMap;
+  llvm::DenseMap<GenericTypeParamType *, Type> subMap;
   llvm::DenseMap<TypeBase *, SmallVector<ProtocolConformanceRef, 1>>
     conformanceMap;
 
@@ -165,7 +166,7 @@ class SubstitutionMap {
   // instead, use GenericSignature::getSubstitutionMap() or
   // GenericEnvironment::getSubstitutionMap().
 
-  void addSubstitution(CanSubstitutableType type, Type replacement);
+  void addSubstitution(CanGenericTypeParamType type, Type replacement);
   void addConformance(CanType type, ProtocolConformanceRef conformance);
 };
 

include/swift/SIL/TypeSubstCloner.h

@@ -41,8 +41,9 @@ class TypeSubstCloner : public SILClonerWithScopes<ImplClass> {
   }
 
   void computeSubsMap() {
-    if (auto *env = Original.getGenericEnvironment()) {
-      SubsMap = env->getSubstitutionMap(ApplySubs);
+    if (auto genericSig = Original.getLoweredFunctionType()
+          ->getGenericSignature()) {
+      SubsMap = genericSig->getSubstitutionMap(ApplySubs);
     }
   }
 

lib/AST/GenericEnvironment.cpp

@@ -349,38 +349,6 @@ GenericEnvironment::getForwardingSubstitutions() const {
   return getGenericSignature()->getASTContext().AllocateCopy(result);
 }
 
-SubstitutionMap GenericEnvironment::
-getSubstitutionMap(SubstitutionList subs) const {
-  SubstitutionMap result(const_cast<GenericEnvironment *>(this));
-
-  getGenericSignature()->enumeratePairedRequirements(
-    [&](Type depTy, ArrayRef<Requirement> reqts) -> bool {
-      // Map the interface type to a context type.
-      auto contextTy = depTy.subst(QueryInterfaceTypeSubstitutions(this),
-                                   MakeAbstractConformanceForGenericType());
-
-      auto sub = subs.front();
-      subs = subs.slice(1);
-
-      // Record the replacement type and its conformances.
-      if (auto *archetype = contextTy->getAs<ArchetypeType>()) {
-        result.addSubstitution(CanArchetypeType(archetype), sub.getReplacement());
-        assert(reqts.size() == sub.getConformances().size());
-        for (auto conformance : sub.getConformances())
-          result.addConformance(CanType(archetype), conformance);
-        return false;
-      }
-
-      assert(contextTy->hasError());
-      return false;
-    });
-
-  assert(subs.empty() && "did not use all substitutions?!");
-
-  result.verify();
-  return result;
-}
-
 SubstitutionMap
 GenericEnvironment::
 getSubstitutionMap(TypeSubstitutionFn subs,
@@ -389,6 +357,7 @@ getSubstitutionMap(TypeSubstitutionFn subs,
 
   getGenericSignature()->enumeratePairedRequirements(
     [&](Type depTy, ArrayRef<Requirement> reqs) -> bool {
+      auto canTy = depTy->getCanonicalType();
 
       // Map the interface type to a context type.
       auto contextTy = depTy.subst(QueryInterfaceTypeSubstitutions(this),
@@ -397,20 +366,19 @@ getSubstitutionMap(TypeSubstitutionFn subs,
       // Compute the replacement type.
       Type currentReplacement = contextTy.subst(subs, lookupConformance,
                                                 SubstFlags::UseErrorType);
-      if (auto archetypeTy = contextTy->getAs<ArchetypeType>()) {
-        subMap.addSubstitution(CanArchetypeType(archetypeTy),
-                               currentReplacement);
-
-        // Collect the conformances.
-        for (auto req: reqs) {
-          assert(req.getKind() == RequirementKind::Conformance);
-          auto protoType = req.getSecondType()->castTo<ProtocolType>();
-          auto conformance = lookupConformance(CanArchetypeType(archetypeTy),
-                                               currentReplacement,
-                                               protoType);
-          if (conformance)
-            subMap.addConformance(CanArchetypeType(archetypeTy), *conformance);
-        }
+
+      if (auto paramTy = dyn_cast<GenericTypeParamType>(canTy))
+        subMap.addSubstitution(paramTy, currentReplacement);
+
+      // Collect the conformances.
+      for (auto req: reqs) {
+        assert(req.getKind() == RequirementKind::Conformance);
+        auto protoType = req.getSecondType()->castTo<ProtocolType>();
+        auto conformance = lookupConformance(canTy,
+                                             currentReplacement,
+                                             protoType);
+        if (conformance)
+          subMap.addConformance(canTy, *conformance);
       }
 
       return false;

lib/AST/GenericSignature.cpp

@@ -387,8 +387,8 @@ GenericSignature::getSubstitutionMap(SubstitutionList subs) const {
       subs = subs.slice(1);
 
       auto canTy = depTy->getCanonicalType();
-      if (isa<SubstitutableType>(canTy))
-        result.addSubstitution(cast<SubstitutableType>(canTy),
+      if (auto paramTy = dyn_cast<GenericTypeParamType>(canTy))
+        result.addSubstitution(paramTy,
                                sub.getReplacement());
       assert(reqts.size() == sub.getConformances().size());
       for (auto conformance : sub.getConformances())
@@ -415,8 +415,8 @@ getSubstitutionMap(TypeSubstitutionFn subs,
     // Compute the replacement type.
     Type currentReplacement = depTy.subst(subs, lookupConformance,
                                           SubstFlags::UseErrorType);
-    if (auto substTy = dyn_cast<SubstitutableType>(canTy))
-      subMap.addSubstitution(substTy, currentReplacement);
+    if (auto paramTy = dyn_cast<GenericTypeParamType>(canTy))
+      subMap.addSubstitution(paramTy, currentReplacement);
 
     // Collect the conformances.
     for (auto req: reqs) {

lib/AST/ProtocolConformance.cpp

@@ -696,10 +696,10 @@ SpecializedProtocolConformance::getTypeWitnessAndDecl(
   }
 
   // Otherwise, perform substitutions to create this witness now.
-  auto *genericEnv = GenericConformance->getGenericEnvironment();
+  auto *genericSig = GenericConformance->getGenericSignature();
 
   auto substitutionMap =
-      genericEnv->getSubstitutionMap(GenericSubstitutions);
+      genericSig->getSubstitutionMap(GenericSubstitutions);
 
   auto genericWitnessAndDecl
     = GenericConformance->getTypeWitnessAndDecl(assocType, resolver);
@@ -730,8 +730,8 @@ SpecializedProtocolConformance::getAssociatedConformance(Type assocType,
   ProtocolConformanceRef conformance =
     GenericConformance->getAssociatedConformance(assocType, protocol, resolver);
 
-  auto genericEnv = GenericConformance->getGenericEnvironment();
-  auto subMap = genericEnv->getSubstitutionMap(GenericSubstitutions);
+  auto genericSig = GenericConformance->getGenericSignature();
+  auto subMap = genericSig->getSubstitutionMap(GenericSubstitutions);
 
   Type origType =
     (conformance.isConcrete()
@@ -811,30 +811,10 @@ ProtocolConformance::subst(Type substType,
                == substType->getNominalOrBoundGenericNominal()
              && "substitution mapped to different nominal?!");
 
-      // Since this is a normal conformance, the substitution maps archetypes
-      // in the environment of the conformance to types containing archetypes
-      // of some other generic environment.
-      //
-      // ASTContext::getSpecializedConformance() wants a substitution map
-      // with interface types as keys, so do the mapping here.
-      //
-      // Once the type of a normal conformance becomes an interface type,
-      // we can remove this.
       SubstitutionMap subMap;
       if (auto *genericSig = getGenericSignature()) {
         auto *genericEnv = getGenericEnvironment();
-        subMap = genericSig->getSubstitutionMap(
-          [&](SubstitutableType *t) -> Type {
-            return genericEnv->mapTypeIntoContext(
-              t).subst(subs, conformances, SubstFlags::UseErrorType);
-          },
-          [&](CanType origType, Type substType, ProtocolType *protoType)
-            -> Optional<ProtocolConformanceRef> {
-            origType = CanType(
-              genericEnv->mapTypeIntoContext(
-                origType)->castTo<ArchetypeType>());
-            return conformances(origType, substType, protoType);
-          });
+        subMap = genericEnv->getSubstitutionMap(subs, conformances);
       }
 
       return substType->getASTContext()

lib/AST/SubstitutionMap.cpp

@@ -58,7 +58,19 @@ bool SubstitutionMap::hasDynamicSelf() const {
 }
 
 Type SubstitutionMap::lookupSubstitution(CanSubstitutableType type) const {
-  auto known = subMap.find(type);
+  // If we have an archetype, map out of the context so we can compute a
+  // conformance access path.
+  if (auto archetype = dyn_cast<ArchetypeType>(type)) {
+    if (archetype->isOpenedExistential() ||
+        archetype->getParent() != nullptr)
+      return Type();
+
+    auto *genericEnv = archetype->getGenericEnvironment();
+    type = cast<GenericTypeParamType>(
+      genericEnv->mapTypeOutOfContext(archetype)->getCanonicalType());
+  }
+
+  auto known = subMap.find(cast<GenericTypeParamType>(type));
   if (known != subMap.end() && known->second)
     return known->second;
 
@@ -67,23 +79,20 @@ Type SubstitutionMap::lookupSubstitution(CanSubstitutableType type) const {
 }
 
 void SubstitutionMap::
-addSubstitution(CanSubstitutableType type, Type replacement) {
-  assert(!(type->isTypeParameter() && !getGenericSignature()) &&
-         "type parameter substitution map without generic signature");
+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;
 }
 
 Optional<ProtocolConformanceRef>
-SubstitutionMap::lookupConformance(CanType origType, ProtocolDecl *proto) const {
-  CanType type = origType;
-
+SubstitutionMap::lookupConformance(CanType type, ProtocolDecl *proto) const {
   // If we have an archetype, map out of the context so we can compute a
   // conformance access path.
-  GenericEnvironment *genericEnv = nullptr;
-  if (auto archetype = type->getAs<ArchetypeType>()) {
-    genericEnv = archetype->getGenericEnvironment();
+  if (auto archetype = dyn_cast<ArchetypeType>(type)) {
+    auto *genericEnv = archetype->getGenericEnvironment();
     type = genericEnv->mapTypeOutOfContext(type)->getCanonicalType();
   }
 
@@ -95,11 +104,6 @@ SubstitutionMap::lookupConformance(CanType origType, ProtocolDecl *proto) const
   // Retrieve the starting conformance from the conformance map.
   auto getInitialConformance =
     [&](Type type, ProtocolDecl *proto) -> Optional<ProtocolConformanceRef> {
-      // We're working relative to a generic environment, map into that
-      // context before looking into the conformance map.
-      if (genericEnv)
-        type = genericEnv->mapTypeIntoContext(type);
-
       auto known = conformanceMap.find(type->getCanonicalType().getPointer());
       if (known == conformanceMap.end())
         return None;
@@ -148,7 +152,7 @@ SubstitutionMap::lookupConformance(CanType origType, ProtocolDecl *proto) const
       // FIXME: Rip this out once we can get a concrete conformance from
       // an archetype.
       auto *M = proto->getParentModule();
-      auto substType = origType.subst(*this);
+      auto substType = type.subst(*this);
       if (substType &&
           !substType->is<ArchetypeType>() &&
           !substType->isTypeParameter() &&
@@ -185,6 +189,7 @@ SubstitutionMap::lookupConformance(CanType origType, ProtocolDecl *proto) const
 
 void SubstitutionMap::
 addConformance(CanType type, ProtocolConformanceRef conformance) {
+  assert(!isa<ArchetypeType>(type));
   conformanceMap[type.getPointer()].push_back(conformance);
 }
 

lib/SILGen/SILGenPoly.cpp

@@ -2646,7 +2646,7 @@ buildThunkSignature(SILGenFunction &gen,
     genericEnv = gen.F.getGenericEnvironment();
     auto subsArray = gen.F.getForwardingSubstitutions();
     interfaceSubs = genericSig->getSubstitutionMap(subsArray);
-    contextSubs = genericEnv->getSubstitutionMap(subsArray);
+    contextSubs = interfaceSubs;
     return genericSig;
   }
 
@@ -2681,11 +2681,11 @@ buildThunkSignature(SILGenFunction &gen,
 
   // Calculate substitutions to map the caller's archetypes to the thunk's
   // archetypes.
-  if (auto *calleeGenericEnv = gen.F.getGenericEnvironment()) {
-    contextSubs = calleeGenericEnv->getSubstitutionMap(
+  if (auto calleeGenericSig = gen.F.getLoweredFunctionType()
+          ->getGenericSignature()) {
+    contextSubs = calleeGenericSig->getSubstitutionMap(
       [&](SubstitutableType *type) -> Type {
-        auto depTy = calleeGenericEnv->mapTypeOutOfContext(type);
-        return genericEnv->mapTypeIntoContext(depTy);
+        return genericEnv->mapTypeIntoContext(type);
       },
       MakeAbstractConformanceForGenericType());
   }

lib/SILOptimizer/Transforms/PerformanceInliner.cpp

@@ -194,7 +194,8 @@ static bool calleeHasPartialApplyWithOpenedExistentials(FullApplySite AI) {
   if (HasNoOpenedExistentials)
     return false;
 
-  auto SubsMap = Callee->getGenericEnvironment()->getSubstitutionMap(Subs);
+  auto SubsMap = Callee->getLoweredFunctionType()
+    ->getGenericSignature()->getSubstitutionMap(Subs);
 
   for (auto &BB : *Callee) {
     for (auto &I : BB) {
@@ -355,7 +356,8 @@ bool SILPerformanceInliner::isProfitableToInline(FullApplySite AI,
 
   SubstitutionMap CalleeSubstMap;
   if (IsGeneric) {
-    CalleeSubstMap = Callee->getGenericEnvironment()
+    CalleeSubstMap = Callee->getLoweredFunctionType()
+      ->getGenericSignature()
       ->getSubstitutionMap(AI.getSubstitutions());
   }