IRGen: Categories on class stubs should go in a separate list and section

slavapestov · slavapestov · commit d6868acf8a07 · 2019-04-19T16:50:29.000-04:00
diff --git a/lib/IRGen/GenCast.cpp b/lib/IRGen/GenCast.cpp
@@ -107,8 +107,7 @@ FailableCastResult irgen::emitClassIdenticalCast(IRGenFunction &IGF,
   llvm::Value *targetMetadata;
   if ((targetMetadata =
            tryEmitConstantHeapMetadataRef(IGF.IGM, toType.getASTType(),
-                                          /*allowUninitialized*/ false,
-                                          /*allowStub*/ false))) {
+                                          /*allowUninitialized*/ false))) {
     // ok
   } else {
     targetMetadata
diff --git a/lib/IRGen/GenClass.cpp b/lib/IRGen/GenClass.cpp
@@ -1194,7 +1194,32 @@ namespace {
       if (categoryCount > 0)
         os << categoryCount;
     }
-    
+
+    llvm::Constant *getClassMetadataRef() {
+      auto *theClass = getClass();
+
+      if (theClass->hasClangNode())
+        return IGM.getAddrOfObjCClass(theClass, NotForDefinition);
+
+      // Note that getClassMetadataStrategy() will return
+      // ClassMetadataStrategy::Resilient if the class is
+      // from another resilience domain, even if inside that
+      // resilience domain the class has fixed metadata
+      // layout.
+      //
+      // Since a class only has a class stub if its class
+      // hierarchy crosses resilience domains, we use a
+      // slightly different query here.
+      if (theClass->checkAncestry(AncestryFlags::ResilientOther)) {
+        assert(IGM.Context.LangOpts.EnableObjCResilientClassStubs);
+        return IGM.getAddrOfObjCResilientClassStub(theClass, NotForDefinition,
+                                             TypeMetadataAddress::AddressPoint);
+      }
+
+      auto type = getSelfType(theClass).getASTType();
+      return tryEmitConstantHeapMetadataRef(IGM, type, /*allowUninit*/ true);
+    }
+
   public:
     llvm::Constant *emitCategory() {
       assert(TheExtension && "can't emit category data for a class");
@@ -1205,18 +1230,7 @@ namespace {
       //   char const *name;
       fields.add(IGM.getAddrOfGlobalString(CategoryName));
       //   const class_t *theClass;
-      if (getClass()->hasClangNode())
-        fields.add(IGM.getAddrOfObjCClass(getClass(), NotForDefinition));
-      else {
-        auto type = getSelfType(getClass()).getASTType();
-        llvm::Constant *metadata =
-          tryEmitConstantHeapMetadataRef(IGM, type,
-                                         /*allowUninit*/ true,
-                                         /*allowStub*/ true);
-        assert(metadata &&
-               "extended objc class doesn't have constant metadata?");
-        fields.add(metadata);
-      }
+      fields.add(getClassMetadataRef());
       //   const method_list_t *instanceMethods;
       fields.add(buildInstanceMethodList());
       //   const method_list_t *classMethods;
diff --git a/lib/IRGen/GenDecl.cpp b/lib/IRGen/GenDecl.cpp
@@ -928,6 +928,12 @@ void IRGenModule::emitGlobalLists() {
                                       "regular,no_dead_strip"),
                    llvm::GlobalValue::InternalLinkage, Int8PtrTy, false);
 
+    // And categories on class stubs.
+    emitGlobalList(*this, ObjCCategoriesOnStubs, "objc_categories_stubs",
+                   GetObjCSectionName("__objc_catlist2",
+                                      "regular,no_dead_strip"),
+                   llvm::GlobalValue::InternalLinkage, Int8PtrTy, false);
+
     // Emit nonlazily realized class references in a second magic section to make
     // sure they are realized by the Objective-C runtime before any instances
     // are allocated.
@@ -3802,7 +3808,15 @@ void IRGenModule::emitExtension(ExtensionDecl *ext) {
            "foreign types cannot have categories emitted");
     llvm::Constant *category = emitCategoryData(*this, ext);
     category = llvm::ConstantExpr::getBitCast(category, Int8PtrTy);
-    ObjCCategories.push_back(category);
+
+    auto *theClass = ext->getSelfClassDecl();
+
+    // Categories on class stubs are added to a separate list.
+    if (theClass->checkAncestry(AncestryFlags::ResilientOther))
+      ObjCCategoriesOnStubs.push_back(category);
+    else
+      ObjCCategories.push_back(category);
+
     ObjCCategoryDecls.push_back(ext);
   }
 }
