[CodeGen] Make sure the EH cleanup for block captures is conditional …

clang/lib/CodeGen/CGDecl.cpp

@@ -2094,21 +2094,47 @@ void CodeGenFunction::pushStackRestore(CleanupKind Kind, Address SPMem) {
   EHStack.pushCleanup<CallStackRestore>(Kind, SPMem);
 }
 
-void CodeGenFunction::pushLifetimeExtendedDestroy(
-    CleanupKind cleanupKind, Address addr, QualType type,
-    Destroyer *destroyer, bool useEHCleanupForArray) {
-  // Push an EH-only cleanup for the object now.
-  // FIXME: When popping normal cleanups, we need to keep this EH cleanup
-  // around in case a temporary's destructor throws an exception.
-  if (cleanupKind & EHCleanup)
-    EHStack.pushCleanup<DestroyObject>(
-        static_cast<CleanupKind>(cleanupKind & ~NormalCleanup), addr, type,
+void CodeGenFunction::pushLifetimeExtendedDestroy(CleanupKind cleanupKind,
+                                                  Address addr, QualType type,
+                                                  Destroyer *destroyer,
+                                                  bool useEHCleanupForArray) {
+  // If we're not in a conditional branch, we don't need to bother generating a
+  // conditional cleanup.
+  if (!isInConditionalBranch()) {
+    // Push an EH-only cleanup for the object now.
+    // FIXME: When popping normal cleanups, we need to keep this EH cleanup
+    // around in case a temporary's destructor throws an exception.
+    if (cleanupKind & EHCleanup)
+      EHStack.pushCleanup<DestroyObject>(
+          static_cast<CleanupKind>(cleanupKind & ~NormalCleanup), addr, type,
+          destroyer, useEHCleanupForArray);
+
+    return pushCleanupAfterFullExprWithActiveFlag<DestroyObject>(
+        cleanupKind, Address::invalid(), addr, type, destroyer, useEHCleanupForArray);
+  }
+
+  // Otherwise, we should only destroy the object if it's been initialized.
+  // Re-use the active flag and saved address across both the EH and end of
+  // scope cleanups.
+
+  using SavedType = typename DominatingValue<Address>::saved_type;
+  using ConditionalCleanupType =
+      EHScopeStack::ConditionalCleanup<DestroyObject, Address, QualType,
+                                       Destroyer *, bool>;
+
+  Address ActiveFlag = createCleanupActiveFlag();
+  SavedType SavedAddr = saveValueInCond(addr);
+
+  if (cleanupKind & EHCleanup) {
+    EHStack.pushCleanup<ConditionalCleanupType>(
+        static_cast<CleanupKind>(cleanupKind & ~NormalCleanup), SavedAddr, type,
         destroyer, useEHCleanupForArray);
+    initFullExprCleanupWithFlag(ActiveFlag);
+  }
 
-  // Remember that we need to push a full cleanup for the object at the
-  // end of the full-expression.
-  pushCleanupAfterFullExpr<DestroyObject>(
-      cleanupKind, addr, type, destroyer, useEHCleanupForArray);
+  pushCleanupAfterFullExprWithActiveFlag<ConditionalCleanupType>(
+      cleanupKind, ActiveFlag, SavedAddr, type, destroyer,
+      useEHCleanupForArray);
 }
 
 /// emitDestroy - Immediately perform the destruction of the given

clang/lib/CodeGen/CodeGenFunction.h

@@ -672,12 +672,13 @@ class CodeGenFunction : public CodeGenTypeCache {
     initFullExprCleanup();
   }
 
-  /// Queue a cleanup to be pushed after finishing the current
-  /// full-expression.
+  /// Queue a cleanup to be pushed after finishing the current full-expression,
+  /// potentially with an active flag.
   template <class T, class... As>
   void pushCleanupAfterFullExpr(CleanupKind Kind, As... A) {
     if (!isInConditionalBranch())
-      return pushCleanupAfterFullExprImpl<T>(Kind, Address::invalid(), A...);
+      return pushCleanupAfterFullExprWithActiveFlag<T>(Kind, Address::invalid(),
+                                                       A...);
 
     Address ActiveFlag = createCleanupActiveFlag();
     assert(!DominatingValue<Address>::needsSaving(ActiveFlag) &&
@@ -687,12 +688,12 @@ class CodeGenFunction : public CodeGenTypeCache {
     SavedTuple Saved{saveValueInCond(A)...};
 
     typedef EHScopeStack::ConditionalCleanup<T, As...> CleanupType;
-    pushCleanupAfterFullExprImpl<CleanupType>(Kind, ActiveFlag, Saved);
+    pushCleanupAfterFullExprWithActiveFlag<CleanupType>(Kind, ActiveFlag, Saved);
   }
 
   template <class T, class... As>
-  void pushCleanupAfterFullExprImpl(CleanupKind Kind, Address ActiveFlag,
-                                    As... A) {
+  void pushCleanupAfterFullExprWithActiveFlag(CleanupKind Kind,
+                                              Address ActiveFlag, As... A) {
     LifetimeExtendedCleanupHeader Header = {sizeof(T), Kind,
                                             ActiveFlag.isValid()};
 

clang/test/CodeGenObjC/arc-blocks-exceptions.m

@@ -0,0 +1,35 @@
+// RUN: %clang_cc1 -triple x86_64-apple-darwin10 -emit-llvm -fblocks -fobjc-arc -fobjc-runtime-has-weak -fexceptions -disable-llvm-passes -o - %s | FileCheck %s
+
+void test1(_Bool c) {
+  void test1_fn(void (^blk)());
+  __weak id weakId = 0;
+  test1_fn(c ? ^{ (void)weakId; } : 0);
+
+  // CHECK: [[CLEANUP_COND:%.*]] = alloca i1
+  // CHECK-NEXT: [[CLEANUP_SAVE:%.*]] = alloca i8**
+
+  // CHECK: store i1 true, i1* [[CLEANUP_COND]]
+  // CHECK-NEXT: store i8** {{.*}}, i8*** [[CLEANUP_SAVE]]
+
+  // CHECK: invoke void @test1_fn(
+  // CHECK-NEXT: to label %[[INVOKE_CONT:.*]] unwind label %[[LANDING_PAD_LAB:.*]]
+
+  // CHECK: [[INVOKE_CONT]]:
+  // CHECK-NEXT: [[LOAD:%.*]] = load i1, i1* [[CLEANUP_COND]]
+  // CHECK-NEXT: br i1 [[LOAD]], label %[[END_OF_SCOPE_LAB:.*]], label
+
+  // CHECK: [[END_OF_SCOPE_LAB]]:
+  // CHECK-NEXT: [[LOAD:%.*]] = load i8**, i8*** [[CLEANUP_SAVE]]
+  // CHECK-NEXT: call void @llvm.objc.destroyWeak(i8** [[LOAD]])
+  // CHECK-NEXT: br label
+
+  // CHECK: [[LANDING_PAD_LAB]]:
+  //   /* some EH stuff */
+  // CHECK:      [[LOAD:%.*]] = load i1, i1* [[CLEANUP_COND]]
+  // CHECK-NEXT: br i1 [[LOAD]], label %[[EH_CLEANUP_LAB:.*]], label
+
+  // CHECK: [[EH_CLEANUP_LAB]]:
+  // CHECK-NEXT: [[LOAD:%.*]] = load i8**, i8*** [[CLEANUP_SAVE]]
+  // CHECK-NEXT: call void @llvm.objc.destroyWeak(i8** [[LOAD]])
+  // CHECK-NEXT: br label
+}