Merge pull request #76536 from slavapestov/small-subst-cleanups

Tiny optimization and cleanups

Author: slavapestov

include/swift/AST/Types.h

@@ -7158,22 +7158,6 @@ class DependentMemberType : public TypeBase {
   AssociatedTypeDecl *getAssocType() const {
     return NameOrAssocType.dyn_cast<AssociatedTypeDecl *>();
   }
-  
-  /// Substitute the base type, looking up our associated type in it if it is
-  /// non-dependent. Returns null if the member could not be found in the new
-  /// base.
-  Type substBaseType(Type base);
-
-  /// Substitute the base type, looking up our associated type in it if it is
-  /// non-dependent. Returns null if the member could not be found in the new
-  /// base.
-  Type substBaseType(Type base, LookupConformanceFn lookupConformance,
-                     SubstOptions options);
-
-  /// Substitute the root generic type, looking up the chain of associated types.
-  /// Returns null if the member could not be found in the new root.
-  Type substRootParam(Type newRoot, LookupConformanceFn lookupConformance,
-                      SubstOptions options);
 
   // Implement isa/cast/dyncast/etc.
   static bool classof(const TypeBase *T) {

lib/AST/ASTDemangler.cpp

@@ -305,29 +305,29 @@ Type ASTBuilder::createBoundGenericType(GenericTypeDecl *decl,
   if (auto *nominalDecl = dyn_cast<NominalTypeDecl>(decl))
     return BoundGenericType::get(nominalDecl, parent, args);
 
-  // Combine the substitutions from our parent type with our generic
-  // arguments.
-  TypeSubstitutionMap subs;
-  if (parent)
-    subs = parent->getContextSubstitutions(decl->getDeclContext());
-
   auto *aliasDecl = cast<TypeAliasDecl>(decl);
+  auto *dc = aliasDecl->getDeclContext();
 
-  auto genericSig = aliasDecl->getGenericSignature();
-  for (unsigned i = 0, e = args.size(); i < e; ++i) {
-    auto origTy = genericSig.getInnermostGenericParams()[i];
-    auto substTy = args[i];
+  SmallVector<Type, 2> subs;
 
-    subs[origTy->getCanonicalType()->castTo<GenericTypeParamType>()] =
-      substTy;
+  // Combine the substitutions from our parent type with our generic
+  // arguments.
+  if (dc->isLocalContext()) {
+    for (auto *param : dc->getGenericSignatureOfContext().getGenericParams()) {
+      subs.push_back(param);
+    }
+  } else if (parent) {
+    auto parentSubs = parent->getContextSubstitutionMap(
+        dc).getReplacementTypes();
+    subs.append(parentSubs.begin(), parentSubs.end());
   }
 
-  auto subMap = SubstitutionMap::get(genericSig,
-                                     QueryTypeSubstitutionMap{subs},
-                                     LookUpConformanceInModule());
-  if (!subMap)
-    return Type();
+  auto genericSig = aliasDecl->getGenericSignature();
+  ASSERT(genericSig.getInnermostGenericParams().size() == args.size());
+  subs.append(args.begin(), args.end());
 
+  auto subMap = SubstitutionMap::get(genericSig, subs,
+                                     LookUpConformanceInModule());
   return aliasDecl->getDeclaredInterfaceType().subst(subMap);
 }
 

lib/AST/ProtocolConformanceRef.cpp

@@ -198,28 +198,40 @@ Type ProtocolConformanceRef::getAssociatedType(Type conformingType,
                                                Type assocType) const {
   if (isPack()) {
     auto *pack = getPack();
-    assert(conformingType->isEqual(pack->getType()));
+    ASSERT(conformingType->isEqual(pack->getType()));
     return pack->getAssociatedType(assocType);
   }
 
-  assert(!isConcrete() || getConcrete()->getType()->isEqual(conformingType));
-
   auto type = assocType->getCanonicalType();
-  auto proto = getRequirement();
 
   // Fast path for generic parameters.
-  if (isa<GenericTypeParamType>(type)) {
-    assert(type->isEqual(proto->getSelfInterfaceType()) &&
+  if (auto paramTy = dyn_cast<GenericTypeParamType>(type)) {
+    ASSERT(paramTy->getDepth() == 0 && paramTy->getIndex() == 0 &&
            "type parameter in protocol was not Self");
     return 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 &&
-      isConcrete())
-    return getConcrete()->getTypeWitness(memberType->getAssocType());
+  if (isInvalid())
+    return ErrorType::get(assocType->getASTContext());
+
+  auto proto = getRequirement();
+
+  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;
+    }
+  }
 
   // General case: consult the substitution map.
   auto substMap =

lib/AST/RequirementMachine/ConcreteContraction.cpp

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

lib/AST/RequirementMachine/GenericSignatureQueries.cpp

@@ -33,6 +33,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "swift/AST/ASTContext.h"
+#include "swift/AST/ConformanceLookup.h"
 #include "swift/AST/Decl.h"
 #include "swift/AST/GenericEnvironment.h"
 #include "swift/AST/GenericSignature.h"
@@ -332,6 +333,8 @@ bool RequirementMachine::isReducedType(Type type) const {
 /// Given a type parameter 'T.A1.A2...An', a suffix length m where m <= n,
 /// and a replacement type U, produce the type 'U.A(n-m)...An' by replacing
 /// 'T.A1...A(n-m-1)' with 'U'.
+///
+/// FIXME: Remove this.
 static Type substPrefixType(Type type, unsigned suffixLength, Type prefixType,
                             GenericSignature sig) {
   if (suffixLength == 0)
@@ -340,9 +343,12 @@ static Type substPrefixType(Type type, unsigned suffixLength, Type prefixType,
   auto *memberType = type->castTo<DependentMemberType>();
   auto substBaseType = substPrefixType(memberType->getBase(), suffixLength - 1,
                                        prefixType, sig);
-  return memberType->substBaseType(
-      substBaseType, LookUpConformanceInModule(),
-      std::nullopt);
+  auto *assocDecl = memberType->getAssocType();
+  auto *proto = assocDecl->getProtocol();
+  auto conformance = lookupConformance(substBaseType, proto);
+  return conformance.getAssociatedType(
+      substBaseType,
+      assocDecl->getDeclaredInterfaceType());
 }
 
 Type RequirementMachine::getReducedTypeParameter(
@@ -380,6 +386,8 @@ Type RequirementMachine::getReducedTypeParameter(
   //
   // Note that V can be empty if T is fully valid; we expect this to be
   // true most of the time.
+  //
+  // FIXME: Remove all of this.
   auto prefix = getLongestValidPrefix(term);
 
   // Get a type (concrete or dependent) for U.

lib/AST/RequirementMachine/RequirementLowering.cpp

@@ -150,6 +150,7 @@
 
 #include "RequirementLowering.h"
 #include "swift/AST/ASTContext.h"
+#include "swift/AST/ConformanceLookup.h"
 #include "swift/AST/Decl.h"
 #include "swift/AST/DiagnosticsSema.h"
 #include "swift/AST/Requirement.h"
@@ -663,9 +664,9 @@ struct InferRequirementsWalker : public TypeWalker {
       if (differentiableProtocol && fnTy->isDifferentiable()) {
         auto addSameTypeConstraint = [&](Type firstType,
                                          AssociatedTypeDecl *assocType) {
-          auto secondType = assocType->getDeclaredInterfaceType()
-              ->castTo<DependentMemberType>()
-              ->substBaseType(firstType);
+          auto conformance = lookupConformance(firstType, differentiableProtocol);
+          auto secondType = conformance.getAssociatedType(
+              firstType, assocType->getDeclaredInterfaceType());
           reqs.push_back({Requirement(RequirementKind::SameType,
                                       firstType, secondType),
                           SourceLoc()});

lib/AST/Type.cpp

@@ -3746,12 +3746,13 @@ void ParameterizedProtocolType::getRequirements(
   auto argTypes = getArgs();
   assert(argTypes.size() <= assocTypes.size());
 
+  auto conformance = lookupConformance(baseType, protoDecl);
+
   for (unsigned i : indices(argTypes)) {
     auto argType = argTypes[i];
     auto *assocType = assocTypes[i];
-    auto subjectType = assocType->getDeclaredInterfaceType()
-        ->castTo<DependentMemberType>()
-        ->substBaseType(baseType);
+    auto subjectType = conformance.getAssociatedType(
+        baseType, assocType->getDeclaredInterfaceType());
     reqs.emplace_back(RequirementKind::SameType, subjectType, argType);
   }
 }
@@ -4457,14 +4458,13 @@ TypeBase::getAutoDiffTangentSpace(LookupConformanceFn lookupConformance) {
   auto associatedTypeLookup =
       differentiableProtocol->lookupDirect(ctx.Id_TangentVector);
   assert(associatedTypeLookup.size() == 1);
-  auto *dependentType = DependentMemberType::get(
-      differentiableProtocol->getDeclaredInterfaceType(),
-      cast<AssociatedTypeDecl>(associatedTypeLookup[0]));
+  auto *assocDecl = cast<AssociatedTypeDecl>(associatedTypeLookup[0]);
 
   // Try to get the `TangentVector` associated type of `base`.
   // Return the associated type if it is valid.
-  auto assocTy =
-      dependentType->substBaseType(this, lookupConformance, std::nullopt);
+  auto conformance = swift::lookupConformance(this, differentiableProtocol);
+  auto assocTy = conformance.getAssociatedType(
+      this, assocDecl->getDeclaredInterfaceType());
   if (!assocTy->hasError())
     return cache(TangentSpace::getTangentVector(assocTy));