AST: Replace calls to substBaseType() with getAssociatedType()

Author: slavapestov

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));