[clang] Fix self-capturing __block variables

clang/include/clang/AST/Decl.h

@@ -1103,6 +1103,11 @@ class VarDecl : public DeclaratorDecl, public Redeclarable<VarDecl> {
 
     LLVM_PREFERRED_TYPE(bool)
     unsigned IsCXXCondDecl : 1;
+
+    /// Whether this is a __block variable that is captured by a block
+    /// referenced in its own initializer.
+    LLVM_PREFERRED_TYPE(bool)
+    unsigned CapturedByOwnInit : 1;
   };
 
   union {
@@ -1588,10 +1593,23 @@ class VarDecl : public DeclaratorDecl, public Redeclarable<VarDecl> {
   /// escaping block.
   bool isNonEscapingByref() const;
 
-  void setEscapingByref() {
-    NonParmVarDeclBits.EscapingByref = true;
+  void setEscapingByref();
+
+  /// Indicates this is a __block variable that is captured by a block
+  /// referenced in its own initializer.
+  bool isCapturedByOwnInit() const {
+    return !isa<ParmVarDecl>(this) && NonParmVarDeclBits.CapturedByOwnInit;
   }
 
+  void setCapturedByOwnInit() {
+    assert(!isa<ParmVarDecl>(this));
+    NonParmVarDeclBits.CapturedByOwnInit = true;
+  }
+
+  /// Check if this is a __block variable which needs to be initialized
+  /// directly on the heap.
+  bool needsInitOnHeap() const;
+
   bool isCXXCondDecl() const {
     return isa<ParmVarDecl>(this) ? false : NonParmVarDeclBits.IsCXXCondDecl;
   }

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -2358,6 +2358,12 @@ def note_var_fixit_add_initialization : Note<
 def note_uninit_fixit_remove_cond : Note<
   "remove the %select{'%1' if its condition|condition if it}0 "
   "is always %select{false|true}2">;
+def warn_implicit_block_var_alloc : Warning<
+  "variable %q0 will be initialized in a heap allocation">,
+  InGroup<DiagGroup<"implicit-block-var-alloc">>, DefaultIgnore;
+def note_because_captured_by_block : Note<
+  "because it is captured by a block used in its own initializer">;
+
 def err_init_incomplete_type : Error<"initialization of incomplete type %0">;
 def err_list_init_in_parens : Error<
   "cannot initialize %select{non-class|reference}0 type %1 with a "

clang/lib/AST/Decl.cpp

@@ -2680,6 +2680,19 @@ bool VarDecl::isNonEscapingByref() const {
   return hasAttr<BlocksAttr>() && !NonParmVarDeclBits.EscapingByref;
 }
 
+void VarDecl::setEscapingByref() {
+  assert(!isa<ParmVarDecl>(this) && hasAttr<BlocksAttr>());
+  NonParmVarDeclBits.EscapingByref = true;
+}
+
+bool VarDecl::needsInitOnHeap() const {
+  // We need direct initialization on the heap for self-capturing __block
+  // variables of types that cause CodeGenFunction::getEvaluationKind to return
+  // TEK_Aggregate.  The only such types that can be captured are records.
+  return isCapturedByOwnInit() &&
+         getType().getAtomicUnqualifiedType()->isRecordType();
+}
+
 bool VarDecl::hasDependentAlignment() const {
   QualType T = getType();
   return T->isDependentType() || T->isUndeducedType() ||

clang/lib/CodeGen/CGBlocks.cpp

@@ -2336,7 +2336,7 @@ generateByrefCopyHelper(CodeGenFunction &CGF, const BlockByrefInfo &byrefInfo,
     // Create a scope with an artificial location for the body of this function.
   auto AL = ApplyDebugLocation::CreateArtificial(CGF);
 
-  if (generator.needsCopy()) {
+  if (generator.needsCopy() && !byrefInfo.ForInitOnHeap) {
     // dst->x
     Address destField = CGF.GetAddrOfLocalVar(&Dst);
     destField = Address(CGF.Builder.CreateLoad(destField), byrefInfo.Type,
@@ -2405,9 +2405,32 @@ generateByrefDisposeHelper(CodeGenFunction &CGF,
     Address addr = CGF.GetAddrOfLocalVar(&Src);
     addr = Address(CGF.Builder.CreateLoad(addr), byrefInfo.Type,
                    byrefInfo.ByrefAlignment);
-    addr = CGF.emitBlockByrefAddress(addr, byrefInfo, false, "object");
 
+    llvm::BasicBlock *skipBB = nullptr;
+    if (byrefInfo.ForInitOnHeap) {
+      // For objects initialized on the heap, a separate field tracks whether
+      // the object has completed initialization.  If it's still false here,
+      // then initialization threw an exception, so we don't want to run the
+      // destructor.
+      skipBB = CGF.createBasicBlock("skip");
+      llvm::BasicBlock *runBB = CGF.createBasicBlock("run");
+
+      Address initializedAddr = CGF.Builder.CreateStructGEP(
+          addr, byrefInfo.IndexOfInitializedFlag, "byref.initialized");
+      // This is safe to load non-atomically because the store of true (if any)
+      // happens while the creating function still holds a reference.
+      llvm::Value *initialized =
+          CGF.Builder.CreateFlagLoad(initializedAddr.emitRawPointer(CGF));
+
+      CGF.Builder.CreateCondBr(initialized, runBB, skipBB);
+      CGF.EmitBlock(runBB);
+    }
+
+    addr = CGF.emitBlockByrefAddress(addr, byrefInfo, false, "object");
     generator.emitDispose(CGF, addr);
+
+    if (skipBB)
+      CGF.EmitBlock(skipBB);
   }
 
   CGF.FinishFunction();
@@ -2503,15 +2526,14 @@ CodeGenFunction::buildByrefHelpers(llvm::StructType &byrefType,
 
     // ARC __strong __block variables need to be retained.
     case Qualifiers::OCL_Strong:
-      // Block pointers need to be copied, and there's no direct
-      // transfer possible.
       if (type->isBlockPointerType()) {
+        // Block pointers need to be copied, and there's no direct
+        // transfer possible.
         return ::buildByrefHelpers(CGM, byrefInfo,
                                    ARCStrongBlockByrefHelpers(valueAlignment));
-
-      // Otherwise, we transfer ownership of the retain from the stack
-      // to the heap.
       } else {
+        // Otherwise, we transfer ownership of the retain from the stack
+        // to the heap.
         return ::buildByrefHelpers(CGM, byrefInfo,
                                    ARCStrongByrefHelpers(valueAlignment));
       }
@@ -2623,6 +2645,14 @@ const BlockByrefInfo &CodeGenFunction::getBlockByrefInfo(const VarDecl *D) {
     size += CharUnits::fromQuantity(PointerSizeInBytes);
   }
 
+  unsigned indexOfInitializedFlag = UINT_MAX;
+  if (D->needsInitOnHeap()) {
+    /// uint8_t initialized;
+    types.push_back(Int8Ty);
+    size += CharUnits::fromQuantity(1);
+    indexOfInitializedFlag = types.size() - 1;
+  }
+
   // T x;
   llvm::Type *varTy = ConvertTypeForMem(Ty);
 
@@ -2652,12 +2682,56 @@ const BlockByrefInfo &CodeGenFunction::getBlockByrefInfo(const VarDecl *D) {
   info.FieldIndex = types.size() - 1;
   info.FieldOffset = varOffset;
   info.ByrefAlignment = std::max(varAlign, getPointerAlign());
+  info.ForInitOnHeap = D->needsInitOnHeap();
+  info.IndexOfInitializedFlag = indexOfInitializedFlag;
 
   auto pair = BlockByrefInfos.insert({D, info});
   assert(pair.second && "info was inserted recursively?");
   return pair.first->second;
 }
 
+/// Allocate a byref on the heap (but don't initialize it yet).
+void CodeGenFunction::emitByrefHeapAlloc(llvm::Value *stackByref,
+                                         QualType declType) {
+  // The object itself is initialized directly on the heap.  But for ABI
+  // backwards compatibility reasons, we have to set up a fake byref struct on
+  // the stack (with the structural components initialized but not the object
+  // itself), then call _Block_object_assign to move it to the heap (which is
+  // safe because we forced a no-op copy helper), then call
+  // _Block_object_dispose to release the extra ref from _Block_object_assign.
+  BlockFieldFlags flags = BLOCK_FIELD_IS_BYREF;
+  RawAddress heapByref =
+      CreateDefaultAlignTempAlloca(stackByref->getType(), "initOnHeap.ptr");
+
+  llvm::Value *args[] = {heapByref.getPointer(), stackByref,
+                         llvm::ConstantInt::get(Int32Ty, flags.getBitMask())};
+  EmitNounwindRuntimeCall(CGM.getBlockObjectAssign(), args);
+
+  BuildBlockRelease(stackByref, flags, false);
+
+  // We need to add a cleanup as soon as the allocation happens, to ensure the
+  // allocation is freed if initialization throws.
+  if (CGM.getLangOpts().getGC() != LangOptions::GCOnly) {
+    enterByrefCleanup(NormalAndEHCleanup, heapByref, BLOCK_FIELD_IS_BYREF,
+                      /*LoadBlockVarAddr*/ true,
+                      cxxDestructorCanThrow(declType));
+  }
+}
+
+/// Set the initialized flag of a heap-initialized byref.
+void CodeGenFunction::emitByrefMarkInitialized(
+    const AutoVarEmission &emission) {
+  auto &info = getBlockByrefInfo(emission.Variable);
+  Address forwardingAddr =
+      Builder.CreateStructGEP(emission.Addr, 1, "forwarding");
+  Address heapByref(Builder.CreateLoad(forwardingAddr), info.Type,
+                    info.ByrefAlignment);
+
+  Address initializedAddr = Builder.CreateStructGEP(
+      heapByref, info.IndexOfInitializedFlag, "byref.initialized");
+  Builder.CreateFlagStore(true, initializedAddr.emitRawPointer(*this));
+}
+
 /// Initialize the structural components of a __block variable, i.e.
 /// everything but the actual object.
 void CodeGenFunction::emitByrefStructureInit(const AutoVarEmission &emission) {
@@ -2699,8 +2773,8 @@ void CodeGenFunction::emitByrefStructureInit(const AutoVarEmission &emission) {
   storeHeaderField(V, getPointerSize(), "byref.isa");
 
   // Store the address of the variable into its own forwarding pointer.
-  storeHeaderField(addr.emitRawPointer(*this), getPointerSize(),
-                   "byref.forwarding");
+  llvm::Value *pointer = addr.emitRawPointer(*this);
+  storeHeaderField(pointer, getPointerSize(), "byref.forwarding");
 
   // Blocks ABI:
   //   c) the flags field is set to either 0 if no helper functions are
@@ -2764,6 +2838,12 @@ void CodeGenFunction::emitByrefStructureInit(const AutoVarEmission &emission) {
     auto layoutInfo = CGM.getObjCRuntime().BuildByrefLayout(CGM, type);
     storeHeaderField(layoutInfo, getPointerSize(), "byref.layout");
   }
+
+  if (emission.NeedsInitOnHeap) {
+    V = llvm::ConstantInt::get(Int8Ty, 0);
+    storeHeaderField(V, CharUnits::fromQuantity(1), "byref.initialized");
+    emitByrefHeapAlloc(pointer, type);
+  }
 }
 
 void CodeGenFunction::BuildBlockRelease(llvm::Value *V, BlockFieldFlags flags,

clang/lib/CodeGen/CGBlocks.h

@@ -136,9 +136,14 @@ inline BlockFieldFlags operator|(BlockFieldFlag_t l, BlockFieldFlag_t r) {
 class BlockByrefInfo {
 public:
   llvm::StructType *Type;
+  /// Index of the field for the variable itself.
   unsigned FieldIndex;
   CharUnits ByrefAlignment;
   CharUnits FieldOffset;
+  /// Whether this will be initialized directly on the heap.
+  bool ForInitOnHeap;
+  /// If ForInitOnHeap is true, index of the 'initialized' field.
+  unsigned IndexOfInitializedFlag;
 };
 
 /// Represents a type of copy/destroy operation that should be performed for an