diff --git a/lib/IRGen/GenMeta.cpp b/lib/IRGen/GenMeta.cpp
@@ -2649,8 +2649,7 @@ namespace {
       Type type = Target->mapTypeIntoContext(Target->getSuperclass());
       auto *metadata = tryEmitConstantHeapMetadataRef(
           IGM, type->getCanonicalType(),
-          /*allowUninit*/ false,
-          /*allowStub*/ false);
+          /*allowUninit*/ false);
       assert(metadata != nullptr);
       B.add(metadata);
     }
diff --git a/lib/IRGen/IRGenModule.h b/lib/IRGen/IRGenModule.h
@@ -987,6 +987,8 @@ class IRGenModule {
   SmallVector<llvm::WeakTrackingVH, 4> ObjCNonLazyClasses;
   /// List of Objective-C categories, bitcast to i8*.
   SmallVector<llvm::WeakTrackingVH, 4> ObjCCategories;
+  /// List of Objective-C categories on class stubs, bitcast to i8*.
+  SmallVector<llvm::WeakTrackingVH, 4> ObjCCategoriesOnStubs;
   /// List of non-ObjC protocols described by this module.
   SmallVector<ProtocolDecl *, 4> SwiftProtocols;
   /// List of protocol conformances to generate descriptors for.
diff --git a/lib/IRGen/MetadataRequest.cpp b/lib/IRGen/MetadataRequest.cpp
@@ -517,29 +517,10 @@ irgen::getRuntimeReifiedType(IRGenModule &IGM, CanType type) {
 llvm::Constant *
 irgen::tryEmitConstantHeapMetadataRef(IRGenModule &IGM,
                                       CanType type,
-                                      bool allowDynamicUninitialized,
-                                      bool allowStub) {
+                                      bool allowDynamicUninitialized) {
   auto theDecl = type->getClassOrBoundGenericClass();
   assert(theDecl && "emitting constant heap metadata ref for non-class type?");
 
-  // Note that getClassMetadataStrategy() will return
-  // ClassMetadataStrategy::Resilient if the class is
-  // from another resilience domain, even if inside that
-  // resilience domain the class has fixed metadata
-  // layout.
-  //
-  // Since a class only has a class stub if its class
-  // hierarchy crosses resilience domains, we use a
-  // slightly different query here.
-  if (theDecl->checkAncestry(AncestryFlags::ResilientOther)) {
-    if (allowStub && IGM.Context.LangOpts.EnableObjCResilientClassStubs) {
-      return IGM.getAddrOfObjCResilientClassStub(theDecl, NotForDefinition,
-                                            TypeMetadataAddress::AddressPoint);
-    }
-
-    return nullptr;
-  }
-
   switch (IGM.getClassMetadataStrategy(theDecl)) {
   case ClassMetadataStrategy::Resilient:
   case ClassMetadataStrategy::Singleton:
diff --git a/lib/IRGen/MetadataRequest.h b/lib/IRGen/MetadataRequest.h
@@ -592,8 +592,7 @@ CanType getRuntimeReifiedType(IRGenModule &IGM, CanType type);
 /// by a constant.
 llvm::Constant *tryEmitConstantHeapMetadataRef(IRGenModule &IGM,
                                                CanType type,
-                                               bool allowUninitialized,
-                                               bool allowStub);
+                                               bool allowUninitialized);
 
 enum class MetadataValueType { ObjCClass, TypeMetadata };
 
diff --git a/test/IRGen/class_update_callback_with_stub.swift b/test/IRGen/class_update_callback_with_stub.swift
@@ -104,11 +104,14 @@ import resilient_objc_class
 // -- The category list
 
 // CHECK-LABEL: @objc_categories = internal global
+// CHECK-SAME: @"_CATEGORY__TtC15resilient_class22ResilientOutsideParent_$_class_update_callback_with_stub"
+// CHECK-SAME: , section "__DATA,__objc_catlist,regular,no_dead_strip"
+
+// CHECK-LABEL: @objc_categories_stubs = internal global
 // CHECK-SAME: @"_CATEGORY__TtC31class_update_callback_with_stub17ResilientSubclass_$_class_update_callback_with_stub"
 // CHECK-SAME: @"_CATEGORY__TtC31class_update_callback_with_stub25ResilientNSObjectSubclass_$_class_update_callback_with_stub"
 // CHECK-SAME: @"_CATEGORY__TtC31class_update_callback_with_stub27FixedLayoutNSObjectSubclass_$_class_update_callback_with_stub"
-// CHECK-SAME: @"_CATEGORY__TtC15resilient_class22ResilientOutsideParent_$_class_update_callback_with_stub"
-// CHECK-SAME: , section "__DATA,__objc_catlist,regular,no_dead_strip"
+// CHECK-SAME: , section "__DATA,__objc_catlist2,regular,no_dead_strip"
 
 
 // -- Address point for class stubs
