Support class extensions in objcImpl

include/swift/AST/Decl.h

@@ -355,7 +355,7 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl> {
   // for the inline bitfields.
   union { uint64_t OpaqueBits;
 
-  SWIFT_INLINE_BITFIELD_BASE(Decl, bitmax(NumDeclKindBits,8)+1+1+1+1+1+1,
+  SWIFT_INLINE_BITFIELD_BASE(Decl, bitmax(NumDeclKindBits,8)+1+1+1+1+1+1+1,
     Kind : bitmax(NumDeclKindBits,8),
 
     /// Whether this declaration is invalid.
@@ -381,7 +381,12 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl> {
     Hoisted : 1,
 
     /// Whether the set of semantic attributes has been computed.
-    SemanticAttrsComputed : 1
+    SemanticAttrsComputed : 1,
+
+    /// True if \c ObjCInterfaceAndImplementationRequest has been computed
+    /// and did \em not find anything. This is the fast path where we can bail
+    /// out without checking other caches or computing anything.
+    LacksObjCInterfaceOrImplementation : 1
   );
 
   SWIFT_INLINE_BITFIELD_FULL(PatternBindingDecl, Decl, 1+1+2+16,
@@ -816,13 +821,6 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl> {
 private:
   llvm::PointerUnion<DeclContext *, ASTContext *> Context;
 
-  /// The imported Clang declaration representing the \c @_objcInterface for
-  /// this declaration (or vice versa), or \c nullptr if there is none.
-  ///
-  /// If \c this (an otherwise nonsensical value), the value has not yet been
-  /// computed.
-  Decl *CachedObjCImplementationDecl;
-
   Decl(const Decl&) = delete;
   void operator=(const Decl&) = delete;
   SourceLoc getLocFromSource() const;
@@ -840,14 +838,15 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl> {
 protected:
 
   Decl(DeclKind kind, llvm::PointerUnion<DeclContext *, ASTContext *> context)
-    : Context(context), CachedObjCImplementationDecl(this) {
+    : Context(context) {
     Bits.OpaqueBits = 0;
     Bits.Decl.Kind = unsigned(kind);
     Bits.Decl.Invalid = false;
     Bits.Decl.Implicit = false;
     Bits.Decl.FromClang = false;
     Bits.Decl.EscapedFromIfConfig = false;
     Bits.Decl.Hoisted = false;
+    Bits.Decl.LacksObjCInterfaceOrImplementation = false;
   }
 
   /// Get the Clang node associated with this declaration.
@@ -1168,29 +1167,36 @@ 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 {
+    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()
-  Decl *getImplementedObjCDecl() const;
+  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.
   ///
   /// \seeAlso ExtensionDecl::isObjCInterface()
   Decl *getObjCImplementationDecl() const;
 
-  llvm::Optional<Decl *> getCachedObjCImplementationDecl() const {
-    if (CachedObjCImplementationDecl == this)
-      return llvm::None;
-    return CachedObjCImplementationDecl;
+  bool getCachedLacksObjCInterfaceOrImplementation() const {
+    return Bits.Decl.LacksObjCInterfaceOrImplementation;
   }
 
-  void setCachedObjCImplementationDecl(Decl *decl) {
-    assert((CachedObjCImplementationDecl == this
-              || CachedObjCImplementationDecl == decl)
-               && "can't change CachedObjCInterfaceDecl once it's computed");
-    assert(decl != this && "can't form circular reference");
-    CachedObjCImplementationDecl = decl;
+  void setCachedLacksObjCInterfaceOrImplementation(bool value) {
+    Bits.Decl.LacksObjCInterfaceOrImplementation = value;
   }
 
   /// Return the GenericContext if the Decl has one.
@@ -1842,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
@@ -4927,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/ClangImporter.h

@@ -132,6 +132,12 @@ typedef llvm::PointerUnion<const clang::Decl *, const clang::MacroInfo *,
 class ClangImporter final : public ClangModuleLoader {
   friend class ClangModuleUnit;
 
+  // Make requests in the ClangImporter zone friends so they can access `Impl`.
+#define SWIFT_REQUEST(Zone, Name, Sig, Caching, LocOptions)                    \
+  friend class Name;
+#include "swift/ClangImporter/ClangImporterTypeIDZone.def"
+#undef SWIFT_REQUEST
+
 public:
   class Implementation;
 

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;