Merge pull request swiftlang#76585 from slavapestov/small-subst-cleanups-2

More type substitution cleanups

Author: slavapestov

include/swift/AST/PackConformance.h

@@ -74,7 +74,8 @@ class alignas(1 << DeclAlignInBits) PackConformance final
 
   PackConformance *getCanonicalConformance() const;
 
-  PackType *getAssociatedType(Type assocType) const;
+  PackType *getTypeWitness(AssociatedTypeDecl *assocType,
+                           SubstOptions options=std::nullopt) const;
 
   PackConformance *
   getAssociatedConformance(Type assocType, ProtocolDecl *protocol) const;

include/swift/AST/ProtocolConformanceRef.h

@@ -170,6 +170,11 @@ class ProtocolConformanceRef {
   /// Map contextual types to interface types in the conformance.
   ProtocolConformanceRef mapConformanceOutOfContext() const;
 
+  /// Look up the type witness for an associated type declaration in this
+  /// conformance.
+  Type getTypeWitness(Type origType, AssociatedTypeDecl *assocType,
+                      SubstOptions options = std::nullopt) const;
+
   /// Given a dependent type (expressed in terms of this conformance's
   /// protocol), follow it from the conforming type.
   Type getAssociatedType(Type origType, Type dependentType) const;

lib/AST/GenericSignature.cpp

@@ -412,7 +412,8 @@ bool GenericSignatureImpl::isRequirementSatisfied(
 
     requirement = requirement.subst(
         QueryInterfaceTypeSubstitutions{genericEnv},
-        LookUpConformanceInModule());
+        LookUpConformanceInModule(),
+        SubstFlags::PreservePackExpansionLevel);
   }
 
   SmallVector<Requirement, 2> subReqs;
