Support class extensions in objcImpl

Their requirements are now included when typechecking the main body extension.

Fixes rdar://118535473.

Author: beccadax

lib/ClangImporter/ClangImporter.cpp

@@ -5885,6 +5885,10 @@ Identifier ExtensionDecl::getObjCCategoryName() const {
 
   // We'll look for an implementation with this category name.
   auto clangCategoryName = category->getName();
+  if (clangCategoryName.empty())
+    // Class extension (has an empty name).
+    return Identifier();
+  
   return getASTContext().getIdentifier(clangCategoryName);
 }
 
@@ -6109,6 +6113,13 @@ Decl *Decl::getObjCImplementationDecl() const {
     // This *is* the implementation, if it has one.
     return nullptr;
 
+  // If this is an imported class extension, it's implemented by the main body's
+  // implementation. (But the relationship is not reciprocal, and hence not
+  // modeled by `ObjCInterfaceAndImplementationRequest`.)
+  if (auto ext = dyn_cast<ExtensionDecl>(this))
+    if (ext->getObjCCategoryName().empty())
+      return ext->getSelfNominalTypeDecl()->getObjCImplementationDecl();
+
   ObjCInterfaceAndImplementationRequest req{const_cast<Decl *>(this)};
   return evaluateOrDefault(getASTContext().evaluator, req, {})
              .implementationDecl;

lib/Sema/TypeCheckDeclObjC.cpp

@@ -3068,6 +3068,17 @@ class ObjCImplementationChecker {
 
       addRequirement(member);
     }
+
+    if (auto classDecl = dyn_cast<ClassDecl>(idc)) {
+      // The extension for the main class body should also implement the
+      // contents of class extensions. Find all class extensions (imported
+      // categories with no name) and add their requirements too.
+      for (auto ext : classDecl->getExtensions()) {
+        if (ext->hasClangNode() && ext->getObjCCategoryName().empty()) {
+          addRequirements(ext);
+        }
+      }
+    }
   }
 
   void addCandidates(ExtensionDecl *ext) {
@@ -3422,9 +3433,9 @@ class ObjCImplementationChecker {
 
     // Check only applies to members of implementations, not implementations in
     // their own right.
-    if (!cand->isObjCImplementation()
-          && cand->getDeclContext()->getImplementedObjCContext()
-                 != req->getDeclContext())
+    if (!cand->isObjCImplementation() &&
+          getCategoryName(cand->getDeclContext()->getImplementedObjCContext())
+               != getCategoryName(req->getDeclContext()))
       return MatchOutcome::WrongCategory;
 
     if (cand->getKind() != req->getKind())

test/decl/ext/Inputs/objc_implementation.h

@@ -59,6 +59,13 @@
 
 @end
 
+@interface ObjCClass ()
+
+- (void)extensionMethodFromHeader1:(int)param;
+- (void)extensionMethodFromHeader2:(int)param;
+
+@end
+
 @interface ObjCClass (PresentAdditions)
 
 - (void)categoryMethodFromHeader1:(int)param;

test/decl/ext/objc_implementation.swift

@@ -12,6 +12,7 @@ protocol EmptySwiftProto {}
   // FIXME: give better diagnostic expected-warning@-6 {{extension for main class interface should provide implementation for instance method 'method(fromHeader3:)'}}
   // expected-warning@-7 {{'@_objcImplementation' extension cannot add conformance to 'EmptySwiftProto'; add this conformance with an ordinary extension}}
   // expected-warning@-8 {{'@_objcImplementation' extension cannot add conformance to 'EmptyObjCProto'; add this conformance in the Objective-C header}}
+  // expected-warning@-9 {{extension for main class interface should provide implementation for instance method 'extensionMethod(fromHeader2:)'}}
 
   func method(fromHeader1: CInt) {
     // OK, provides an implementation for the header's method.
@@ -215,6 +216,10 @@ protocol EmptySwiftProto {}
   @nonobjc public convenience init(notFromHeader6: CInt) {
     // OK
   }
+
+  @objc func extensionMethod(fromHeader1: CInt) {
+    // OK
+  }
 }
 
 @_objcImplementation(PresentAdditions) extension ObjCClass {