[Clang importer / Interface printing] Print extensions in the appropriate submodule.

When printing the interface for a (sub)module, make sure that we only
print those extensions that were created to hold that submodule's
globals that were imported as members.

lib/ClangImporter/ClangImporter.cpp

@@ -3286,7 +3286,8 @@ void ClangModuleUnit::getTopLevelDecls(SmallVectorImpl<Decl*> &results) const {
   DarwinBlacklistDeclConsumer blacklistConsumer(filterConsumer,
                                                 getClangASTContext());
 
-  const clang::Module *topLevelModule = clangModule->getTopLevelModule();
+  const clang::Module *topLevelModule =
+    clangModule ? clangModule->getTopLevelModule() : nullptr;
 
   swift::VisibleDeclConsumer *actualConsumer = &filterConsumer;
   if (DarwinBlacklistDeclConsumer::needsBlacklist(topLevelModule))
@@ -3304,14 +3305,22 @@ void ClangModuleUnit::getTopLevelDecls(SmallVectorImpl<Decl*> &results) const {
         results.push_back(extension);
     }
 
-    // Add any extensions created during this import
-    for (auto mapItr : owner.Impl.extensionPoints) {
-      // TODO: Any other entry points where we need to add these? Is there any
-      // way to to instead marshal the extension through the consumer more
-      // directly?
-      results.push_back(mapItr.getSecond());
-    }
+    // Retrieve all of the globals that will be mapped to members.
 
+    // FIXME: Since we don't represent Clang submodules as Swift
+    // modules, we're getting everything.
+    llvm::SmallPtrSet<ExtensionDecl *, 8> knownExtensions;
+    for (auto entry : lookupTable->allGlobalsAsMembers()) {
+      auto decl = entry.get<clang::NamedDecl *>();
+      auto importedDecl = owner.Impl.importDecl(decl);
+      if (!importedDecl) continue;
+
+      auto ext = dyn_cast<ExtensionDecl>(importedDecl->getDeclContext());
+      if (!ext) continue;
+
+      if (knownExtensions.insert(ext).second)
+        results.push_back(ext);
+    }
   }
 }
 

lib/ClangImporter/ImportDecl.cpp

@@ -5770,7 +5770,7 @@ ClangImporter::Implementation::importDeclContextOf(
   // Create a new extension for this nominal type/Clang submodule pair.
   auto swiftTyLoc = TypeLoc::withoutLoc(nominal->getDeclaredType());
   auto ext = ExtensionDecl::create(SwiftContext, SourceLoc(), swiftTyLoc, {},
-                                   importDeclContextOf(decl), nullptr);
+                                   getClangModuleForDecl(decl), nullptr);
   ext->setValidated();
   ext->setCheckedInheritanceClause();
   ext->setMemberLoader(this, reinterpret_cast<uintptr_t>(declSubmodule));

lib/ClangImporter/SwiftLookupTable.cpp

@@ -356,6 +356,31 @@ SwiftLookupTable::lookupGlobalsAsMembers(EffectiveClangContext context) {
   return lookupGlobalsAsMembers(*storedContext);
 }
 
