[CodeGenObjCXX] Fix handling of blocks in lambda.

This fixes a crash that occurs when a block captures a reference that is
captured by its enclosing lambda.

rdar://problem/18586651

Differential Revision: http://reviews.llvm.org/D19536

llvm-svn: 268532

2 files changed, 47 insertions, 26 deletions

diff --git a/clang/lib/CodeGen/CGBlocks.cpp b/clang/lib/CodeGen/CGBlocks.cpp
index d603d21bbe0..1389749195c 100644
--- a/clang/lib/CodeGen/CGBlocks.cpp
+++ b/clang/lib/CodeGen/CGBlocks.cpp
@@ -780,35 +780,34 @@ llvm::Value *CodeGenFunction::EmitBlockLiteral(const CGBlockInfo &blockInfo) {
     // Compute the address of the thing we're going to move into the
     // block literal.
     Address src = Address::invalid();
-    if (BlockInfo && CI.isNested()) {
-      // We need to use the capture from the enclosing block.
-      const CGBlockInfo::Capture &enclosingCapture =
-        BlockInfo->getCapture(variable);
-
-      // This is a [[type]]*, except that a byref entry wil just be an i8**.
-      src = Builder.CreateStructGEP(LoadBlockStruct(),
-                                    enclosingCapture.getIndex(),
-                                    enclosingCapture.getOffset(),
-                                    "block.capture.addr");
-    } else if (blockDecl->isConversionFromLambda()) {
+
+    if (blockDecl->isConversionFromLambda()) {
       // The lambda capture in a lambda's conversion-to-block-pointer is
       // special; we'll simply emit it directly.
       src = Address::invalid();
-    } else {
-      // Just look it up in the locals map, which will give us back a
-      // [[type]]*.  If that doesn't work, do the more elaborate DRE
-      // emission.
-      auto it = LocalDeclMap.find(variable);
-      if (it != LocalDeclMap.end()) {
-        src = it->second;
+    } else if (CI.isByRef()) {
+      if (BlockInfo && CI.isNested()) {
+        // We need to use the capture from the enclosing block.
+        const CGBlockInfo::Capture &enclosingCapture =
+            BlockInfo->getCapture(variable);
+
+        // This is a [[type]]*, except that a byref entry wil just be an i8**.
+        src = Builder.CreateStructGEP(LoadBlockStruct(),
+                                      enclosingCapture.getIndex(),
+                                      enclosingCapture.getOffset(),
+                                      "block.capture.addr");
       } else {
-        DeclRefExpr declRef(
-            const_cast<VarDecl *>(variable),
-            /*RefersToEnclosingVariableOrCapture*/ CI.isNested(), type,
-            VK_LValue, SourceLocation());
-        src = EmitDeclRefLValue(&declRef).getAddress();
+        auto I = LocalDeclMap.find(variable);
+        assert(I != LocalDeclMap.end());
+        src = I->second;
       }
-    }
+    } else {
+      DeclRefExpr declRef(const_cast<VarDecl *>(variable),
+                          /*RefersToEnclosingVariableOrCapture*/ CI.isNested(),
+                          type.getNonReferenceType(), VK_LValue,
+                          SourceLocation());
+      src = EmitDeclRefLValue(&declRef).getAddress();
+    };
 
     // For byrefs, we just write the pointer to the byref struct into
     // the block field.  There's no need to chase the forwarding
@@ -842,8 +841,7 @@ llvm::Value *CodeGenFunction::EmitBlockLiteral(const CGBlockInfo &blockInfo) {
 
     // If it's a reference variable, copy the reference into the block field.
     } else if (type->isReferenceType()) {
-      llvm::Value *ref = Builder.CreateLoad(src, "ref.val");
-      Builder.CreateStore(ref, blockField);
+      Builder.CreateStore(src.getPointer(), blockField);
 
     // If this is an ARC __strong block-pointer variable, don't do a
     // block copy.
diff --git a/clang/test/CodeGenObjCXX/block-nested-in-lambda.cpp b/clang/test/CodeGenObjCXX/block-nested-in-lambda.cpp
new file mode 100644
index 00000000000..51b7abf034e
--- /dev/null
+++ b/clang/test/CodeGenObjCXX/block-nested-in-lambda.cpp
@@ -0,0 +1,23 @@
+// 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();
+}