[PrintAsObjC] Handle generic parameters in extensions. (#7818)

Most of the time, "generics" means "cannot be exposed to Objective-C"
and certainly "cannot be exposed in the generated header", but there
is one exception: imported Objective-C parameterized types, and their
extensions. We were previously dropping this on the floor and printing
`Foo</* BarType */>` in the generated header, which is nonsense.

https://bugs.swift.org/browse/SR-3480

Author: jrose-apple

lib/PrintAsObjC/PrintAsObjC.cpp

@@ -15,6 +15,7 @@
 #include "swift/AST/AST.h"
 #include "swift/AST/ASTVisitor.h"
 #include "swift/AST/ForeignErrorConvention.h"
+#include "swift/AST/GenericEnvironment.h"
 #include "swift/AST/NameLookup.h"
 #include "swift/AST/PrettyStackTrace.h"
 #include "swift/AST/ProtocolConformance.h"
@@ -172,6 +173,23 @@ class ObjCPrinter : private DeclVisitor<ObjCPrinter>,
     ASTVisitor::visit(const_cast<Decl *>(D));
   }
 
+  void maybePrintObjCGenericParameters(const ClassDecl *importedClass) {
+    auto *clangDecl = importedClass->getClangDecl();
+    auto *objcClass = dyn_cast_or_null<clang::ObjCInterfaceDecl>(clangDecl);
+    if (!objcClass)
+      return;
+    if (!objcClass->getTypeParamList())
+      return;
+    assert(objcClass->getTypeParamList()->size() != 0);
+    os << "<";
+    interleave(*objcClass->getTypeParamList(),
+               [this](const clang::ObjCTypeParamDecl *param) {
+                 os << param->getName();
+               },
+               [this] { os << ", "; });
+    os << ">";
+  }
+
   void printAdHocCategory(iterator_range<const ValueDecl * const *> members) {
     assert(members.begin() != members.end());
 
@@ -182,8 +200,10 @@ class ObjCPrinter : private DeclVisitor<ObjCPrinter>,
       baseClass = extendedTy->getClassOrBoundGenericClass();
     }
 
-    os << "@interface " << getNameForObjC(baseClass)
-    << " (SWIFT_EXTENSION(" << origDC->getParentModule()->getName() << "))\n";
+    os << "@interface " << getNameForObjC(baseClass);
+    maybePrintObjCGenericParameters(baseClass);
+    os << " (SWIFT_EXTENSION(" << origDC->getParentModule()->getName()
+       << "))\n";
     printMembers</*allowDelayed*/true>(members);
     os << "@end\n\n";
   }
@@ -312,8 +332,9 @@ class ObjCPrinter : private DeclVisitor<ObjCPrinter>,
   void visitExtensionDecl(ExtensionDecl *ED) {
     auto baseClass = ED->getExtendedType()->getClassOrBoundGenericClass();
 
-    os << "@interface " << getNameForObjC(baseClass)
-       << " (SWIFT_EXTENSION(" << ED->getModuleContext()->getName() << "))";
+    os << "@interface " << getNameForObjC(baseClass);
+    maybePrintObjCGenericParameters(baseClass);
+    os << " (SWIFT_EXTENSION(" << ED->getModuleContext()->getName() << "))";
     printProtocols(ED->getLocalProtocols(ConformanceLookupKind::OnlyExplicit));
     os << "\n";
     printMembers(ED->getMembers());
@@ -1574,6 +1595,30 @@ class ObjCPrinter : private DeclVisitor<ObjCPrinter>,
       visitType(MT, optionalKind);
     }
   }
+
+  void visitGenericTypeParamType(GenericTypeParamType *type,
+                                 Optional<OptionalTypeKind> optionalKind) {
+    const GenericTypeParamDecl *decl = type->getDecl();
+    assert(decl && "can't print canonicalized GenericTypeParamType");
+
+    if (auto *extension = dyn_cast<ExtensionDecl>(decl->getDeclContext())) {
+      const ClassDecl *extendedClass =
+          extension->getAsClassOrClassExtensionContext();
+      assert(extendedClass->isGeneric());
+      assert(extension->getGenericParams()->size() ==
+             extendedClass->getGenericParams()->size() &&
+             "extensions with custom generic parameters?");
+      assert(extension->getGenericSignature()->getCanonicalSignature() ==
+             extendedClass->getGenericSignature()->getCanonicalSignature() &&
+             "constrained extensions or custom generic parameters?");
+      type = extendedClass->getGenericEnvironment()->getSugaredType(type);
+      decl = type->getDecl();
+    }
+
+    assert(decl->getClangDecl() && "can only handle imported ObjC generics");
+    os << cast<clang::ObjCTypeParamDecl>(decl->getClangDecl())->getName();
+    printNullability(optionalKind);
+  }
 
   void printFunctionType(FunctionType *FT, char pointerSigil,
                          Optional<OptionalTypeKind> optionalKind) {

test/Inputs/clang-importer-sdk/usr/include/objc_generics.h

@@ -67,6 +67,9 @@ void takeGenericClass(__nullable GenericClass<NSString *> *thing);
 
 @end
 
+@interface PettableOverextendedMetaphor: NSObject <Pettable>
+@end
+
 @protocol Fungible
 @end
 

test/PrintAsObjC/extensions.swift

@@ -84,7 +84,7 @@ extension CGColor {
   func anyOldMethod() {}
 }
 
-// CHECK-LABEL: @interface GenericClass (SWIFT_EXTENSION(extensions))
+// CHECK-LABEL: @interface GenericClass<T> (SWIFT_EXTENSION(extensions))
 // CHECK-NEXT: - (void)bar;
 // CHECK-NEXT: @end
 extension GenericClass {
@@ -115,3 +115,17 @@ extension NSString {
   class func fromColor(_ color: NSColor) -> NSString? { return nil; }
 }
 
+// CHECK-LABEL: @interface PettableContainer<T> (SWIFT_EXTENSION(extensions))
+// CHECK-NEXT: - (PettableContainer<T> * _Nonnull)duplicate SWIFT_WARN_UNUSED_RESULT;
+// CHECK-NEXT: - (PettableContainer<T> * _Nonnull)duplicate2 SWIFT_WARN_UNUSED_RESULT;
+// CHECK-NEXT: - (PettableContainer<PettableOverextendedMetaphor *> * _Nonnull)duplicate3 SWIFT_WARN_UNUSED_RESULT;
+// CHECK-NEXT: - (T _Nonnull)extract SWIFT_WARN_UNUSED_RESULT;
+// CHECK-NEXT: - (T _Nullable)extract2 SWIFT_WARN_UNUSED_RESULT;
+// CHECK-NEXT: @end
+extension PettableContainer {
+  func duplicate() -> PettableContainer { fatalError() }
+  func duplicate2() -> PettableContainer<T> { fatalError() }
+  func duplicate3() -> PettableContainer<PettableOverextendedMetaphor> { fatalError() }
+  func extract() -> T { fatalError() }
+  func extract2() -> T? { fatalError() }
+}