[NFC] Make objcImpl request one-to-many

An @_objcImpl extension with no category name *should* implement not only the class’s main @interface, but also any class extension @interfaces. Start making this true by making ObjCInterfaceAndImplementationRequest return all of these decls as the interfaces for such an implementation.

This commit doesn’t actually change Sema or IRGen to process the extra interfaces, so it’s NFC.

Author: beccadax

include/swift/AST/Decl.h

@@ -1167,10 +1167,23 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl> {
   }
 
   /// If this is the Swift implementation of a declaration imported from ObjC,
-  /// returns the imported declaration. Otherwise return \c nullptr.
+  /// returns the imported declaration. (If there are several, only the main
+  /// class body will be returned.) Otherwise return \c nullptr.
   ///
   /// \seeAlso ExtensionDecl::isObjCInterface()
-  Decl *getImplementedObjCDecl() const;
+  Decl *getImplementedObjCDecl() const {
+    auto impls = getAllImplementedObjCDecls();
+    if (impls.empty())
+      return nullptr;
+    return impls.front();
+  }
+
+  /// If this is the Swift implementation of a declaration imported from ObjC,
+  /// returns the imported declarations. (There may be several for a main class
+  /// body; if so, the first will be the class itself.) Otherwise return an empty list.
+  ///
+  /// \seeAlso ExtensionDecl::isObjCInterface()
+  llvm::TinyPtrVector<Decl *> getAllImplementedObjCDecls() const;
 
   /// If this is the ObjC interface of a declaration implemented in Swift,
   /// returns the implementating declaration. Otherwise return \c nullptr.
@@ -1835,8 +1848,8 @@ class ExtensionDecl final : public GenericContext, public Decl,
   bool isEquivalentToExtendedContext() const;
 
   /// Returns the name of the category specified by the \c \@_objcImplementation
-  /// attribute, or \c None if the name is invalid. Do not call unless
-  /// \c isObjCImplementation() returns \c true.
+  /// attribute, or \c None if the name is invalid or
+  /// \c isObjCImplementation() is false.
   llvm::Optional<Identifier> getCategoryNameForObjCImplementation() const;
 
   /// If this extension represents an imported Objective-C category, returns the
