[ClangImporter] Support lazy member loading for import-as-member globals in extensions

lib/AST/NameLookup.cpp

@@ -1277,6 +1277,10 @@ class swift::MemberLookupTable : public ASTAllocated<swift::MemberLookupTable> {
   /// Add the given members to the lookup table.
   void addMembers(DeclRange members);
 
+  /// Add the members of the extension to the lookup table, if necessary
+  /// registering it for future lazy member loading.
+  void addExtension(ExtensionDecl *ext);
+
   void addExtensionWithLazyMembers(ExtensionDecl *ext) {
     ExtensionsWithLazyMembers.push_back(ext);
   }
@@ -1467,23 +1471,29 @@ void MemberLookupTable::addMembers(DeclRange members) {
   }
 }
 
+void MemberLookupTable::addExtension(ExtensionDecl *ext) {
+  // If we can lazy-load this extension, only take the members we've loaded
+  // so far.
+  if (ext->hasLazyMembers()) {
+    addMembers(ext->getCurrentMembersWithoutLoading());
+    clearLazilyCompleteCache();
+    clearLazilyCompleteForMacroExpansionCache();
+    addExtensionWithLazyMembers(ext);
+  } else {
+    // Else, load all the members into the table.
+    addMembers(ext->getMembers());
+  }
+  addContainerWithMacroExpansions(ext);
+}
+
 void NominalTypeDecl::addedExtension(ExtensionDecl *ext) {
   if (!LookupTable.getInt())
     return;
 
   auto *table = LookupTable.getPointer();
   assert(table);
 
-  if (ext->wasDeserialized() || ext->hasClangNode()) {
-    table->addMembers(ext->getCurrentMembersWithoutLoading());
-    table->clearLazilyCompleteCache();
-    table->clearLazilyCompleteForMacroExpansionCache();
-    table->addExtensionWithLazyMembers(ext);
-  } else {
-    table->addMembers(ext->getMembers());
-  }
-
-  table->addContainerWithMacroExpansions(ext);
+  table->addExtension(ext);
 }
 
 void NominalTypeDecl::addedMember(Decl *member) {
@@ -1603,8 +1613,6 @@ populateLookupTableEntryFromExtensions(ASTContext &ctx,
       continue;
     }
 
-    assert(e->wasDeserialized() || e->hasClangNode() &&
-           "Extension without deserializable content has lazy members!");
     assert(!e->hasUnparsedMembers());
 
     populateLookupTableEntryFromLazyIDCLoader(ctx, table, name, e);
@@ -1952,21 +1960,7 @@ void NominalTypeDecl::prepareLookupTable() {
 
   // Note: this calls prepareExtensions()
   for (auto e : getExtensions()) {
-    // If we can lazy-load this extension, only take the members we've loaded
-    // so far.
-    //
-    // FIXME: This should be 'e->hasLazyMembers()' but that crashes` because
-    // some imported extensions don't have a Clang node, and only support
-    // LazyMemberLoader::loadAllMembers() and not
-    // LazyMemberLoader::loadNamedMembers().
-    if (e->wasDeserialized() || e->hasClangNode()) {
-      table->addMembers(e->getCurrentMembersWithoutLoading());
-      table->addExtensionWithLazyMembers(e);
-      continue;
-    }
-
-    // Else, load all the members into the table.
-    table->addMembers(e->getMembers());
+    table->addExtension(e);
   }
 
   // Any extensions added after this point will add their members to the

lib/ClangImporter/ClangImporter.cpp

@@ -6222,13 +6222,11 @@ swift::extractNearestSourceLoc(const ClangCategoryLookupDescriptor &desc) {
 
 TinyPtrVector<ValueDecl *>
 ClangImporter::Implementation::loadNamedMembers(
-    const IterableDeclContext *IDC, DeclBaseName N, uint64_t contextData) {
-
+    const IterableDeclContext *IDC, DeclBaseName N, uint64_t extra) {
   auto *D = IDC->getDecl();
   auto *DC = D->getInnermostDeclContext();
   auto *CD = D->getClangDecl();
-  auto *CDC = cast<clang::DeclContext>(CD);
-  assert(CD && "loadNamedMembers on a Decl without a clangDecl");
+  auto *CDC = cast_or_null<clang::DeclContext>(CD);
 
   auto *nominal = DC->getSelfNominalTypeDecl();
   auto effectiveClangContext = getEffectiveClangContext(nominal);
@@ -6248,15 +6246,22 @@ ClangImporter::Implementation::loadNamedMembers(
   // findLookupTable, below, handles the first two cases; we assert on the
   // third.
 
-  auto CMO = getClangSubmoduleForDecl(CD);
+  llvm::Optional<clang::Module *> CMO;
+  if (CD)
+    CMO = getClangSubmoduleForDecl(CD);
+  else {
+    // IDC is an extension containing globals imported as members, so it doesn't
+    // have a clang node but the submodule pointer has been stashed in `extra`.
+    CMO = reinterpret_cast<clang::Module *>(static_cast<uintptr_t>(extra));
+  }
   assert(CMO && "loadNamedMembers on a forward-declared Decl");
 
   auto table = findLookupTable(*CMO);
   assert(table && "clang module without lookup table");
 
-  assert(!isa<clang::NamespaceDecl>(CD) && "Namespace members should be loaded"
-                                           "via a request.");
-  assert(isa<clang::ObjCContainerDecl>(CD));
+  assert(!isa_and_nonnull<clang::NamespaceDecl>(CD)
+            && "Namespace members should be loaded via a request.");
+  assert(!CD || isa<clang::ObjCContainerDecl>(CD));
 
   // Force the members of the entire inheritance hierarchy to be loaded and
   // deserialized before loading the named member of a class. This warms up
@@ -6271,32 +6276,37 @@ ClangImporter::Implementation::loadNamedMembers(
 
   // TODO: update this to use the requestified lookup.
   TinyPtrVector<ValueDecl *> Members;
-  for (auto entry : table->lookup(SerializedSwiftName(N),
-                                  effectiveClangContext)) {
-    if (!entry.is<clang::NamedDecl *>()) continue;
-    auto member = entry.get<clang::NamedDecl *>();
-    if (!isVisibleClangEntry(member)) continue;
 
-    // Skip Decls from different clang::DeclContexts
-    if (member->getDeclContext() != CDC) continue;
-
-    SmallVector<Decl*, 4> tmp;
-    insertMembersAndAlternates(member, tmp);
-    for (auto *TD : tmp) {
-      if (auto *V = dyn_cast<ValueDecl>(TD)) {
-        // Skip ValueDecls if they import under different names.
-        if (V->getBaseName() == N) {
-          Members.push_back(V);
+  // Lookup actual, factual clang-side members of the context. No need to do
+  // this if we're handling an import-as-member extension.
+  if (CD) {
+    for (auto entry : table->lookup(SerializedSwiftName(N),
+                                    effectiveClangContext)) {
+      if (!entry.is<clang::NamedDecl *>()) continue;
+      auto member = entry.get<clang::NamedDecl *>();
+      if (!isVisibleClangEntry(member)) continue;
+
+      // Skip Decls from different clang::DeclContexts
+      if (member->getDeclContext() != CDC) continue;
+
+      SmallVector<Decl*, 4> tmp;
+      insertMembersAndAlternates(member, tmp, DC);
+      for (auto *TD : tmp) {
+        if (auto *V = dyn_cast<ValueDecl>(TD)) {
+          // Skip ValueDecls if they import under different names.
+          if (V->getBaseName() == N) {
+            Members.push_back(V);
+          }
         }
-      }
 
-      // If the property's accessors have alternate decls, we might have
-      // to import those too.
-      if (auto *ASD = dyn_cast<AbstractStorageDecl>(TD)) {
-        for (auto *AD : ASD->getAllAccessors()) {
-          for (auto *D : getAlternateDecls(AD)) {
-            if (D->getBaseName() == N)
-              Members.push_back(D);
+        // If the property's accessors have alternate decls, we might have
+        // to import those too.
+        if (auto *ASD = dyn_cast<AbstractStorageDecl>(TD)) {
+          for (auto *AD : ASD->getAllAccessors()) {
+            for (auto *D : getAlternateDecls(AD)) {
+              if (D->getBaseName() == N)
+                Members.push_back(D);
+            }
           }
         }
       }
@@ -6309,11 +6319,14 @@ ClangImporter::Implementation::loadNamedMembers(
     auto member = entry.get<clang::NamedDecl *>();
     if (!isVisibleClangEntry(member)) continue;
 
-    // Skip Decls from different clang::DeclContexts
-    if (member->getDeclContext() != CDC) continue;
+    // Skip Decls from different clang::DeclContexts. We don't do this for
+    // import-as-member extensions because we don't know what decl context to
+    // expect; for instance, an enum constant is inside the enum decl, not in
+    // the translation unit.
+    if (CDC && member->getDeclContext() != CDC) continue;
 
     SmallVector<Decl*, 4> tmp;
-    insertMembersAndAlternates(member, tmp);
+    insertMembersAndAlternates(member, tmp, DC);
     for (auto *TD : tmp) {
       if (auto *V = dyn_cast<ValueDecl>(TD)) {
         // Skip ValueDecls if they import under different names.
@@ -6324,7 +6337,7 @@ ClangImporter::Implementation::loadNamedMembers(
     }
   }
 
-  if (N.isConstructor()) {
+  if (CD && N.isConstructor()) {
     if (auto *classDecl = dyn_cast<ClassDecl>(D)) {
       SmallVector<Decl *, 4> ctors;
       importInheritedConstructors(cast<clang::ObjCInterfaceDecl>(CD),
@@ -6334,7 +6347,7 @@ ClangImporter::Implementation::loadNamedMembers(
     }
   }
 
-  if (!isa<ProtocolDecl>(D)) {
+  if (CD && !isa<ProtocolDecl>(D)) {
     if (auto *OCD = dyn_cast<clang::ObjCContainerDecl>(CD)) {
       SmallVector<Decl *, 1> newMembers;
       importMirroredProtocolMembers(OCD, DC, N, newMembers);

lib/ClangImporter/ImportDecl.cpp

@@ -9136,6 +9136,8 @@ void ClangImporter::Implementation::loadAllMembersOfRecordDecl(
         (isa<clang::FieldDecl>(nd) || nd == nd->getCanonicalDecl());
     if (isCanonicalInContext && nd->getDeclContext() == clangRecord &&
         isVisibleClangEntry(nd))
+      // We don't pass `swiftDecl` as `expectedDC` because we might be in a
+      // recursive call that adds base class members to a derived class.
       insertMembersAndAlternates(nd, members);
   }
 
@@ -9376,7 +9378,9 @@ void ClangImporter::Implementation::loadAllMembersOfObjcContainer(
 }
 
 void ClangImporter::Implementation::insertMembersAndAlternates(
-    const clang::NamedDecl *nd, SmallVectorImpl<Decl *> &members) {
+    const clang::NamedDecl *nd,
+    SmallVectorImpl<Decl *> &members,
+    DeclContext *expectedDC) {
 
   size_t start = members.size();
   llvm::SmallPtrSet<Decl *, 4> knownAlternateMembers;
@@ -9391,9 +9395,13 @@ void ClangImporter::Implementation::insertMembersAndAlternates(
       return false;
     }
 
+    // If no DC was provided, use wherever the primary decl was imported into.
+    if (!expectedDC)
+      expectedDC = member->getDeclContext();
+
     // If there are alternate declarations for this member, add them.
     for (auto alternate : getAlternateDecls(member)) {
-      if (alternate->getDeclContext() == member->getDeclContext() &&
+      if (alternate->getDeclContext() == expectedDC &&
           knownAlternateMembers.insert(alternate).second) {
         members.push_back(alternate);
       }
@@ -9404,7 +9412,8 @@ void ClangImporter::Implementation::insertMembersAndAlternates(
     if (shouldSuppressDeclImport(nd))
       return true;
 
-    members.push_back(member);
+    if (member->getDeclContext() == expectedDC)
+      members.push_back(member);
     if (nameVersion.supportsConcurrency()) {
       assert(!asyncImport &&
              "Should only have a single version with concurrency enabled");
@@ -9436,7 +9445,7 @@ void ClangImporter::Implementation::collectMembersToAdd(
     if (nd && nd == nd->getCanonicalDecl() &&
         nd->getDeclContext() == objcContainer &&
         isVisibleClangEntry(nd))
-      insertMembersAndAlternates(nd, members);
+      insertMembersAndAlternates(nd, members, DC);
   }
 
   // Objective-C protocols don't require any special handling.

lib/ClangImporter/ImporterImpl.h

@@ -1621,7 +1621,8 @@ class LLVM_LIBRARY_VISIBILITY ClangImporter::Implementation
                            Decl *D, DeclContext *DC,
                            SmallVectorImpl<Decl *> &members);
   void insertMembersAndAlternates(const clang::NamedDecl *nd,
-                                  SmallVectorImpl<Decl *> &members);
+                                  SmallVectorImpl<Decl *> &members,
+                                  DeclContext *expectedDC = nullptr);
   void loadAllMembersIntoExtension(Decl *D, uint64_t extra);
 
   /// Imports \p decl under \p nameVersion with the name \p newName, and adds