RequirementMachine: Generalize hack that allows associated type inheritance clauses to reference protocol typealiases

slavapestov · slavapestov · commit bc43cdd104cb · 2022-04-05T18:42:35.000-04:00
We want to allow this for all conformance requirements written in protocols
or the `where` clause of protocol extensions.

Fixes rdar://problem/91304291.
diff --git a/include/swift/AST/TypeCheckRequests.h b/include/swift/AST/TypeCheckRequests.h
@@ -604,7 +604,6 @@ struct WhereClauseOwner {
     return !(lhs == rhs);
   }
 
-public:
   /// Retrieve the array of requirements.
   MutableArrayRef<RequirementRepr> getRequirements() const;
 
@@ -1903,8 +1902,7 @@ class InferredGenericSignatureRequest :
 /// InferredGenericSignatureRequest.
 class InferredGenericSignatureRequestRQM :
     public SimpleRequest<InferredGenericSignatureRequestRQM,
-                         GenericSignatureWithError (ModuleDecl *,
-                                                    const GenericSignatureImpl *,
+                         GenericSignatureWithError (const GenericSignatureImpl *,
                                                     GenericParamList *,
                                                     WhereClauseOwner,
                                                     SmallVector<Requirement, 2>,
@@ -1920,7 +1918,6 @@ class InferredGenericSignatureRequestRQM :
   // Evaluation.
   GenericSignatureWithError
   evaluate(Evaluator &evaluator,
-           ModuleDecl *parentModule,
            const GenericSignatureImpl *baseSignature,
            GenericParamList *genericParams,
            WhereClauseOwner whereClause,
diff --git a/lib/AST/GenericSignatureBuilder.cpp b/lib/AST/GenericSignatureBuilder.cpp
@@ -8464,7 +8464,6 @@ InferredGenericSignatureRequest::evaluate(
     return evaluateOrDefault(
         ctx.evaluator,
         InferredGenericSignatureRequestRQM{
-          parentModule,
           parentSig,
           genericParams,
           whereClause,
diff --git a/lib/AST/RequirementMachine/RequirementLowering.cpp b/lib/AST/RequirementMachine/RequirementLowering.cpp
@@ -160,6 +160,7 @@
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/SetVector.h"
 #include "RewriteContext.h"
+#include "NameLookup.h"
 
 using namespace swift;
 using namespace rewriting;
@@ -384,12 +385,40 @@ swift::rewriting::desugarRequirement(Requirement req, SourceLoc loc,
 // Requirement realization and inference.
 //
 
-static void realizeTypeRequirement(Type subjectType, Type constraintType,
+static void realizeTypeRequirement(DeclContext *dc,
+                                   Type subjectType, Type constraintType,
                                    SourceLoc loc,
                                    SmallVectorImpl<StructuralRequirement> &result,
                                    SmallVectorImpl<RequirementError> &errors) {
   SmallVector<Requirement, 2> reqs;
 
+  // The GenericSignatureBuilder allowed the right hand side of a
+  // conformance or superclass requirement to reference a protocol
+  // typealias whose underlying type was a protocol or class.
+  //
+  // Since protocol typealiases resolve to DependentMemberTypes in
+  // ::Structural mode, this relied on the GSB's "delayed requirements"
+  // mechanism.
+  //
+  // The RequirementMachine does not have an equivalent, and cannot really
+  // support that because we need to collect the protocols mentioned on
+  // the right hand sides of conformance requirements ahead of time.
+  //
+  // However, we can support it in simple cases where the typealias is
+  // defined in the protocol itself and is accessed as a member of 'Self'.
+  if (auto *proto = dc->getSelfProtocolDecl()) {
+    if (auto memberType = constraintType->getAs<DependentMemberType>()) {
+      if (memberType->getBase()->isEqual(proto->getSelfInterfaceType())) {
+        SmallVector<TypeDecl *, 1> result;
+        lookupConcreteNestedType(proto, memberType->getName(), result);
+        auto *typeDecl = findBestConcreteNestedType(result);
+        if (auto *aliasDecl = dyn_cast_or_null<TypeAliasDecl>(typeDecl)) {
+          constraintType = aliasDecl->getUnderlyingType();
+        }
+      }
+    }
+  }
+
   if (constraintType->isConstraintType()) {
     // Handle conformance requirements.
     desugarConformanceRequirement(subjectType, constraintType, loc, reqs, errors);
@@ -534,18 +563,20 @@ void swift::rewriting::inferRequirements(
 /// Desugar a requirement and perform requirement inference if requested
 /// to obtain zero or more structural requirements.
 void swift::rewriting::realizeRequirement(
+    DeclContext *dc,
     Requirement req, RequirementRepr *reqRepr,
-    ModuleDecl *moduleForInference,
+    bool shouldInferRequirements,
     SmallVectorImpl<StructuralRequirement> &result,
     SmallVectorImpl<RequirementError> &errors) {
   auto firstType = req.getFirstType();
   auto loc = (reqRepr ? reqRepr->getSeparatorLoc() : SourceLoc());
+  auto *moduleForInference = dc->getParentModule();
 
   switch (req.getKind()) {
   case RequirementKind::Superclass:
   case RequirementKind::Conformance: {
     auto secondType = req.getSecondType();
-    if (moduleForInference) {
+    if (shouldInferRequirements) {
       auto firstLoc = (reqRepr ? reqRepr->getSubjectRepr()->getStartLoc()
                                : SourceLoc());
       inferRequirements(firstType, firstLoc, moduleForInference, result);
@@ -555,12 +586,12 @@ void swift::rewriting::realizeRequirement(
       inferRequirements(secondType, secondLoc, moduleForInference, result);
     }
 
-    realizeTypeRequirement(firstType, secondType, loc, result, errors);
+    realizeTypeRequirement(dc, firstType, secondType, loc, result, errors);
     break;
   }
 
   case RequirementKind::Layout: {
-    if (moduleForInference) {
+    if (shouldInferRequirements) {
       auto firstLoc = (reqRepr ? reqRepr->getSubjectRepr()->getStartLoc()
                                : SourceLoc());
       inferRequirements(firstType, firstLoc, moduleForInference, result);
@@ -578,7 +609,7 @@ void swift::rewriting::realizeRequirement(
 
   case RequirementKind::SameType: {
     auto secondType = req.getSecondType();
-    if (moduleForInference) {
+    if (shouldInferRequirements) {
       auto firstLoc = (reqRepr ? reqRepr->getFirstTypeRepr()->getStartLoc()
                                : SourceLoc());
       inferRequirements(firstType, firstLoc, moduleForInference, result);
@@ -602,11 +633,13 @@ void swift::rewriting::realizeRequirement(
 /// Collect structural requirements written in the inheritance clause of an
 /// AssociatedTypeDecl or GenericTypeParamDecl.
 void swift::rewriting::realizeInheritedRequirements(
-    TypeDecl *decl, Type type, ModuleDecl *moduleForInference,
+    TypeDecl *decl, Type type, bool shouldInferRequirements,
     SmallVectorImpl<StructuralRequirement> &result,
     SmallVectorImpl<RequirementError> &errors) {
   auto &ctx = decl->getASTContext();
   auto inheritedTypes = decl->getInherited();
+  auto *dc = decl->getInnermostDeclContext();
+  auto *moduleForInference = dc->getParentModule();
 
   for (unsigned index : indices(inheritedTypes)) {
     Type inheritedType
@@ -616,41 +649,13 @@ void swift::rewriting::realizeInheritedRequirements(
                           Type());
     if (!inheritedType) continue;
 
-    // The GenericSignatureBuilder allowed an associated type's inheritance
-    // clause to reference a protocol typealias whose underlying type was a
-    // protocol or class.
-    //
-    // Since protocol typealiases resolve to DependentMemberTypes in
-    // ::Structural mode, this relied on the GSB's "delayed requirements"
-    // mechanism.
-    //
-    // The RequirementMachine does not have an equivalent, and cannot really
-    // support that because we need to collect the protocols mentioned on
-    // the right hand sides of conformance requirements ahead of time.
-    //
-    // However, we can support it in simple cases where the typealias is
-    // defined in the protocol itself and is accessed as a member of 'Self'.
-    if (auto *assocTypeDecl = dyn_cast<AssociatedTypeDecl>(decl)) {
-      if (auto memberType = inheritedType->getAs<DependentMemberType>()) {
-        if (memberType->getBase()->isEqual(
-            assocTypeDecl->getProtocol()->getSelfInterfaceType())) {
-          inheritedType
-            = evaluateOrDefault(ctx.evaluator,
-                                InheritedTypeRequest{decl, index,
-                                TypeResolutionStage::Interface},
-                                Type());
-          if (!inheritedType) continue;
-        }
-      }
-    }
-
     auto *typeRepr = inheritedTypes[index].getTypeRepr();
     SourceLoc loc = (typeRepr ? typeRepr->getStartLoc() : SourceLoc());
-    if (moduleForInference) {
+    if (shouldInferRequirements) {
       inferRequirements(inheritedType, loc, moduleForInference, result);
     }
 
-    realizeTypeRequirement(type, inheritedType, loc, result, errors);
+    realizeTypeRequirement(dc, type, inheritedType, loc, result, errors);
   }
 }
 
@@ -823,14 +828,14 @@ StructuralRequirementsRequest::evaluate(Evaluator &evaluator,
   auto selfTy = proto->getSelfInterfaceType();
 
   realizeInheritedRequirements(proto, selfTy,
-                               /*moduleForInference=*/nullptr,
+                               /*inferRequirements=*/false,
                                result, errors);
 
   // Add requirements from the protocol's own 'where' clause.
   WhereClauseOwner(proto).visitRequirements(TypeResolutionStage::Structural,
       [&](const Requirement &req, RequirementRepr *reqRepr) {
-        realizeRequirement(req, reqRepr,
-                           /*moduleForInference=*/nullptr,
+        realizeRequirement(proto, req, reqRepr,
+                           /*inferRequirements=*/false,
                            result, errors);
         return false;
       });
@@ -855,15 +860,15 @@ StructuralRequirementsRequest::evaluate(Evaluator &evaluator,
     // Add requirements placed directly on this associated type.
     auto assocType = assocTypeDecl->getDeclaredInterfaceType();
     realizeInheritedRequirements(assocTypeDecl, assocType,
-                                 /*moduleForInference=*/nullptr,
+                                 /*inferRequirements=*/false,
                                  result, errors);
 
     // Add requirements from this associated type's where clause.
     WhereClauseOwner(assocTypeDecl).visitRequirements(
         TypeResolutionStage::Structural,
         [&](const Requirement &req, RequirementRepr *reqRepr) {
-          realizeRequirement(req, reqRepr,
-                             /*moduleForInference=*/nullptr,
+          realizeRequirement(proto, req, reqRepr,
+                             /*inferRequirements=*/false,
                              result, errors);
           return false;
         });
diff --git a/lib/AST/RequirementMachine/RequirementLowering.h b/lib/AST/RequirementMachine/RequirementLowering.h
@@ -45,13 +45,14 @@ void desugarRequirement(Requirement req, SourceLoc loc,
 void inferRequirements(Type type, SourceLoc loc, ModuleDecl *module,
                        SmallVectorImpl<StructuralRequirement> &result);
 
-void realizeRequirement(Requirement req, RequirementRepr *reqRepr,
-                        ModuleDecl *moduleForInference,
+void realizeRequirement(DeclContext *dc,
+                        Requirement req, RequirementRepr *reqRepr,
+                        bool shouldInferRequirements,
                         SmallVectorImpl<StructuralRequirement> &result,
                         SmallVectorImpl<RequirementError> &errors);
 
 void realizeInheritedRequirements(TypeDecl *decl, Type type,
-                                  ModuleDecl *moduleForInference,
+                                  bool shouldInferRequirements,
                                   SmallVectorImpl<StructuralRequirement> &result,
                                   SmallVectorImpl<RequirementError> &errors);
 
diff --git a/lib/AST/RequirementMachine/RequirementMachineRequests.cpp b/lib/AST/RequirementMachine/RequirementMachineRequests.cpp
@@ -685,7 +685,6 @@ AbstractGenericSignatureRequestRQM::evaluate(
 GenericSignatureWithError
 InferredGenericSignatureRequestRQM::evaluate(
         Evaluator &evaluator,
-        ModuleDecl *parentModule,
         const GenericSignatureImpl *parentSigImpl,
         GenericParamList *genericParamList,
         WhereClauseOwner whereClause,
@@ -694,8 +693,6 @@ InferredGenericSignatureRequestRQM::evaluate(
         bool allowConcreteGenericParams) const {
   GenericSignature parentSig(parentSigImpl);
 
-  auto &ctx = parentModule->getASTContext();
-
   SmallVector<GenericTypeParamType *, 4> genericParams(
       parentSig.getGenericParams().begin(),
       parentSig.getGenericParams().end());
@@ -705,9 +702,12 @@ InferredGenericSignatureRequestRQM::evaluate(
   for (const auto &req : parentSig.getRequirements())
     requirements.push_back({req, SourceLoc(), /*wasInferred=*/false});
 
+  DeclContext *lookupDC = nullptr;
+
   const auto visitRequirement = [&](const Requirement &req,
                                     RequirementRepr *reqRepr) {
-    realizeRequirement(req, reqRepr, parentModule, requirements, errors);
+    realizeRequirement(lookupDC, req, reqRepr, /*inferRequirements=*/true,
+                       requirements, errors);
     return false;
   };
 
@@ -738,11 +738,12 @@ InferredGenericSignatureRequestRQM::evaluate(
                              ->castTo<GenericTypeParamType>();
         genericParams.push_back(gpType);
 
-        realizeInheritedRequirements(gpDecl, gpType, parentModule,
+        realizeInheritedRequirements(gpDecl, gpType,
+                                     /*inferRequirements=*/true,
                                      requirements, errors);
       }
 
-      auto *lookupDC = (*gpList->begin())->getDeclContext();
+      lookupDC = (*gpList->begin())->getDeclContext();
 
       // Add the generic parameter list's 'where' clause to the builder.
       //
@@ -758,6 +759,8 @@ InferredGenericSignatureRequestRQM::evaluate(
   // Realize all requirements in the free-standing 'where' clause, if there
   // is one.
   if (whereClause) {
+    lookupDC = whereClause.dc;
+
     if (loc.isInvalid())
       loc = whereClause.getLoc();
 
@@ -766,13 +769,16 @@ InferredGenericSignatureRequestRQM::evaluate(
         visitRequirement);
   }
 
+  auto *moduleForInference = lookupDC->getParentModule();
+
   // Perform requirement inference from function parameter and result
   // types and such.
   for (auto sourcePair : inferenceSources) {
     auto *typeRepr = sourcePair.getTypeRepr();
     auto loc = typeRepr ? typeRepr->getStartLoc() : SourceLoc();
 
-    inferRequirements(sourcePair.getType(), loc, parentModule, requirements);
+    inferRequirements(sourcePair.getType(), loc, moduleForInference,
+                      requirements);
   }
 
   // Finish by adding any remaining requirements. This is used to introduce
@@ -781,6 +787,7 @@ InferredGenericSignatureRequestRQM::evaluate(
   for (const auto &req : addedRequirements)
     requirements.push_back({req, SourceLoc(), /*wasInferred=*/true});
 
+  auto &ctx = moduleForInference->getASTContext();
   auto &rewriteCtx = ctx.getRewriteContext();
 
   if (rewriteCtx.getDebugOptions().contains(DebugFlags::Timers)) {
diff --git a/test/Generics/conformance_requirement_with_protocol_typealias.swift b/test/Generics/conformance_requirement_with_protocol_typealias.swift
@@ -1,5 +1,5 @@
-// RUN: %target-typecheck-verify-swift -requirement-machine-protocol-signatures=on
-// RUN: not %target-swift-frontend -typecheck %s -debug-generic-signatures -requirement-machine-protocol-signatures=on 2>&1 | %FileCheck %s
+// RUN: %target-typecheck-verify-swift -requirement-machine-protocol-signatures=on -requirement-machine-inferred-signatures=on
+// RUN: not %target-swift-frontend -typecheck %s -debug-generic-signatures -requirement-machine-protocol-signatures=on -requirement-machine-inferred-signatures=on 2>&1 | %FileCheck %s
 
 // CHECK-LABEL: .P@
 // CHECK-NEXT: Requirement signature: <Self where Self.[P]T : C>
@@ -20,17 +20,21 @@ protocol PBad {
   // expected-error@-1 {{type 'Self.T1' constrained to non-protocol, non-class type 'Self.B.A'}}
 
   associatedtype T2 where T2 : A
-  // expected-error@-1 {{type 'Self.T2' constrained to non-protocol, non-class type 'Self.A'}}
 }
 
-// FIXME: Terrible diagnostics.
-
 protocol PWorse {
-// expected-error@-1 5{{circular reference}}
-// expected-note@-2 9{{through reference here}}
   typealias A = C
 
   associatedtype T : Self.A
-// expected-note@-1 5{{while resolving type 'Self.A'}}
-// expected-note@-2 5{{through reference here}}
 }
+
+protocol Q1 {}
+protocol Q2 {}
+
+extension P {
+  typealias B = (Q1 & Q2)
+}
+
+// CHECK-LABEL: ExtensionDecl line={{.*}} base=P
+// CHECK-NEXT: Generic signature: <Self where Self : P, Self : Q1, Self : Q2>
+extension P where Self: B {}
