[interop][SwiftToCxx] emit members in extensions for same-type

Author: hyp

lib/PrintAsClang/DeclAndTypePrinter.cpp

@@ -348,8 +348,17 @@ class DeclAndTypePrinter::Implementation
     // FIXME: Print struct's availability.
     ClangValueTypePrinter printer(os, owningPrinter.prologueOS,
                                   owningPrinter.interopContext);
-    printer.printValueTypeDecl(
-        SD, /*bodyPrinter=*/[&]() { printMembers(SD->getMembers()); });
+    printer.printValueTypeDecl(SD, /*bodyPrinter=*/[&]() {
+      printMembers(SD->getMembers());
+      for (const auto *ed :
+           owningPrinter.interopContext.getExtensionsForNominalType(SD)) {
+        auto sign = ed->getGenericSignature();
+        // FIXME: support requirements.
+        if (!sign.getRequirements().empty())
+          continue;
+        printMembers(ed->getMembers());
+      }
+    });
   }
 
   void visitExtensionDecl(ExtensionDecl *ED) {
@@ -881,7 +890,11 @@ class DeclAndTypePrinter::Implementation
       if (isClassMethod)
         return;
       assert(!AFD->isStatic());
-      auto *typeDeclContext = cast<NominalTypeDecl>(AFD->getParent());
+      auto *typeDeclContext = dyn_cast<NominalTypeDecl>(AFD->getParent());
+      if (!typeDeclContext) {
+        typeDeclContext =
+            cast<ExtensionDecl>(AFD->getParent())->getExtendedNominal();
+      }
 
       std::string cFuncDecl;
       llvm::raw_string_ostream cFuncPrologueOS(cFuncDecl);

lib/PrintAsClang/ModuleContentsWriter.cpp

@@ -452,8 +452,13 @@ class ModuleWriter {
   bool writeStruct(const StructDecl *SD) {
     if (addImport(SD))
       return true;
-    if (outputLangMode == OutputLanguageMode::Cxx)
+    if (outputLangMode == OutputLanguageMode::Cxx) {
       (void)forwardDeclareMemberTypes(SD->getMembers(), SD);
+      for (const auto *ed :
+           printer.getInteropContext().getExtensionsForNominalType(SD)) {
+        (void)forwardDeclareMemberTypes(ed->getMembers(), SD);
+      }
+    }
     printer.print(SD);
     return true;
   }
@@ -555,6 +560,8 @@ class ModuleWriter {
         return !printer.shouldInclude(VD);
 
       if (auto ED = dyn_cast<ExtensionDecl>(D)) {
+        if (outputLangMode == OutputLanguageMode::Cxx)
+          return false;
         auto baseClass = ED->getSelfClassDecl();
         return !baseClass || !printer.shouldInclude(baseClass) ||
                baseClass->isForeign();
@@ -580,6 +587,8 @@ class ModuleWriter {
 
         if (auto ED = dyn_cast<ExtensionDecl>(D)) {
           auto baseClass = ED->getSelfClassDecl();
+          if (!baseClass)
+              return ED->getExtendedNominal()->getName().str();
           return baseClass->getName().str();
         }
         llvm_unreachable("unknown top-level ObjC decl");
@@ -632,6 +641,16 @@ class ModuleWriter {
     assert(declsToWrite.empty());
     declsToWrite.assign(decls.begin(), decls.end());
 
+    if (outputLangMode == OutputLanguageMode::Cxx) {
+      for (const Decl *D : declsToWrite) {
+        if (auto *ED = dyn_cast<ExtensionDecl>(D)) {
+          const auto *type = ED->getExtendedNominal();
+          if (isa<StructDecl>(type))
+            printer.getInteropContext().recordExtensions(type, ED);
+        }
+      }
+    }
+
     while (!declsToWrite.empty()) {
       const Decl *D = declsToWrite.back();
       bool success = true;

lib/PrintAsClang/SwiftToClangInteropContext.cpp

@@ -40,3 +40,19 @@ void SwiftToClangInteropContext::recordEmittedClangTypeDecl(
   assert(typeDecl->hasClangNode());
   referencedClangTypeDecls.insert(typeDecl);
 }
+
+void SwiftToClangInteropContext::recordExtensions(
+    const NominalTypeDecl *typeDecl, const ExtensionDecl *ext) {
+  auto it = extensions.insert(
+      std::make_pair(typeDecl, std::vector<const ExtensionDecl *>()));
+  it.first->second.push_back(ext);
+}
+
+llvm::ArrayRef<const ExtensionDecl *>
+SwiftToClangInteropContext::getExtensionsForNominalType(
+    const NominalTypeDecl *typeDecl) const {
+  auto exts = extensions.find(typeDecl);
+  if (exts != extensions.end())
+    return exts->getSecond();
+  return {};
+}

lib/PrintAsClang/SwiftToClangInteropContext.h

@@ -14,6 +14,8 @@
 #define SWIFT_PRINTASCLANG_SWIFTTOCLANGINTEROPCONTEXT_H
 
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/Optional.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/StringSet.h"
@@ -25,6 +27,7 @@ class Decl;
 class IRABIDetailsProvider;
 class IRGenOptions;
 class ModuleDecl;
+class ExtensionDecl;
 class NominalTypeDecl;
 
 /// The \c SwiftToClangInteropContext class is responsible for providing
@@ -52,12 +55,20 @@ class SwiftToClangInteropContext {
     return referencedClangTypeDecls;
   }
 
+  void recordExtensions(const NominalTypeDecl *typeDecl,
+                        const ExtensionDecl *ext);
+
+  llvm::ArrayRef<const ExtensionDecl *>
+  getExtensionsForNominalType(const NominalTypeDecl *typeDecl) const;
+
 private:
   ModuleDecl &mod;
   const IRGenOptions &irGenOpts;
   std::unique_ptr<IRABIDetailsProvider> irABIDetails;
   llvm::StringSet<> emittedStubs;
   llvm::SetVector<const NominalTypeDecl *> referencedClangTypeDecls;
+  llvm::DenseMap<const NominalTypeDecl *, std::vector<const ExtensionDecl *>>
+      extensions;
 };
 
 } // end namespace swift

test/Interop/SwiftToCxx/extension/struct-extension-in-cxx.swift

@@ -0,0 +1,26 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend %s -typecheck -module-name Structs -clang-header-expose-public-decls -emit-clang-header-path %t/structs.h
+// RUN: %FileCheck %s < %t/structs.h
+
+// RUN: %check-interop-cxx-header-in-clang(%t/structs.h -Wno-unused-private-field -Wno-unused-function)
+
+
+public struct TypeAfterArray {
+    var x: Int16
+}
+
+public struct Array {
+    public var x: Int
+}
+
+extension Array {
+    public var val: Structs.TypeAfterArray {
+        return TypeAfterArray(x: 42)
+    }
+}
+
+// CHECK class TypeAfterArray;
+// CHECK: class Array final {
+// CHECK:        swift::Int getX() const;
+// CHECK-NEXT:   inline void setX(swift::Int value);
+// CHECK-NEXT:   TypeAfterArray getVal() const;