AST: Refactor unqualified lookup a little bit

slavapestov · calda · commit 628d6241009c · 2024-03-18T16:09:01.000-07:00
diff --git a/include/swift/AST/ASTScope.h b/include/swift/AST/ASTScope.h
@@ -482,45 +482,33 @@ class Portion : public ASTAllocated<ASTScopeImpl> {
 
   virtual NullablePtr<ASTScopeImpl>
   insertionPointForDeferredExpansion(IterableTypeScope *) const = 0;
-  };
-
-  // For the whole Decl scope of a GenericType or an Extension
-  class GenericTypeOrExtensionWholePortion final : public Portion {
-  public:
-    GenericTypeOrExtensionWholePortion() : Portion("Decl") {}
-    virtual ~GenericTypeOrExtensionWholePortion() {}
-
-    // Just for TypeAlias
-    ASTScopeImpl *expandScope(GenericTypeOrExtensionScope *,
-                              ScopeCreator &) const override;
+};
 
-    SourceRange getChildlessSourceRangeOf(const GenericTypeOrExtensionScope *,
-                                          bool omitAssertions) const override;
+// For the whole Decl scope of a GenericType or an Extension
+class GenericTypeOrExtensionWholePortion final : public Portion {
+public:
+  GenericTypeOrExtensionWholePortion() : Portion("Decl") {}
+  virtual ~GenericTypeOrExtensionWholePortion() {}
 
-    NullablePtr<const ASTScopeImpl>
-    getLookupLimitFor(const GenericTypeOrExtensionScope *) const override;
+  // Just for TypeAlias
+  ASTScopeImpl *expandScope(GenericTypeOrExtensionScope *,
+                            ScopeCreator &) const override;
 
-    NullablePtr<ASTScopeImpl>
-    insertionPointForDeferredExpansion(IterableTypeScope *) const override;
-  };
+  SourceRange getChildlessSourceRangeOf(const GenericTypeOrExtensionScope *,
+                                        bool omitAssertions) const override;
 
-  /// GenericTypeOrExtension = GenericType or Extension
-  class GenericTypeOrExtensionWhereOrBodyPortion : public Portion {
-  public:
-    GenericTypeOrExtensionWhereOrBodyPortion(const char *n) : Portion(n) {}
-    virtual ~GenericTypeOrExtensionWhereOrBodyPortion() {}
+  NullablePtr<const ASTScopeImpl>
+  getLookupLimitFor(const GenericTypeOrExtensionScope *) const override;
 
-    bool lookupMembersOf(const GenericTypeOrExtensionScope *scope,
-                         ASTScopeImpl::DeclConsumer consumer) const override;
+  NullablePtr<ASTScopeImpl>
+  insertionPointForDeferredExpansion(IterableTypeScope *) const override;
 };
 
 /// Behavior specific to representing the trailing where clause of a
 /// GenericTypeDecl or ExtensionDecl scope.
