[CodeGen][ObjC] Convert '[self alloc]' in a class method to 'objc_alloc(self)'

[ahatanak]

Summary

In a class method, 'self' points to the class object, so it should be safe to pass it to objc_alloc. Also, convert '[[self alloc] init]' in a class method to 'objc_alloc_init(self)'.

rdar://problem/50855121

[rjmccall]

lib/CodeGen/CGObjC.cpp

376 ↗

(On Diff #202113)

I think it's fine to just call this a class message, unless I'm missing a reason to distinguish them in this function.

[ahatanak]

Just set isClassMessage instead of passing a separate flag to tryGenerateSpecializedMessageSend.

[erik.pilkington]

I thinks this looks good, aside from a null concern.

lib/CodeGen/CGObjC.cpp

460 ↗

(On Diff #202821)

I think (?) that CurFuncDecl is nullable, might want to dyn_cast_or_null here.

528 ↗

(On Diff #202821)

(ditto)

cfe/trunk/lib/CodeGen/CGObjC.cpp

if (!CGM.getCodeGenOpts().ObjCConvertMessagesToRuntimeCalls)
     return None;
 
   auto &Runtime = CGM.getLangOpts().ObjCRuntime;
   switch (Sel.getMethodFamily()) {
   case OMF_alloc:
     if (isClassMessage &&
         Runtime.shouldUseRuntimeFunctionsForAlloc() &&
         ResultType->isObjCObjectPointerType()) {
-        // [Foo alloc] -> objc_alloc(Foo)
+        // [Foo alloc] -> objc_alloc(Foo) or
+        // [self alloc] -> objc_alloc(self)
         if (Sel.isUnarySelector() && Sel.getNameForSlot(0) == "alloc")
           return CGF.EmitObjCAlloc(Receiver, CGF.ConvertType(ResultType));
-        // [Foo allocWithZone:nil] -> objc_allocWithZone(Foo)
+        // [Foo allocWithZone:nil] -> objc_allocWithZone(Foo) or
+        // [self allocWithZone:nil] -> objc_allocWithZone(self)
         if (Sel.isKeywordSelector() && Sel.getNumArgs() == 1 &&
             Args.size() == 1 && Args.front().getType()->isPointerType() &&
             Sel.getNameForSlot(0) == "allocWithZone") {
           const llvm::Value* arg = Args.front().getKnownRValue().getScalarVal();
           if (isa<llvm::ConstantPointerNull>(arg))
             return CGF.EmitObjCAllocWithZone(Receiver,
                                              CGF.ConvertType(ResultType));
           return None;
tryEmitSpecializedAllocInit(CodeGenFunction &CGF, const ObjCMessageExpr *OME) {
 
   // Match the exact pattern '[[MyClass alloc] init]'.
   Selector Sel = OME->getSelector();
   if (OME->getReceiverKind() != ObjCMessageExpr::Instance ||
       !OME->getType()->isObjCObjectPointerType() || !Sel.isUnarySelector() ||
       Sel.getNameForSlot(0) != "init")
     return None;
 
-  // Okay, this is '[receiver init]', check if 'receiver' is '[cls alloc]'.
+  // Okay, this is '[receiver init]', check if 'receiver' is '[cls alloc]' or
+  // we are in an ObjC class method and 'receiver' is '[self alloc]'.
   auto *SubOME =
-      dyn_cast<ObjCMessageExpr>(OME->getInstanceReceiver()->IgnoreParens());
+      dyn_cast<ObjCMessageExpr>(OME->getInstanceReceiver()->IgnoreParenCasts());
   if (!SubOME)
     return None;
   Selector SubSel = SubOME->getSelector();
-  if (SubOME->getReceiverKind() != ObjCMessageExpr::Class ||
-      !SubOME->getType()->isObjCObjectPointerType() ||
+
+  // Check if we are in an ObjC class method and the receiver expression is
+  // 'self'.
+  const Expr *SelfInClassMethod = nullptr;
+  if (const auto *CurMD = dyn_cast_or_null<ObjCMethodDecl>(CGF.CurFuncDecl))
+    if (CurMD->isClassMethod())
+      if ((SelfInClassMethod = SubOME->getInstanceReceiver()))
+        if (!SelfInClassMethod->isObjCSelfExpr())
+          SelfInClassMethod = nullptr;
+
+  if ((SubOME->getReceiverKind() != ObjCMessageExpr::Class &&
+       !SelfInClassMethod) || !SubOME->getType()->isObjCObjectPointerType() ||
       !SubSel.isUnarySelector() || SubSel.getNameForSlot(0) != "alloc")
     return None;
 
+  llvm::Value *Receiver;
+  if (SelfInClassMethod) {
+    Receiver = CGF.EmitScalarExpr(SelfInClassMethod);
+  } else {
     QualType ReceiverType = SubOME->getClassReceiver();
     const ObjCObjectType *ObjTy = ReceiverType->getAs<ObjCObjectType>();
     const ObjCInterfaceDecl *ID = ObjTy->getInterface();
     assert(ID && "null interface should be impossible here");
-  llvm::Value *Receiver = CGF.CGM.getObjCRuntime().GetClass(CGF, ID);
+    Receiver = CGF.CGM.getObjCRuntime().GetClass(CGF, ID);
+  }
   return CGF.EmitObjCAllocInit(Receiver, CGF.ConvertType(OME->getType()));
 }
 
 RValue CodeGenFunction::EmitObjCMessageExpr(const ObjCMessageExpr *E,
                                             ReturnValueSlot Return) {
   // Only the lookup mechanism and first two arguments of the method
   // implementation vary between runtimes.  We can get the receiver and
   // arguments in generic code.
RValue CodeGenFunction::EmitObjCMessageExpr(const ObjCMessageExpr *E,
   bool isClassMessage = false;
   ObjCInterfaceDecl *OID = nullptr;
   // Find the receiver
   QualType ReceiverType;
   llvm::Value *Receiver = nullptr;
   switch (E->getReceiverKind()) {
   case ObjCMessageExpr::Instance:
     ReceiverType = E->getInstanceReceiver()->getType();
+    if (auto *OMD = dyn_cast_or_null<ObjCMethodDecl>(CurFuncDecl))
+      if (OMD->isClassMethod())
+        if (E->getInstanceReceiver()->isObjCSelfExpr())
+          isClassMessage = true;
     if (retainSelf) {
       TryEmitResult ter = tryEmitARCRetainScalarExpr(*this,
                                                    E->getInstanceReceiver());
       Receiver = ter.getPointer();
       if (ter.getInt()) retainSelf = false;
     } else
       Receiver = EmitScalarExpr(E->getInstanceReceiver());
     break;

cfe/trunk/test/CodeGenObjC/convert-messages-to-runtime-calls.m

 // not a pointer type.
 // CHECK-LABEL: define {{.*}}void @test_cannot_message_return_float
 float test_cannot_message_return_float(C *c) {
   // MSGS: {{call.*@objc_msgSend}}
   // CALLS: {{call.*@objc_msgSend}}
   return [c retain];
 }
 
+@interface TestSelf
++ (instancetype)alloc;
++ (instancetype)allocWithZone:(void*)zone;
++ (id)classMeth;
+- (id)instanceMeth;
+@end
+
+@implementation TestSelf
+// CHECK-LABEL: define internal i8* @"\01+[TestSelf classMeth]"(
++ (id)classMeth {
+  // MSGS: {{call.*@objc_msgSend}}
+  // MSGS: {{call.*@objc_msgSend}}
+  // CALLS: {{call.*@objc_allocWithZone\(}}
+  // CALLS: {{call.*@objc_alloc\(}}
+  [self allocWithZone:nil];
+  return [self alloc];
+}
+// CHECK-LABEL: define internal i8* @"\01-[TestSelf instanceMeth]"(
+- (id)instanceMeth {
+  // MSGS: {{call.*@objc_msgSend}}
+  // MSGS: {{call.*@objc_msgSend}}
+  // CALLS: {{call.*@objc_msgSend}}
+  // CALLS: {{call.*@objc_msgSend}}
+  [self allocWithZone:nil];
+  return [self alloc];
+}
+@end
+
 @interface NSString : NSObject
 + (void)retain_self;
 - (void)retain_super;
 @end
 
 @implementation NSString
 
 // Make sure we can convert a message to a dynamic receiver to a call

cfe/trunk/test/CodeGenObjC/objc-alloc-init.m

void f() {
   } @catch (X *x) {
   }
   // OPTIMIZED: invoke i8* @objc_alloc_init(
   // NOT_OPTIMIZED: invoke i8* @objc_alloc(
 }
 
 @interface Y : X
 +(void)meth;
+-(void)instanceMeth;
 @end
 
 @implementation Y
 +(void)meth {
   [[self alloc] init];
+  // OPTIMIZED: call i8* @objc_alloc_init(
+  // NOT_OPTIMIZED: call i8* @objc_alloc(
+}
+-(void)instanceMeth {
   // EITHER-NOT: call i8* @objc_alloc
   // EITHER: call {{.*}} @objc_msgSend
   // EITHER: call {{.*}} @objc_msgSend
+  [[self alloc] init];
 }
 @end
 
 // rdar://48247290
 @interface Base
 -(instancetype)init;
 @end
 
 @interface Derived : Base
 @end
 @implementation Derived
 -(void)meth {
   [super init];
 }
 @end