[CodeGen][ObjC] Block captures should inherit the type of the captured

field in the enclosing lambda or block.

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

rdar://problem/18586386

Differential Revision: https://reviews.llvm.org/D21104

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

diff --git a/lib/CodeGen/CGBlocks.cpp b/lib/CodeGen/CGBlocks.cpp
index e3658ab9b7624487548d51fc2d0ea9312896a938..94ab61d0e64c9f246cb774080a38511cdfc3c344 100644 (file)
--- a/lib/CodeGen/CGBlocks.cpp
+++ b/lib/CodeGen/CGBlocks.cpp
@@ -199,13 +199,14 @@ namespace {
     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) {
@@ -213,8 +214,8 @@ namespace {
         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});
       }
     }
   };
@@ -363,7 +364,7 @@ static void computeBlockInfo(CodeGenModule &CGM, CodeGenFunction *CGF,
 
     layout.push_back(BlockLayoutChunk(tinfo.second, tinfo.first,
                                       Qualifiers::OCL_None,
-                                      nullptr, llvmType));
+                                      nullptr, llvmType, thisType));
   }
 
   // Next, all the block captures.
@@ -380,7 +381,7 @@ static void computeBlockInfo(CodeGenModule &CGM, CodeGenFunction *CGF,
 
       layout.push_back(BlockLayoutChunk(align, CGM.getPointerSize(),
                                         Qualifiers::OCL_None, &CI,
-                                        CGM.VoidPtrTy));
+                                        CGM.VoidPtrTy, variable->getType()));
       continue;
     }
 
@@ -436,6 +437,14 @@ static void computeBlockInfo(CodeGenModule &CGM, CodeGenFunction *CGF,
     }
 
     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);
     
@@ -444,7 +453,8 @@ static void computeBlockInfo(CodeGenModule &CGM, CodeGenFunction *CGF,
     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.
@@ -775,7 +785,7 @@ llvm::Value *CodeGenFunction::EmitBlockLiteral(const CGBlockInfo &blockInfo) {
     // 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**.
@@ -1033,9 +1043,8 @@ Address CodeGenFunction::GetAddrOfBlockDecl(const VarDecl *variable,
                                  variable->getName());
   }
 
-  if (auto refType = variable->getType()->getAs<ReferenceType>()) {
+  if (auto refType = capture.fieldType()->getAs<ReferenceType>())
     addr = EmitLoadOfReference(addr, refType);
-  }
 
   return addr;
 }

diff --git a/lib/CodeGen/CGBlocks.h b/lib/CodeGen/CGBlocks.h
index aaf06792df97ca25cfd24c95dbeb93c9f538e2a7..80e255f754170cb168e8a9a56f7c8db1bcdbffec 100644 (file)
--- a/lib/CodeGen/CGBlocks.h
+++ b/lib/CodeGen/CGBlocks.h
@@ -159,6 +159,11 @@ public:
     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(); }
@@ -185,10 +190,16 @@ public:
       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;
     }
 

diff --git a/test/CodeGenObjCXX/lambda-expressions.mm b/test/CodeGenObjCXX/lambda-expressions.mm
index 0c309c0a4e615844f46079129ea7b3fdc9f442f4..d987083e7b1c099a8fa1705bc9797b87ef8f5841 100644 (file)
--- a/test/CodeGenObjCXX/lambda-expressions.mm
+++ b/test/CodeGenObjCXX/lambda-expressions.mm
@@ -4,6 +4,8 @@
 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(
@@ -60,6 +62,40 @@ void take_block(void (^block)()) { block(); }
 }
 @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