[CodeGen][ObjC] Block captures should inherit the type of the captured field in the enclosing lambda or block

[ahatanak]

Summary

This patch fixes a bug in code-gen where it uses the type of the declared variable rather than the type of the capture of the enclosing lambda or block for the block capture. For example, in the following function, code-gen currently uses i32* for the block capture "a" because "a" is passed to foo1 as a reference, but it should use i32 since the enclosing lambda captures "a" by value.

void foo1(int &a) {
  auto lambda = [a]{
    auto block1 = ^{
      i = a;
    };
    block1();
  };
  lambda();
}

[ahatanak]

Rebase.

[rjmccall]

LGTM.

cfe/trunk/lib/CodeGen/CGBlocks.h

public:
   /// The field index of 'this' within the block, if there is one.
   unsigned CXXThisIndex;
 
   class Capture {
     uintptr_t Data;
     EHScopeStack::stable_iterator Cleanup;
     CharUnits::QuantityType Offset;
 
+    /// Type of the capture field. Normally, this is identical to the type of
+    /// the capture's VarDecl, but can be different if there is an enclosing
+    /// lambda.
+    QualType FieldType;
+
   public:
     bool isIndex() const { return (Data & 1) != 0; }
     bool isConstant() const { return !isIndex(); }
 
     unsigned getIndex() const {
       assert(isIndex());
       return Data >> 1;
     }
void setCleanup(EHScopeStack::stable_iterator cleanup) {
       Cleanup = cleanup;
     }
 
     llvm::Value *getConstant() const {
       assert(isConstant());
       return reinterpret_cast<llvm::Value*>(Data);
     }
 
-    static Capture makeIndex(unsigned index, CharUnits offset) {
+    QualType fieldType() const {
+      return FieldType;
+    }
+
+    static Capture makeIndex(unsigned index, CharUnits offset,
+                             QualType FieldType) {
       Capture v;
       v.Data = (index << 1) | 1;
       v.Offset = offset.getQuantity();
+      v.FieldType = FieldType;
       return v;
     }
 
     static Capture makeConstant(llvm::Value *value) {
       Capture v;
       v.Data = reinterpret_cast<uintptr_t>(value);
       return v;
     }    

cfe/trunk/lib/CodeGen/CGBlocks.cpp

 namespace {
   /// A chunk of data that we actually have to capture in the block.
   struct BlockLayoutChunk {
     CharUnits Alignment;
     CharUnits Size;
     Qualifiers::ObjCLifetime Lifetime;
     const BlockDecl::Capture *Capture; // null for 'this'
     llvm::Type *Type;
+    QualType FieldType;
 
     BlockLayoutChunk(CharUnits align, CharUnits size,
                      Qualifiers::ObjCLifetime lifetime,
                      const BlockDecl::Capture *capture,
-                     llvm::Type *type)
+                     llvm::Type *type, QualType fieldType)
       : Alignment(align), Size(size), Lifetime(lifetime),
-        Capture(capture), Type(type) {}
+        Capture(capture), Type(type), FieldType(fieldType) {}
 
     /// Tell the block info that this chunk has the given field index.
     void setIndex(CGBlockInfo &info, unsigned index, CharUnits offset) {
       if (!Capture) {
         info.CXXThisIndex = index;
         info.CXXThisOffset = offset;
       } else {
-        info.Captures.insert({Capture->getVariable(),
-                              CGBlockInfo::Capture::makeIndex(index, offset)});
+        auto C = CGBlockInfo::Capture::makeIndex(index, offset, FieldType);
+        info.Captures.insert({Capture->getVariable(), C});
       }
     }
   };
 
   /// Order by 1) all __strong together 2) next, all byfref together 3) next,
   /// all __weak together. Preserve descending alignment in all situations.
   bool operator<(const BlockLayoutChunk &left, const BlockLayoutChunk &right) {
     if (left.Alignment != right.Alignment)
if (block->capturesCXXThis()) {
     // don't assume standard pointer size/align.
     llvm::Type *llvmType = CGM.getTypes().ConvertType(thisType);
     std::pair<CharUnits,CharUnits> tinfo
       = CGM.getContext().getTypeInfoInChars(thisType);
     maxFieldAlign = std::max(maxFieldAlign, tinfo.second);
 
     layout.push_back(BlockLayoutChunk(tinfo.second, tinfo.first,
                                       Qualifiers::OCL_None,
-                                      nullptr, llvmType));
+                                      nullptr, llvmType, thisType));
   }
 
   // Next, all the block captures.
   for (const auto &CI : block->captures()) {
     const VarDecl *variable = CI.getVariable();
 
     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);
 
       layout.push_back(BlockLayoutChunk(align, CGM.getPointerSize(),
                                         Qualifiers::OCL_None, &CI,
-                                        CGM.VoidPtrTy));
+                                        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;
// And so do types with destructors.
         if (!record->hasTrivialDestructor()) {
           info.HasCXXObject = true;
           info.NeedsCopyDispose = true;
         }
       }
     }
 
     QualType VT = variable->getType();
