[CodeGen] Simplify the way lifetime of block captures is extended

Rather than pushing inactive cleanups for the block captures at the
entry of a full expression and activating them during the creation of
the block literal, just call pushLifetimeExtendedDestroy to ensure the
cleanups are popped at the end of the scope enclosing the block
expression.

rdar://problem/63996471

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

Author: ahatanaka

clang/lib/CodeGen/CGBlocks.cpp

@@ -36,7 +36,7 @@ CGBlockInfo::CGBlockInfo(const BlockDecl *block, StringRef name)
   : Name(name), CXXThisIndex(0), CanBeGlobal(false), NeedsCopyDispose(false),
     HasCXXObject(false), UsesStret(false), HasCapturedVariableLayout(false),
     CapturesNonExternalType(false), LocalAddress(Address::invalid()),
-    StructureType(nullptr), Block(block), DominatingIP(nullptr) {
+    StructureType(nullptr), Block(block) {
 
   // Skip asm prefix, if any.  'name' is usually taken directly from
   // the mangled name of the enclosing function.
@@ -775,151 +775,23 @@ static void computeBlockInfo(CodeGenModule &CGM, CodeGenFunction *CGF,
     llvm::StructType::get(CGM.getLLVMContext(), elementTypes, true);
 }
 
-/// Enter the scope of a block.  This should be run at the entrance to
-/// a full-expression so that the block's cleanups are pushed at the
-/// right place in the stack.
-static void enterBlockScope(CodeGenFunction &CGF, BlockDecl *block) {
-  assert(CGF.HaveInsertPoint());
-
-  // Allocate the block info and place it at the head of the list.
-  CGBlockInfo &blockInfo =
-    *new CGBlockInfo(block, CGF.CurFn->getName());
-  blockInfo.NextBlockInfo = CGF.FirstBlockInfo;
-  CGF.FirstBlockInfo = &blockInfo;
-
-  // Compute information about the layout, etc., of this block,
-  // pushing cleanups as necessary.
-  computeBlockInfo(CGF.CGM, &CGF, blockInfo);
-
-  // Nothing else to do if it can be global.
-  if (blockInfo.CanBeGlobal) return;
-
-  // Make the allocation for the block.
-  blockInfo.LocalAddress = CGF.CreateTempAlloca(blockInfo.StructureType,
-                                                blockInfo.BlockAlign, "block");
-
-  // If there are cleanups to emit, enter them (but inactive).
-  if (!blockInfo.NeedsCopyDispose) return;
-
-  // Walk through the captures (in order) and find the ones not
-  // captured by constant.
-  for (const auto &CI : block->captures()) {
-    // Ignore __block captures; there's nothing special in the
-    // on-stack block that we need to do for them.
-    if (CI.isByRef()) continue;
-
-    // Ignore variables that are constant-captured.
-    const VarDecl *variable = CI.getVariable();
-    CGBlockInfo::Capture &capture = blockInfo.getCapture(variable);
-    if (capture.isConstant()) continue;
-
-    // Ignore objects that aren't destructed.
-    QualType VT = getCaptureFieldType(CGF, CI);
-    QualType::DestructionKind dtorKind = VT.isDestructedType();
-    if (dtorKind == QualType::DK_none) continue;
-
-    CodeGenFunction::Destroyer *destroyer;
-
-    // Block captures count as local values and have imprecise semantics.
-    // They also can't be arrays, so need to worry about that.
-    //
-    // For const-qualified captures, emit clang.arc.use to ensure the captured
-    // object doesn't get released while we are still depending on its validity
-    // within the block.
-    if (VT.isConstQualified() &&
-        VT.getObjCLifetime() == Qualifiers::OCL_Strong &&
-        CGF.CGM.getCodeGenOpts().OptimizationLevel != 0) {
-      assert(CGF.CGM.getLangOpts().ObjCAutoRefCount &&
-             "expected ObjC ARC to be enabled");
-      destroyer = CodeGenFunction::emitARCIntrinsicUse;
-    } else if (dtorKind == QualType::DK_objc_strong_lifetime) {
-      destroyer = CodeGenFunction::destroyARCStrongImprecise;
-    } else {
-      destroyer = CGF.getDestroyer(dtorKind);
-    }
-
-    // GEP down to the address.
-    Address addr =
-        CGF.Builder.CreateStructGEP(blockInfo.LocalAddress, capture.getIndex());
-
-    // We can use that GEP as the dominating IP.
-    if (!blockInfo.DominatingIP)
-      blockInfo.DominatingIP = cast<llvm::Instruction>(addr.getPointer());
-
-    CleanupKind cleanupKind = InactiveNormalCleanup;
-    bool useArrayEHCleanup = CGF.needsEHCleanup(dtorKind);
-    if (useArrayEHCleanup)
-      cleanupKind = InactiveNormalAndEHCleanup;
-
-    CGF.pushDestroy(cleanupKind, addr, VT,
-                    destroyer, useArrayEHCleanup);
-
-    // Remember where that cleanup was.
-    capture.setCleanup(CGF.EHStack.stable_begin());
-  }
-}
-
-/// Enter a full-expression with a non-trivial number of objects to
-/// clean up.
-void CodeGenFunction::enterNonTrivialFullExpression(const FullExpr *E) {
-  if (const auto EWC = dyn_cast<ExprWithCleanups>(E)) {
-    assert(EWC->getNumObjects() != 0);
-    for (const ExprWithCleanups::CleanupObject &C : EWC->getObjects())
-      if (auto *BD = C.dyn_cast<BlockDecl *>())
-        enterBlockScope(*this, BD);
-  }
-}
-
-/// Find the layout for the given block in a linked list and remove it.
-static CGBlockInfo *findAndRemoveBlockInfo(CGBlockInfo **head,
-                                           const BlockDecl *block) {
-  while (true) {
-    assert(head && *head);
-    CGBlockInfo *cur = *head;
-
-    // If this is the block we're looking for, splice it out of the list.
-    if (cur->getBlockDecl() == block) {
-      *head = cur->NextBlockInfo;
-      return cur;
-    }
-
-    head = &cur->NextBlockInfo;
-  }
-}
-
-/// Destroy a chain of block layouts.
-void CodeGenFunction::destroyBlockInfos(CGBlockInfo *head) {
-  assert(head && "destroying an empty chain");
-  do {
-    CGBlockInfo *cur = head;
-    head = cur->NextBlockInfo;
-    delete cur;
-  } while (head != nullptr);
-}
-
 /// Emit a block literal expression in the current function.
 llvm::Value *CodeGenFunction::EmitBlockLiteral(const BlockExpr *blockExpr) {
   // If the block has no captures, we won't have a pre-computed
   // layout for it.
-  if (!blockExpr->getBlockDecl()->hasCaptures()) {
+  if (!blockExpr->getBlockDecl()->hasCaptures())
     // The block literal is emitted as a global variable, and the block invoke
     // function has to be extracted from its initializer.
-    if (llvm::Constant *Block = CGM.getAddrOfGlobalBlockIfEmitted(blockExpr)) {
+    if (llvm::Constant *Block = CGM.getAddrOfGlobalBlockIfEmitted(blockExpr))
       return Block;
-    }
-    CGBlockInfo blockInfo(blockExpr->getBlockDecl(), CurFn->getName());
-    computeBlockInfo(CGM, this, blockInfo);
-    blockInfo.BlockExpression = blockExpr;
-    return EmitBlockLiteral(blockInfo);
-  }
-
-  // Find the block info for this block and take ownership of it.
-  std::unique_ptr<CGBlockInfo> blockInfo;
-  blockInfo.reset(findAndRemoveBlockInfo(&FirstBlockInfo,
-                                         blockExpr->getBlockDecl()));
 
-  blockInfo->BlockExpression = blockExpr;
-  return EmitBlockLiteral(*blockInfo);
+  CGBlockInfo blockInfo(blockExpr->getBlockDecl(), CurFn->getName());
+  computeBlockInfo(CGM, this, blockInfo);
+  blockInfo.BlockExpression = blockExpr;
+  if (!blockInfo.CanBeGlobal)
+    blockInfo.LocalAddress = CreateTempAlloca(blockInfo.StructureType,
+                                              blockInfo.BlockAlign, "block");
+  return EmitBlockLiteral(blockInfo);
 }
 
 llvm::Value *CodeGenFunction::EmitBlockLiteral(const CGBlockInfo &blockInfo) {
@@ -1161,12 +1033,64 @@ llvm::Value *CodeGenFunction::EmitBlockLiteral(const CGBlockInfo &blockInfo) {
                      /*captured by init*/ false);
     }
 
-    // Activate the cleanup if layout pushed one.
-    if (!CI.isByRef()) {
-      EHScopeStack::stable_iterator cleanup = capture.getCleanup();
-      if (cleanup.isValid())
-        ActivateCleanupBlock(cleanup, blockInfo.DominatingIP);
+    // Push a cleanup for the capture if necessary.
+    if (!blockInfo.NeedsCopyDispose)
+      continue;
+
+    // Ignore __block captures; there's nothing special in the on-stack block
+    // that we need to do for them.
+    if (CI.isByRef())
+      continue;
+
+    // Ignore objects that aren't destructed.
+    QualType::DestructionKind dtorKind = type.isDestructedType();
+    if (dtorKind == QualType::DK_none)
+      continue;
+
+    CodeGenFunction::Destroyer *destroyer;
+
+    // Block captures count as local values and have imprecise semantics.
+    // They also can't be arrays, so need to worry about that.
+    //
+    // For const-qualified captures, emit clang.arc.use to ensure the captured
+    // object doesn't get released while we are still depending on its validity
+    // within the block.
+    if (type.isConstQualified() &&
+        type.getObjCLifetime() == Qualifiers::OCL_Strong &&
+        CGM.getCodeGenOpts().OptimizationLevel != 0) {
+      assert(CGM.getLangOpts().ObjCAutoRefCount &&
+             "expected ObjC ARC to be enabled");
+      destroyer = emitARCIntrinsicUse;
+    } else if (dtorKind == QualType::DK_objc_strong_lifetime) {
+      destroyer = destroyARCStrongImprecise;
+    } else {
+      destroyer = getDestroyer(dtorKind);
     }
+
+    CleanupKind cleanupKind = NormalCleanup;
+    bool useArrayEHCleanup = needsEHCleanup(dtorKind);
+    if (useArrayEHCleanup)
+      cleanupKind = NormalAndEHCleanup;
+
+    // Extend the lifetime of the capture to the end of the scope enclosing the
+    // block expression except when the block decl is in the list of RetExpr's
+    // cleanup objects, in which case its lifetime ends after the full
+    // expression.
+    auto IsBlockDeclInRetExpr = [&]() {
+      auto *EWC = llvm::dyn_cast_or_null<ExprWithCleanups>(RetExpr);
+      if (EWC)
+        for (auto &C : EWC->getObjects())
+          if (auto *BD = C.dyn_cast<BlockDecl *>())
+            if (BD == blockDecl)
+              return true;
+      return false;
+    };
+
+    if (IsBlockDeclInRetExpr())
+      pushDestroy(cleanupKind, blockField, type, destroyer, useArrayEHCleanup);
+    else
+      pushLifetimeExtendedDestroy(cleanupKind, blockField, type, destroyer,
+                                  useArrayEHCleanup);
   }
 
   // Cast to the converted block-pointer type, which happens (somewhat

clang/lib/CodeGen/CGBlocks.h

@@ -257,10 +257,6 @@ class CGBlockInfo {
   // This could be zero if no forced alignment is required.
   CharUnits BlockHeaderForcedGapSize;
 
-  /// An instruction which dominates the full-expression that the
-  /// block is inside.
-  llvm::Instruction *DominatingIP;
-
   /// The next block in the block-info chain.  Invalid if this block
   /// info is not part of the CGF's block-info chain, which is true
   /// if it corresponds to a global block or a block whose expression

clang/lib/CodeGen/CGCleanup.cpp

@@ -179,12 +179,10 @@ void *EHScopeStack::pushCleanup(CleanupKind Kind, size_t Size) {
   char *Buffer = allocate(EHCleanupScope::getSizeForCleanupSize(Size));
   bool IsNormalCleanup = Kind & NormalCleanup;
   bool IsEHCleanup = Kind & EHCleanup;
-  bool IsActive = !(Kind & InactiveCleanup);
   bool IsLifetimeMarker = Kind & LifetimeMarker;
   EHCleanupScope *Scope =
     new (Buffer) EHCleanupScope(IsNormalCleanup,
                                 IsEHCleanup,
-                                IsActive,
                                 Size,
                                 BranchFixups.size(),
                                 InnermostNormalCleanup,

clang/lib/CodeGen/CGCleanup.h

@@ -284,16 +284,16 @@ class alignas(8) EHCleanupScope : public EHScope {
     return sizeof(EHCleanupScope) + CleanupBits.CleanupSize;
   }
 
-  EHCleanupScope(bool isNormal, bool isEH, bool isActive,
-                 unsigned cleanupSize, unsigned fixupDepth,
+  EHCleanupScope(bool isNormal, bool isEH, unsigned cleanupSize,
+                 unsigned fixupDepth,
                  EHScopeStack::stable_iterator enclosingNormal,
                  EHScopeStack::stable_iterator enclosingEH)
       : EHScope(EHScope::Cleanup, enclosingEH),
         EnclosingNormal(enclosingNormal), NormalBlock(nullptr),
         ActiveFlag(nullptr), ExtInfo(nullptr), FixupDepth(fixupDepth) {
     CleanupBits.IsNormalCleanup = isNormal;
     CleanupBits.IsEHCleanup = isEH;
-    CleanupBits.IsActive = isActive;
+    CleanupBits.IsActive = true;
     CleanupBits.IsLifetimeMarker = false;
     CleanupBits.TestFlagInNormalCleanup = false;
     CleanupBits.TestFlagInEHCleanup = false;

clang/lib/CodeGen/CGDecl.cpp

@@ -762,10 +762,9 @@ void CodeGenFunction::EmitScalarInit(const Expr *init, const ValueDecl *D,
 
   // If we're emitting a value with lifetime, we have to do the
   // initialization *before* we leave the cleanup scopes.
-  if (const FullExpr *fe = dyn_cast<FullExpr>(init)) {
-    enterFullExpression(fe);
+  if (const FullExpr *fe = dyn_cast<FullExpr>(init))
     init = fe->getSubExpr();
-  }
+
   CodeGenFunction::RunCleanupsScope Scope(*this);
 
   // We have to maintain the illusion that the variable is

clang/lib/CodeGen/CGExpr.cpp

@@ -1330,7 +1330,6 @@ LValue CodeGenFunction::EmitLValue(const Expr *E) {
 
   case Expr::ExprWithCleanupsClass: {
     const auto *cleanups = cast<ExprWithCleanups>(E);
-    enterFullExpression(cleanups);
     RunCleanupsScope Scope(*this);
     LValue LV = EmitLValue(cleanups->getSubExpr());
     if (LV.isSimple()) {

clang/lib/CodeGen/CGExprAgg.cpp

@@ -1349,7 +1349,6 @@ AggExprEmitter::VisitLambdaExpr(LambdaExpr *E) {
 }
 
 void AggExprEmitter::VisitExprWithCleanups(ExprWithCleanups *E) {
-  CGF.enterFullExpression(E);
   CodeGenFunction::RunCleanupsScope cleanups(CGF);
   Visit(E->getSubExpr());
 }

clang/lib/CodeGen/CGExprComplex.cpp

@@ -222,7 +222,6 @@ class ComplexExprEmitter
     return Visit(DIE->getExpr());
   }
   ComplexPairTy VisitExprWithCleanups(ExprWithCleanups *E) {
-    CGF.enterFullExpression(E);
     CodeGenFunction::RunCleanupsScope Scope(CGF);
     ComplexPairTy Vals = Visit(E->getSubExpr());
     // Defend against dominance problems caused by jumps out of expression

clang/lib/CodeGen/CGExprScalar.cpp

@@ -2342,7 +2342,6 @@ Value *ScalarExprEmitter::VisitStmtExpr(const StmtExpr *E) {
 }
 
 Value *ScalarExprEmitter::VisitExprWithCleanups(ExprWithCleanups *E) {
-  CGF.enterFullExpression(E);
   CodeGenFunction::RunCleanupsScope Scope(CGF);
   Value *V = Visit(E->getSubExpr());
   // Defend against dominance problems caused by jumps out of expression

clang/lib/CodeGen/CGObjC.cpp

@@ -3256,7 +3256,6 @@ static llvm::Value *emitARCRetainLoadOfScalar(CodeGenFunction &CGF,
 llvm::Value *CodeGenFunction::EmitARCRetainScalarExpr(const Expr *e) {
   // The retain needs to happen within the full-expression.
   if (const ExprWithCleanups *cleanups = dyn_cast<ExprWithCleanups>(e)) {
-    enterFullExpression(cleanups);
     RunCleanupsScope scope(*this);
     return EmitARCRetainScalarExpr(cleanups->getSubExpr());
   }
@@ -3272,7 +3271,6 @@ llvm::Value *
 CodeGenFunction::EmitARCRetainAutoreleaseScalarExpr(const Expr *e) {
   // The retain needs to happen within the full-expression.
   if (const ExprWithCleanups *cleanups = dyn_cast<ExprWithCleanups>(e)) {
-    enterFullExpression(cleanups);
     RunCleanupsScope scope(*this);
     return EmitARCRetainAutoreleaseScalarExpr(cleanups->getSubExpr());
   }
@@ -3383,7 +3381,6 @@ static llvm::Value *emitARCUnsafeUnretainedScalarExpr(CodeGenFunction &CGF,
 llvm::Value *CodeGenFunction::EmitARCUnsafeUnretainedScalarExpr(const Expr *e) {
   // Look through full-expressions.
   if (const ExprWithCleanups *cleanups = dyn_cast<ExprWithCleanups>(e)) {
-    enterFullExpression(cleanups);
     RunCleanupsScope scope(*this);
     return emitARCUnsafeUnretainedScalarExpr(*this, cleanups->getSubExpr());
   }

clang/lib/CodeGen/CGStmt.cpp

@@ -1070,6 +1070,19 @@ void CodeGenFunction::EmitReturnOfRValue(RValue RV, QualType Ty) {
   EmitBranchThroughCleanup(ReturnBlock);
 }
 
+namespace {
+// RAII struct used to save and restore a return statment's result expression.
+struct SaveRetExprRAII {
+  SaveRetExprRAII(const Expr *RetExpr, CodeGenFunction &CGF)
+      : OldRetExpr(CGF.RetExpr), CGF(CGF) {
+    CGF.RetExpr = RetExpr;
+  }
+  ~SaveRetExprRAII() { CGF.RetExpr = OldRetExpr; }
+  const Expr *OldRetExpr;
+  CodeGenFunction &CGF;
+};
+} // namespace
+
 /// EmitReturnStmt - Note that due to GCC extensions, this can have an operand
 /// if the function returns void, or may be missing one if the function returns
 /// non-void.  Fun stuff :).
@@ -1095,15 +1108,19 @@ void CodeGenFunction::EmitReturnStmt(const ReturnStmt &S) {
   // Emit the result value, even if unused, to evaluate the side effects.
   const Expr *RV = S.getRetValue();
 
-  // Treat block literals in a return expression as if they appeared
-  // in their own scope.  This permits a small, easily-implemented
-  // exception to our over-conservative rules about not jumping to
-  // statements following block literals with non-trivial cleanups.
+  // Record the result expression of the return statement. The recorded
+  // expression is used to determine whether a block capture's lifetime should
+  // end at the end of the full expression as opposed to the end of the scope
+  // enclosing the block expression.
+  //
+  // This permits a small, easily-implemented exception to our over-conservative
+  // rules about not jumping to statements following block literals with
+  // non-trivial cleanups.
+  SaveRetExprRAII SaveRetExpr(RV, *this);
+
   RunCleanupsScope cleanupScope(*this);
-  if (const FullExpr *fe = dyn_cast_or_null<FullExpr>(RV)) {
-    enterFullExpression(fe);
+  if (const FullExpr *fe = dyn_cast_or_null<FullExpr>(RV))
     RV = fe->getSubExpr();
-  }
 
   // FIXME: Clean this up by using an LValue for ReturnTemp,
   // EmitStoreThroughLValue, and EmitAnyExpr.