Support @objc(CustomName) on extensions

This now specifies a category name that’s used in TBDGen, IRGen, and PrintAsClang. There are also now category name conflict diagnostics; these subsume some @implementation diagnostics.

(It turns out there was already a check for @objc(CustomName) to make sure it wasn’t a selector!)

Author: beccadax

include/swift/AST/Decl.h

@@ -5064,6 +5064,10 @@ class ClassDecl final : public NominalTypeDecl {
   llvm::TinyPtrVector<Decl *>
   getImportedObjCCategory(Identifier name) const;
 
+  /// Record the presence of an extension that may become an Objective-C
+  /// category.
+  void recordObjCCategory(ExtensionDecl *ext);
+
   // Implement isa/cast/dyncast/etc.
   static bool classof(const Decl *D) {
     return D->getKind() == DeclKind::Class;

include/swift/AST/DiagnosticsClangImporter.def

@@ -136,12 +136,10 @@ WARNING(api_pattern_attr_ignored, none,
         (StringRef, StringRef))
 
 ERROR(objc_implementation_two_impls, none,
-      "duplicate implementation of Objective-C %select{|category %0 on }0"
-      "%kind1",
-      (Identifier, Decl *))
-
+      "duplicate implementation of imported %kind0",
+      (Decl *))
 NOTE(previous_objc_implementation, none,
-     "previously implemented by extension here", ())
+     "previously implemented here", ())
 
 NOTE(macro_not_imported_unsupported_operator, none, "operator not supported in macro arithmetic", ())
 NOTE(macro_not_imported_unsupported_named_operator, none, "operator '%0' not supported in macro arithmetic", (StringRef))

include/swift/AST/DiagnosticsSema.def

@@ -6450,6 +6450,12 @@ ERROR(objc_redecl_same,none,
       "previous declaration with the same Objective-C selector",
       (unsigned, DeclName, unsigned, DeclName, ObjCSelector))
 
+ERROR(objc_redecl_category_name,none,
+      "%select{|imported }0extension with Objective-C category name %2 "
+      "conflicts with previous %select{|imported }1extension with the same "
+      "category name",
+      (bool, bool, Identifier))
+
 ERROR(objc_override_other,none,
       OBJC_DIAG_SELECT " with Objective-C selector %4 conflicts with "
       OBJC_DIAG_SELECT_2 " from superclass %5 with the same Objective-C "

include/swift/AST/SourceFile.h

@@ -359,6 +359,16 @@ class SourceFile final : public FileUnit {
   /// List of Objective-C member conflicts we have found during type checking.
   llvm::SetVector<ObjCMethodConflict> ObjCMethodConflicts;
 
+  struct ObjCCategoryConflict {
+    NominalTypeDecl *typeDecl;
+    Identifier name;
+
+    ObjCCategoryConflict(NominalTypeDecl *typeDecl, Identifier name)
+        : typeDecl(typeDecl), name(name)
+    {}
+  };
+  llvm::SetVector<ObjCCategoryConflict> ObjCCategoryConflicts;
+
   /// List of attributes added by access notes, used to emit remarks for valid
   /// ones.
   llvm::DenseMap<ValueDecl *, std::vector<DeclAttribute *>>
@@ -876,6 +886,25 @@ struct DenseMapInfo<swift::SourceFile::ObjCMethodConflict> {
   }
 };
 
+template<>
+struct DenseMapInfo<swift::SourceFile::ObjCCategoryConflict> {
+  using ObjCCategoryConflict = swift::SourceFile::ObjCCategoryConflict;
+
+  static inline ObjCCategoryConflict getEmptyKey() {
+    return { nullptr, DenseMapInfo<swift::Identifier>::getEmptyKey() };
+  }
+  static inline ObjCCategoryConflict getTombstoneKey() {
+    return { nullptr, DenseMapInfo<swift::Identifier>::getTombstoneKey() };
+  }
+  static inline unsigned getHashValue(ObjCCategoryConflict a) {
+    return hash_combine(hash_value(a.typeDecl),
+                  DenseMapInfo<swift::Identifier>::getHashValue(a.name));
+  }
+  static bool isEqual(ObjCCategoryConflict a, ObjCCategoryConflict b) {
+    return a.typeDecl == b.typeDecl && a.name == b.name;
+  }
+};
+
 }
 
 

lib/AST/Decl.cpp

@@ -3718,12 +3718,9 @@ Identifier ExtensionDecl::getObjCCategoryName() const {
   }
 
   // Fall back to @_objcImplementation attribute.
