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

@@ -4199,6 +4199,12 @@ class NominalTypeDecl : public GenericTypeDecl, public IterableDeclContext {
   /// Retrieve the set of extensions of this type.
   ExtensionRange getExtensions();
 
+  /// Retrieve the extension most recently added to this type. Helpful to
+  /// determine if an extension has been added.
+  ExtensionDecl *getLastExtension() const {
+    return LastExtension;
+  }
+
   /// Special-behaviour flags passed to lookupDirect()
   enum class LookupDirectFlags {
     /// Whether to include @_implements members.
@@ -5064,6 +5070,11 @@ class ClassDecl final : public NominalTypeDecl {
   llvm::TinyPtrVector<Decl *>
   getImportedObjCCategory(Identifier name) const;
 
+  /// Return a map of category names to extensions with that category name,
+  /// whether imported or otherwise. 
+  llvm::DenseMap<Identifier, llvm::TinyPtrVector<ExtensionDecl *>>
+  getObjCCategoryNameMap();
+
   // 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/NameLookupRequests.h

@@ -970,6 +970,39 @@ class LookupIntrinsicRequest
   bool isCached() const { return true; }
 };
 
+using ObjCCategoryNameMap =
+  llvm::DenseMap<Identifier, llvm::TinyPtrVector<ExtensionDecl *>>;
+
+/// Generate a map of all known extensions of the given class that have an
+/// explicit category name. This request does not force clang categories that
+/// haven't been imported already, but it will generate a new map if new
+/// categories have been imported since the cached value was generated.
+///
+/// \seeAlso ClassDecl::getObjCCategoryNameMap()
+class ObjCCategoryNameMapRequest
+    : public SimpleRequest<ObjCCategoryNameMapRequest,
+                           ObjCCategoryNameMap(ClassDecl *, ExtensionDecl *),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+  // Convenience to automatically extract `lastExtension`.
+  ObjCCategoryNameMapRequest(ClassDecl *classDecl)
+    : ObjCCategoryNameMapRequest(classDecl, classDecl->getLastExtension())
+  {}
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  ObjCCategoryNameMap evaluate(Evaluator &evaluator,
+                               ClassDecl *classDecl,
+                               ExtensionDecl *lastExtension) const;
+
+public:
+  bool isCached() const { return true; }
+};
+
 #define SWIFT_TYPEID_ZONE NameLookup
 #define SWIFT_TYPEID_HEADER "swift/AST/NameLookupTypeIDZone.def"
 #include "swift/Basic/DefineTypeIDZone.h"

include/swift/AST/NameLookupTypeIDZone.def

@@ -116,3 +116,5 @@ SWIFT_REQUEST(NameLookup, ImplementsAttrProtocolRequest,
               ProtocolDecl *(const ImplementsAttr *, DeclContext *), Cached, NoLocationInfo)
 SWIFT_REQUEST(NameLookup, LookupIntrinsicRequest,
               FuncDecl *(ModuleDecl *, Identifier), Cached, NoLocationInfo)
+SWIFT_REQUEST(NameLookup, ObjCCategoryNameMapRequest,
+              ObjCCategoryNameMap(ClassDecl *, ExtensionDecl *), Cached, NoLocationInfo)

include/swift/AST/SourceFile.h

@@ -359,6 +359,10 @@ class SourceFile final : public FileUnit {
   /// List of Objective-C member conflicts we have found during type checking.
   llvm::SetVector<ObjCMethodConflict> ObjCMethodConflicts;
 
+  /// Categories (extensions with explicit @objc names) declared in this
+  /// source file. They need to be checked for conflicts after type checking.
+  llvm::TinyPtrVector<ExtensionDecl *> ObjCCategories;
+
   /// List of attributes added by access notes, used to emit remarks for valid
   /// ones.
   llvm::DenseMap<ValueDecl *, std::vector<DeclAttribute *>>

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

@@ -2262,6 +2262,35 @@ void NominalTypeDecl::recordObjCMethod(AbstractFunctionDecl *method,
   vec.push_back(method);
 }
 
+ObjCCategoryNameMap
+ObjCCategoryNameMapRequest::evaluate(Evaluator &evaluator,
+                                     ClassDecl *classDecl,
+                                     ExtensionDecl *lastExtension) const {
+  // The purpose of the `lastExtension` parameter is to bake something into the
+  // request that will change when another extension is added. This ensures that
+  // if more extensions are added later, we don't return an outdated version of
+  // the extension map by accident.
+
+  ObjCCategoryNameMap results;
+
+  for (auto ext : classDecl->getExtensions()) {
+    auto catName = ext->getObjCCategoryName();
+    if (!catName.empty())
+      results[catName].push_back(ext);
+
+    if (ext == lastExtension)
+      break;
+  }
+
+  return results;
+}
+
+ObjCCategoryNameMap ClassDecl::getObjCCategoryNameMap() {
+  return evaluateOrDefault(getASTContext().evaluator,
+                           ObjCCategoryNameMapRequest{this},
+                           ObjCCategoryNameMap());
+}
+
 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)