+SmallVector<SwiftLookupTable::SingleEntry, 4>
+SwiftLookupTable::allGlobalsAsMembers() {
+  // If we have a reader, deserialize all of the globals-as-members data.
+  if (Reader) {
+    for (auto context : Reader->getGlobalsAsMembersContexts()) {
+      (void)lookupGlobalsAsMembers(context);
+    }
+  }
+
+  // Collect all of the keys and sort them.
+  SmallVector<StoredContext, 8> contexts;
+  for (const auto &globalAsMember : GlobalsAsMembers) {
+    contexts.push_back(globalAsMember.first);
+  }
+  llvm::array_pod_sort(contexts.begin(), contexts.end());
+
+  // Collect all of the results in order.
+  SmallVector<SwiftLookupTable::SingleEntry, 4> results;
+  for (const auto &context : contexts) {
+    for (auto &entry : GlobalsAsMembers[context])
+      results.push_back(mapStored(entry));
+  }
+  return results;
+}
+
 SmallVector<SwiftLookupTable::SingleEntry, 4>
 SwiftLookupTable::lookup(StringRef baseName,
                          EffectiveClangContext searchContext) {

lib/ClangImporter/SwiftLookupTable.h

@@ -351,6 +351,10 @@ public:
   SmallVector<SingleEntry, 4>
   lookupGlobalsAsMembers(EffectiveClangContext context);
 
+  /// Retrieve the set of global declarations that are going to be
+  /// imported as members.
+  SmallVector<SingleEntry, 4> allGlobalsAsMembers();
+
   /// Deserialize all entries.
   void deserializeAll();
 

lib/IDE/ModuleInterfacePrinting.cpp

@@ -201,6 +201,30 @@ swift::ide::collectModuleGroups(Module *M, std::vector<StringRef> &Scratch) {
   return llvm::makeArrayRef(Scratch);
 }
 
+/// Determine whether the given extension has a Clang node that
+/// created it (vs. being a Swift extension).
+static bool extensionHasClangNode(ExtensionDecl *ext) {
+  // If it has a Clang node (directly), 
+  if (ext->hasClangNode()) return true;
+
+  // If it has a global imported as a member.
+  auto members = ext->getMembers();
+  if (members.empty()) return false;
+  return members.front()->hasClangNode();
+}
+
+/// Retrieve the Clang node for the given extension, if it has one.
+/// created it (vs. being a Swift extension).
+static ClangNode extensionGetClangNode(ExtensionDecl *ext) {
+  // If it has a Clang node (directly), 
+  if (ext->hasClangNode()) return ext->getClangNode();
+
+  // If it has a global imported as a member.
+  auto members = ext->getMembers();
+  if (members.empty()) return ClangNode();
+  return members.front()->getClangNode();
+}
+
 void swift::ide::printSubmoduleInterface(
        Module *M,
        ArrayRef<StringRef> FullModuleName,
@@ -329,9 +353,8 @@ void swift::ide::printSubmoduleInterface(
       continue;
     }
 
-    auto addToClangDecls = [&](Decl *D) {
-      assert(D->hasClangNode());
-      auto CN = D->getClangNode();
+    auto addToClangDecls = [&](Decl *D, ClangNode CN) {
+      assert(CN && "No Clang node here");
       clang::SourceLocation Loc = CN.getLocation();
 
       auto *OwningModule = Importer.getClangOwningModule(CN);
@@ -342,9 +365,19 @@ void swift::ide::printSubmoduleInterface(
     };
 
     if (D->hasClangNode()) {
-      addToClangDecls(D);
+      addToClangDecls(D, D->getClangNode());
       continue;
     }
+
+    // If we have an extension containing globals imported as members,
+    // use the first member as the Clang node.
+    if (auto Ext = dyn_cast<ExtensionDecl>(D)) {
+      if (extensionHasClangNode(Ext)) {
+        addToClangDecls(Ext, extensionGetClangNode(Ext));
+        continue;
+      }
+    }
+
     if (FullModuleName.empty()) {
       // If group name is given and the decl does not belong to the group, skip it.
       if (!GroupNames.empty()){
@@ -428,7 +461,7 @@ void swift::ide::printSubmoduleInterface(
       // Clang extensions (categories) are always printed in source order.
       // Swift extensions are printed with their associated type unless it's
       // a cross-module extension.
-      if (!Ext->hasClangNode()) {
+      if (!extensionHasClangNode(Ext)) {
         auto ExtendedNominal = Ext->getExtendedType()->getAnyNominal();
         if (Ext->getModuleContext() == ExtendedNominal->getModuleContext())
           return false;
@@ -455,7 +488,7 @@ void swift::ide::printSubmoduleInterface(
               Printer.callAvoidPrintDeclPost(Ext);
               continue;
             }
-            if (Ext->hasClangNode())
+            if (extensionHasClangNode(Ext))
               continue; // will be printed in its source location, see above.
             Printer << "\n";
             Ext->print(Printer, AdjustedOptions);

test/IDE/Inputs/custom-modules/module.map

@@ -5,7 +5,10 @@ module OmitNeedlessWords {
 
 module ImportAsMember {
   export *
-  header "ImportAsMember.h"
+
+  module A {
+    header "ImportAsMember.h"
+  }
 
   module B {
     header "ImportAsMemberB.h"

test/IDE/import_as_member.swift

@@ -1,11 +1,9 @@
-// RUN: %target-swift-ide-test(mock-sdk: -sdk %S/../Inputs/clang-importer-sdk -I %t -I %S/Inputs/custom-modules) -print-module -source-filename %s -module-to-print=ImportAsMember -enable-omit-needless-words -always-argument-labels > %t.printed.txt
-// RUN: FileCheck %s -check-prefix=PRINT -strict-whitespace < %t.printed.txt
+// RUN: %target-swift-ide-test(mock-sdk: -sdk %S/../Inputs/clang-importer-sdk -I %t -I %S/Inputs/custom-modules) -print-module -source-filename %s -module-to-print=ImportAsMember.A -enable-omit-needless-words -always-argument-labels > %t.printed.A.txt
 
 // RUN: %target-swift-ide-test(mock-sdk: -sdk %S/../Inputs/clang-importer-sdk -I %t -I %S/Inputs/custom-modules) -print-module -source-filename %s -module-to-print=ImportAsMember.B -enable-omit-needless-words -always-argument-labels > %t.printed.B.txt
 
-// RUN: FileCheck %s -check-prefix=PRINTB -strict-whitespace < %t.printed.txt
-
-// FIXME: The extensions are getting printed in multiple submodules.
+// RUN: FileCheck %s -check-prefix=PRINT -strict-whitespace < %t.printed.A.txt
+// RUN: FileCheck %s -check-prefix=PRINTB -strict-whitespace < %t.printed.B.txt
 
 // PRINT: struct Struct1 {
 // PRINT-NEXT:   var x: Double
@@ -14,6 +12,10 @@
 // PRINT-NEXT:   init()
 // PRINT-NEXT:   init(x x: Double, y y: Double, z z: Double)
 // PRINT-NEXT: }
+
+// Make sure the other extension isn't here.
+// PRINT-NOT: static var static1: Double
+
 // PRINT:      extension Struct1 {
 // PRINT-NEXT:   static var globalVar: Double
 // PRINT-NEXT:   init(value value: Double)
@@ -22,15 +24,21 @@
 // PRINT-NEXT:   func selfComesLast(x x: Double)
 // PRINT-NEXT:   func selfComesThird(a a: Int32, b b: Float, x x: Double)
 // PRINT-NEXT: }
+// PRINT-NOT: static var static1: Double
 
 // RUN: %target-parse-verify-swift -I %S/Inputs/custom-modules
 
+// Make sure the other extension isn't here.
+// PRINTB-NOT: static var globalVar: Double
+
 // PRINTB:      extension Struct1 {
-// PRINTB:   static var static1: Double
+// PRINTB:        static var static1: Double
 // PRINTB-NEXT:   static var static2: Float
 // PRINTB-NEXT:   init(float value: Float)
 // PRINTB-NEXT: }
 
+// PRINTB-NOT: static var globalVar: Double
+
 // RUN: %target-swift-frontend %s -parse -I %S/Inputs/custom-modules -verify
 
 import ImportAsMember