Merge pull request swiftlang#28922 from CodaFi/neo-compl-cache

[NFC] One-Shot Name Lookup + Lazily-Complete Base Name Cache

Author: CodaFi

include/swift/AST/Decl.h

@@ -3300,6 +3300,10 @@ class NominalTypeDecl : public GenericTypeDecl, public IterableDeclContext {
   /// so that it can also be added to the lookup table (if needed).
   void addedMember(Decl *member);
 
+  /// Note that we have added an extension into the nominal type,
+  /// so that its members can eventually be added to the lookup table.
+  void addedExtension(ExtensionDecl *ext);
+
   /// A lookup table used to find the protocol conformances of
   /// a given nominal type.
   mutable ConformanceLookupTable *ConformanceTable = nullptr;

lib/AST/Decl.cpp

@@ -3664,12 +3664,16 @@ void NominalTypeDecl::addExtension(ExtensionDecl *extension) {
   if (!FirstExtension) {
     FirstExtension = extension;
     LastExtension = extension;
+
+    addedExtension(extension);
     return;
   }
 
   // Add to the end of the list.
   LastExtension->NextExtension.setPointer(extension);
   LastExtension = extension;
+
+  addedExtension(extension);
 }
 
 ArrayRef<VarDecl *> NominalTypeDecl::getStoredProperties() const {

lib/AST/NameLookup.cpp

@@ -880,6 +880,11 @@ class swift::MemberLookupTable {
   /// Lookup table mapping names to the set of declarations with that name.
   LookupTable Lookup;
 
+  /// The set of names of lazily-loaded members that the lookup table has a
+  /// complete accounting of with respect to all known extensions of its
+  /// parent nominal type.
+  llvm::DenseSet<DeclBaseName> LazilyCompleteNames;
+
 public:
   /// Create a new member lookup table.
   explicit MemberLookupTable(ASTContext &ctx);
@@ -893,6 +898,24 @@ class swift::MemberLookupTable {
   /// Add the given members to the lookup table.
   void addMembers(DeclRange members);
 
+  /// Returns \c true if the lookup table has a complete accounting of the
+  /// given name.
+  bool isLazilyComplete(DeclBaseName name) const {
+    return LazilyCompleteNames.find(name) != LazilyCompleteNames.end();
+  }
+
+  /// Mark a given lazily-loaded name as being complete.
+  void markLazilyComplete(DeclBaseName name) {
+    LazilyCompleteNames.insert(name);
+  }
+
+  /// Clears the cache of lazily-complete names.  This _must_ be called when
+  /// new extensions with lazy members are added to the type, or direct lookup
+  /// will return inconsistent or stale results.
+  void clearLazilyCompleteCache() {
+    LazilyCompleteNames.clear();
+  }
+
   /// Iterator into the lookup table.
   typedef LookupTable::iterator iterator;
 
@@ -912,7 +935,11 @@ class swift::MemberLookupTable {
 
     os << "Lookup:\n  ";
     for (auto &pair : Lookup) {
-      pair.getFirst().print(os) << ":\n  ";
+      pair.getFirst().print(os);
+      if (isLazilyComplete(pair.getFirst().getBaseName())) {
+        os << " (lazily complete)";
+      }
+      os << ":\n  ";
       for (auto &decl : pair.getSecond()) {
         os << "- ";
         decl->dumpRef(os);
@@ -926,21 +953,6 @@ class swift::MemberLookupTable {
     dump(llvm::errs());
   }
 
-  // Mark all Decls in this table as not-resident in a table, drop
-  // references to them. Should only be called when this was not fully-populated
-  // from an IterableDeclContext.
-  void clear() {
-    // LastExtensionIncluded would only be non-null if this was populated from
-    // an IterableDeclContext (though it might still be null in that case).
-    assert(LastExtensionIncluded == nullptr);
-    for (auto const &i : Lookup) {
-      for (auto d : i.getSecond()) {
-        d->setAlreadyInLookupTable(false);
-      }
-    }
-    Lookup.clear();
-  }
-
   // Only allow allocation of member lookup tables using the allocator in
   // ASTContext or by doing a placement new.
   void *operator new(size_t Bytes, ASTContext &C,
@@ -1043,25 +1055,26 @@ void MemberLookupTable::updateLookupTable(NominalTypeDecl *nominal) {
   }
 }
 
-void NominalTypeDecl::addedMember(Decl *member) {
-  auto *vd = dyn_cast<ValueDecl>(member);
-  if (!vd)
-    return;
+void NominalTypeDecl::addedExtension(ExtensionDecl *ext) {
+  if (!LookupTable) return;
 
+  if (ext->hasLazyMembers()) {
+    LookupTable->addMembers(ext->getCurrentMembersWithoutLoading());
+    LookupTable->clearLazilyCompleteCache();
+  } else {
+    LookupTable->addMembers(ext->getMembers());
+  }
+}
+
+void NominalTypeDecl::addedMember(Decl *member) {
   // If we have a lookup table, add the new member to it. If not, we'll pick up
   // this member when we first create the table.
-  if (auto *lookup = LookupTable) {
-    if (hasLazyMembers()) {
-      // If we have lazy members, only add the new member to the lookup
-      // table if we already have another member with the same name.
-      // The presence of a lookup table entry indicates that the
-      // nominal as well as all extensions have already been searched.
-      if (lookup->find(vd->getBaseName()) == lookup->end())
-        return;
-    }
+  auto *vd = dyn_cast<ValueDecl>(member);
+  auto *lookup = LookupTable;
+  if (!vd || !lookup)
+    return;
 
-    lookup->addMember(vd);
-  }
+  lookup->addMember(vd);
 }
 
 void ExtensionDecl::addedMember(Decl *member) {
@@ -1094,8 +1107,8 @@ void ExtensionDecl::addedMember(Decl *member) {
 // │ExtensionDecl *LastExtension ─┼───────┐│      │                   └───┐
 // │                              │       ││      └──────────────────────┐│
 // │MemberLookupTable *LookupTable├─┐     ││                             ││
-// │bool LookupTableComplete      │ │     ││     ┌─────────────────┐     ││
-// └──────────────────────────────┘ │     ││     │ExtensionDecl    │     ││
+// └──────────────────────────────┘ │     ││     ┌─────────────────┐     ││
+//                                  │     ││     │ExtensionDecl    │     ││
 //                                  │     ││     │-------------    │     ││
 //                    ┌─────────────┘     │└────▶│ExtensionDecl    │     ││
 //                    │                   │      │  *NextExtension ├──┐  ││
@@ -1156,34 +1169,30 @@ populateLookupTableEntryFromLazyIDCLoader(ASTContext &ctx,
   }
 }
 
-static void populateLookupTableEntryFromCurrentMembers(
-    ASTContext &ctx, MemberLookupTable &LookupTable, DeclBaseName name,
-    IterableDeclContext *IDC) {
-  for (auto m : IDC->getMembers()) {
-    if (auto v = dyn_cast<ValueDecl>(m)) {
-      if (v->getBaseName() == name) {
-        LookupTable.addMember(m);
-      }
-    }
-  }
-}
-
 static void
 populateLookupTableEntryFromExtensions(ASTContext &ctx,
                                        MemberLookupTable &table,
                                        NominalTypeDecl *nominal,
                                        DeclBaseName name) {
+  assert(!table.isLazilyComplete(name) &&
+         "Should not be searching extensions for complete name!");
+
   for (auto e : nominal->getExtensions()) {
-    // If we can retrieve the members of this extension without deserializing
-    // anything, do so now.
-    if (!e->wasDeserialized() && !e->hasClangNode()) {
-      populateLookupTableEntryFromCurrentMembers(ctx, table, name, e);
+    // If there's no lazy members to look at, all the members of this extension
+    // are present in the lookup table.
+    if (!e->hasLazyMembers()) {
       continue;
     }
 
+    assert(e->wasDeserialized() || e->hasClangNode() &&
+           "Extension without deserializable content has lazy members!");
     assert(!e->hasUnparsedMembers());
+
+    // Try lazy loading. If that fails, then we fall back by loading the
+    // entire extension. FIXME: It's rather unfortunate that we fall off the
+    // happy path because the Clang Importer can't handle lazy import-as-member.
     if (populateLookupTableEntryFromLazyIDCLoader(ctx, table, name, e)) {
-      populateLookupTableEntryFromCurrentMembers(ctx, table, name, e);
+      e->loadAllMembers();
     }
   }
 }
@@ -1204,19 +1213,18 @@ void NominalTypeDecl::prepareLookupTable() {
 
     // 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, then update table entries from the extensions that
-    // have the same names as any such initial-population members.
+    // globals-as-members.
     LookupTable->addMembers(getCurrentMembersWithoutLoading());
-
-    llvm::SmallSet<DeclBaseName, 4> baseNamesPresent;
-    for (auto entry : *LookupTable) {
-      auto baseName = entry.getFirst().getBaseName();
-      if (!baseNamesPresent.insert(baseName).second)
+    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;
+      }
 
-      populateLookupTableEntryFromExtensions(getASTContext(),
-                                             *LookupTable,
-                                             this, baseName);
+      // Else, load all the members into the table.
+      LookupTable->addMembers(e->getMembers());
     }
   } else {
     LookupTable->addMembers(getMembers());
@@ -1276,8 +1284,9 @@ NominalTypeDecl::lookupDirect(DeclName name,
           DeclName name) -> Optional<TinyPtrVector<ValueDecl *>> {
     // Look for a declaration with this name.
     auto known = table->find(name);
-    if (known == table->end())
+    if (known == table->end()) {
       return None;
+    }
 
     // We found something; return it.
     return maybeFilterOutAttrImplements(known->second, name,
@@ -1297,29 +1306,22 @@ NominalTypeDecl::lookupDirect(DeclName name,
 
   if (!useNamedLazyMemberLoading) {
     updateLookupTable(LookupTable);
+  } else if (!LookupTable->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);
+    }
+    Table.markLazilyComplete(baseName);
   }
 
   // Look for a declaration with this name.
-  if (auto lookup = tryCacheLookup(LookupTable, name))
-    return lookup.getValue();
-
-  if (!useNamedLazyMemberLoading) {
-    return { };
-  }
-
-  // If we get here, we had a cache-miss and _are_ using
-  // NamedLazyMemberLoading. Try to populate a _single_ entry in the
-  // MemberLookupTable from both this nominal and all of its extensions, and
-  // retry.
-  auto &Table = *LookupTable;
-  if (populateLookupTableEntryFromLazyIDCLoader(ctx, Table,
-                                                name.getBaseName(), this)) {
-    updateLookupTable(LookupTable);
-  } else {
-    populateLookupTableEntryFromExtensions(ctx, Table, this,
-                                           name.getBaseName());
-  }
-
   return tryCacheLookup(LookupTable, name)
             .getValueOr(TinyPtrVector<ValueDecl *>());
 }

test/NameBinding/named_lazy_member_loading_objc_category.swift

@@ -8,7 +8,7 @@
 // Check that named-lazy-member-loading reduces the number of Decls deserialized
 // RUN: %target-swift-frontend -typecheck -I %S/Inputs/NamedLazyMembers -disable-named-lazy-member-loading -stats-output-dir %t/stats-pre %s
 // RUN: %target-swift-frontend -typecheck -I %S/Inputs/NamedLazyMembers -stats-output-dir %t/stats-post %s
-// RUN: %{python} %utils/process-stats-dir.py --evaluate-delta 'NumTotalClangImportedEntities < -10' %t/stats-pre %t/stats-post
+// RUN: %{python} %utils/process-stats-dir.py --evaluate-delta 'NumTotalClangImportedEntities < -30' %t/stats-pre %t/stats-post
 
 import NamedLazyMembers
 

test/NameBinding/named_lazy_member_loading_objc_protocol.swift

@@ -8,7 +8,7 @@
 // Check that named-lazy-member-loading reduces the number of Decls deserialized
 // RUN: %target-swift-frontend -typecheck -I %S/Inputs/NamedLazyMembers -disable-named-lazy-member-loading -stats-output-dir %t/stats-pre %s
 // RUN: %target-swift-frontend -typecheck -I %S/Inputs/NamedLazyMembers -stats-output-dir %t/stats-post %s
-// RUN: %{python} %utils/process-stats-dir.py --evaluate-delta 'NumTotalClangImportedEntities < -9' %t/stats-pre %t/stats-post
+// RUN: %{python} %utils/process-stats-dir.py --evaluate-delta 'NumTotalClangImportedEntities < -40' %t/stats-pre %t/stats-post
 
 import NamedLazyMembers