+
+    // If the variable is captured by an enclosing block or lambda expression,
+    // use the type of the capture field.
+    if (CGF->BlockInfo && CI.isNested())
+      VT = CGF->BlockInfo->getCapture(variable).fieldType();
+    else if (auto *FD = CGF->LambdaCaptureFields.lookup(variable))
+      VT = FD->getType();
+
     CharUnits size = C.getTypeSizeInChars(VT);
     CharUnits align = C.getDeclAlign(variable);
     
     maxFieldAlign = std::max(maxFieldAlign, align);
 
     llvm::Type *llvmType =
       CGM.getTypes().ConvertTypeForMem(VT);
     
-    layout.push_back(BlockLayoutChunk(align, size, lifetime, &CI, llvmType));
+    layout.push_back(
+        BlockLayoutChunk(align, size, lifetime, &CI, llvmType, VT));
   }
 
   // If that was everything, we're done here.
   if (layout.empty()) {
     info.StructureType =
       llvm::StructType::get(CGM.getLLVMContext(), elementTypes, true);
     info.CanBeGlobal = true;
     return;
llvm::Value *CodeGenFunction::EmitBlockLiteral(const CGBlockInfo &blockInfo) {
   // Next, captured variables.
   for (const auto &CI : blockDecl->captures()) {
     const VarDecl *variable = CI.getVariable();
     const CGBlockInfo::Capture &capture = blockInfo.getCapture(variable);
 
     // Ignore constant captures.
     if (capture.isConstant()) continue;
 
-    QualType type = variable->getType();
+    QualType type = capture.fieldType();
 
     // 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.
if (isByRef) {
 
     auto byrefPointerType = llvm::PointerType::get(byrefInfo.Type, 0);
     addr = Builder.CreateBitCast(addr, byrefPointerType, "byref.addr");
 
     addr = emitBlockByrefAddress(addr, byrefInfo, /*follow*/ true,
                                  variable->getName());
   }
 
-  if (auto refType = variable->getType()->getAs<ReferenceType>()) {
+  if (auto refType = capture.fieldType()->getAs<ReferenceType>())
     addr = EmitLoadOfReference(addr, refType);
-  }
 
   return addr;
 }
 
 llvm::Constant *
 CodeGenModule::GetAddrOfGlobalBlock(const BlockExpr *blockExpr,
                                     const char *name) {
   CGBlockInfo blockInfo(blockExpr->getBlockDecl(), name);

cfe/trunk/test/CodeGenObjCXX/lambda-expressions.mm

 // RUN: %clang_cc1 -triple x86_64-apple-darwin10.0.0 -emit-llvm -o - %s -fexceptions -std=c++11 -fblocks -fobjc-arc | FileCheck -check-prefix=ARC %s
 // RUN: %clang_cc1 -triple x86_64-apple-darwin10.0.0 -emit-llvm -o - %s -fexceptions -std=c++11 -fblocks | FileCheck -check-prefix=MRC %s
 
 typedef int (^fp)();
 fp f() { auto x = []{ return 3; }; return x; }
 
+// ARC: %[[LAMBDACLASS:.*]] = type { i32 }
+
 // MRC: @OBJC_METH_VAR_NAME{{.*}} = private global [5 x i8] c"copy\00"
 // MRC: @OBJC_METH_VAR_NAME{{.*}} = private global [12 x i8] c"autorelease\00"
 // MRC-LABEL: define i32 ()* @_Z1fv(
 // MRC-LABEL: define internal i32 ()* @"_ZZ1fvENK3$_0cvU13block_pointerFivEEv"
 // MRC: store i8* bitcast (i8** @_NSConcreteStackBlock to i8*)
 // MRC: store i8* bitcast (i32 (i8*)* @"___ZZ1fvENK3$_0cvU13block_pointerFivEEv_block_invoke" to i8*)
 // MRC: call i32 ()* bitcast (i8* (i8*, i8*, ...)* @objc_msgSend to i32 ()* (i8*, i8*)*)
 // MRC: call i32 ()* bitcast (i8* (i8*, i8*, ...)* @objc_msgSend to i32 ()* (i8*, i8*)*)
take_block(^{
                   [super test];
               });
           });
       });
    });
 }
 @end
 
