[CodeGen] Fix handling of variables captured by a block that is nested inside a lambda

[ahatanak]

Summary

Currently IRGen doesn't handle variables captured by a block correctly when the variable is captured by reference by a lambda that encloses the block.

For example, in the following code, the type of capture 'x' in the block literal is a reference to 'id' because 'x' is captured by reference by the enclosing lambda. In this case, copy/dispose functions shouldn't be needed, but currently IRGen emits them.

void test() {
  id x;
  [&]{ ^{ (void)x; }(); }();
}

This happens because there are a few places in CGBlocks.cpp that use the variable's type ('id' in the example above) instead of the capture type ('id &' in the example above).

[rjmccall]

LGTM.

cfe/trunk/lib/CodeGen/CGBlocks.cpp

for (const auto &CI : block->captures()) {
     if (CI.isByRef()) {
       // We have to copy/dispose of the __block reference.
       info.NeedsCopyDispose = true;
 
       // Just use void* instead of a pointer to the byref type.
       CharUnits align = CGM.getPointerAlign();
       maxFieldAlign = std::max(maxFieldAlign, align);
 
+      // Since a __block variable cannot be captured by lambdas, its type and
+      // the capture field type should always match.
+      assert(getCaptureFieldType(*CGF, CI) == variable->getType() &&
+             "capture type differs from the variable type");
       layout.push_back(BlockLayoutChunk(align, CGM.getPointerSize(),
                                         Qualifiers::OCL_None, &CI,
                                         CGM.VoidPtrTy, variable->getType()));
       continue;
     }
 
     // Otherwise, build a layout chunk with the size and alignment of
     // the declaration.
     if (llvm::Constant *constant = tryCaptureAsConstant(CGM, CGF, variable)) {
       info.Captures[variable] = CGBlockInfo::Capture::makeConstant(constant);
       continue;
     }
 
+    QualType VT = getCaptureFieldType(*CGF, CI);
+
     // If we have a lifetime qualifier, honor it for capture purposes.
     // That includes *not* copying it if it's __unsafe_unretained.
-    Qualifiers::ObjCLifetime lifetime =
-      variable->getType().getObjCLifetime();
+    Qualifiers::ObjCLifetime lifetime = VT.getObjCLifetime();
     if (lifetime) {
       switch (lifetime) {
       case Qualifiers::OCL_None: llvm_unreachable("impossible");
       case Qualifiers::OCL_ExplicitNone:
       case Qualifiers::OCL_Autoreleasing:
         break;
 
       case Qualifiers::OCL_Strong:
       case Qualifiers::OCL_Weak:
         info.NeedsCopyDispose = true;
       }
 
     // Block pointers require copy/dispose.  So do Objective-C pointers.
-    } else if (variable->getType()->isObjCRetainableType()) {
+    } else if (VT->isObjCRetainableType()) {
       // But honor the inert __unsafe_unretained qualifier, which doesn't
       // actually make it into the type system.
-       if (variable->getType()->isObjCInertUnsafeUnretainedType()) {
+       if (VT->isObjCInertUnsafeUnretainedType()) {
         lifetime = Qualifiers::OCL_ExplicitNone;
       } else {
         info.NeedsCopyDispose = true;
         // used for mrr below.
         lifetime = Qualifiers::OCL_Strong;
       }
 
     // So do types that require non-trivial copy construction.
     } else if (CI.hasCopyExpr()) {
       info.NeedsCopyDispose = true;
       info.HasCXXObject = true;
-      if (!variable->getType()->getAsCXXRecordDecl()->isExternallyVisible())
+      if (!VT->getAsCXXRecordDecl()->isExternallyVisible())
         info.CapturesNonExternalType = true;
 
     // So do C structs that require non-trivial copy construction or
     // destruction.
-    } else if (variable->getType().isNonTrivialToPrimitiveCopy() ==
-                   QualType::PCK_Struct ||
-               variable->getType().isDestructedType() ==
-                   QualType::DK_nontrivial_c_struct) {
+    } else if (VT.isNonTrivialToPrimitiveCopy() == QualType::PCK_Struct ||
+               VT.isDestructedType() == QualType::DK_nontrivial_c_struct) {
       info.NeedsCopyDispose = true;
 
     // And so do types with destructors.
     } else if (CGM.getLangOpts().CPlusPlus) {
-      if (const CXXRecordDecl *record =
-            variable->getType()->getAsCXXRecordDecl()) {
+      if (const CXXRecordDecl *record = VT->getAsCXXRecordDecl()) {
         if (!record->hasTrivialDestructor()) {
           info.HasCXXObject = true;
           info.NeedsCopyDispose = true;
           if (!record->isExternallyVisible())
             info.CapturesNonExternalType = true;
         }
       }
     }
 
-    QualType VT = getCaptureFieldType(*CGF, CI);
     CharUnits size = C.getTypeSizeInChars(VT);
     CharUnits align = C.getDeclAlign(variable);
 
     maxFieldAlign = std::max(maxFieldAlign, align);
 
     llvm::Type *llvmType =
       CGM.getTypes().ConvertTypeForMem(VT);
 
static void findBlockCapturedManagedEntities(
     const CGBlockInfo &BlockInfo, const LangOptions &LangOpts,
     SmallVectorImpl<BlockCaptureManagedEntity> &ManagedCaptures) {
   for (const auto &CI : BlockInfo.getBlockDecl()->captures()) {
     const VarDecl *Variable = CI.getVariable();
     const CGBlockInfo::Capture &Capture = BlockInfo.getCapture(Variable);
     if (Capture.isConstant())
       continue;
 
-    auto CopyInfo =
-       computeCopyInfoForBlockCapture(CI, Variable->getType(), LangOpts);
-    auto DisposeInfo =
-       computeDestroyInfoForBlockCapture(CI, Variable->getType(), LangOpts);
+    QualType VT = Capture.fieldType();
+    auto CopyInfo = computeCopyInfoForBlockCapture(CI, VT, LangOpts);
+    auto DisposeInfo = computeDestroyInfoForBlockCapture(CI, VT, LangOpts);
     if (CopyInfo.first != BlockCaptureEntityKind::None ||
         DisposeInfo.first != BlockCaptureEntityKind::None)
       ManagedCaptures.emplace_back(CopyInfo.first, DisposeInfo.first,
                                    CopyInfo.second, DisposeInfo.second, CI,
                                    Capture);
   }
 
   // Sort the captures by offset.

cfe/trunk/test/CodeGenObjCXX/block-nested-in-lambda.cpp

-// RUN: %clang_cc1 -triple=x86_64-apple-darwin10 -emit-llvm -std=c++11 -fblocks -o - %s | FileCheck %s
-
-// CHECK: %[[BLOCK_CAPTURED0:.*]] = getelementptr inbounds <{ i8*, i32, i32, i8*, %struct.__block_descriptor*, i32*, i32* }>, <{ i8*, i32, i32, i8*, %struct.__block_descriptor*, i32*, i32* }>* %[[BLOCK:.*]], i32 0, i32 5
-// CHECK: %[[V0:.*]] = getelementptr inbounds %[[LAMBDA_CLASS:.*]], %[[LAMBDA_CLASS]]* %[[THIS:.*]], i32 0, i32 0
-// CHECK: %[[V1:.*]] = load i32*, i32** %[[V0]], align 8
-// CHECK: store i32* %[[V1]], i32** %[[BLOCK_CAPTURED0]], align 8
-// CHECK: %[[BLOCK_CAPTURED1:.*]] = getelementptr inbounds <{ i8*, i32, i32, i8*, %struct.__block_descriptor*, i32*, i32* }>, <{ i8*, i32, i32, i8*, %struct.__block_descriptor*, i32*, i32* }>* %[[BLOCK]], i32 0, i32 6
-// CHECK: %[[V2:.*]] = getelementptr inbounds %[[LAMBDA_CLASS]], %[[LAMBDA_CLASS]]* %[[THIS]], i32 0, i32 1
-// CHECK: %[[V3:.*]] = load i32*, i32** %[[V2]], align 8
-// CHECK: store i32* %[[V3]], i32** %[[BLOCK_CAPTURED1]], align 8
-
-void foo1(int &, int &);
-
-void block_in_lambda(int &s1, int &s2) {
-  auto lambda = [&s1, &s2]() {
-    auto block = ^{
-      foo1(s1, s2);
-    };
-    block();
-  };
-
-  lambda();
-}

cfe/trunk/test/CodeGenObjCXX/block-nested-in-lambda.mm

+// RUN: %clang_cc1 -triple=x86_64-apple-darwin10 -emit-llvm -std=c++11 -fblocks -fobjc-arc -o - %s | FileCheck %s
+
+// CHECK: %[[STRUCT_BLOCK_DESCRIPTOR:.*]] = type { i64, i64 }
+
+// CHECK: %[[BLOCK_CAPTURED0:.*]] = getelementptr inbounds <{ i8*, i32, i32, i8*, %struct.__block_descriptor*, i32*, i32* }>, <{ i8*, i32, i32, i8*, %struct.__block_descriptor*, i32*, i32* }>* %[[BLOCK:.*]], i32 0, i32 5
+// CHECK: %[[V0:.*]] = getelementptr inbounds %[[LAMBDA_CLASS:.*]], %[[LAMBDA_CLASS]]* %[[THIS:.*]], i32 0, i32 0
+// CHECK: %[[V1:.*]] = load i32*, i32** %[[V0]], align 8
+// CHECK: store i32* %[[V1]], i32** %[[BLOCK_CAPTURED0]], align 8
+// CHECK: %[[BLOCK_CAPTURED1:.*]] = getelementptr inbounds <{ i8*, i32, i32, i8*, %struct.__block_descriptor*, i32*, i32* }>, <{ i8*, i32, i32, i8*, %struct.__block_descriptor*, i32*, i32* }>* %[[BLOCK]], i32 0, i32 6
+// CHECK: %[[V2:.*]] = getelementptr inbounds %[[LAMBDA_CLASS]], %[[LAMBDA_CLASS]]* %[[THIS]], i32 0, i32 1
+// CHECK: %[[V3:.*]] = load i32*, i32** %[[V2]], align 8
+// CHECK: store i32* %[[V3]], i32** %[[BLOCK_CAPTURED1]], align 8
+
+void foo1(int &, int &);
+
+void block_in_lambda(int &s1, int &s2) {
+  auto lambda = [&s1, &s2]() {
+    auto block = ^{
+      foo1(s1, s2);
+    };
+    block();
+  };
+
+  lambda();
+}
+
+namespace CaptureByReference {
+
+id getObj();
+void use(id);
+
+// Block copy/dispose helpers aren't needed because 'a' is captured by
+// reference.
+
+// CHECK-LABEL: define void @_ZN18CaptureByReference5test0Ev(
+// CHECK-LABEL: define internal void @"_ZZN18CaptureByReference5test0EvENK3$_1clEv"(
+// CHECK: %[[BLOCK_DESCRIPTOR:.*]] = getelementptr inbounds <{ i8*, i32, i32, i8*, %[[STRUCT_BLOCK_DESCRIPTOR]]*, i8** }>, <{ i8*, i32, i32, i8*, %[[STRUCT_BLOCK_DESCRIPTOR]]*, i8** }>* %{{.*}}, i32 0, i32 4
+// CHECK: store %[[STRUCT_BLOCK_DESCRIPTOR]]* bitcast ({ i64, i64, i8*, i64 }* @"__block_descriptor_40_e5_v8@?0ls32l8" to %[[STRUCT_BLOCK_DESCRIPTOR]]*), %[[STRUCT_BLOCK_DESCRIPTOR]]** %[[BLOCK_DESCRIPTOR]], align 8
+
+void test0() {
+  id a = getObj();
+  [&]{ ^{ a = 0; }(); }();
+}
+
+// Block copy/dispose helpers shouldn't have to retain/release 'a' because it
+// is captured by reference.
+
+// CHECK-LABEL: define void @_ZN18CaptureByReference5test1Ev(
+// CHECK-LABEL: define internal void @"_ZZN18CaptureByReference5test1EvENK3$_2clEv"(
+// CHECK: %[[BLOCK_DESCRIPTOR:.*]] = getelementptr inbounds <{ i8*, i32, i32, i8*, %[[STRUCT_BLOCK_DESCRIPTOR]]*, i8*, i8*, i8** }>, <{ i8*, i32, i32, i8*, %[[STRUCT_BLOCK_DESCRIPTOR]]*, i8*, i8*, i8** }>* %{{.*}}, i32 0, i32 4
+// CHECK: store %[[STRUCT_BLOCK_DESCRIPTOR]]* bitcast ({ i64, i64, i8*, i8*, i8*, i64 }* @"__block_descriptor_56_8_32s40s_e5_v8@?0l" to %[[STRUCT_BLOCK_DESCRIPTOR]]*), %[[STRUCT_BLOCK_DESCRIPTOR]]** %[[BLOCK_DESCRIPTOR]], align 8
+
+// CHECK-LABEL: define linkonce_odr hidden void @__copy_helper_block_8_32s40s(
+// CHECK-NOT: call void @objc_storeStrong(
+// CHECK: %[[V4:.*]] = getelementptr inbounds <{ i8*, i32, i32, i8*, %[[STRUCT_BLOCK_DESCRIPTOR]]*, i8*, i8*, i8** }>, <{ i8*, i32, i32, i8*, %[[STRUCT_BLOCK_DESCRIPTOR]]*, i8*, i8*, i8** }>* %{{.*}}, i32 0, i32 5
+// CHECK: %[[V5:.*]] = getelementptr inbounds <{ i8*, i32, i32, i8*, %[[STRUCT_BLOCK_DESCRIPTOR]]*, i8*, i8*, i8** }>, <{ i8*, i32, i32, i8*, %[[STRUCT_BLOCK_DESCRIPTOR]]*, i8*, i8*, i8** }>* %{{.*}}, i32 0, i32 5
+// CHECK: %[[BLOCKCOPY_SRC:.*]] = load i8*, i8** %[[V4]], align 8
+// CHECK: store i8* null, i8** %[[V5]], align 8
+// CHECK: call void @objc_storeStrong(i8** %[[V5]], i8* %[[BLOCKCOPY_SRC]])
+// CHECK: %[[V6:.*]] = getelementptr inbounds <{ i8*, i32, i32, i8*, %[[STRUCT_BLOCK_DESCRIPTOR]]*, i8*, i8*, i8** }>, <{ i8*, i32, i32, i8*, %[[STRUCT_BLOCK_DESCRIPTOR]]*, i8*, i8*, i8** }>* %{{.*}}, i32 0, i32 6
+// CHECK: %[[V7:.*]] = getelementptr inbounds <{ i8*, i32, i32, i8*, %[[STRUCT_BLOCK_DESCRIPTOR]]*, i8*, i8*, i8** }>, <{ i8*, i32, i32, i8*, %[[STRUCT_BLOCK_DESCRIPTOR]]*, i8*, i8*, i8** }>* %{{.*}}, i32 0, i32 6
+// CHECK: %[[BLOCKCOPY_SRC2:.*]] = load i8*, i8** %[[V6]], align 8
+// CHECK: store i8* null, i8** %[[V7]], align 8
+// CHECK: call void @objc_storeStrong(i8** %[[V7]], i8* %[[BLOCKCOPY_SRC2]])
+// CHECK-NOT: call void @objc_storeStrong(
+// CHECK: ret void
+
+// CHECK-LABEL: define linkonce_odr hidden void @__destroy_helper_block_8_32s40s(
+// CHECK: %[[V2:.*]] = getelementptr inbounds <{ i8*, i32, i32, i8*, %[[STRUCT_BLOCK_DESCRIPTOR]]*, i8*, i8*, i8** }>, <{ i8*, i32, i32, i8*, %[[STRUCT_BLOCK_DESCRIPTOR]]*, i8*, i8*, i8** }>* %{{.*}}, i32 0, i32 5
+// CHECK: %[[V3:.*]] = getelementptr inbounds <{ i8*, i32, i32, i8*, %[[STRUCT_BLOCK_DESCRIPTOR]]*, i8*, i8*, i8** }>, <{ i8*, i32, i32, i8*, %[[STRUCT_BLOCK_DESCRIPTOR]]*, i8*, i8*, i8** }>* %{{.*}}, i32 0, i32 6
+// CHECK-NOT: call void @objc_storeStrong(
+// CHECK: call void @objc_storeStrong(i8** %[[V3]], i8* null)
+// CHECK: call void @objc_storeStrong(i8** %[[V2]], i8* null)
+// CHECK-NOT: call void @objc_storeStrong(
+// CHECK: ret void
+
+void test1() {
+  id a = getObj(), b = getObj(), c = getObj();
+  [&a, b, c]{ ^{ a = 0; use(b); use(c); }(); }();
+}
+
+}