Add DirectLookupRequest

include/swift/AST/Decl.h

@@ -3326,6 +3326,7 @@ class NominalTypeDecl : public GenericTypeDecl, public IterableDeclContext {
   friend class ExtensionDecl;
   friend class DeclContext;
   friend class IterableDeclContext;
+  friend class DirectLookupRequest;
   friend ArrayRef<ValueDecl *>
   ValueDecl::getSatisfiedProtocolRequirements(bool Sorted) const;
 
@@ -3400,6 +3401,12 @@ class NominalTypeDecl : public GenericTypeDecl, public IterableDeclContext {
     /// Whether to include @_implements members.
     /// Used by conformance-checking to find special @_implements members.
     IncludeAttrImplements = 1 << 0,
+    /// Whether to avoid loading lazy members from any new extensions that would otherwise be found
+    /// by deserialization.
+    ///
+    /// Used by the module loader to break recursion and as an optimization e.g. when it is known that a
+    /// particular member declaration will never appear in an extension.
+    IgnoreNewExtensions = 1 << 1,
   };
 
   /// Find all of the declarations with the given name within this nominal type
@@ -7413,6 +7420,9 @@ ParameterList *getParameterList(ValueDecl *source);
 /// nullptr if the source does not have a parameter list.
 const ParamDecl *getParameterAt(const ValueDecl *source, unsigned index);
 
+void simple_display(llvm::raw_ostream &out,
+                    OptionSet<NominalTypeDecl::LookupDirectFlags> options);
+
 /// Display Decl subclasses.
 void simple_display(llvm::raw_ostream &out, const Decl *decl);
 

include/swift/AST/DeclContext.h

@@ -681,18 +681,9 @@ enum class IterableDeclContextKind : uint8_t {
 /// Note that an iterable declaration context must inherit from both
 /// \c IterableDeclContext and \c DeclContext.
 class IterableDeclContext {
-  enum LazyMembers : unsigned {
-    Present = 1 << 0,
-
-    /// Lazy member loading has a variety of feedback loops that need to
-    /// switch to pseudo-empty-member behaviour to avoid infinite recursion;
-    /// we use this flag to control them.
-    InProgress = 1 << 1,
-  };
-
   /// The first declaration in this context along with a bit indicating whether
   /// the members of this context will be lazily produced.
-  mutable llvm::PointerIntPair<Decl *, 2, LazyMembers> FirstDeclAndLazyMembers;
+  mutable llvm::PointerIntPair<Decl *, 1, bool> FirstDeclAndLazyMembers;
 
   /// The last declaration in this context, used for efficient insertion,
   /// along with the kind of iterable declaration context.
@@ -780,20 +771,7 @@ class IterableDeclContext {
 
   /// Check whether there are lazily-loaded members.
   bool hasLazyMembers() const {
-    return FirstDeclAndLazyMembers.getInt() & LazyMembers::Present;
-  }
-
-  bool isLoadingLazyMembers() {
-    return FirstDeclAndLazyMembers.getInt() & LazyMembers::InProgress;
-  }
-
-  void setLoadingLazyMembers(bool inProgress) {
-    LazyMembers status = FirstDeclAndLazyMembers.getInt();
-    if (inProgress)
-      status = LazyMembers(status | LazyMembers::InProgress);
-    else
-      status = LazyMembers(status & ~LazyMembers::InProgress);
-    FirstDeclAndLazyMembers.setInt(status);
+    return FirstDeclAndLazyMembers.getInt();
   }
 
   /// Setup the loader for lazily-loaded members.

include/swift/AST/NameLookupRequests.h

@@ -467,6 +467,54 @@ class QualifiedLookupRequest
            NLOptions opts) const;
 };
 
+/// The input type for a direct lookup request.
+class DirectLookupDescriptor final {
+  using LookupOptions = OptionSet<NominalTypeDecl::LookupDirectFlags>;
+
+public:
+  NominalTypeDecl *DC;
+  DeclName Name;
+  LookupOptions Options;
+
+  DirectLookupDescriptor(NominalTypeDecl *dc, DeclName name,
+                         LookupOptions options = {})
+      : DC(dc), Name(name), Options(options) {}
+
+  friend llvm::hash_code hash_value(const DirectLookupDescriptor &desc) {
+    return llvm::hash_combine(desc.Name, desc.DC, desc.Options.toRaw());
+  }
+
+  friend bool operator==(const DirectLookupDescriptor &lhs,
+                         const DirectLookupDescriptor &rhs) {
+    return lhs.Name == rhs.Name && lhs.DC == rhs.DC &&
+           lhs.Options.toRaw() == rhs.Options.toRaw();
+  }
+
+  friend bool operator!=(const DirectLookupDescriptor &lhs,
+                         const DirectLookupDescriptor &rhs) {
+    return !(lhs == rhs);
+  }
+};
+
+void simple_display(llvm::raw_ostream &out, const DirectLookupDescriptor &desc);
+
+SourceLoc extractNearestSourceLoc(const DirectLookupDescriptor &desc);
+
+class DirectLookupRequest
+    : public SimpleRequest<DirectLookupRequest,
+                           TinyPtrVector<ValueDecl *>(DirectLookupDescriptor),
+                           CacheKind::Uncached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  llvm::Expected<TinyPtrVector<ValueDecl *>>
+  evaluate(Evaluator &evaluator, DirectLookupDescriptor desc) const;
+};
+
 #define SWIFT_TYPEID_ZONE NameLookup
 #define SWIFT_TYPEID_HEADER "swift/AST/NameLookupTypeIDZone.def"
 #include "swift/Basic/DefineTypeIDZone.h"

include/swift/AST/NameLookupTypeIDZone.def

@@ -21,6 +21,9 @@ SWIFT_REQUEST(NameLookup, AnyObjectLookupRequest,
 SWIFT_REQUEST(NameLookup, CustomAttrNominalRequest,
               NominalTypeDecl *(CustomAttr *, DeclContext *), Cached,
               NoLocationInfo)
+SWIFT_REQUEST(NameLookup, DirectLookupRequest,
+              TinyPtrVector<ValueDecl *>(DirectLookupDescriptor), Uncached,
+              NoLocationInfo)
 SWIFT_REQUEST(NameLookup, ExpandASTScopeRequest,
               ast_scope::ASTScopeImpl* (ast_scope::ASTScopeImpl*, ast_scope::ScopeCreator*),
               SeparatelyCached,

include/swift/Basic/Statistics.def

@@ -240,9 +240,6 @@ FRONTEND_STATISTIC(Sema, NumLazyRequirementSignaturesLoaded)
 /// Number of lazy iterable declaration contexts constructed.
 FRONTEND_STATISTIC(Sema, NumLazyIterableDeclContexts)
 
-/// Number of direct member-name lookups performed on nominal types.
-FRONTEND_STATISTIC(Sema, NominalTypeLookupDirectCount)
-
 /// Number of member-name lookups that avoided loading all members.
 FRONTEND_STATISTIC(Sema, NamedLazyMemberLoadSuccessCount)
 

lib/AST/Decl.cpp

@@ -4648,7 +4648,8 @@ ValueDecl *ProtocolDecl::getSingleRequirement(DeclName name) const {
 }
 
 AssociatedTypeDecl *ProtocolDecl::getAssociatedType(Identifier name) const {
-  auto results = const_cast<ProtocolDecl *>(this)->lookupDirect(name);
+  const auto flags = NominalTypeDecl::LookupDirectFlags::IgnoreNewExtensions;
+  auto results = const_cast<ProtocolDecl *>(this)->lookupDirect(name, flags);
   for (auto candidate : results) {
     if (candidate->getDeclContext() == this &&
         isa<AssociatedTypeDecl>(candidate)) {
@@ -7917,6 +7918,15 @@ void swift::simple_display(llvm::raw_ostream &out, const Decl *decl) {
   }
 }
 
+void swift::simple_display(llvm::raw_ostream &out,
+                           OptionSet<NominalTypeDecl::LookupDirectFlags> opts) {
+  out << "{ ";
+  using LookupFlags = NominalTypeDecl::LookupDirectFlags;
+  if (opts.contains(LookupFlags::IncludeAttrImplements))
+    out << "IncludeAttrImplements";
+  out << " }";
+}
+
 void swift::simple_display(llvm::raw_ostream &out, const ValueDecl *decl) {
   if (decl) decl->dumpRef(out);
   else out << "(null)";

lib/AST/DeclContext.cpp

@@ -798,8 +798,7 @@ void IterableDeclContext::setMemberLoader(LazyMemberLoader *loader,
 
   ASTContext &ctx = getASTContext();
   auto contextInfo = ctx.getOrCreateLazyIterableContextData(this, loader);
-  auto lazyMembers = FirstDeclAndLazyMembers.getInt() | LazyMembers::Present;
-  FirstDeclAndLazyMembers.setInt(LazyMembers(lazyMembers));
+  FirstDeclAndLazyMembers.setInt(true);
   contextInfo->memberData = contextData;
 
   ++NumLazyIterableDeclContexts;
@@ -855,12 +854,11 @@ void IterableDeclContext::loadAllMembers() const {
     return;
 
   // Don't try to load all members re-entrant-ly.
-  auto contextInfo = ctx.getOrCreateLazyIterableContextData(this,
-    /*lazyLoader=*/nullptr);
-  auto lazyMembers = FirstDeclAndLazyMembers.getInt() & ~LazyMembers::Present;
-  FirstDeclAndLazyMembers.setInt(LazyMembers(lazyMembers));
+  FirstDeclAndLazyMembers.setInt(false);
 
   const Decl *container = getDecl();
+  auto contextInfo = ctx.getOrCreateLazyIterableContextData(this,
+    /*lazyLoader=*/nullptr);
   contextInfo->loader->loadAllMembers(const_cast<Decl *>(container),
                                       contextInfo->memberData);
 

lib/AST/NameLookup.cpp

@@ -1148,11 +1148,9 @@ populateLookupTableEntryFromLazyIDCLoader(ASTContext &ctx,
                                           MemberLookupTable &LookupTable,
                                           DeclBaseName name,
                                           IterableDeclContext *IDC) {
-  IDC->setLoadingLazyMembers(true);
   auto ci = ctx.getOrCreateLazyIterableContextData(IDC,
                                                    /*lazyLoader=*/nullptr);
   if (auto res = ci->loader->loadNamedMembers(IDC, name, ci->memberData)) {
-    IDC->setLoadingLazyMembers(false);
     if (auto s = ctx.Stats) {
       ++s->getFrontendCounters().NamedLazyMemberLoadSuccessCount;
     }
@@ -1161,7 +1159,6 @@ populateLookupTableEntryFromLazyIDCLoader(ASTContext &ctx,
     }
     return false;
   } else {
-    IDC->setLoadingLazyMembers(false);
     if (auto s = ctx.Stats) {
       ++s->getFrontendCounters().NamedLazyMemberLoadFailureCount;
     }
@@ -1210,25 +1207,21 @@ void NominalTypeDecl::prepareLookupTable() {
 
   if (hasLazyMembers()) {
     assert(!hasUnparsedMembers());
-
-    // Lazy members: if the table needs population, populate the table _only
-    // from those members already in the IDC member list_ such as implicits or
-    // globals-as-members.
     LookupTable->addMembers(getCurrentMembersWithoutLoading());
-    for (auto e : getExtensions()) {
-      // If we can lazy-load this extension, only take the members we've loaded
-      // so far.
-      if (e->wasDeserialized() || e->hasClangNode()) {
-        LookupTable->addMembers(e->getCurrentMembersWithoutLoading());
-        continue;
-      }
-
-      // Else, load all the members into the table.
-      LookupTable->addMembers(e->getMembers());
-    }
   } else {
     LookupTable->addMembers(getMembers());
-    LookupTable->updateLookupTable(this);
+  }
+
+  for (auto e : getExtensions()) {
+    // If we can lazy-load this extension, only take the members we've loaded
+    // so far.
+    if (e->wasDeserialized() || e->hasClangNode()) {
+      LookupTable->addMembers(e->getCurrentMembersWithoutLoading());
+      continue;
+    }
+
+    // Else, load all the members into the table.
+    LookupTable->addMembers(e->getMembers());
   }
 }
 
@@ -1257,34 +1250,42 @@ maybeFilterOutAttrImplements(TinyPtrVector<ValueDecl *> decls,
 TinyPtrVector<ValueDecl *>
 NominalTypeDecl::lookupDirect(DeclName name,
                               OptionSet<LookupDirectFlags> flags) {
-  ASTContext &ctx = getASTContext();
-  if (auto s = ctx.Stats) {
-    ++s->getFrontendCounters().NominalTypeLookupDirectCount;
-  }
+  return evaluateOrDefault(getASTContext().evaluator,
+                           DirectLookupRequest({this, name, flags}), {});
+}
+
+llvm::Expected<TinyPtrVector<ValueDecl *>>
+DirectLookupRequest::evaluate(Evaluator &evaluator,
+                              DirectLookupDescriptor desc) const {
+  const auto &name = desc.Name;
+  const auto flags = desc.Options;
+  auto *decl = desc.DC;
 
   // We only use NamedLazyMemberLoading when a user opts-in and we have
   // not yet loaded all the members into the IDC list in the first place.
+  ASTContext &ctx = decl->getASTContext();
   const bool useNamedLazyMemberLoading = (ctx.LangOpts.NamedLazyMemberLoading &&
-                                          hasLazyMembers());
-
+                                          decl->hasLazyMembers());
+  const bool disableAdditionalExtensionLoading =
+      flags.contains(NominalTypeDecl::LookupDirectFlags::IgnoreNewExtensions);
   const bool includeAttrImplements =
-      flags.contains(LookupDirectFlags::IncludeAttrImplements);
+      flags.contains(NominalTypeDecl::LookupDirectFlags::IncludeAttrImplements);
 
-  LLVM_DEBUG(llvm::dbgs() << getNameStr() << ".lookupDirect("
+  LLVM_DEBUG(llvm::dbgs() << decl->getNameStr() << ".lookupDirect("
                           << name << ")"
-                          << ", hasLazyMembers()=" << hasLazyMembers()
-                          << ", useNamedLazyMemberLoading=" << useNamedLazyMemberLoading
+                          << ", hasLazyMembers()=" << decl->hasLazyMembers()
+                          << ", useNamedLazyMemberLoading="
+                          << useNamedLazyMemberLoading
                           << "\n");
 
-
-  prepareLookupTable();
+  decl->prepareLookupTable();
 
   auto tryCacheLookup =
-      [=](MemberLookupTable *table,
+      [=](MemberLookupTable &table,
           DeclName name) -> Optional<TinyPtrVector<ValueDecl *>> {
     // Look for a declaration with this name.
-    auto known = table->find(name);
-    if (known == table->end()) {
+    auto known = table.find(name);
+    if (known == table.end()) {
       return None;
     }
 
@@ -1293,36 +1294,40 @@ NominalTypeDecl::lookupDirect(DeclName name,
                                         includeAttrImplements);
   };
 
-  auto updateLookupTable = [this](MemberLookupTable *table) {
+  auto updateLookupTable = [&decl](MemberLookupTable &table,
+                                   bool noExtensions) {
     // Make sure we have the complete list of members (in this nominal and in
     // all extensions).
-    (void)getMembers();
+    (void)decl->getMembers();
+
+    if (noExtensions)
+      return;
 
-    for (auto E : getExtensions())
+    for (auto E : decl->getExtensions())
       (void)E->getMembers();
 
-    LookupTable->updateLookupTable(this);
+    table.updateLookupTable(decl);
   };
 
+  auto &Table = *decl->LookupTable;
   if (!useNamedLazyMemberLoading) {
-    updateLookupTable(LookupTable);
-  } else if (!LookupTable->isLazilyComplete(name.getBaseName())) {
+    updateLookupTable(Table, disableAdditionalExtensionLoading);
+  } else if (!Table.isLazilyComplete(name.getBaseName())) {
     // The lookup table believes it doesn't have a complete accounting of this
     // name - either because we're never seen it before, or another extension
     // was registered since the last time we searched. Ask the loaders to give
     // us a hand.
-    auto &Table = *LookupTable;
     DeclBaseName baseName(name.getBaseName());
-    if (populateLookupTableEntryFromLazyIDCLoader(ctx, Table, baseName, this)) {
-      updateLookupTable(LookupTable);
-    } else {
-      populateLookupTableEntryFromExtensions(ctx, Table, this, baseName);
+    if (populateLookupTableEntryFromLazyIDCLoader(ctx, Table, baseName, decl)) {
+      updateLookupTable(Table, disableAdditionalExtensionLoading);
+    } else if (!disableAdditionalExtensionLoading) {
+      populateLookupTableEntryFromExtensions(ctx, Table, decl, baseName);
     }
     Table.markLazilyComplete(baseName);
   }
 
   // Look for a declaration with this name.
-  return tryCacheLookup(LookupTable, name)
+  return tryCacheLookup(Table, name)
             .getValueOr(TinyPtrVector<ValueDecl *>());
 }
 

lib/AST/NameLookupRequests.cpp

@@ -188,6 +188,24 @@ swift::extractNearestSourceLoc(const UnqualifiedLookupDescriptor &desc) {
   return extractNearestSourceLoc(desc.DC);
 }
 
+//----------------------------------------------------------------------------//
+// DirectLookupRequest computation.
+//----------------------------------------------------------------------------//
+
+void swift::simple_display(llvm::raw_ostream &out,
+                           const DirectLookupDescriptor &desc) {
+  out << "directly looking up ";
+  simple_display(out, desc.Name);
+  out << " on ";
+  simple_display(out, desc.DC);
+  out << " with options ";
+  simple_display(out, desc.Options);
+}
+
+SourceLoc swift::extractNearestSourceLoc(const DirectLookupDescriptor &desc) {
+  return extractNearestSourceLoc(desc.DC);
+}
+
 // Define request evaluation functions for each of the name lookup requests.
 static AbstractRequestFunction *nameLookupRequestFunctions[] = {
 #define SWIFT_REQUEST(Zone, Name, Sig, Caching, LocOptions)                    \