+// ARC: define void @_ZN13LambdaCapture4foo1ERi(i32* dereferenceable(4) %{{.*}})
+// ARC:   %[[CAPTURE0:.*]] = getelementptr inbounds %[[LAMBDACLASS]], %[[LAMBDACLASS]]* %{{.*}}, i32 0, i32 0
+// ARC:   store i32 %{{.*}}, i32* %[[CAPTURE0]]
+
+// ARC: define internal void @"_ZZN13LambdaCapture4foo1ERiENK3$_3clEv"(%[[LAMBDACLASS]]* %{{.*}})
+// ARC:   %[[BLOCK:.*]] = alloca <{ i8*, i32, i32, i8*, %struct.__block_descriptor*, i32 }>
+// ARC:   %[[CAPTURE1:.*]] = getelementptr inbounds <{ i8*, i32, i32, i8*, %struct.__block_descriptor*, i32 }>, <{ i8*, i32, i32, i8*, %struct.__block_descriptor*, i32 }>* %[[BLOCK]], i32 0, i32 5
+// ARC:   store i32 %{{.*}}, i32* %[[CAPTURE1]]
+
+// ARC: define internal void @"___ZZN13LambdaCapture4foo1ERiENK3$_3clEv_block_invoke"
+// ARC:   %[[CAPTURE2:.*]] = getelementptr inbounds <{ i8*, i32, i32, i8*, %struct.__block_descriptor*, i32 }>, <{ i8*, i32, i32, i8*, %struct.__block_descriptor*, i32 }>* %{{.*}}, i32 0, i32 5
+// ARC:   store i32 %{{.*}}, i32* %[[CAPTURE2]]
+
+// ARC: define internal void @"___ZZN13LambdaCapture4foo1ERiENK3$_3clEv_block_invoke_2"(i8* %{{.*}})
+// ARC:   %[[CAPTURE3:.*]] = getelementptr inbounds <{ i8*, i32, i32, i8*, %struct.__block_descriptor*, i32 }>, <{ i8*, i32, i32, i8*, %struct.__block_descriptor*, i32 }>* %{{.*}}, i32 0, i32 5
+// ARC:   %[[V1:.*]] = load i32, i32* %[[CAPTURE3]]
+// ARC:   store i32 %[[V1]], i32* @_ZN13LambdaCapture1iE
+
+namespace LambdaCapture {
+  int i;
+  void foo1(int &a) {
+    auto lambda = [a]{
+      auto block1 = ^{
+        auto block2 = ^{
+          i = a;
+        };
+        block2();
+      };
+      block1();
+    };
+    lambda();
+  }
+}
+
 // ARC-LABEL: define linkonce_odr i32 ()* @_ZZNK13StaticMembersIfE1fMUlvE_clEvENKUlvE_cvU13block_pointerFivEEv
 
 // Check lines for BlockInLambda test below
 // ARC-LABEL: define internal i32 @___ZZN13BlockInLambda1X1fEvENKUlvE_clEv_block_invoke
 // ARC: [[Y:%.*]] = getelementptr inbounds %"struct.BlockInLambda::X", %"struct.BlockInLambda::X"* {{.*}}, i32 0, i32 1
 // ARC-NEXT: [[YVAL:%.*]] = load i32, i32* [[Y]], align 4
 // ARC-NEXT: ret i32 [[YVAL]]