-  if (auto attr =
+  if (auto implAttr =
        getAttrs().getAttribute<ObjCImplementationAttr>(/*AllowInvalid=*/true)) {
-    if (!attr->isCategoryNameInvalid())
-      return attr->CategoryName;
-
-    return Identifier();
+    return implAttr->CategoryName;
   }
 
   // Not a category, evidently.

lib/AST/NameLookup.cpp

@@ -1524,6 +1524,8 @@ void NominalTypeDecl::addedExtension(ExtensionDecl *ext) {
   assert(table);
 
   table->addExtension(ext);
+  if (auto CD = dyn_cast<ClassDecl>(this))
+    CD->recordObjCCategory(ext);
 }
 
 void NominalTypeDecl::addedMember(Decl *member) {
@@ -2012,8 +2014,11 @@ void NominalTypeDecl::prepareLookupTable() {
   }
 
   // Note: this calls prepareExtensions()
+  auto CD = dyn_cast<ClassDecl>(this);
   for (auto e : getExtensions()) {
     table->addExtension(e);
+    if (CD)
+      CD->recordObjCCategory(e);
   }
 
   // Any extensions added after this point will add their members to the
@@ -2259,6 +2264,27 @@ void NominalTypeDecl::recordObjCMethod(AbstractFunctionDecl *method,
   vec.push_back(method);
 }
 
+void ClassDecl::recordObjCCategory(ExtensionDecl *ext) {
+  Identifier name = ext->getObjCCategoryName();
+  if (name.empty())
+    return;
+
+  // Use a linear scan on the assumption that there aren't very many extensions.
+  for (auto *other : getExtensions()) {
+    if (other != ext && other->getObjCCategoryName() == name
+        // @implementations don't conflict with their interfaces
+        && !((ext->hasClangNode() && other->isObjCImplementation())
+             || (ext->isObjCImplementation() && other->hasClangNode()))) {
+
+      // Collision.
+      if (auto *sf = ext->getParentSourceFile()) {
+        sf->ObjCCategoryConflicts.insert({ this, name });
+      }
+      break;
+    }
+  }
+}
+
 static bool missingExplicitImportForMemberDecl(const DeclContext *dc,
                                                ValueDecl *decl) {
   // Only require explicit imports for members when MemberImportVisibility is

lib/ClangImporter/ClangImporter.cpp

@@ -5984,10 +5984,15 @@ constructResult(const llvm::TinyPtrVector<Decl *> &interfaces,
       auto attr = extraImpl->getAttrs().getAttribute<ObjCImplementationAttr>();
       attr->setCategoryNameInvalid();
 
-      diags.diagnose(attr->getLocation(), diag::objc_implementation_two_impls,
-                     categoryName, diagnoseOn)
-        .fixItRemove(attr->getRangeWithAt());
-      diags.diagnose(impls.front(), diag::previous_objc_implementation);
+      // @objc @implementations for categories are diagnosed as category
+      // conflicts, so we're only concerned with main class bodies and
+      // non-category implementations here.
+      if (categoryName.empty() || !isa<ExtensionDecl>(impls.front())) {
+        diags.diagnose(attr->getLocation(), diag::objc_implementation_two_impls,
+                       diagnoseOn)
+          .fixItRemove(attr->getRangeWithAt());
+        diags.diagnose(impls.front(), diag::previous_objc_implementation);
+      }
     }
   }
 

lib/IRGen/GenClass.cpp

@@ -1138,8 +1138,8 @@ namespace {
       return pair.second;
     }
 
-    std::optional<StringRef> getObjCImplCategoryName() const {
-      if (!TheExtension || !TheExtension->isObjCImplementation())
+    std::optional<StringRef> getCustomCategoryName() const {
+      if (!TheExtension)
         return std::nullopt;
       assert(!TheExtension->hasClangNode());
       auto ident = TheExtension->getObjCCategoryName();
@@ -1426,8 +1426,8 @@ namespace {
     void buildCategoryName(SmallVectorImpl<char> &s) {
       llvm::raw_svector_ostream os(s);
 
-      if (auto implementationCategoryName = getObjCImplCategoryName()) {
-        os << *implementationCategoryName;
+      if (auto customCategoryName = getCustomCategoryName()) {
+        os << *customCategoryName;
         return;
       }
 
@@ -2449,8 +2449,8 @@ namespace {
         os << "@";
         TheExtension->getLoc().print(os, IGM.Context.SourceMgr);
       }
-      if (auto name = getObjCImplCategoryName()) {
-        os << " objc_impl_category_name=" << *name;
+      if (auto name = getCustomCategoryName()) {
+        os << " custom_category_name=" << *name;
       }
 
       if (HasNonTrivialConstructor)

lib/IRGen/TBDGen.cpp

@@ -824,6 +824,10 @@ class APIGenRecorder final : public APIRecorder {
   void buildCategoryName(const ExtensionDecl *ext, const ClassDecl *cls,
                          SmallVectorImpl<char> &s) {
     llvm::raw_svector_ostream os(s);
+    if (!ext->getObjCCategoryName().empty()) {
+      os << ext->getObjCCategoryName();
+      return;
+    }
     ModuleDecl *module = ext->getParentModule();
     os << module->getName();
     unsigned categoryCount = CategoryCounts[{cls, module}]++;

lib/PrintAsClang/DeclAndTypePrinter.cpp

@@ -403,7 +403,10 @@ class DeclAndTypePrinter::Implementation
       os << "\n";
     os << "@interface " << getNameForObjC(baseClass);
     maybePrintObjCGenericParameters(baseClass);
-    os << " (SWIFT_EXTENSION(" << ED->getModuleContext()->getName() << "))";
+    if (ED->getObjCCategoryName().empty())
+      os << " (SWIFT_EXTENSION(" << ED->getModuleContext()->getName() << "))";
+    else
+      os << " (" << ED->getObjCCategoryName() << ")";
     printProtocols(ED->getLocalProtocols(ConformanceLookupKind::OnlyExplicit));
     os << "\n";
     printMembers(ED->getMembers());