[CodeGenObjCXX] Fix handling of blocks in lambda

[ahatanak]

Summary

This fixes a crash that occurs when a block nested in a c++ lambda captures a reference that is captured by the enclosing lambda.

rdar://problem/18586651

[rjmccall]

lib/CodeGen/CGBlocks.cpp

806 ↗

(On Diff #55022)

Hmm. It's become increasingly clear that my original decision to expand the decl-ref logic here was a mistake. Please just turn this entire block of conditions into:

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 {
  DeclRefExpr declRef(
          const_cast<VarDecl *>(variable),
          /*RefersToEnclosingVariableOrCapture*/ CI.isNested(), type,
          VK_LValue, SourceLocation());
  src = EmitDeclRefLValue(&declRef).getAddress();
}

[ahatanak]

lib/CodeGen/CGBlocks.cpp

806 ↗

(On Diff #55022)

Seems like a lot of regression tests fail if I try the code above. Currently investigating.

[ahatanak]

Apply John's patch with a few modifications.

I had to special case __block variables and remove the line creating a load from a reference to fix the failing regression tests.

[rjmccall]

Yes, that looks good, thanks.

cfe/trunk/lib/CodeGen/CGBlocks.cpp

for (const auto &CI : blockDecl->captures()) {
     // This will be a [[type]]*, except that a byref entry will just be
     // an i8**.
     Address blockField =
       projectField(capture.getIndex(), capture.getOffset(), "block.captured");
 
     // Compute the address of the thing we're going to move into the
     // block literal.
     Address src = Address::invalid();
+
+    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 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 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 {
-        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;
+      }
+      DeclRefExpr declRef(const_cast<VarDecl *>(variable),
+                          /*RefersToEnclosingVariableOrCapture*/ CI.isNested(),
+                          type.getNonReferenceType(), VK_LValue,
+                          SourceLocation());
+    };
 
     // For byrefs, we just write the pointer to the byref struct into
     // the block field.  There's no need to chase the forwarding
     // pointer at this point, since we're building something that will
     // live a shorter life than the stack byref anyway.
     if (CI.isByRef()) {
       // Get a void* that points to the byref struct.
       llvm::Value *byrefPointer;
// If we have a copy constructor, evaluate that into the block field.
                                   AggValueSlot::IsNotAliased);
         EmitAggExpr(copyExpr, Slot);
       } else {
         EmitSynthesizedCXXCopyCtor(blockField, src, copyExpr);
       }
 
     // 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.
     //
     // TODO: this can be generalized into the normal initialization logic:
     // we should never need to do a block-copy when initializing a local
     // variable, because the local variable's lifetime should be strictly
     // contained within the stack block's.

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();
+}