[CodeGen] Avoid destructing a callee-destructued struct type in a

function if a function delegates to another function.

Fix a bug introduced in r328731, which caused a struct with ObjC __weak
fields that was passed to a function to be destructed twice, once in the
callee function and once in another function the callee function
delegates to. To prevent this, keep track of the callee-destructed
structs passed to a function and disable their cleanups at the point of
the call to the delegated function.

rdar://problem/39194693

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

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

7 files changed, 143 insertions, 5 deletions

diff --git a/lib/CodeGen/CGCall.cpp b/lib/CodeGen/CGCall.cpp
index 392c7b1ded..d4374ff6e0 100644
--- a/lib/CodeGen/CGCall.cpp
+++ b/lib/CodeGen/CGCall.cpp
@@ -3063,6 +3063,22 @@ void CodeGenFunction::EmitDelegateCallArg(CallArgList &args,
   } else {
     args.add(convertTempToRValue(local, type, loc), type);
   }
+
+  // Deactivate the cleanup for the callee-destructed param that was pushed.
+  if (hasAggregateEvaluationKind(type) &&
+      getContext().isParamDestroyedInCallee(type)) {
+    EHScopeStack::stable_iterator cleanup =
+        CalleeDestructedParamCleanups.lookup(cast<ParmVarDecl>(param));
+    if (cleanup.isValid()) {
+      // This unreachable is a temporary marker which will be removed later.
+      llvm::Instruction *isActive = Builder.CreateUnreachable();
+      args.addArgCleanupDeactivation(cleanup, isActive);
+    } else
+      // A param cleanup should have been pushed unless we are code-generating
+      // a thunk.
+      assert(CurFuncIsThunk &&
+             "cleanup for callee-destructed param not recorded");
+  }
 }
 
 static bool isProvablyNull(llvm::Value *addr) {
diff --git a/lib/CodeGen/CGCleanup.cpp b/lib/CodeGen/CGCleanup.cpp
index 526def2cc2..c5f935bdef 100644
--- a/lib/CodeGen/CGCleanup.cpp
+++ b/lib/CodeGen/CGCleanup.cpp
@@ -1233,8 +1233,10 @@ void CodeGenFunction::DeactivateCleanupBlock(EHScopeStack::stable_iterator C,
   EHCleanupScope &Scope = cast<EHCleanupScope>(*EHStack.find(C));
   assert(Scope.isActive() && "double deactivation");
 
-  // If it's the top of the stack, just pop it.
-  if (C == EHStack.stable_begin()) {
+  // If it's the top of the stack, just pop it, but do so only if it belongs
+  // to the current RunCleanupsScope.
+  if (C == EHStack.stable_begin() &&
+      CurrentCleanupScopeDepth.strictlyEncloses(C)) {
     // If it's a normal cleanup, we need to pretend that the
     // fallthrough is unreachable.
     CGBuilderTy::InsertPoint SavedIP = Builder.saveAndClearIP();
diff --git a/lib/CodeGen/CGDecl.cpp b/lib/CodeGen/CGDecl.cpp
index c9b80e38d4..75251df55b 100644
--- a/lib/CodeGen/CGDecl.cpp
+++ b/lib/CodeGen/CGDecl.cpp
@@ -1962,6 +1962,8 @@ void CodeGenFunction::EmitParmDecl(const VarDecl &D, ParamValue Arg,
                 DtorKind == QualType::DK_nontrivial_c_struct) &&
                "unexpected destructor type");
         pushDestroy(DtorKind, DeclPtr, Ty);
+        CalleeDestructedParamCleanups[cast<ParmVarDecl>(&D)] =
+            EHStack.stable_begin();
       }
     }
   } else {
diff --git a/lib/CodeGen/CodeGenFunction.h b/lib/CodeGen/CodeGenFunction.h
index 091fdc7ab7..7889d08907 100644
--- a/lib/CodeGen/CodeGenFunction.h
+++ b/lib/CodeGen/CodeGenFunction.h
@@ -587,7 +587,7 @@ public:
   /// \brief Enters a new scope for capturing cleanups, all of which
   /// will be executed once the scope is exited.
   class RunCleanupsScope {
-    EHScopeStack::stable_iterator CleanupStackDepth;
+    EHScopeStack::stable_iterator CleanupStackDepth, OldCleanupScopeDepth;
     size_t LifetimeExtendedCleanupStackSize;
     bool OldDidCallStackSave;
   protected:
@@ -610,6 +610,8 @@ public:
           CGF.LifetimeExtendedCleanupStack.size();
       OldDidCallStackSave = CGF.DidCallStackSave;
       CGF.DidCallStackSave = false;
+      OldCleanupScopeDepth = CGF.CurrentCleanupScopeDepth;
+      CGF.CurrentCleanupScopeDepth = CleanupStackDepth;
     }
 
     /// \brief Exit this cleanup scope, emitting any accumulated cleanups.
@@ -635,9 +637,14 @@ public:
       CGF.PopCleanupBlocks(CleanupStackDepth, LifetimeExtendedCleanupStackSize,
                            ValuesToReload);
       PerformCleanup = false;
+      CGF.CurrentCleanupScopeDepth = OldCleanupScopeDepth;
     }
   };
 
