AST: Refine name lookup rule for protocol extension 'where' clauses

slavapestov · slavapestov · commit e46a0c56a8f5 · 2024-03-16T08:34:42.000-04:00
We have two "levels" of name lookup, and the more primitive level is
used by name lookup itself to avoid certain cycles. For example,
extension binding, resolution of inheritance clauses, etc.

One interesting case is that a protocol extension can impose additional
requiremnts on `Self`, and members of the right-hand side type are
visible to unqualified lookup.

The right-hand side of a `Self` requirement in this case is always a
protocol type or class type canonically, but it might be written to
refer to a protocol type alias.

Before some changes for noncopyable generics, the primitive name
lookup mechanism, implemented in directReferencesForTypeRepr() and
such, would check if the TypeRepr had already been resolved by
resolveType(). If so, it would immediately return the decl.

This masked an issue, where the right-hand side of a `Self` requirement
was resolved in the parent DeclContext. A more subtle rule is needed;
for a protocol extension, we must resolve the right-hand side in the
protocol, but disregard the protocol extension's `Self` requirements,
because doing so would recursively trigger the same lookup again.

Fixes rdar://problem/124498054.
diff --git a/include/swift/AST/NameLookup.h b/include/swift/AST/NameLookup.h
@@ -247,6 +247,9 @@ enum class UnqualifiedLookupFlags {
   MacroLookup            = 1 << 7,
   /// This lookup should only return modules
   ModuleLookup           = 1 << 8,
+  /// This lookup should discard 'Self' requirements in protocol extension
+  /// 'where' clauses.
+  DisregardSelfBounds    = 1 << 9
 };
 
 using UnqualifiedLookupOptions = OptionSet<UnqualifiedLookupFlags>;
diff --git a/lib/AST/NameLookup.cpp b/lib/AST/NameLookup.cpp
@@ -1086,7 +1086,8 @@ namespace {
 static DirectlyReferencedTypeDecls
 directReferencesForTypeRepr(Evaluator &evaluator, ASTContext &ctx,
                             TypeRepr *typeRepr, DeclContext *dc,
-                            bool allowUsableFromInline=false);
+                            bool allowUsableFromInline,
+                            bool rhsOfSelfRequirement);
 
 /// Retrieve the set of type declarations that are directly referenced from
 /// the given type.
