[5.1] [IRGen] Mangle Swift @objc(renamed) protocols as Objective-C in metadata

include/swift/AST/ASTMangler.h

@@ -38,6 +38,10 @@ class ASTMangler : public Mangler {
   /// Optimize out protocol names if a type only conforms to one protocol.
   bool OptimizeProtocolNames = true;
 
+  /// If enabled, use Objective-C runtime names when mangling @objc Swift
+  /// protocols.
+  bool UseObjCProtocolNames = false;
+
   /// If enabled, non-canonical types are allowed and type alias types get a
   /// special mangling.
   bool DWARFMangling;
@@ -187,7 +191,7 @@ class ASTMangler : public Mangler {
   };
 
   static Optional<SpecialContext>
-  getSpecialManglingContext(const ValueDecl *decl);
+  getSpecialManglingContext(const ValueDecl *decl, bool useObjCProtocolNames);
 
   static const clang::NamedDecl *
   getClangDeclForMangling(const ValueDecl *decl);

lib/AST/ASTMangler.cpp

@@ -367,7 +367,7 @@ std::string ASTMangler::mangleTypeForDebugger(Type Ty, const DeclContext *DC) {
   DWARFMangling = true;
   OptimizeProtocolNames = false;
   beginMangling();
-  
+
   if (DC)
     bindGenericParameters(DC);
 
@@ -601,6 +601,33 @@ static StringRef getPrivateDiscriminatorIfNecessary(const ValueDecl *decl) {
   return discriminator.str();
 }
 
+/// If the declaration is an @objc protocol defined in Swift and the
+/// Objective-C name has been overrridden from the default, return the
+/// specified name.
+///
+/// \param useObjCProtocolNames When false, always returns \c None.
+static Optional<std::string> getOverriddenSwiftProtocolObjCName(
+                                                  const ValueDecl *decl,
+                                                  bool useObjCProtocolNames) {
+  if (!useObjCProtocolNames)
+    return None;
+
+  auto proto = dyn_cast<ProtocolDecl>(decl);
+  if (!proto) return None;
+
+  if (!proto->isObjC()) return None;
+
+  // If there is an 'objc' attribute with a name, use that name.
+  if (auto objc = proto->getAttrs().getAttribute<ObjCAttr>()) {
+    if (auto name = objc->getName()) {
+      llvm::SmallString<4> buffer;
+      return std::string(name->getString(buffer));
+    }
+  }
+
+  return None;
+}
+
 void ASTMangler::appendDeclName(const ValueDecl *decl) {
   DeclBaseName name = decl->getBaseName();
   assert(!name.isSpecial() && "Cannot print special names");
@@ -625,6 +652,11 @@ void ASTMangler::appendDeclName(const ValueDecl *decl) {
         appendOperator("oi");
         break;
     }
+  } else if (auto objCName =
+               getOverriddenSwiftProtocolObjCName(decl, UseObjCProtocolNames)) {
+    // @objc Swift protocols should be mangled as Objective-C protocols,
+    // so append the Objective-C runtime name.
+    appendIdentifier(*objCName);
   } else if (!name.empty()) {
     appendIdentifier(name.getIdentifier().str());
   } else {
@@ -1335,7 +1367,8 @@ void ASTMangler::appendImplFunctionType(SILFunctionType *fn) {
 }
 
 Optional<ASTMangler::SpecialContext>
-ASTMangler::getSpecialManglingContext(const ValueDecl *decl) {
+ASTMangler::getSpecialManglingContext(const ValueDecl *decl,
+                                      bool useObjCProtocolNames) {
   // Declarations provided by a C module have a special context mangling.
   //   known-context ::= 'So'
   //
@@ -1351,6 +1384,11 @@ ASTMangler::getSpecialManglingContext(const ValueDecl *decl) {
     }
   }
 
+  // If @objc Swift protocols should be mangled as Objective-C protocols,
+  // they are defined in the Objective-C context.
+  if (getOverriddenSwiftProtocolObjCName(decl, useObjCProtocolNames))
+    return ASTMangler::ObjCContext;
+
   // Nested types imported from C should also get use the special "So" context.
   if (isa<TypeDecl>(decl)) {
     if (auto *clangDecl = cast_or_null<clang::NamedDecl>(decl->getClangDecl())){
@@ -1376,7 +1414,7 @@ ASTMangler::getSpecialManglingContext(const ValueDecl *decl) {
 /// This is the top-level entrypoint for mangling <context>.
 void ASTMangler::appendContextOf(const ValueDecl *decl) {
   // Check for a special mangling context.
-  if (auto context = getSpecialManglingContext(decl)) {
+  if (auto context = getSpecialManglingContext(decl, UseObjCProtocolNames)) {
     switch (*context) {
     case ClangImporterContext:
       return appendOperator("SC");

lib/IRGen/GenDecl.cpp

@@ -747,7 +747,9 @@ IRGenModule::getAddrOfParentContextDescriptor(DeclContext *from,
   // Some types get special treatment.
   if (auto Type = dyn_cast<NominalTypeDecl>(from)) {
     // Use a special module context if we have one.
-    if (auto context = Mangle::ASTMangler::getSpecialManglingContext(Type)) {
+    if (auto context =
+            Mangle::ASTMangler::getSpecialManglingContext(
+              Type, /*UseObjCProtocolNames=*/false)) {
       switch (*context) {
       case Mangle::ASTMangler::ObjCContext:
         return {getAddrOfObjCModuleContextDescriptor(),

lib/IRGen/IRGenMangler.cpp

@@ -83,6 +83,7 @@ IRGenMangler::withSymbolicReferences(IRGenModule &IGM,
                                   llvm::function_ref<void ()> body) {
   Mod = IGM.getSwiftModule();
   OptimizeProtocolNames = false;
+  UseObjCProtocolNames = true;
 
   llvm::SaveAndRestore<bool>
     AllowSymbolicReferencesLocally(AllowSymbolicReferences);

test/Runtime/associated_type_demangle_objc.swift

@@ -0,0 +1,40 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift %s -module-name main -o %t/a.out
+// RUN: %target-codesign %t/a.out
+// RUN: %target-run %t/a.out
+// REQUIRES: executable_test
+// REQUIRES: objc_interop
+
+import Foundation
+import StdlibUnittest
+
+@objc(MyObjCProtocol) public protocol MyProtocol { }
+
+public protocol OtherProtocol {
+  associatedtype AssocType
+}
+
+public struct MyThing: OtherProtocol {
+  public typealias AssocType = MyProtocol
+}
+
+func getAssocType<T: OtherProtocol>(_: T.Type) -> Any.Type {
+  return T.AssocType.self
+}
+let RenamedObjCDemangleTests = TestSuite("RenamedObjCDemangle")
+
+RenamedObjCDemangleTests.test("@objc protocols") {
+  expectEqual(getAssocType(MyThing.self), MyProtocol.self)
+}
+
+@objc(MyObjCClass) class MyClass: NSObject { }
+
+struct MyClassWrapper: OtherProtocol {
+  typealias AssocType = MyClass
+}
+
+RenamedObjCDemangleTests.test("@objc classes") {
+  expectEqual(getAssocType(MyClassWrapper.self), MyClass.self)
+}
+
+runAllTests()

test/Runtime/demangleToMetadataObjC.swift

@@ -11,14 +11,24 @@ import Dispatch
 let DemangleToMetadataTests = TestSuite("DemangleToMetadataObjC")
 
 @objc class C : NSObject { }
+@objc(CRenamedInObjC) class CRenamed : NSObject { }
 @objc enum E: Int { case a }
 @objc protocol P1 { }
 protocol P2 { }
 @objc protocol P3: P1 { }
 @objc protocol mainP4 { }
 
+@objc(P5RenamedInObjC) protocol P5 { }
+
 DemangleToMetadataTests.test("@objc classes") {
   expectEqual(type(of: C()), _typeByName("4main1CC")!)
+
+  // @objc class that's been renamed, which can be found by its Objective-C
+  // name...
+  expectEqual(type(of: CRenamed()), _typeByName("So14CRenamedInObjCC")!)
+
+  // ... but not by it's Swift name.
+  expectNil(_typeByName("4main8CRenamed"))
 }
 
 DemangleToMetadataTests.test("@objc enums") {
@@ -40,9 +50,16 @@ DemangleToMetadataTests.test("Objective-C classes") {
 }
 
 func f1_composition_NSCoding(_: NSCoding) { }
+func f1_composition_P5(_: P5) { }
 
 DemangleToMetadataTests.test("Objective-C protocols") {
   expectEqual(type(of: f1_composition_NSCoding), _typeByName("yySo8NSCoding_pc")!)
+
+  // @objc Swift protocols can be found by their Objective-C names...
+  expectEqual(type(of: f1_composition_P5), _typeByName("yySo15P5RenamedInObjC_pc")!)
+
+  // ... but not their Swift names.
+  expectNil(_typeByName("yy4main2P5_pc"))
 }
 
 DemangleToMetadataTests.test("Classes that don't exist") {