+  // Cleanup stack depth of the RunCleanupsScope that was pushed most recently.
+  EHScopeStack::stable_iterator CurrentCleanupScopeDepth =
+      EHScopeStack::stable_end();
+
   class LexicalScope : public RunCleanupsScope {
     SourceRange Range;
     SmallVector<const LabelDecl*, 4> Labels;
@@ -1095,6 +1102,11 @@ private:
   /// decls.
   DeclMapTy LocalDeclMap;
 
+  // Keep track of the cleanups for callee-destructed parameters pushed to the
+  // cleanup stack so that they can be deactivated later.
+  llvm::DenseMap<const ParmVarDecl *, EHScopeStack::stable_iterator>
+      CalleeDestructedParamCleanups;
+
   /// SizeArguments - If a ParmVarDecl had the pass_object_size attribute, this
   /// will contain a mapping from said ParmVarDecl to its implicit "object_size"
   /// parameter.
diff --git a/test/CodeGenObjCXX/arc-forwarded-lambda-call.mm b/test/CodeGenObjCXX/arc-forwarded-lambda-call.mm
index 35abe19183..0a575c7697 100644
--- a/test/CodeGenObjCXX/arc-forwarded-lambda-call.mm
+++ b/test/CodeGenObjCXX/arc-forwarded-lambda-call.mm
@@ -10,6 +10,17 @@ void test0(id x) {
   // CHECK-NEXT: ret i8* [[T2]]
 }
 
+// Check that the delegating block invoke function doesn't destruct the Weak
+// object that is passed.
+
+// CHECK-LABEL: define internal void @___Z8testWeakv_block_invoke(
+// CHECK: call void @"_ZZ8testWeakvENK3$_2clE4Weak"(
+// CHECK-NEXT: ret void
+
+// CHECK-LABEL: define internal void @"_ZZ8testWeakvENK3$_2clE4Weak"(
+// CHECK: call void @_ZN4WeakD1Ev(
+// CHECK-NEXT: ret void
+
 id test1_rv;
 
 void test1() {
@@ -21,3 +32,12 @@ void test1() {
   // CHECK-NEXT: [[T2:%.*]] = tail call i8* @objc_autoreleaseReturnValue(i8* [[T1]])
   // CHECK-NEXT: ret i8* [[T2]]
 }
+
+struct Weak {
+  __weak id x;
+};
+
+void testWeak() {
+  extern void testWeak_helper(void (^)(Weak));
+  testWeak_helper([](Weak){});
+}
diff --git a/test/CodeGenObjCXX/arc-special-member-functions.mm b/test/CodeGenObjCXX/arc-special-member-functions.mm
index 278c779481..d620187d70 100644
--- a/test/CodeGenObjCXX/arc-special-member-functions.mm
+++ b/test/CodeGenObjCXX/arc-special-member-functions.mm
@@ -1,4 +1,4 @@
-// RUN: %clang_cc1 -fobjc-arc -fblocks -triple x86_64-apple-darwin10.0.0 -emit-llvm -o - %s | FileCheck %s
+// RUN: %clang_cc1 -std=c++11 -fobjc-arc -fblocks -triple x86_64-apple-darwin10.0.0 -fobjc-runtime-has-weak -emit-llvm -o - %s | FileCheck %s
 
 struct ObjCMember {
   id member;
@@ -12,6 +12,59 @@ struct ObjCBlockMember {
   int (^bp)(int);
 };
 
+// CHECK: %[[STRUCT_CONTAINSWEAK:.*]] = type { %[[STRUCT_WEAK:.*]] }
+// CHECK: %[[STRUCT_WEAK]] = type { i8* }
+
+// The Weak object that is passed is destructed in this constructor.
+
+// CHECK: define void @_ZN12ContainsWeakC2E4Weak(
+// CHECK: call void @_ZN4WeakC1ERKS_(
+// CHECK: call void @_ZN4WeakD1Ev(
+
+// Check that the Weak object passed to this constructor is not destructed after
+// the delegate constructor is called.
+
+// CHECK: define void @_ZN12ContainsWeakC1E4Weak(
+// CHECK: call void @_ZN12ContainsWeakC2E4Weak(
+// CHECK-NEXT: ret void
+
+struct Weak {
+  Weak(id);
+  __weak id x;
+};
+
+struct ContainsWeak {
+  ContainsWeak(Weak);
+  Weak w;
+};
+
+ContainsWeak::ContainsWeak(Weak a) : w(a) {}
+
+// The Weak object that is passed is destructed in this constructor.
+
+// CHECK: define void @_ZN4BaseC2E4Weak(
+// CHECK: call void @_ZN4WeakD1Ev(
+// CHECK: ret void
+
+// Check that the Weak object passed to this constructor is not destructed after
+// the delegate constructor is called.
+
+// CHECK: define linkonce_odr void @_ZN7DerivedCI14BaseE4Weak(
+// CHECK: call void @_ZN7DerivedCI24BaseE4Weak(
+// CHECK-NEXT: ret void
+
+struct Base {
+  Base(Weak);
+};
+
+Base::Base(Weak a) {}
+
+struct Derived : Base {
+  using Base::Base;
+};
+
+Derived d(Weak(0));
+
 // CHECK-LABEL: define void @_Z42test_ObjCMember_default_construct_destructv(
 void test_ObjCMember_default_construct_destruct() {
   // CHECK: call void @_ZN10ObjCMemberC1Ev
@@ -111,6 +164,13 @@ void test_ObjCBlockMember_copy_assign(ObjCBlockMember m1, ObjCBlockMember m2) {
 // CHECK-NEXT: call void @objc_release(i8* [[T7]])
 // CHECK-NEXT: ret
 
+// Check that the Weak object passed to this constructor is not destructed after
+// the delegate constructor is called.
+
+// CHECK: define linkonce_odr void @_ZN7DerivedCI24BaseE4Weak(
+// CHECK: call void @_ZN4BaseC2E4Weak(
+// CHECK-NEXT: ret void
+
 // Implicitly-generated default constructor for ObjCMember
 // CHECK-LABEL: define linkonce_odr void @_ZN10ObjCMemberC2Ev
 // CHECK-NOT: objc_release
diff --git a/test/CodeGenObjCXX/lambda-expressions.mm b/test/CodeGenObjCXX/lambda-expressions.mm
index c8247e2e0a..f60655c61b 100644
--- a/test/CodeGenObjCXX/lambda-expressions.mm
+++ b/test/CodeGenObjCXX/lambda-expressions.mm
@@ -1,4 +1,4 @@
-// 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 -fobjc-arc -fobjc-runtime-has-weak -DWEAK_SUPPORTED | 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)();
@@ -138,5 +138,31 @@ namespace BlockInLambda {
 }
 @end
 
+// Check that the delegating invoke function doesn't destruct the Weak object
+// that is passed.
+
+// ARC-LABEL: define internal void @"_ZZN14LambdaDelegate4testEvEN3$_58__invokeENS_4WeakE"(
+// ARC: call void @"_ZZN14LambdaDelegate4testEvENK3$_5clENS_4WeakE"(
+// ARC-NEXT: ret void
+
+// ARC-LABEL: define internal void @"_ZZN14LambdaDelegate4testEvENK3$_5clENS_4WeakE"(
+// ARC: call void @_ZN14LambdaDelegate4WeakD1Ev(
+
+#ifdef WEAK_SUPPORTED
+
+namespace LambdaDelegate {
+
+struct Weak {
+  __weak id x;
+};
+
+void test() {
+  void (*p)(Weak) = [](Weak a) { };
+}
+
+};
+
+#endif
+
 // ARC: attributes [[NUW]] = { noinline nounwind{{.*}} }
 // MRC: attributes [[NUW]] = { noinline nounwind{{.*}} }