Fix a couple of problems with initialization and assignment to

__block variables where the act of initialization/assignment
itself causes the __block variable to be copied to the heap
because the variable is of block type and is being assigned
a block literal which captures the variable.

rdar://problem/9814099



git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@136337 91177308-0d34-0410-b5e6-96231b3b80d8

4 files changed, 92 insertions, 8 deletions

diff --git a/lib/CodeGen/CGDecl.cpp b/lib/CodeGen/CGDecl.cpp
index a372d7d696..870fe811f1 100644
--- a/lib/CodeGen/CGDecl.cpp
+++ b/lib/CodeGen/CGDecl.cpp
@@ -507,7 +507,7 @@ void CodeGenFunction::EmitScalarInit(const Expr *init,
   // actually perform the initialization with an assign.
   bool accessedByInit = false;
   if (lifetime != Qualifiers::OCL_ExplicitNone)
-    accessedByInit = isAccessedBy(D, init);
+    accessedByInit = (capturedByInit || isAccessedBy(D, init));
   if (accessedByInit) {
     LValue tempLV = lvalue;
     // Drill down to the __block object if necessary.
diff --git a/lib/CodeGen/CGObjC.cpp b/lib/CodeGen/CGObjC.cpp
index f81b6b75e3..2f9ff22495 100644
--- a/lib/CodeGen/CGObjC.cpp
+++ b/lib/CodeGen/CGObjC.cpp
@@ -2144,10 +2144,20 @@ CodeGenFunction::EmitARCStoreStrong(const BinaryOperator *e,
   TryEmitResult result = tryEmitARCRetainScalarExpr(*this, e->getRHS());
   llvm::Value *value = result.getPointer();
 
+  bool hasImmediateRetain = result.getInt();
+
+  // If we didn't emit a retained object, and the l-value is of block
+  // type, then we need to emit the block-retain immediately in case
+  // it invalidates the l-value.
+  if (!hasImmediateRetain && e->getType()->isBlockPointerType()) {
+    value = EmitARCRetainBlock(value);
+    hasImmediateRetain = true;
+  }
+
   LValue lvalue = EmitLValue(e->getLHS());
 
   // If the RHS was emitted retained, expand this.
-  if (result.getInt()) {
+  if (hasImmediateRetain) {
     llvm::Value *oldValue =
       EmitLoadOfScalar(lvalue.getAddress(), lvalue.isVolatileQualified(),
                        lvalue.getAlignment(), e->getType(),
diff --git a/test/CodeGenObjC/arc.m b/test/CodeGenObjC/arc.m
index 3d0abc5c30..109e360b79 100644
--- a/test/CodeGenObjC/arc.m
+++ b/test/CodeGenObjC/arc.m
@@ -1684,3 +1684,73 @@ void test59(void) {
   // CHECK-NEXT: call void @objc_release(i8* [[T1]])
   // CHECK-NEXT: ret void
 }
+
+// rdar://problem/9814099
+// Test that we correctly initialize __block variables
+// when the initialization captures the variable.
+void test60a(void) {
+  __block void (^block)(void) = ^{ block(); };
+  // CHECK:    define void @test60a()
+  // CHECK:      [[BYREF:%.*]] = alloca [[BYREF_T:%.*]],
+
+  // Zero-initialization before running the initializer.
+  // CHECK:      [[T0:%.*]] = getelementptr inbounds [[BYREF_T]]* [[BYREF]], i32 0, i32 6
+  // CHECK-NEXT: store void ()* null, void ()** [[T0]], align 8
+
+  // Run the initializer as an assignment.
+  // CHECK:      [[T0:%.*]] = bitcast void ()* {{%.*}} to i8*
+  // CHECK-NEXT: [[T1:%.*]] = call i8* @objc_retainBlock(i8* [[T0]])
+  // CHECK-NEXT: [[T2:%.*]] = bitcast i8* [[T1]] to void ()*
+  // CHECK-NEXT: [[T3:%.*]] = getelementptr inbounds [[BYREF_T]]* [[BYREF]], i32 0, i32 1
+  // CHECK-NEXT: [[T4:%.*]] = load [[BYREF_T]]** [[T3]]
+  // CHECK-NEXT: [[T5:%.*]] = getelementptr inbounds [[BYREF_T]]* [[T4]], i32 0, i32 6
+  // CHECK-NEXT: [[T6:%.*]] = load void ()** [[T5]], align 8
+  // CHECK-NEXT: store void ()* {{%.*}}, void ()** [[T5]], align 8
+  // CHECK-NEXT: [[T7:%.*]] = bitcast void ()* [[T6]] to i8*
+  // CHECK-NEXT: call void @objc_release(i8* [[T7]])
+
+  // Destroy at end of function.
+  // CHECK-NEXT: [[SLOT:%.*]] = getelementptr inbounds [[BYREF_T]]* [[BYREF]], i32 0, i32 6
+  // CHECK-NEXT: [[T0:%.*]] = bitcast [[BYREF_T]]* [[BYREF]] to i8*
+  // CHECK-NEXT: call void @_Block_object_dispose(i8* [[T0]], i32 8)
+  // CHECK-NEXT: [[T1:%.*]] = load void ()** [[SLOT]]
+  // CHECK-NEXT: [[T2:%.*]] = bitcast void ()* [[T1]] to i8*
+  // CHECK-NEXT: call void @objc_release(i8* [[T2]])
+  // CHECK-NEXT: ret void
+}
+
+// Test that we correctly assign to __block variables when the
+// assignment captures the variable.
+void test60b(void) {
+  __block void (^block)(void);
+  block = ^{ block(); };
+
+  // CHECK:    define void @test60b()
+  // CHECK:      [[BYREF:%.*]] = alloca [[BYREF_T:%.*]],
+
+  // Zero-initialize.
+  // CHECK:      [[T0:%.*]] = getelementptr inbounds [[BYREF_T]]* [[BYREF]], i32 0, i32 6
+  // CHECK-NEXT: store void ()* null, void ()** [[T0]], align 8
+
+  // CHECK-NEXT: [[SLOT:%.*]] = getelementptr inbounds [[BYREF_T]]* [[BYREF]], i32 0, i32 6
+
+  // The assignment.
+  // CHECK:      [[T0:%.*]] = bitcast void ()* {{%.*}} to i8*
+  // CHECK-NEXT: [[T1:%.*]] = call i8* @objc_retainBlock(i8* [[T0]])
+  // CHECK-NEXT: [[T2:%.*]] = bitcast i8* [[T1]] to void ()*
+  // CHECK-NEXT: [[T3:%.*]] = getelementptr inbounds [[BYREF_T]]* [[BYREF]], i32 0, i32 1
+  // CHECK-NEXT: [[T4:%.*]] = load [[BYREF_T]]** [[T3]]
+  // CHECK-NEXT: [[T5:%.*]] = getelementptr inbounds [[BYREF_T]]* [[T4]], i32 0, i32 6
+  // CHECK-NEXT: [[T6:%.*]] = load void ()** [[T5]], align 8
+  // CHECK-NEXT: store void ()* {{%.*}}, void ()** [[T5]], align 8
+  // CHECK-NEXT: [[T7:%.*]] = bitcast void ()* [[T6]] to i8*
+  // CHECK-NEXT: call void @objc_release(i8* [[T7]])
+
+  // Destroy at end of function.
+  // CHECK-NEXT: [[T0:%.*]] = bitcast [[BYREF_T]]* [[BYREF]] to i8*
+  // CHECK-NEXT: call void @_Block_object_dispose(i8* [[T0]], i32 8)
+  // CHECK-NEXT: [[T1:%.*]] = load void ()** [[SLOT]]
+  // CHECK-NEXT: [[T2:%.*]] = bitcast void ()* [[T1]] to i8*
+  // CHECK-NEXT: call void @objc_release(i8* [[T2]])
+  // CHECK-NEXT: ret void
+}
diff --git a/test/CodeGenObjCXX/arc-special-member-functions.mm b/test/CodeGenObjCXX/arc-special-member-functions.mm
index d88a2bd62f..34d43250a2 100644
--- a/test/CodeGenObjCXX/arc-special-member-functions.mm
+++ b/test/CodeGenObjCXX/arc-special-member-functions.mm
@@ -92,12 +92,16 @@ void test_ObjCBlockMember_copy_assign(ObjCBlockMember m1, ObjCBlockMember m2) {
 
 // Implicitly-generated copy assignment operator for ObjCBlockMember
 // CHECK:    define linkonce_odr {{%.*}}* @_ZN15ObjCBlockMemberaSERKS_(
-// CHECK:      [[T0:%.*]] = call i8* @objc_retainBlock(
-// CHECK-NEXT: [[T1:%.*]] = bitcast i8* [[T0]] to i32 (i32)*
-// CHECK-NEXT: [[T2:%.*]] = load {{.*}} [[SLOT:%.*]],
-// CHECK: store
-// CHECK-NEXT: [[T3:%.*]] = bitcast 
-// CHECK-NEXT: call void @objc_release(i8* [[T3]])
+// CHECK:      [[T0:%.*]] = getelementptr inbounds [[T:%.*]]* {{%.*}}, i32 0, i32 0
+// CHECK-NEXT: [[T1:%.*]] = load i32 (i32)** [[T0]], align 8
+// CHECK-NEXT: [[T2:%.*]] = bitcast i32 (i32)* [[T1]] to i8*
+// CHECK-NEXT: [[T3:%.*]] = call i8* @objc_retainBlock(i8* [[T2]])
+// CHECK-NEXT: [[T4:%.*]] = bitcast i8* [[T3]] to i32 (i32)*
+// CHECK-NEXT: [[T5:%.*]] = getelementptr inbounds [[T]]* {{%.*}}, i32 0, i32 0
+// CHECK-NEXT: [[T6:%.*]] = load i32 (i32)** [[T5]], align 8
+// CHECK-NEXT: store i32 (i32)* [[T4]], i32 (i32)** [[T5]]
+// CHECK-NEXT: [[T7:%.*]] = bitcast i32 (i32)* [[T6]] to i8*
+// CHECK-NEXT: call void @objc_release(i8* [[T7]])
 // CHECK-NEXT: ret
 
 // Implicitly-generated copy constructor for ObjCBlockMember