@@ -4920,11 +4933,14 @@ class ClassDecl final : public NominalTypeDecl {
   /// the Objective-C runtime.
   StringRef getObjCRuntimeName(llvm::SmallVectorImpl<char> &buffer) const;
 
-  /// Return the imported declaration for the category with the given name; this
-  /// will always be an Objective-C-backed \c ExtensionDecl or, if \p name is
-  /// empty, \c ClassDecl. Returns \c nullptr if the class was not imported from
-  /// Objective-C or does not have an imported category by that name.
-  IterableDeclContext *getImportedObjCCategory(Identifier name) const;
+  /// Return the imported declaration(s) for the category with the given name; this
+  /// will either be a single imported \c ExtensionDecl, an imported
+  /// \c ClassDecl followed by zero or more imported \c ExtensionDecl s (if
+  /// \p name is empty; the extensions are for any class extensions), or empty
+  /// if the class was not imported from Objective-C or does not have a
+  /// category by that name.
+  llvm::TinyPtrVector<Decl *>
+  getImportedObjCCategory(Identifier name) const;
 
   // Implement isa/cast/dyncast/etc.
   static bool classof(const Decl *D) {

include/swift/ClangImporter/ClangImporterRequests.h

@@ -207,17 +207,21 @@ void simple_display(llvm::raw_ostream &out,
                     const ClangCategoryLookupDescriptor &desc);
 SourceLoc extractNearestSourceLoc(const ClangCategoryLookupDescriptor &desc);
 
-/// Given a Swift class, find the imported Swift decl representing the
-/// \c \@interface with the given category name. That is, this will return an
-/// \c swift::ExtensionDecl backed by a \c clang::ObjCCategoryDecl, or a
-/// \c swift::ClassDecl backed by a \c clang::ObjCInterfaceDecl, or \c nullptr
-/// if the class is not imported from Clang or it does not have a category by
-/// that name.
+/// Given a Swift class, find the imported Swift decl(s) representing the
+/// \c \@interface with the given category name. An empty \c categoryName
+/// represents the main interface for the class.
 ///
-/// An empty/invalid \c categoryName requests the main interface for the class.
+/// That is, this request will return one of:
+///
+/// \li a single \c swift::ExtensionDecl backed by a \c clang::ObjCCategoryDecl
+/// \li a \c swift::ClassDecl backed by a \c clang::ObjCInterfaceDecl, plus
+///     zero or more \c swift::ExtensionDecl s backed by 
+///     \c clang::ObjCCategoryDecl s (representing ObjC class extensions).
+/// \li an empty list if the class is not imported from Clang or it does not
+///     have a category by that name.
 class ClangCategoryLookupRequest
     : public SimpleRequest<ClangCategoryLookupRequest,
-                           IterableDeclContext *(ClangCategoryLookupDescriptor),
+                           TinyPtrVector<Decl *>(ClangCategoryLookupDescriptor),
                            RequestFlags::Uncached> {
 public:
   using SimpleRequest::SimpleRequest;
@@ -226,39 +230,46 @@ class ClangCategoryLookupRequest
   friend SimpleRequest;
 
   // Evaluation.
-  IterableDeclContext *evaluate(Evaluator &evaluator,
+ TinyPtrVector<Decl *> evaluate(Evaluator &evaluator,
                                 ClangCategoryLookupDescriptor desc) const;
 };
 
-/// Links an imported Clang decl to the native Swift decl(s) that implement it
-/// using \c \@_objcImplementation.
+/// Links an \c \@_objcImplementation decl to the imported declaration(s) that
+/// it implements.
+///
+/// There is usually a 1:1 correspondence between interfaces and
+/// implementations, except that a class's main implementation implements
+/// both its main interface and any class extension interfaces. In this
+/// situation, the main class is always the first decl in \c interfaceDecls.
 struct ObjCInterfaceAndImplementation final {
-  Decl *interfaceDecl;
+  llvm::TinyPtrVector<Decl *> interfaceDecls;
   Decl *implementationDecl;
 
-  ObjCInterfaceAndImplementation(Decl *interfaceDecl,
+  ObjCInterfaceAndImplementation(llvm::TinyPtrVector<Decl *> interfaceDecls,
                                  Decl *implementationDecl)
-      : interfaceDecl(interfaceDecl), implementationDecl(implementationDecl)
+      : interfaceDecls(interfaceDecls), implementationDecl(implementationDecl)
   {
-    assert(interfaceDecl && implementationDecl &&
+    assert(!interfaceDecls.empty() && implementationDecl &&
            "interface and implementation are both non-null");
   }
 
   ObjCInterfaceAndImplementation()
-      : interfaceDecl(nullptr), implementationDecl(nullptr) {}
+      : interfaceDecls(), implementationDecl(nullptr) {}
 
   operator bool() const {
-    return interfaceDecl;
+    return interfaceDecls.empty();
   }
 
   friend llvm::hash_code
   hash_value(const ObjCInterfaceAndImplementation &pair) {
-    return llvm::hash_combine(pair.interfaceDecl, pair.implementationDecl);
+    return hash_combine(llvm::hash_combine_range(pair.interfaceDecls.begin(),
+                                                 pair.interfaceDecls.end()),
+                        pair.implementationDecl);
   }
 
   friend bool operator==(const ObjCInterfaceAndImplementation &lhs,
                          const ObjCInterfaceAndImplementation &rhs) {
-    return lhs.interfaceDecl == rhs.interfaceDecl
+    return lhs.interfaceDecls == rhs.interfaceDecls
                && lhs.implementationDecl == rhs.implementationDecl;
   }
 
@@ -272,13 +283,12 @@ void simple_display(llvm::raw_ostream &out,
                     const ObjCInterfaceAndImplementation &desc);
 SourceLoc extractNearestSourceLoc(const ObjCInterfaceAndImplementation &desc);
 
-/// Given a \c Decl whose declaration is imported from ObjC but whose
-/// implementation is provided by a Swift \c \@_objcImplementation
-/// \c extension , return both decls, with the imported interface first.
-/// Otherwise return \c {nullptr,nullptr} .
+/// Given a \c Decl , determine if it is an implementation with separate
+/// interfaces imported from ObjC (or vice versa) and if so, return all of the
+/// declarations involved in this relationship. Otherwise return an empty value.
 ///
-/// We retrieve both in a single request because we want to cache the
-/// relationship on both sides to avoid duplicating work.
+/// We perform this lookup in both directions using a single request because
+/// we want to cache the relationship on both sides to avoid duplicating work.
 class ObjCInterfaceAndImplementationRequest
     : public SimpleRequest<ObjCInterfaceAndImplementationRequest,
                            ObjCInterfaceAndImplementation(Decl *),

include/swift/ClangImporter/ClangImporterTypeIDZone.def

@@ -25,7 +25,7 @@ SWIFT_REQUEST(ClangImporter, ClangRecordMemberLookup,
               Decl *(ClangRecordMemberLookupDescriptor), Uncached,
               NoLocationInfo)
 SWIFT_REQUEST(ClangImporter, ClangCategoryLookupRequest,
-              IterableDeclContext *(ClangCategoryLookupDescriptor), Uncached,
+              llvm::TinyPtrVector<Decl *>(ClangCategoryLookupDescriptor), Uncached,
               NoLocationInfo)
 SWIFT_REQUEST(ClangImporter, ObjCInterfaceAndImplementationRequest,
               ObjCInterfaceAndImplementation(Decl *), SeparatelyCached,

lib/AST/Decl.cpp

@@ -1854,11 +1854,9 @@ bool Decl::isObjCImplementation() const {
 
 llvm::Optional<Identifier>
 ExtensionDecl::getCategoryNameForObjCImplementation() const {
-  assert(isObjCImplementation());
-
   auto attr = getAttrs()
                   .getAttribute<ObjCImplementationAttr>(/*AllowInvalid=*/true);
-  if (attr->isCategoryNameInvalid())
+  if (!attr || attr->isCategoryNameInvalid())
     return llvm::None;
 
   return attr->CategoryName;