[CodeGen][ObjC] Fix broken IR generated when there is a nil receiver check

[ahatanak]

Summary

This patch fixes a bug in emitARCOperationAfterCall where it inserts the fall-back call after a bitcast instruction and then replaces the bitcast's operand with the result of the fall-back call. The generated IR without this patch looks like this:

msgSend.call:                                     ; preds = %entry
  %call = call i8* bitcast (i8* (i8*, i8*, ...)* @objc_msgSend to i8* (i8*, i8*, i8*)*)(i8* %6, i8* %5, i8* %4) #4, !clang.arc.no_objc_arc_exceptions !11
  br label %msgSend.cont

msgSend.null-receiver:                            ; preds = %entry
  call void @llvm.objc.release(i8* %4) #1, !clang.imprecise_release !11
  br label %msgSend.cont

msgSend.cont:                                     ; preds = %msgSend.null-receiver, %msgSend.call
  %8 = phi i8* [ %call, %msgSend.call ], [ null, %msgSend.null-receiver ]
  %9 = bitcast i8* %10 to %0*
  %10 = call i8* @llvm.objc.retain(i8* %8) #1

Notice that %9 = bitcast i8* %10 to %0* is taking operand %10 which is defined after it.

To fix the bug, this patch modifies the insert point to point to the bitcast instruction so that the fall-back call is inserted before the bitcast. In addition, it teaches the function to look at phi instructions that are generated when there is a check for a null receiver and insert the retainRV/claimRV instruction right after the call instead of inserting a fall-back call right after the phi instruction.

rdar://73360225

[rjmccall]

This looks good, but in the long term it would be great if we could eliminate the need for emitARCOperationAfterCall and just emit the call properly in the first place.

Please add a test where we have a null check around an invoke.

[ahatanak]

Add a test for invoke and a FIXME for eliminating the need for emitARCOperationAfterCall.

clang/lib/CodeGen/CGObjC.cpp

typedef llvm::function_ref<llvm::Value *(CodeGenFunction &CGF,
                                          llvm::Value *value)>
   ValueTransform;
 
 /// Insert code immediately after a call.
 static llvm::Value *emitARCOperationAfterCall(CodeGenFunction &CGF,
                                               llvm::Value *value,
                                               ValueTransform doAfterCall,
                                               ValueTransform doFallback) {
-  if (llvm::CallInst *call = dyn_cast<llvm::CallInst>(value)) {
-    CGBuilderTy::InsertPoint ip = CGF.Builder.saveIP();
 
+  if (llvm::CallInst *call = dyn_cast<llvm::CallInst>(value)) {
     // Place the retain immediately following the call.
     CGF.Builder.SetInsertPoint(call->getParent(),
                                ++llvm::BasicBlock::iterator(call));
     value = doAfterCall(CGF, value);
-
-    CGF.Builder.restoreIP(ip);
-    return value;
   } else if (llvm::InvokeInst *invoke = dyn_cast<llvm::InvokeInst>(value)) {
-    CGBuilderTy::InsertPoint ip = CGF.Builder.saveIP();
-
     // Place the retain at the beginning of the normal destination block.
     llvm::BasicBlock *BB = invoke->getNormalDest();
     CGF.Builder.SetInsertPoint(BB, BB->begin());
     value = doAfterCall(CGF, value);
 
-    CGF.Builder.restoreIP(ip);
-    return value;
-
   // Bitcasts can arise because of related-result returns.  Rewrite
   // the operand.
   } else if (llvm::BitCastInst *bitcast = dyn_cast<llvm::BitCastInst>(value)) {
+    // Change the insert point to avoid emitting the fall-back call after the
+    // bitcast.
+    CGF.Builder.SetInsertPoint(bitcast->getParent(), bitcast->getIterator());
     llvm::Value *operand = bitcast->getOperand(0);
     operand = emitARCOperationAfterCall(CGF, operand, doAfterCall, doFallback);
     bitcast->setOperand(0, operand);
-    return bitcast;
-
-  // Generic fall-back case.
+    value = bitcast;
   } else {
+    auto *phi = dyn_cast<llvm::PHINode>(value);
+    if (phi && phi->getNumIncomingValues() == 2 &&
+        isa<llvm::ConstantPointerNull>(phi->getIncomingValue(1)) &&
+        isa<llvm::CallBase>(phi->getIncomingValue(0))) {
+      // Handle phi instructions that are generated when it's necessary to check
+      // whether the receiver of a message is null.
+      llvm::Value *inVal = phi->getIncomingValue(0);
+      inVal = emitARCOperationAfterCall(CGF, inVal, doAfterCall, doFallback);
+      phi->setIncomingValue(0, inVal);
+      value = phi;
+    } else {
+      // Generic fall-back case.
       // Retain using the non-block variant: we never need to do a copy
       // of a block that's been returned to us.
-    return doFallback(CGF, value);
+      value = doFallback(CGF, value);
     }
   }
 
+  CGF.Builder.restoreIP(ip);
+  return value;
+}
+
 /// Given that the given expression is some sort of call (which does
 /// not return retained), emit a retain following it.
 static llvm::Value *emitARCRetainCallResult(CodeGenFunction &CGF,
                                             const Expr *e) {
   llvm::Value *value = CGF.EmitScalarExpr(e);
   return emitARCOperationAfterCall(CGF, value,
            [](CodeGenFunction &CGF, llvm::Value *value) {
              return CGF.EmitARCRetainAutoreleasedReturnValue(value);

clang/test/CodeGenObjC/ns_consume_null_check.m

 // RUN: %clang_cc1 -triple x86_64-apple-darwin10 -emit-llvm -fobjc-arc -fobjc-dispatch-method=mixed -fobjc-runtime-has-weak -fexceptions -o - %s | FileCheck %s
 
 @interface NSObject
 - (id) new;
 @end
 
 @interface MyObject : NSObject
 - (char)isEqual:(id) __attribute__((ns_consumed)) object;
 - (_Complex float) asComplexWithArg: (id) __attribute__((ns_consumed)) object;
++(instancetype)m0:(id) __attribute__((ns_consumed)) object;
 @end
 
 MyObject *x;
 
 // rdar://10444476
 void test0(void) {
   id obj = [NSObject new];
   [x isEqual : obj];

 //   Epilogue.
 // CHECK-NEXT: call void @llvm.objc.destroyWeak(i8** [[WEAKOBJ]]) [[NUW]]
 // CHECK-NEXT: call void @llvm.objc.storeStrong(i8** [[OBJ]], i8* null) [[NUW]]
 // CHECK-NEXT: ret void
 //   Cleanup.
 // CHECK:      landingpad
 // CHECK:      call void @llvm.objc.destroyWeak(i8** [[WEAKOBJ]]) [[NUW]]
 
+void test2(id a) {
+  id obj = [MyObject m0:a];
+}
+
+// CHECK: define{{.*}} void @test2(i8* %[[A:.*]]) #0 {
+// CHECK: %[[CALL:.*]] = call i8* bitcast (i8* (i8*, i8*, ...)* @objc_msgSend
+// CHECK-NEXT: %[[V6:.*]] = {{.*}}call i8* @llvm.objc.retainAutoreleasedReturnValue(i8* %[[CALL]])
+
+// CHECK: phi i8* [ %[[V6]], %{{.*}} ], [ null, %{{.*}} ]
+
+
 // CHECK: attributes [[NUW]] = { nounwind }