@@ -1118,13 +1119,6 @@ SelfBounds SelfBoundsFromWhereClauseRequest::evaluate(
   const DeclContext *dc =
       protoDecl ? (const DeclContext *)protoDecl : (const DeclContext *)extDecl;
 
-  // A protocol or extension 'where' clause can reference associated types of
-  // the protocol itself, so we have to start unqualified lookup from 'dc'.
-  //
-  // However, the right hand side of a 'Self' conformance constraint must be
-  // resolved before unqualified lookup into 'dc' can work, so we make an
-  // exception here and begin lookup from the parent context instead.
-  auto *lookupDC = dc->getParent();
   auto requirements = protoDecl ? protoDecl->getTrailingWhereClause()
                                 : extDecl->getTrailingWhereClause();
 
@@ -1151,7 +1145,10 @@ SelfBounds SelfBoundsFromWhereClauseRequest::evaluate(
     // Resolve the right-hand side.
     DirectlyReferencedTypeDecls rhsDecls;
     if (auto typeRepr = req.getConstraintRepr()) {
-      rhsDecls = directReferencesForTypeRepr(evaluator, ctx, typeRepr, lookupDC);
+      rhsDecls = directReferencesForTypeRepr(evaluator, ctx, typeRepr,
+                                             const_cast<DeclContext *>(dc),
+                                             /*allowUsableFromInline=*/false,
+                                             /*rhsOfSelfRequirement=*/true);
     }
 
     SmallVector<ModuleDecl *, 2> modulesFound;
@@ -1213,7 +1210,9 @@ TypeDeclsFromWhereClauseRequest::evaluate(Evaluator &evaluator,
 
   DirectlyReferencedTypeDecls result;
   auto resolve = [&](TypeRepr *typeRepr) {
-    auto decls = directReferencesForTypeRepr(evaluator, ctx, typeRepr, ext);
+    auto decls = directReferencesForTypeRepr(evaluator, ctx, typeRepr, ext,
+                                             /*allowUsableFromInline=*/false,
+                                             /*rhsOfSelfRequirement=*/false);
     result.first.insert(result.first.end(),
                         decls.first.begin(),
                         decls.first.end());
@@ -2834,7 +2833,8 @@ static DirectlyReferencedTypeDecls
 directReferencesForUnqualifiedTypeLookup(DeclNameRef name,
                                          SourceLoc loc, DeclContext *dc,
                                          LookupOuterResults lookupOuter,
-                                         bool allowUsableFromInline=false) {
+                                         bool allowUsableFromInline,
+                                         bool rhsOfSelfRequirement) {
   UnqualifiedLookupOptions options =
       UnqualifiedLookupFlags::TypeLookup |
       UnqualifiedLookupFlags::AllowProtocolMembers;
@@ -2856,23 +2856,17 @@ directReferencesForUnqualifiedTypeLookup(DeclNameRef name,
   //
   // extension MyProto where AssocType == Int { ... }
   //
-  // For this reason, ASTScope maps source locations inside the 'where'
-  // clause to a scope that performs the lookup into the protocol or
-  // protocol extension.
-  //
-  // However, protocol and protocol extensions can also put bounds on 'Self',
-  // for example:
-  //
-  // protocol MyProto where Self : MyClass { ... }
-  //
-  // We must start searching for 'MyClass' at the top level, otherwise
-  // we end up with a cycle, because qualified lookup wants to resolve
-  // 'Self' bounds to build the set of declarations to search inside of.
-  //
-  // To make this work, we handle the top-level lookup case explicitly
-  // here, bypassing unqualified lookup and ASTScope altogether.
-  if (dc->isModuleScopeContext())
-    loc = SourceLoc();
+  // To avoid cycles when resolving the right-hand side, we perform the
+  // lookup in the parent context (for a protocol), or a special mode where
+  // we disregard 'Self' requirements (for a protocol extension).
+  if (rhsOfSelfRequirement) {
+    if (dc->getExtendedProtocolDecl())
+      options |= UnqualifiedLookupFlags::DisregardSelfBounds;
+    else {
+      dc = dc->getModuleScopeContext();
+      loc = SourceLoc();
+    }
+  }
 
   DirectlyReferencedTypeDecls results;
 
@@ -2956,10 +2950,12 @@ directReferencesForQualifiedTypeLookup(Evaluator &evaluator,
 static DirectlyReferencedTypeDecls
 directReferencesForDeclRefTypeRepr(Evaluator &evaluator, ASTContext &ctx,
                                    DeclRefTypeRepr *repr, DeclContext *dc,
-                                   bool allowUsableFromInline) {
+                                   bool allowUsableFromInline,
+                                   bool rhsOfSelfRequirement) {
   if (auto *qualIdentTR = dyn_cast<QualifiedIdentTypeRepr>(repr)) {
     auto result = directReferencesForTypeRepr(
-        evaluator, ctx, qualIdentTR->getBase(), dc, allowUsableFromInline);
+        evaluator, ctx, qualIdentTR->getBase(), dc,
+        allowUsableFromInline, rhsOfSelfRequirement);
 
     // For a qualified identifier, perform qualified name lookup.
     result.first = directReferencesForQualifiedTypeLookup(
@@ -2972,13 +2968,14 @@ directReferencesForDeclRefTypeRepr(Evaluator &evaluator, ASTContext &ctx,
   // For an unqualified identifier, perform unqualified name lookup.
   return directReferencesForUnqualifiedTypeLookup(
       repr->getNameRef(), repr->getLoc(), dc, LookupOuterResults::Excluded,
-      allowUsableFromInline);
+      allowUsableFromInline, rhsOfSelfRequirement);
 }
 
 static DirectlyReferencedTypeDecls
 directReferencesForTypeRepr(Evaluator &evaluator,
                             ASTContext &ctx, TypeRepr *typeRepr,
-                            DeclContext *dc, bool allowUsableFromInline) {
+                            DeclContext *dc, bool allowUsableFromInline,
+                            bool rhsOfSelfRequirement) {
   DirectlyReferencedTypeDecls result;
 
   switch (typeRepr->getKind()) {
@@ -2990,15 +2987,17 @@ directReferencesForTypeRepr(Evaluator &evaluator,
     auto attributed = cast<AttributedTypeRepr>(typeRepr);
     return directReferencesForTypeRepr(evaluator, ctx,
                                        attributed->getTypeRepr(), dc,
-                                       allowUsableFromInline);
+                                       allowUsableFromInline,
+                                       rhsOfSelfRequirement);
   }
 
   case TypeReprKind::Composition: {
     auto composition = cast<CompositionTypeRepr>(typeRepr);
     for (auto component : composition->getTypes()) {
       auto componentResult =
           directReferencesForTypeRepr(evaluator, ctx, component, dc,
-                                      allowUsableFromInline);
+                                      allowUsableFromInline,
+                                      rhsOfSelfRequirement);
       result.first.insert(result.first.end(),
                           componentResult.first.begin(),
                           componentResult.first.end());
@@ -3013,7 +3012,8 @@ directReferencesForTypeRepr(Evaluator &evaluator,
   case TypeReprKind::UnqualifiedIdent:
     return directReferencesForDeclRefTypeRepr(evaluator, ctx,
                                               cast<DeclRefTypeRepr>(typeRepr),
-                                              dc, allowUsableFromInline);
+                                              dc, allowUsableFromInline,
+                                              rhsOfSelfRequirement);
 
   case TypeReprKind::Dictionary:
     result.first.push_back(ctx.getDictionaryDecl());
@@ -3024,7 +3024,8 @@ directReferencesForTypeRepr(Evaluator &evaluator,
     if (tupleRepr->isParenType()) {
       result = directReferencesForTypeRepr(evaluator, ctx,
                                            tupleRepr->getElementType(0), dc,
-                                           allowUsableFromInline);
+                                           allowUsableFromInline,
+                                           rhsOfSelfRequirement);
     } else {
       result.first.push_back(ctx.getBuiltinTupleDecl());
     }
@@ -3035,21 +3036,24 @@ directReferencesForTypeRepr(Evaluator &evaluator,
     auto packExpansionRepr = cast<VarargTypeRepr>(typeRepr);
     return directReferencesForTypeRepr(evaluator, ctx,
                                        packExpansionRepr->getElementType(), dc,
-                                       allowUsableFromInline);
+                                       allowUsableFromInline,
+                                       rhsOfSelfRequirement);
   }
 
   case TypeReprKind::PackExpansion: {
     auto packExpansionRepr = cast<PackExpansionTypeRepr>(typeRepr);
     return directReferencesForTypeRepr(evaluator, ctx,
                                        packExpansionRepr->getPatternType(), dc,
-                                       allowUsableFromInline);
+                                       allowUsableFromInline,
+                                       rhsOfSelfRequirement);
   }
 
   case TypeReprKind::PackElement: {
     auto packReferenceRepr = cast<PackElementTypeRepr>(typeRepr);
     return directReferencesForTypeRepr(evaluator, ctx,
                                        packReferenceRepr->getPackType(), dc,
-                                       allowUsableFromInline);
+                                       allowUsableFromInline,
+                                       rhsOfSelfRequirement);
   }
 
   case TypeReprKind::Inverse: {
@@ -3058,7 +3062,8 @@ directReferencesForTypeRepr(Evaluator &evaluator,
     auto *inverseRepr = cast<InverseTypeRepr>(typeRepr);
     auto innerResult = directReferencesForTypeRepr(evaluator, ctx,
                                                    inverseRepr->getConstraint(), dc,
-                                                   allowUsableFromInline);
+                                                   allowUsableFromInline,
+                                                   rhsOfSelfRequirement);
     if (innerResult.first.size() == 1) {
       if (auto *proto = dyn_cast<ProtocolDecl>(innerResult.first[0])) {
         if (auto ip = proto->getInvertibleProtocolKind()) {
@@ -3169,7 +3174,9 @@ DirectlyReferencedTypeDecls InheritedDeclsReferencedRequest::evaluate(
       dc = (DeclContext *)decl.get<const ExtensionDecl *>();
 
     return directReferencesForTypeRepr(evaluator, dc->getASTContext(), typeRepr,
-                                       const_cast<DeclContext *>(dc));
+                                       const_cast<DeclContext *>(dc),
+                                       /*allowUsableFromInline=*/false,
+                                       /*rhsOfSelfRequirement=*/false);
   }
 
   // Fall back to semantic types.
@@ -3188,7 +3195,9 @@ DirectlyReferencedTypeDecls UnderlyingTypeDeclsReferencedRequest::evaluate(
   // Prefer syntactic information when we have it.
   if (auto typeRepr = typealias->getUnderlyingTypeRepr()) {
     return directReferencesForTypeRepr(evaluator, typealias->getASTContext(),
-                                       typeRepr, typealias);
+                                       typeRepr, typealias,
+                                       /*allowUsableFromInline=*/false,
+                                       /*rhsOfSelfRequirement=*/false);
   }
 
   // Fall back to semantic types.
@@ -3354,7 +3363,8 @@ ExtendedNominalRequest::evaluate(Evaluator &evaluator,
   ASTContext &ctx = ext->getASTContext();
   DirectlyReferencedTypeDecls referenced =
     directReferencesForTypeRepr(evaluator, ctx, typeRepr, ext->getParent(),
-                                ext->isInSpecializeExtensionContext());
+                                ext->isInSpecializeExtensionContext(),
+                                /*rhsOfSelfRequirement=*/false);
 
   // Resolve those type declarations to nominal type declarations.
   SmallVector<ModuleDecl *, 2> modulesFound;
@@ -3402,7 +3412,9 @@ static bool declsAreProtocols(ArrayRef<TypeDecl *> decls) {
 
 bool TypeRepr::isProtocolOrProtocolComposition(DeclContext *dc) {
   auto &ctx = dc->getASTContext();
-  auto references = directReferencesForTypeRepr(ctx.evaluator, ctx, this, dc);
+  auto references = directReferencesForTypeRepr(ctx.evaluator, ctx, this, dc,
+                                                /*allowUsableFromInline=*/false,
+                                                /*rhsOfSelfRequirement=*/false);
   return declsAreProtocols(references.first);
 }
 
@@ -3435,7 +3447,9 @@ createTupleExtensionGenericParams(ASTContext &ctx,
   DirectlyReferencedTypeDecls referenced =
     directReferencesForTypeRepr(ctx.evaluator, ctx,
                                 extendedTypeRepr,
-                                ext->getParent());
+                                ext->getParent(),
+                                /*allowUsableFromInline=*/false,
+                                /*rhsOfSelfRequirement=*/false);
   assert(referenced.second.empty() && "Implement me");
   if (referenced.first.size() != 1 || !isa<TypeAliasDecl>(referenced.first[0]))
     return nullptr;
@@ -3708,7 +3722,9 @@ CustomAttrNominalRequest::evaluate(Evaluator &evaluator,
   DirectlyReferencedTypeDecls decls;
   if (auto *typeRepr = attr->getTypeRepr()) {
     decls = directReferencesForTypeRepr(
-        evaluator, ctx, typeRepr, dc);
+        evaluator, ctx, typeRepr, dc,
+        /*allowUsableFromInline=*/false,
+        /*rhsOfSelfRequirement=*/false);
   } else if (Type type = attr->getType()) {
     decls = directReferencesForType(type);
   }
@@ -3737,7 +3753,9 @@ CustomAttrNominalRequest::evaluate(Evaluator &evaluator,
         modulesFound.clear();
         anyObject = false;
         decls = directReferencesForUnqualifiedTypeLookup(
-            name, loc, dc, LookupOuterResults::Included);
+            name, loc, dc, LookupOuterResults::Included,
+            /*allowUsableFromInline=*/false,
+            /*rhsOfSelfRequirement=*/false);
         nominals = resolveTypeDeclsToNominal(evaluator, ctx, decls.first,
                                              ResolveToNominalOptions(),
                                              modulesFound, anyObject);
@@ -3970,7 +3988,9 @@ ProtocolDecl *ImplementsAttrProtocolRequest::evaluate(
 
   ASTContext &ctx = dc->getASTContext();
   DirectlyReferencedTypeDecls referenced =
-    directReferencesForTypeRepr(evaluator, ctx, typeRepr, dc);
+    directReferencesForTypeRepr(evaluator, ctx, typeRepr, dc,
+                                /*allowUsableFromInline=*/false,
+                                /*rhsOfSelfRequirement=*/false);
 
   // Resolve those type declarations to nominal type declarations.
   SmallVector<ModuleDecl *, 2> modulesFound;
diff --git a/lib/AST/UnqualifiedLookup.cpp b/lib/AST/UnqualifiedLookup.cpp
@@ -730,31 +730,27 @@ bool ASTScopeDeclGatherer::consume(ArrayRef<ValueDecl *> valuesArg,
   return false;
 }
 
-static SmallVector<NominalTypeDecl *, 2> findSelfBounds(const DeclContext *dc) {
-  auto nominal = dc->getSelfNominalTypeDecl();
+// TODO: in future, migrate this functionality into ASTScopes
+bool ASTScopeDeclConsumerForUnqualifiedLookup::lookInMembers(
+    const DeclContext *scopeDC) const {
+  auto nominal = scopeDC->getSelfNominalTypeDecl();
   if (!nominal)
-    return {};
+    return false;
 
   SmallVector<NominalTypeDecl *, 2> selfBounds;
   selfBounds.push_back(nominal);
 
   // For a protocol extension, check whether there are additional "Self"
   // constraints that can affect name lookup.
-  if (dc->getExtendedProtocolDecl()) {
-    auto ext = cast<ExtensionDecl>(dc);
-    auto bounds = getSelfBoundsFromWhereClause(ext);
-    for (auto bound : bounds.decls)
-      selfBounds.push_back(bound);
+  if (!factory.options.contains(UnqualifiedLookupFlags::DisregardSelfBounds)) {
+    if (scopeDC->getExtendedProtocolDecl()) {
+      auto ext = cast<ExtensionDecl>(scopeDC);
+      auto bounds = getSelfBoundsFromWhereClause(ext);
+      for (auto bound : bounds.decls)
+        selfBounds.push_back(bound);
+    }
   }
 
-  return selfBounds;
-}
-
-// TODO: in future, migrate this functionality into ASTScopes
-bool ASTScopeDeclConsumerForUnqualifiedLookup::lookInMembers(
-    const DeclContext *scopeDC) const {
-  auto selfBounds = findSelfBounds(scopeDC);
-
   // We're looking for members of a type.
   //
   // If we started the looking from inside a scope where a 'self' parameter
diff --git a/test/NameLookup/protocol_extension_where_clause.swift b/test/NameLookup/protocol_extension_where_clause.swift
@@ -0,0 +1,31 @@
+// RUN: %target-typecheck-verify-swift
+
+protocol P1 {
+  typealias A = P2
+}
+
+protocol P2 {
+  associatedtype B
+  typealias A1 = Int
+  func f1()
+}
+
+extension P2 {
+  typealias A2 = String
+  func f2() {}
+}
+
+extension P1 where Self: A {
+  typealias B1 = A1
+  typealias B2 = A2
+
+  func g() {
+    f1()
+    f2()
+  }
+}
+
+// This is terrible and we should ban it some day
+extension P1 where Self: A, B: Hashable {
+  func h(_: Set<B>) {}
+}