@@ -1173,7 +1174,39 @@ void swift::validateGenericSignature(ASTContext &context,
 
     // If the removed requirement is satisfied by the new generic signature,
     // it is redundant. Complain.
-    if (newSig->isRequirementSatisfied(requirements[victimIndex])) {
+    auto satisfied = [&](Requirement victim) {
+      if (!newSig->isValidTypeParameter(victim.getFirstType()))
+        return false;
+
+      switch (victim.getKind()) {
+      case RequirementKind::SameShape:
+        return (newSig->isValidTypeParameter(victim.getSecondType()) &&
+                newSig->haveSameShape(victim.getFirstType(),
+                                      victim.getSecondType()));
+      case RequirementKind::Conformance:
+        return newSig->requiresProtocol(victim.getFirstType(),
+                                        victim.getProtocolDecl());
+      case RequirementKind::Superclass: {
+        auto superclass = newSig->getSuperclassBound(victim.getFirstType());
+        return (superclass && superclass->isEqual(victim.getSecondType()));
+      }
+      case RequirementKind::SameType:
+        if (!victim.getSecondType().findIf([&](Type t) -> bool {
+            return (!t->isTypeParameter() ||
+                    newSig->isValidTypeParameter(t));
+          })) {
+          return false;
+        }
+        return newSig.getReducedType(victim.getFirstType())
+          ->isEqual(newSig.getReducedType(victim.getSecondType()));
+      case RequirementKind::Layout: {
+        auto layout = newSig->getLayoutConstraint(victim.getFirstType());
+        return (layout && layout == victim.getLayoutConstraint());
+      }
+      }
+    };
+
+    if (satisfied(requirements[victimIndex])) {
       SmallString<32> reqString;
       {
         llvm::raw_svector_ostream out(reqString);

lib/AST/PackConformance.cpp

@@ -97,7 +97,9 @@ PackConformance *PackConformance::getCanonicalConformance() const {
 /// Project the corresponding associated type from each pack element
 /// of the conforming type, collecting the results into a new pack type
 /// that has the same pack expansion structure as the conforming type.
-PackType *PackConformance::getAssociatedType(Type assocType) const {
+PackType *PackConformance::getTypeWitness(
+    AssociatedTypeDecl *assocType,
+    SubstOptions options) const {
   SmallVector<Type, 4> packElements;
 
   auto conformances = getPatternConformances();
@@ -110,8 +112,8 @@ PackType *PackConformance::getAssociatedType(Type assocType) const {
     // conformance.
     if (auto *packExpansion = packElement->getAs<PackExpansionType>()) {
       auto assocTypePattern =
-        conformances[i].getAssociatedType(packExpansion->getPatternType(),
-                                          assocType);
+        conformances[i].getTypeWitness(packExpansion->getPatternType(),
+                                       assocType, options);
 
       packElements.push_back(PackExpansionType::get(
           assocTypePattern, packExpansion->getCountType()));
@@ -120,7 +122,7 @@ PackType *PackConformance::getAssociatedType(Type assocType) const {
     // the associated type witness from the pattern conformance.
     } else {
       auto assocTypeScalar =
-        conformances[i].getAssociatedType(packElement, assocType);
+        conformances[i].getTypeWitness(packElement, assocType, options);
       packElements.push_back(assocTypeScalar);
     }
   }

lib/AST/ProtocolConformance.cpp

@@ -147,7 +147,27 @@ ProtocolConformance::getTypeWitnessAndDecl(AssociatedTypeDecl *assocType,
 
 Type ProtocolConformance::getTypeWitness(AssociatedTypeDecl *assocType,
                                          SubstOptions options) const {
-  return getTypeWitnessAndDecl(assocType, options).getWitnessType();
+  auto witness = getTypeWitnessAndDecl(assocType, options);
+  auto witnessTy = witness.getWitnessType();
+  if (!witnessTy)
+    return witnessTy;
+
+  // This is a hacky feature allowing code completion to migrate to
+  // using Type::subst() without changing output.
+  //
+  // FIXME: Remove this hack and do whatever we need to do in the
+  // ASTPrinter instead.
+  if (options & SubstFlags::DesugarMemberTypes) {
+    if (auto *aliasType = dyn_cast<TypeAliasType>(witnessTy.getPointer()))
+      witnessTy = aliasType->getSinglyDesugaredType();
+
+    // Another hack. If the type witness is a opaque result type. They can
+    // only be referred using the name of the associated type.
+    if (witnessTy->is<OpaqueTypeArchetypeType>())
+      witnessTy = witness.getWitnessDecl()->getDeclaredInterfaceType();
+  }
+
+  return witnessTy;
 }
 
 ConcreteDeclRef

lib/AST/ProtocolConformanceRef.cpp

@@ -20,6 +20,7 @@
 #include "swift/AST/Availability.h"
 #include "swift/AST/ConformanceLookup.h"
 #include "swift/AST/Decl.h"
+#include "swift/AST/GenericEnvironment.h"
 #include "swift/AST/InFlightSubstitution.h"
 #include "swift/AST/Module.h"
 #include "swift/AST/PackConformance.h"
@@ -163,8 +164,7 @@ ProtocolConformanceRef::getTypeWitnessByName(Type type, Identifier name) const {
   if (!assocType)
     return ErrorType::get(proto->getASTContext());
 
-  return assocType->getDeclaredInterfaceType().subst(
-    SubstitutionMap::getProtocolSubstitutions(proto, type, *this));
+  return getTypeWitness(type, assocType);
 }
 
 ConcreteDeclRef
@@ -194,78 +194,115 @@ ProtocolConformanceRef::getConditionalRequirements() const {
     return {};
 }
 
-Type ProtocolConformanceRef::getAssociatedType(Type conformingType,
-                                               Type assocType) const {
+Type ProtocolConformanceRef::getTypeWitness(Type conformingType,
+                                            AssociatedTypeDecl *assocType,
+                                            SubstOptions options) const {
   if (isPack()) {
     auto *pack = getPack();
     ASSERT(conformingType->isEqual(pack->getType()));
-    return pack->getAssociatedType(assocType);
+    return pack->getTypeWitness(assocType);
   }
 
-  auto type = assocType->getCanonicalType();
-
-  // Fast path for generic parameters.
-  if (auto paramTy = dyn_cast<GenericTypeParamType>(type)) {
-    ASSERT(paramTy->getDepth() == 0 && paramTy->getIndex() == 0 &&
-           "type parameter in protocol was not Self");
-    return conformingType;
-  }
+  auto failed = [&]() {
+    return DependentMemberType::get(ErrorType::get(conformingType),
+                                    assocType);
+  };
 
   if (isInvalid())
-    return ErrorType::get(assocType->getASTContext());
+    return failed();
 
   auto proto = getRequirement();
+  ASSERT(assocType->getProtocol() == proto);
 
   if (isConcrete()) {
-    if (auto selfType = conformingType->getAs<DynamicSelfType>())
-      conformingType = selfType->getSelfType();
-    ASSERT(getConcrete()->getType()->isEqual(conformingType));
-
-    // Fast path for dependent member types on 'Self' of our associated types.
-    auto memberType = cast<DependentMemberType>(type);
-    if (memberType.getBase()->isEqual(proto->getSelfInterfaceType()) &&
-        memberType->getAssocType()->getProtocol() == proto) {
-      auto witnessType = getConcrete()->getTypeWitness(
-          memberType->getAssocType());
-      if (!witnessType)
-        return ErrorType::get(assocType->getASTContext());
-      return witnessType;
-    }
+    auto witnessType = getConcrete()->getTypeWitness(assocType, options);
+    if (!witnessType || witnessType->is<ErrorType>())
+      return failed();
+    return witnessType;
   }
 
-  // General case: consult the substitution map.
+  ASSERT(isAbstract());
+
+  if (auto *archetypeType = conformingType->getAs<ArchetypeType>()) {
+    return archetypeType->getNestedType(assocType);
+  }
+
+  CONDITIONAL_ASSERT(conformingType->isTypeParameter() ||
+                     conformingType->isTypeVariableOrMember() ||
+                     conformingType->is<UnresolvedType>() ||
+                     conformingType->is<PlaceholderType>());
+
+  return DependentMemberType::get(conformingType, assocType);
+}
+
+Type ProtocolConformanceRef::getAssociatedType(Type conformingType,
+                                               Type assocType) const {
+  if (isInvalid())
+    return ErrorType::get(assocType->getASTContext());
+
+  auto proto = getRequirement();
+
   auto substMap =
     SubstitutionMap::getProtocolSubstitutions(proto, conformingType, *this);
-  return type.subst(substMap);
+  return assocType.subst(substMap);
 }
 
 ProtocolConformanceRef
 ProtocolConformanceRef::getAssociatedConformance(Type conformingType,
                                                  Type assocType,
                                                  ProtocolDecl *protocol) const {
-  // If this is a pack conformance, project the associated conformances.
+  // If this is a pack conformance, project the associated conformances from
+  // each pack element.
   if (isPack()) {
     auto *pack = getPack();
     assert(conformingType->isEqual(pack->getType()));
     return ProtocolConformanceRef(
         pack->getAssociatedConformance(assocType, protocol));
   }
 
-  // If this is a concrete conformance, look up the associated conformance.
+  // If this is a concrete conformance, project the associated conformance.
   if (isConcrete()) {
     auto conformance = getConcrete();
     assert(conformance->getType()->isEqual(conformingType));
     return conformance->getAssociatedConformance(assocType, protocol);
   }
 
-  // Otherwise, apply the substitution {self -> conformingType}
-  // to the abstract conformance requirement laid upon the dependent type
-  // by the protocol.
-  auto subMap =
-    SubstitutionMap::getProtocolSubstitutions(getRequirement(),
-                                              conformingType, *this);
-  auto abstractConf = ProtocolConformanceRef(protocol);
-  return abstractConf.subst(assocType, subMap);
+  // An associated conformance of an archetype might be known to be
+  // a concrete conformance, if the subject type is fixed to a concrete
+  // type in the archetype's generic signature. We don't actually have
+  // any way to recover the conformance in this case, except via global
+  // conformance lookup.
+  //
+  // However, if we move to a first-class representation of abstract
+  // conformances where they store their subject types, we can also
+  // cache the lookups inside the abstract conformance instance too.
+  if (auto archetypeType = conformingType->getAs<ArchetypeType>()) {
+    conformingType = archetypeType->getInterfaceType();
+    auto *genericEnv = archetypeType->getGenericEnvironment();
+
+    auto subjectType = assocType.transformRec(
+      [&](TypeBase *t) -> std::optional<Type> {
+        if (isa<GenericTypeParamType>(t))
+          return conformingType;
+        return std::nullopt;
+      });
+
+    return lookupConformance(
+        genericEnv->mapTypeIntoContext(subjectType),
+        protocol);
+  }
+
+  // Associated conformances of type parameters and type variables
+  // are always abstract, because we don't know the output generic
+  // signature of the substitution (or in the case of type variables,
+  // we have no visibility into constraints). See the parallel hack
+  // to handle this in SubstitutionMap::lookupConformance().
+  CONDITIONAL_ASSERT(conformingType->isTypeParameter() ||
+                     conformingType->isTypeVariableOrMember() ||
+                     conformingType->is<UnresolvedType>() ||
+                     conformingType->is<PlaceholderType>());
+
+  return ProtocolConformanceRef(protocol);
 }
 
 /// Check of all types used by the conformance are canonical.

lib/AST/RequirementMachine/ConcreteContraction.cpp

@@ -302,8 +302,7 @@ ConcreteContraction::substTypeParameterRec(Type type, Position position) const {
         return std::nullopt;
       }
 
-      return conformance.getAssociatedType(
-          *substBaseType, assocType->getDeclaredInterfaceType());
+      return conformance.getTypeWitness(*substBaseType, assocType);
     }
 
     // An unresolved DependentMemberType stores an identifier. Handle this

lib/AST/RequirementMachine/GenericSignatureQueries.cpp

@@ -346,9 +346,7 @@ static Type substPrefixType(Type type, unsigned suffixLength, Type prefixType,
   auto *assocDecl = memberType->getAssocType();
   auto *proto = assocDecl->getProtocol();
   auto conformance = lookupConformance(substBaseType, proto);
-  return conformance.getAssociatedType(
-      substBaseType,
-      assocDecl->getDeclaredInterfaceType());
+  return conformance.getTypeWitness(substBaseType, assocDecl);
 }
 
 Type RequirementMachine::getReducedTypeParameter(

lib/AST/RequirementMachine/RequirementLowering.cpp

@@ -665,8 +665,7 @@ struct InferRequirementsWalker : public TypeWalker {
         auto addSameTypeConstraint = [&](Type firstType,
                                          AssociatedTypeDecl *assocType) {
           auto conformance = lookupConformance(firstType, differentiableProtocol);
-          auto secondType = conformance.getAssociatedType(
-              firstType, assocType->getDeclaredInterfaceType());
+          auto secondType = conformance.getTypeWitness(firstType, assocType);
           reqs.push_back({Requirement(RequirementKind::SameType,
                                       firstType, secondType),
                           SourceLoc()});

lib/AST/Type.cpp

@@ -3747,12 +3747,16 @@ void ParameterizedProtocolType::getRequirements(
   assert(argTypes.size() <= assocTypes.size());
 
   auto conformance = lookupConformance(baseType, protoDecl);
+  auto subMap = SubstitutionMap::getProtocolSubstitutions(
+      protoDecl, baseType, conformance);
 
   for (unsigned i : indices(argTypes)) {
     auto argType = argTypes[i];
     auto *assocType = assocTypes[i];
-    auto subjectType = conformance.getAssociatedType(
-        baseType, assocType->getDeclaredInterfaceType());
+    // Do a general type substitution here because the associated type might be
+    // from an inherited protocol, in which case we will evaluate a non-trivial
+    // conformance path.
+    auto subjectType = assocType->getDeclaredInterfaceType().subst(subMap);
     reqs.emplace_back(RequirementKind::SameType, subjectType, argType);
   }
 }
@@ -4463,8 +4467,7 @@ TypeBase::getAutoDiffTangentSpace(LookupConformanceFn lookupConformance) {
   // Try to get the `TangentVector` associated type of `base`.
   // Return the associated type if it is valid.
   auto conformance = swift::lookupConformance(this, differentiableProtocol);
-  auto assocTy = conformance.getAssociatedType(
-      this, assocDecl->getDeclaredInterfaceType());
+  auto assocTy = conformance.getTypeWitness(this, assocDecl);
   if (!assocTy->hasError())
     return cache(TangentSpace::getTangentVector(assocTy));