-class GenericTypeOrExtensionWherePortion final
-    : public GenericTypeOrExtensionWhereOrBodyPortion {
+class GenericTypeOrExtensionWherePortion final : public Portion {
 public:
-  GenericTypeOrExtensionWherePortion()
-      : GenericTypeOrExtensionWhereOrBodyPortion("Where") {}
+  GenericTypeOrExtensionWherePortion() : Portion("Where") {}
 
   bool lookupMembersOf(const GenericTypeOrExtensionScope *scope,
                        ASTScopeImpl::DeclConsumer consumer) const override;
@@ -537,11 +525,12 @@ class GenericTypeOrExtensionWherePortion final
 
 /// Behavior specific to representing the Body of a NominalTypeDecl or
 /// ExtensionDecl scope
-class IterableTypeBodyPortion final
-    : public GenericTypeOrExtensionWhereOrBodyPortion {
+class IterableTypeBodyPortion final : public Portion {
 public:
-  IterableTypeBodyPortion()
-      : GenericTypeOrExtensionWhereOrBodyPortion("Body") {}
+  IterableTypeBodyPortion() : Portion("Body") {}
+
+  bool lookupMembersOf(const GenericTypeOrExtensionScope *scope,
+                       ASTScopeImpl::DeclConsumer consumer) const override;
 
   ASTScopeImpl *expandScope(GenericTypeOrExtensionScope *,
                             ScopeCreator &) const override;
diff --git a/lib/AST/ASTScopeLookup.cpp b/lib/AST/ASTScopeLookup.cpp
@@ -274,22 +274,24 @@ bool Portion::lookupMembersOf(const GenericTypeOrExtensionScope *,
   return false;
 }
 
-bool GenericTypeOrExtensionWhereOrBodyPortion::lookupMembersOf(
+bool GenericTypeOrExtensionWherePortion::lookupMembersOf(
     const GenericTypeOrExtensionScope *scope,
     ASTScopeImpl::DeclConsumer consumer) const {
   if (scope->getCorrespondingNominalTypeDecl().isNull())
     return false;
+
+  if (!scope->areMembersVisibleFromWhereClause())
+    return false;
+
   return consumer.lookInMembers(scope->getGenericContext());
 }
 
-bool GenericTypeOrExtensionWherePortion::lookupMembersOf(
+bool IterableTypeBodyPortion::lookupMembersOf(
     const GenericTypeOrExtensionScope *scope,
     ASTScopeImpl::DeclConsumer consumer) const {
-  if (!scope->areMembersVisibleFromWhereClause())
+  if (scope->getCorrespondingNominalTypeDecl().isNull())
     return false;
-
-  return GenericTypeOrExtensionWhereOrBodyPortion::lookupMembersOf(
-    scope, consumer);
+  return consumer.lookInMembers(scope->getGenericContext());
 }
 
 bool GenericTypeOrExtensionScope::areMembersVisibleFromWhereClause() const {
diff --git a/lib/AST/UnqualifiedLookup.cpp b/lib/AST/UnqualifiedLookup.cpp
@@ -51,49 +51,6 @@ namespace {
     using ResultsVector = SmallVector<LookupResultEntry, 4>;
     
   private:
-    /// Finds lookup results based on the types that self conforms to.
-    /// For instance, self always conforms to a struct, enum or class.
-    /// But in addition, self could conform to any number of protocols.
-    /// For example, when there's a protocol extension, e.g. extension P where
-    /// self: P2, self also conforms to P2 so P2 must be searched.
-    class ResultFinderForTypeContext {
-      UnqualifiedLookupFactory *const factory;
-      /// Nontypes are formally members of the base type, i.e. the dynamic type
-      /// of the activation record.
-      const DeclContext *const dynamicContext;
-      /// Types are formally members of the metatype, i.e. the static type of the
-      /// activation record.
-      const DeclContext *const staticContext;
-      using SelfBounds = SmallVector<NominalTypeDecl *, 2>;
-      SelfBounds selfBounds;
-      
-    public:
-      /// \p staticContext is also the context from which to derive the self types
-      ResultFinderForTypeContext(UnqualifiedLookupFactory *factory,
-                                 const DeclContext *dynamicContext,
-                                 const DeclContext *staticContext);
-
-      SWIFT_DEBUG_DUMP;
-      
-    private:
-      SelfBounds findSelfBounds(const DeclContext *dc);
-      ValueDecl *lookupBaseDecl(const DeclContext *baseDC) const;
-      ValueDecl *getBaseDeclForResult(const DeclContext *baseDC) const;
-
-      // Classify this declaration.
-      // Types are formally members of the metatype.
-      const DeclContext *
-      whereValueIsMember(const ValueDecl *const member) const {
-        return isa<TypeDecl>(member) ? staticContext : dynamicContext;
-      }
-
-    public:
-      /// Do the lookups and add matches to results.
-      void findResults(const DeclNameRef &Name, NLOptions baseNLOptions,
-                       const DeclContext *contextForLookup,
-                       SmallVectorImpl<LookupResultEntry> &results) const;
-    };
-    
     // Inputs
     const DeclNameRef Name;
     DeclContext *const DC;
@@ -152,6 +109,10 @@ namespace {
 
 #pragma mark context-based lookup declarations
 
+    ValueDecl *lookupBaseDecl(const DeclContext *baseDC) const;
+
+    ValueDecl *getBaseDeclForResult(const DeclContext *baseDC) const;
+
     /// For diagnostic purposes, move aside the unavailables, and put
     /// them back as a last-ditch effort.
     /// Could be cleaner someday with a richer interface to UnqualifiedLookup.
@@ -172,8 +133,9 @@ namespace {
                          const bool isOriginallyMacroLookup);
 
     void findResultsAndSaveUnavailables(
-        const DeclContext *lookupContextForThisContext,
-        ResultFinderForTypeContext &&resultFinderForTypeContext,
+        const DeclContext *dynamicContext,
+        const DeclContext *staticContext,
+        SmallVector<NominalTypeDecl *, 2> selfBounds,
         NLOptions baseNLOptions);
 
   public:
@@ -335,30 +297,8 @@ void UnqualifiedLookupFactory::lookUpTopLevelNamesInModuleScopeContext(
 
 #pragma mark context-based lookup definitions
 
-void UnqualifiedLookupFactory::ResultFinderForTypeContext::findResults(
-    const DeclNameRef &Name, NLOptions baseNLOptions,
-    const DeclContext *contextForLookup,
-    SmallVectorImpl<LookupResultEntry> &results) const {
-  // An optimization:
-  if (selfBounds.empty())
-    return;
-
-  SmallVector<ValueDecl *, 4> Lookup;
-  contextForLookup->lookupQualified(selfBounds, Name, factory->Loc,
-                                    baseNLOptions, Lookup);
-  for (auto Result : Lookup) {
-    auto baseDC = const_cast<DeclContext *>(whereValueIsMember(Result));
-    auto baseDecl = getBaseDeclForResult(baseDC);
-    results.emplace_back(baseDC, baseDecl, Result);
-#ifndef NDEBUG
-    factory->addedResult(results.back());
-#endif
-  }
-}
-
 ValueDecl *
-UnqualifiedLookupFactory::ResultFinderForTypeContext::getBaseDeclForResult(
-    const DeclContext *baseDC) const {
+UnqualifiedLookupFactory::getBaseDeclForResult(const DeclContext *baseDC) const {
   if (baseDC == nullptr) {
     return nullptr;
   }
@@ -393,7 +333,7 @@ UnqualifiedLookupFactory::ResultFinderForTypeContext::getBaseDeclForResult(
 /// unwrapping condition (e.g. `guard let self else { return }`).
 /// If this is true, then we know any implicit self reference in the
 /// following scope is guaranteed to be non-optional.
-bool implicitSelfReferenceIsUnwrapped(const ValueDecl *selfDecl) {
+static bool implicitSelfReferenceIsUnwrapped(const ValueDecl *selfDecl) {
   ASTContext &Ctx = selfDecl->getASTContext();
 
   // Check if the implicit self decl refers to a var in a conditional stmt
@@ -411,25 +351,19 @@ bool implicitSelfReferenceIsUnwrapped(const ValueDecl *selfDecl) {
   return conditionalStmt->rebindsSelf(Ctx);
 }
 
-ValueDecl *UnqualifiedLookupFactory::ResultFinderForTypeContext::lookupBaseDecl(
-    const DeclContext *baseDC) const {
-  auto dc = factory->DC;
-  if (!dc) {
-    return nullptr;
-  }
-
+ValueDecl *UnqualifiedLookupFactory::lookupBaseDecl(const DeclContext *baseDC) const {
   // Perform an unqualified lookup for the base decl of this result. This
   // handles cases where self was rebound (e.g. `guard let self = self`)
   // earlier in this closure or some outer closure.
-  auto closureExpr =
-      dyn_cast_or_null<ClosureExpr>(dc->getInnermostClosureForSelfCapture());
+   auto closureExpr =
+       dyn_cast_or_null<ClosureExpr>(dc->getInnermostClosureForSelfCapture());
+
   if (!closureExpr) {
     return nullptr;
   }
 
   auto selfDecl = ASTScope::lookupSingleLocalDecl(
-      factory->DC->getParentSourceFile(), DeclName(factory->Ctx.Id_self),
-      factory->Loc);
+      DC->getParentSourceFile(), DeclName(Ctx.Id_self), Loc);
 
   if (!selfDecl) {
     return nullptr;
@@ -466,7 +400,7 @@ ValueDecl *UnqualifiedLookupFactory::ResultFinderForTypeContext::lookupBaseDecl(
   // In these cases, using the Swift 6 lookup behavior doesn't affect
   // how the body is type-checked, so it can be used in Swift 5 mode
   // without breaking source compatibility for non-escaping closures.
-  if (capturesSelfWeakly && !factory->Ctx.LangOpts.isSwiftVersionAtLeast(6) &&
+  if (capturesSelfWeakly && !Ctx.LangOpts.isSwiftVersionAtLeast(6) &&
       !implicitSelfReferenceIsUnwrapped(selfDecl)) {
     return nullptr;
   }
@@ -612,12 +546,25 @@ void UnqualifiedLookupFactory::lookForAModuleWithTheGivenName(
 #pragma mark common helper definitions
 
 void UnqualifiedLookupFactory::findResultsAndSaveUnavailables(
-    const DeclContext *lookupContextForThisContext,
-    ResultFinderForTypeContext &&resultFinderForTypeContext,
+    const DeclContext *dynamicContext,
+    const DeclContext *staticContext,
+    SmallVector<NominalTypeDecl *, 2> selfBounds,
     NLOptions baseNLOptions) {
+  if (selfBounds.empty())
+    return;
+
   auto firstPossiblyUnavailableResult = Results.size();
-  resultFinderForTypeContext.findResults(Name, baseNLOptions,
-                                         lookupContextForThisContext, Results);
+  SmallVector<ValueDecl *, 4> Lookup;
+  staticContext->lookupQualified(selfBounds, Name, Loc, baseNLOptions, Lookup);
+  for (auto Result : Lookup) {
+    auto baseDC = isa<TypeDecl>(Result) ? staticContext : dynamicContext;
+    auto baseDecl = getBaseDeclForResult(baseDC);
+    Results.emplace_back(const_cast<DeclContext *>(baseDC), baseDecl, Result);
+#ifndef NDEBUG
+    addedResult(Results.back());
+#endif
+  }
+
   setAsideUnavailableResults(firstPossiblyUnavailableResult);
 }
 
@@ -651,34 +598,6 @@ void UnqualifiedLookupFactory::recordCompletionOfAScope() {
     IndexOfFirstOuterResult = Results.size();
 }
 
-UnqualifiedLookupFactory::ResultFinderForTypeContext::
-    ResultFinderForTypeContext(UnqualifiedLookupFactory *factory,
-                               const DeclContext *dynamicContext,
-                               const DeclContext *staticContext)
-    : factory(factory), dynamicContext(dynamicContext),
-      staticContext(staticContext), selfBounds(findSelfBounds(staticContext)) {}
-
-UnqualifiedLookupFactory::ResultFinderForTypeContext::SelfBounds
-UnqualifiedLookupFactory::ResultFinderForTypeContext::findSelfBounds(
-    const DeclContext *dc) {
-  auto nominal = dc->getSelfNominalTypeDecl();
-  if (!nominal)
-    return {};
-
-  SelfBounds 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);
-  }
-  return selfBounds;
-}
-
 #pragma mark ASTScopeImpl support
 
 void UnqualifiedLookupFactory::lookInASTScopes() {
@@ -799,19 +718,39 @@ bool ASTScopeDeclGatherer::consume(ArrayRef<ValueDecl *> valuesArg,
   return false;
 }
 
+static SmallVector<NominalTypeDecl *, 2> findSelfBounds(const DeclContext *dc) {
+  auto nominal = dc->getSelfNominalTypeDecl();
+  if (!nominal)
+    return {};
+
+  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);
+  }
+
+  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
   // is visible, instance members are returned with the 'self' parameter's
   // DeclContext as the base, which is how the expression checker knows to
   // convert the unqualified reference into a self member access.
-  auto resultFinder = UnqualifiedLookupFactory::ResultFinderForTypeContext(
-      &factory, candidateSelfDC ? candidateSelfDC : scopeDC, scopeDC);
-  factory.findResultsAndSaveUnavailables(scopeDC, std::move(resultFinder),
-                                         factory.baseNLOptions);
+  factory.findResultsAndSaveUnavailables(candidateSelfDC ? candidateSelfDC : scopeDC,
+                                         scopeDC, selfBounds, factory.baseNLOptions);
   factory.recordCompletionOfAScope();
 
   // We're done looking inside a nominal type declaration. It is possible
@@ -837,18 +776,6 @@ UnqualifiedLookupRequest::evaluate(Evaluator &evaluator,
 
 #pragma mark debugging
 
-void UnqualifiedLookupFactory::ResultFinderForTypeContext::dump() const {
-  (void)factory;
-  llvm::errs() << "dynamicContext: ";
-  dynamicContext->dumpContext();
-  llvm::errs() << "staticContext: ";
-  staticContext->dumpContext();
-  llvm::errs() << "selfBounds: ";
-  for (const auto *D : selfBounds)
-    D->dump(llvm::errs(), 1);
-  llvm::errs() << "\n";
-}
-
 void UnqualifiedLookupFactory::dump() const { print(llvm::errs()); }
 void UnqualifiedLookupFactory::dumpScopes() const { printScopes(llvm::errs()); }
 void UnqualifiedLookupFactory::dumpResults() const {
