ahatanak created this revision.
ahatanak added reviewers: rjmccall, erik.pilkington.
ahatanak added a project: clang.
Herald added subscribers: jdoerfert, dexonsmith, jkorous.

This patch avoids copying blocks that initialize or are assigned to local auto 
variables to the heap when the local auto variables do not have their addresses 
taken and are declared in the same scope as the block. We can possibly add back 
the optimization in the ARC optimizer that was reverted in r189869 
(http://lists.llvm.org/pipermail/llvm-commits/Week-of-Mon-20130902/186509.html),
 but I suspect it would be much more complicated and fragile than doing it in 
the front-end.

I'm not 100% sure whether it's necessary to disable this optimization when the 
address of the local variable is taken. We do pass a block on the stack to a 
function when the block is directly passed instead of first being assigned to a 
local variable, but clang currently doesn't copy the passed block to the heap 
in the callee although the block can possibly escape if the address is taken. 
For example:

  __strong id *g0, g1;
  
  void foo0(BlockTy b) {
    g0 = (__strong id *)&b;
    g1 = *g0; // this is just a retain, not a block copy.
  }
  
  void foo1() {
    foo0(^{...}) // block is on the stack.
  }
  
  void foo2() {
    foo1();
    ((BlockTy)g1)(); // this can crash if the block is still on the stack.
  }

rdar://problem/13289333


Repository:
  rC Clang

https://reviews.llvm.org/D58514

Files:
  include/clang/AST/Decl.h
  include/clang/AST/DeclBase.h
  include/clang/Sema/ScopeInfo.h
  lib/AST/Decl.cpp
  lib/CodeGen/CGObjC.cpp
  lib/Sema/ScopeInfo.cpp
  lib/Sema/Sema.cpp
  lib/Sema/SemaDecl.cpp
  lib/Sema/SemaExpr.cpp
  lib/Serialization/ASTReaderDecl.cpp
  lib/Serialization/ASTWriterDecl.cpp
  test/CodeGenObjC/arc-block-copy-escape.m
  test/CodeGenObjC/arc-blocks.m
  test/CodeGenObjCXX/arc-blocks.mm
  test/PCH/arc-blocks.mm

Index: test/PCH/arc-blocks.mm
===================================================================
--- /dev/null
+++ test/PCH/arc-blocks.mm
@@ -0,0 +1,49 @@
+// RUN: %clang_cc1 -fobjc-arc -fblocks -std=c++1y -emit-pch %s -o %t
+// RUN: %clang_cc1 -fobjc-arc -fblocks -std=c++1y -include-pch %t -emit-llvm -o - %s | FileCheck %s
+
+#ifndef HEADER_INCLUDED
+#define HEADER_INCLUDED
+
+namespace test_block_retain {
+  typedef void (^BlockTy)();
+  void foo1(id);
+
+  inline void initialization(id a) {
+    // Call to @llvm.objc.retainBlock isn't needed.
+    BlockTy b0 = ^{ foo1(a); };
+    b0();
+  }
+
+  inline void assignmentConditional(id a, bool c) {
+    BlockTy b0;
+    if (c)
+      // @llvm.objc.retainBlock is called since 'b0' is declared in the outer scope.
+      b0 = ^{ foo1(a); };
+    b0();
+  }
+}
+
+#else
+
+// CHECK: %[[STRUCT_BLOCK_DESCRIPTOR:.*]] = type { i64, i64 }
+
+namespace test_block_retain {
+// CHECK-LABEL: define linkonce_odr void @_ZN17test_block_retain14initializationEP11objc_object(
+// CHECK-NOT: call i8* @llvm.objc.retainBlock(
+
+  void test_initialization(id a) {
+    initialization(a);
+  }
+
+// CHECK-LABEL: define void @_ZN17test_block_retain26test_assignmentConditionalEP11objc_objectb(
+// CHECK: %[[BLOCK:.*]] = alloca <{ i8*, i32, i32, i8*, %[[STRUCT_BLOCK_DESCRIPTOR]]*, i8* }>, align 8
+// CHECK: %[[V4:.*]] = bitcast <{ i8*, i32, i32, i8*, %[[STRUCT_BLOCK_DESCRIPTOR]]*, i8* }>* %[[BLOCK]] to void ()*
+// CHECK: %[[V5:.*]] = bitcast void ()* %[[V4]] to i8*
+// CHECK: call i8* @llvm.objc.retainBlock(i8* %[[V5]])
+
+  void test_assignmentConditional(id a, bool c) {
+    assignmentConditional(a, c);
+  }
+}
+
+#endif
Index: test/CodeGenObjCXX/arc-blocks.mm
===================================================================
--- test/CodeGenObjCXX/arc-blocks.mm
+++ test/CodeGenObjCXX/arc-blocks.mm
@@ -201,3 +201,141 @@
   ^{ (void)t0; (void)t1; (void)t2; (void)t3; (void)t4; (void)t5; };
 }
 }
+
+// Test that calls to @llvm.objc.retainBlock aren't emitted in some cases.
+
+namespace test_block_retain {
+  typedef void (^BlockTy)();
+
+  void foo1(id);
+
+// CHECK-LABEL: define void @_ZN17test_block_retain14initializationEP11objc_object(
+// CHECK-NOT: @llvm.objc.retainBlock(
+  void initialization(id a) {
+    BlockTy b0 = ^{ foo1(a); };
+    b0();
+  }
+
+// CHECK-LABEL: define void @_ZN17test_block_retain20initializationStaticEP11objc_object(
+// CHECK: @llvm.objc.retainBlock(
+  void initializationStatic(id a) {
+    static BlockTy b0 = ^{ foo1(a); };
+    b0();
+  }
+
+// CHECK-LABEL: define void @_ZN17test_block_retain15initialization2EP11objc_object
+// CHECK: %[[B0:.*]] = alloca void ()*, align 8
+// CHECK: %[[B1:.*]] = alloca void ()*, align 8
+// CHECK: load void ()*, void ()** %[[B0]], align 8
+// CHECK-NOT: @llvm.objc.retainBlock
+// CHECK: %[[V9:.*]] = load void ()*, void ()** %[[B0]], align 8
+// CHECK: %[[V10:.*]] = bitcast void ()* %[[V9]] to i8*
+// CHECK: %[[V11:.*]] = call i8* @llvm.objc.retainBlock(i8* %[[V10]])
+// CHECK: %[[V12:.*]] = bitcast i8* %[[V11]] to void ()*
+// CHECK: store void ()* %[[V12]], void ()** %[[B1]], align 8
+  void initialization2(id a) {
+    BlockTy b0 = ^{ foo1(a); };
+    b0();
+    BlockTy b1 = b0; // can't optimize this yet.
+    b1();
+  }
+
+// CHECK-LABEL: define void @_ZN17test_block_retain10assignmentEP11objc_object(
+// CHECK-NOT: @llvm.objc.retainBlock(
+  void assignment(id a) {
+    BlockTy b0;
+    b0 = ^{ foo1(a); };
+    b0();
+    b0 = ^{ foo1(a); };
+    b0();
+  }
+
+// CHECK-LABEL: define void @_ZN17test_block_retain16assignmentStaticEP11objc_object(
+// CHECK: @llvm.objc.retainBlock(
+  void assignmentStatic(id a) {
+    static BlockTy b0;
+    b0 = ^{ foo1(a); };
+    b0();
+  }
+
+// CHECK-LABEL: define void @_ZN17test_block_retain21assignmentConditionalEP11objc_objectb(
+// CHECK: @llvm.objc.retainBlock(
+  void assignmentConditional(id a, bool c) {
+    BlockTy b0;
+    if (c)
+      // can't optimize this since 'b0' is declared in the outer scope.
+      b0 = ^{ foo1(a); };
+    b0();
+  }
+
+// CHECK-LABEL: define void @_ZN17test_block_retain11assignment2EP11objc_object(
+// CHECK: %[[B0:.*]] = alloca void ()*, align 8
+// CHECK: %[[B1:.*]] = alloca void ()*, align 8
+// CHECK-NOT: @llvm.objc.retainBlock
+// CHECK: store void ()* null, void ()** %[[B1]], align 8
+// CHECK: %[[V9:.*]] = load void ()*, void ()** %[[B0]], align 8
+// CHECK: %[[V10:.*]] = bitcast void ()* %[[V9]] to i8*
+// CHECK: %[[V11:.*]] = call i8* @llvm.objc.retainBlock(i8* %[[V10]]
+// CHECK: %[[V12:.*]] = bitcast i8* %[[V11]] to void ()*
+// CHECK: store void ()* %[[V12]], void ()** %[[B1]], align 8
+  void assignment2(id a) {
+    BlockTy b0 = ^{ foo1(a); };
+    b0();
+    BlockTy b1;
+    b1 = b0; // can't optimize this yet.
+    b1();
+  }
+
+// We cannot remove the call to @llvm.objc.retainBlock if the variable is of type id.
+
+// CHECK: define void @_ZN17test_block_retain21initializationObjCPtrEP11objc_object(
+// CHECK: alloca i8*, align 8
+// CHECK: %[[B0:.*]] = alloca i8*, align 8
+// CHECK: %[[BLOCK:.*]] = alloca <{ i8*, i32, i32, i8*, %[[STRUCT_BLOCK_DESCRIPTOR]]*, i8* }>, align 8
+// CHECK: %[[V3:.*]] = bitcast <{ i8*, i32, i32, i8*, %[[STRUCT_BLOCK_DESCRIPTOR]]*, i8* }>* %[[BLOCK]] to void ()*
+// CHECK: %[[V4:.*]] = bitcast void ()* %[[V3]] to i8*
+// CHECK: %[[V5:.*]] = call i8* @llvm.objc.retainBlock(i8* %[[V4]])
+// CHECK: %[[V6:.*]] = bitcast i8* %[[V5]] to void ()*
+// CHECK: %[[V7:.*]] = bitcast void ()* %[[V6]] to i8*
+// CHECK: store i8* %[[V7]], i8** %[[B0]], align 8
+  void initializationObjCPtr(id a) {
+    id b0 = ^{ foo1(a); };
+    ((BlockTy)b0)();
+  }
+
+// CHECK: define void @_ZN17test_block_retain17assignmentObjCPtrEP11objc_object(
+// CHECK: %[[B0:.*]] = alloca void ()*, align 8
+// CHECK: %[[B1:.*]] = alloca i8*, align 8
+// CHECK: %[[V4:.*]] = load void ()*, void ()** %[[B0]], align 8
+// CHECK: %[[V5:.*]] = bitcast void ()* %[[V4]] to i8*
+// CHECK: %[[V6:.*]] = call i8* @llvm.objc.retainBlock(i8* %[[V5]])
+// CHECK: %[[V7:.*]] = bitcast i8* %[[V6]] to void ()*
+// CHECK: %[[V8:.*]] = bitcast void ()* %[[V7]] to i8*
+// CHECK: store i8* %[[V8]], i8** %[[B1]], align 8
+  void assignmentObjCPtr(id a) {
+    BlockTy b0 = ^{ foo1(a); };
+    id b1;
+    b1 = b0;
+    ((BlockTy)b1)();
+  }
+
+// Calls to @llvm.objc.retainBlock cannot be removed if the address of the local
+// variable is taken.
+//
+  __strong id *gp0;
+
+// CHECK: define void @_ZN17test_block_retain26addressTakenInitializationEP11objc_object(
+// CHECK: call i8* @llvm.objc.retainBlock(
+  void addressTakenInitialization(id a) {
+    BlockTy b0 = ^{ foo1(a); };
+    gp0 = (__strong id *)&b0;
+  }
+
+// CHECK: define void @_ZN17test_block_retain22addressTakenAssignmentEP11objc_object(
+// CHECK: call i8* @llvm.objc.retainBlock(
+  void addressTakenAssignment(id a) {
+    BlockTy b0;
+    b0 = ^{ foo1(a); };
+    gp0 = (__strong id *)&b0;
+  }
+}
Index: test/CodeGenObjC/arc-blocks.m
===================================================================
--- test/CodeGenObjC/arc-blocks.m
+++ test/CodeGenObjC/arc-blocks.m
@@ -338,20 +338,19 @@
   __block void (^block)(void) = ^{ block(); };
   // CHECK-LABEL:    define void @test10a()
   // CHECK:      [[BYREF:%.*]] = alloca [[BYREF_T:%.*]],
+  // CHECK:      [[BLOCK1:%.*]] = alloca <{ i8*, i32, i32, i8*, %[[STRUCT_BLOCK_DESCRIPTOR]]*, i8* }>, align 8
 
   // Zero-initialization before running the initializer.
   // CHECK:      [[T0:%.*]] = getelementptr inbounds [[BYREF_T]], [[BYREF_T]]* [[BYREF]], i32 0, i32 6
   // CHECK-NEXT: store void ()* null, void ()** [[T0]], align 8
 
   // Run the initializer as an assignment.
-  // CHECK:      [[T0:%.*]] = bitcast void ()* {{%.*}} to i8*
-  // CHECK-NEXT: [[T1:%.*]] = call i8* @llvm.objc.retainBlock(i8* [[T0]])
-  // CHECK-NEXT: [[T2:%.*]] = bitcast i8* [[T1]] to void ()*
+  // CHECK:      [[T2:%.*]] = bitcast <{ i8*, i32, i32, i8*, %[[STRUCT_BLOCK_DESCRIPTOR]]*, i8* }>* [[BLOCK1]] to void ()*
   // CHECK-NEXT: [[T3:%.*]] = getelementptr inbounds [[BYREF_T]], [[BYREF_T]]* [[BYREF]], i32 0, i32 1
   // CHECK-NEXT: [[T4:%.*]] = load [[BYREF_T]]*, [[BYREF_T]]** [[T3]]
   // CHECK-NEXT: [[T5:%.*]] = getelementptr inbounds [[BYREF_T]], [[BYREF_T]]* [[T4]], i32 0, i32 6
   // CHECK-NEXT: [[T6:%.*]] = load void ()*, void ()** [[T5]], align 8
-  // CHECK-NEXT: store void ()* {{%.*}}, void ()** [[T5]], align 8
+  // CHECK-NEXT: store void ()* [[T2]], void ()** [[T5]], align 8
   // CHECK-NEXT: [[T7:%.*]] = bitcast void ()* [[T6]] to i8*
   // CHECK-NEXT: call void @llvm.objc.release(i8* [[T7]])
 
@@ -401,6 +400,7 @@
 
   // CHECK-LABEL:    define void @test10b()
   // CHECK:      [[BYREF:%.*]] = alloca [[BYREF_T:%.*]],
+  // CHECK:      [[BLOCK3:%.*]] = alloca <{ i8*, i32, i32, i8*, %[[STRUCT_BLOCK_DESCRIPTOR]]*, i8* }>, align 8
 
   // Zero-initialize.
   // CHECK:      [[T0:%.*]] = getelementptr inbounds [[BYREF_T]], [[BYREF_T]]* [[BYREF]], i32 0, i32 6
@@ -409,14 +409,12 @@
   // CHECK-NEXT: [[SLOT:%.*]] = getelementptr inbounds [[BYREF_T]], [[BYREF_T]]* [[BYREF]], i32 0, i32 6
 
   // The assignment.
-  // CHECK:      [[T0:%.*]] = bitcast void ()* {{%.*}} to i8*
-  // CHECK-NEXT: [[T1:%.*]] = call i8* @llvm.objc.retainBlock(i8* [[T0]])
-  // CHECK-NEXT: [[T2:%.*]] = bitcast i8* [[T1]] to void ()*
+  // CHECK:      [[T2:%.*]] = bitcast <{ i8*, i32, i32, i8*, %[[STRUCT_BLOCK_DESCRIPTOR]]*, i8* }>* [[BLOCK3]] to void ()*
   // CHECK-NEXT: [[T3:%.*]] = getelementptr inbounds [[BYREF_T]], [[BYREF_T]]* [[BYREF]], i32 0, i32 1
   // CHECK-NEXT: [[T4:%.*]] = load [[BYREF_T]]*, [[BYREF_T]]** [[T3]]
   // CHECK-NEXT: [[T5:%.*]] = getelementptr inbounds [[BYREF_T]], [[BYREF_T]]* [[T4]], i32 0, i32 6
   // CHECK-NEXT: [[T6:%.*]] = load void ()*, void ()** [[T5]], align 8
-  // CHECK-NEXT: store void ()* {{%.*}}, void ()** [[T5]], align 8
+  // CHECK-NEXT: store void ()* [[T2]], void ()** [[T5]], align 8
   // CHECK-NEXT: [[T7:%.*]] = bitcast void ()* [[T6]] to i8*
   // CHECK-NEXT: call void @llvm.objc.release(i8* [[T7]])
 
Index: test/CodeGenObjC/arc-block-copy-escape.m
===================================================================
--- test/CodeGenObjC/arc-block-copy-escape.m
+++ test/CodeGenObjC/arc-block-copy-escape.m
@@ -9,14 +9,14 @@
 void test0(int i) {
   block_t block = ^{ use_int(i); };
   // CHECK-LABEL:   define {{.*}}void @test0(
-  // CHECK:     call {{.*}}i8* @llvm.objc.retainBlock(i8* {{%.*}}) [[NUW:#[0-9]+]], !clang.arc.copy_on_escape
+  // CHECK-NOT: @llvm.objc.retainBlock(
   // CHECK:     ret void
 }
 
 void test1(int i) {
   id block = ^{ use_int(i); };
   // CHECK-LABEL:   define {{.*}}void @test1(
-  // CHECK:     call {{.*}}i8* @llvm.objc.retainBlock(i8* {{%.*}}) [[NUW]]
+  // CHECK:     call {{.*}}i8* @llvm.objc.retainBlock(i8* {{%.*}}) [[NUW:#[0-9]+]]
   // CHECK-NOT: !clang.arc.copy_on_escape
   // CHECK:     ret void
 }
Index: lib/Serialization/ASTWriterDecl.cpp
===================================================================
--- lib/Serialization/ASTWriterDecl.cpp
+++ lib/Serialization/ASTWriterDecl.cpp
@@ -1110,6 +1110,7 @@
   Record.push_back(D->blockMissingReturnType());
   Record.push_back(D->isConversionFromLambda());
   Record.push_back(D->doesNotEscape());
+  Record.push_back(D->canAvoidCopyToHeap());
   Record.push_back(D->capturesCXXThis());
   Record.push_back(D->getNumCaptures());
   for (const auto &capture : D->captures()) {
Index: lib/Serialization/ASTReaderDecl.cpp
===================================================================
--- lib/Serialization/ASTReaderDecl.cpp
+++ lib/Serialization/ASTReaderDecl.cpp
@@ -1479,6 +1479,7 @@
   BD->setBlockMissingReturnType(Record.readInt());
   BD->setIsConversionFromLambda(Record.readInt());
   BD->setDoesNotEscape(Record.readInt());
+  BD->setCanAvoidCopyToHeap(Record.readInt());
 
   bool capturesCXXThis = Record.readInt();
   unsigned numCaptures = Record.readInt();
Index: lib/Sema/SemaExpr.cpp
===================================================================
--- lib/Sema/SemaExpr.cpp
+++ lib/Sema/SemaExpr.cpp
@@ -12445,6 +12445,27 @@
       DiagnoseSelfMove(LHS.get(), RHS.get(), OpLoc);
     }
     RecordModifiableNonNullParam(*this, LHS.get());
+
+    // Record blocks that are assigned to auto local variables. We only care
+    // about blocks assigned to local variables declared in the same scope as
+    // the block as it isn't safe to avoid copying the block to the heap if the
+    // local variable is declared in an outer scope.
+    //
+    // Example:
+    //
+    // BlockTy b;
+    // {
+    //   b = ^{...};
+    // }
+    // // It is unsafe to invoke the block here if it wasn't copied to the heap.
+    // b();
+
+    if (!LHS.isInvalid() && !RHS.isInvalid())
+      if (auto *DRE = dyn_cast<DeclRefExpr>(LHS.get()))
+        if (auto *VD = dyn_cast<VarDecl>(DRE->getDecl()))
+          if (auto *BE = dyn_cast<BlockExpr>(RHS.get()))
+            if (VD->hasLocalStorage() && getCurScope()->isDeclScope(VD))
+              getCurFunction()->addLocalVarBlockDecl(VD, BE->getBlockDecl());
     break;
   case BO_PtrMemD:
   case BO_PtrMemI:
Index: lib/Sema/SemaDecl.cpp
===================================================================
--- lib/Sema/SemaDecl.cpp
+++ lib/Sema/SemaDecl.cpp
@@ -11257,6 +11257,12 @@
           << Culprit->getSourceRange();
       }
     }
+
+    if (auto *E = dyn_cast<ExprWithCleanups>(Init))
+      if (auto *BE = dyn_cast<BlockExpr>(E->getSubExpr()))
+        if (VDecl->hasLocalStorage())
+          getCurFunction()->addLocalVarBlockDecl(VDecl, BE->getBlockDecl());
+
   } else if (VDecl->isStaticDataMember() && !VDecl->isInline() &&
              VDecl->getLexicalDeclContext()->isRecord()) {
     // This is an in-class initialization for a static data member, e.g.,
Index: lib/Sema/Sema.cpp
===================================================================
--- lib/Sema/Sema.cpp
+++ lib/Sema/Sema.cpp
@@ -19,6 +19,7 @@
 #include "clang/AST/Expr.h"
 #include "clang/AST/ExprCXX.h"
 #include "clang/AST/PrettyDeclStackTrace.h"
+#include "clang/AST/StmtVisitor.h"
 #include "clang/AST/StmtCXX.h"
 #include "clang/Basic/DiagnosticOptions.h"
 #include "clang/Basic/PartialDiagnostic.h"
@@ -1623,6 +1624,84 @@
   }
 }
 
+namespace {
+
+/// This class walks the AST to find variables whose addresses are possibly
+/// taken.
+class FindAddressTakenLocalVars
+    : public ConstStmtVisitor<FindAddressTakenLocalVars> {
+  typedef ConstStmtVisitor<FindAddressTakenLocalVars> Super;
+  FunctionScopeInfo &Scope;
+
+public:
+  FindAddressTakenLocalVars(FunctionScopeInfo &Scope) : Scope(Scope) {}
+
+  void VisitImplicitCastExpr(const ImplicitCastExpr *ICE) {
+    // If a DeclRefExpr is a subexpression of an lvalue-to-rvalue conversion,
+    // the address of the declaration referenced by the DeclRefExpr isn't
+    // taken.
+    if (ICE->getCastKind() == CK_LValueToRValue)
+      if (isa<DeclRefExpr>(ICE->getSubExpr()))
+        return;
+    return Visit(ICE->getSubExpr());
+  }
+
+  void VisitBinaryOperator(const BinaryOperator *BO) {
+    // If a DeclRefExpr is the LHS of a binary assignment operation, the
+    // address of the declaration referenced by the DeclRefExpr isn't taken.
+    if (!BO->isAssignmentOp() || !isa<DeclRefExpr>(BO->getLHS()))
+      Visit(BO->getLHS());
+    Visit(BO->getRHS());
+  }
+
+  void VisitDeclRefExpr(const DeclRefExpr *DRE) {
+    // Every other variable referenced by a DeclRefExpr is considered to have
+    // its addresses taken.
+    if (auto *VD = dyn_cast<VarDecl>(DRE->getDecl()))
+      Scope.varDeclAddressTaken(VD);
+  }
+
+  void VisitStmt(const Stmt *S) {
+    for (const Stmt *SubStmt : S->children()) {
+      if (SubStmt)
+        Visit(SubStmt);
+    }
+  }
+};
+
+} // namespace
+
+static void setCanAvoidCopyToHeap(FunctionScopeInfo &FSI, Stmt *Body) {
+  if (!Body)
+    return;
+
+  // It is not safe to avoid copying a block to the heap if it initializes or is
+  // assigned to a local variable whose address is possibly taken. For example:
+  //
+  // __strong id g0, *g1;
+  //
+  // void foo() {
+  //   BlockTy b = ^{...};
+  //   g1 = (__strong id *)&b;
+  //   g0 = *g1; // this doesn't cause the block to be copied to the heap.
+  // }
+  //
+  // void test() {
+  //   foo();
+  //   // this can crash if the stack block wasn't copied to the heap.
+  //   ((BlockTy)g0)();
+  // }
+
+  FindAddressTakenLocalVars F(FSI);
+  F.Visit(Body);
+
+  // The remaining blocks do not need to be copied to the heap at the point of
+  // initialization or assignment.
+  for (auto &P : FSI.LocalVarToBlocks)
+    for (BlockDecl *BD : P.getSecond())
+      BD->setCanAvoidCopyToHeap();
+}
+
 void Sema::PopFunctionScopeInfo(const AnalysisBasedWarnings::Policy *WP,
                                 const Decl *D, const BlockExpr *blkExpr) {
   assert(!FunctionScopes.empty() && "mismatched push/pop!");
@@ -1635,6 +1714,9 @@
 
   FunctionScopeInfo *Scope = FunctionScopes.pop_back_val();
 
+  if (D)
+    setCanAvoidCopyToHeap(*Scope, D->getBody());
+
   if (LangOpts.OpenMP)
     popOpenMPFunctionRegion(Scope);
 
Index: lib/Sema/ScopeInfo.cpp
===================================================================
--- lib/Sema/ScopeInfo.cpp
+++ lib/Sema/ScopeInfo.cpp
@@ -55,6 +55,7 @@
   ModifiedNonNullParams.clear();
   Blocks.clear();
   ByrefBlockVars.clear();
+  LocalVarToBlocks.clear();
 }
 
 static const NamedDecl *getBestPropertyDecl(const ObjCPropertyRefExpr *PropE) {
Index: lib/CodeGen/CGObjC.cpp
===================================================================
--- lib/CodeGen/CGObjC.cpp
+++ lib/CodeGen/CGObjC.cpp
@@ -3205,9 +3205,16 @@
 
   TryEmitResult result = tryEmitARCRetainScalarExpr(*this, e);
   llvm::Value *value = result.getPointer();
-  if (!result.getInt())
-    value = EmitARCRetain(e->getType(), value);
-  return value;
+
+  if (result.getInt())
+    return value;
+
+  // Return if it's possible to avoid block-copy.
+  if (auto *be = dyn_cast<BlockExpr>(e))
+    if (be->getBlockDecl()->canAvoidCopyToHeap())
+      return value;
+
+  return EmitARCRetain(e->getType(), value);
 }
 
 llvm::Value *
@@ -3365,7 +3372,11 @@
   // type, then we need to emit the block-retain immediately in case
   // it invalidates the l-value.
   if (!hasImmediateRetain && e->getType()->isBlockPointerType()) {
-    value = EmitARCRetainBlock(value, /*mandatory*/ false);
+    // Avoid the block-retain if the RHS is a block literal that doesn't need to
+    // be copied to the heap.
+    auto *be = dyn_cast<BlockExpr>(e->getRHS());
+    if (!be || !be->getBlockDecl()->canAvoidCopyToHeap())
+      value = EmitARCRetainBlock(value, /*mandatory*/ false);
     hasImmediateRetain = true;
   }
 
Index: lib/AST/Decl.cpp
===================================================================
--- lib/AST/Decl.cpp
+++ lib/AST/Decl.cpp
@@ -4265,6 +4265,7 @@
   setBlockMissingReturnType(true);
   setIsConversionFromLambda(false);
   setDoesNotEscape(false);
+  setCanAvoidCopyToHeap(false);
 }
 
 void BlockDecl::setParams(ArrayRef<ParmVarDecl *> NewParamInfo) {
Index: include/clang/Sema/ScopeInfo.h
===================================================================
--- include/clang/Sema/ScopeInfo.h
+++ include/clang/Sema/ScopeInfo.h
@@ -208,6 +208,13 @@
   /// The set of __block variables that are introduced in this function.
   llvm::TinyPtrVector<VarDecl *> ByrefBlockVars;
 
+  /// This map keeps track of block literals that initialize or are assigned to
+  /// local auto variables declared in the same scope as the block literal. We
+  /// use this information to determine whether it's possible to avoid emitting
+  /// a call to objc_retainBlock at the point of initialization or assignment.
+  llvm::DenseMap<const VarDecl *, llvm::SmallVector<BlockDecl *, 1>>
+      LocalVarToBlocks;
+
   /// A list of PartialDiagnostics created but delayed within the
   /// current function scope.  These diagnostics are vetted for reachability
   /// prior to being emitted.
@@ -442,6 +449,14 @@
     ByrefBlockVars.push_back(VD);
   }
 
+  void addLocalVarBlockDecl(const VarDecl *VD, BlockDecl *BD) {
+    LocalVarToBlocks[VD].push_back(BD);
+  }
+
+  void varDeclAddressTaken(const VarDecl *VD) {
+    LocalVarToBlocks.erase(VD);
+  }
+
   bool isCoroutine() const { return !FirstCoroutineStmtLoc.isInvalid(); }
 
   void setFirstCoroutineStmt(SourceLocation Loc, StringRef Keyword) {
Index: include/clang/AST/DeclBase.h
===================================================================
--- include/clang/AST/DeclBase.h
+++ include/clang/AST/DeclBase.h
@@ -1665,6 +1665,11 @@
     /// A bit that indicates this block is passed directly to a function as a
     /// non-escaping parameter.
     uint64_t DoesNotEscape : 1;
+
+    /// A bit that indicates whether it's possible to avoid coying this block to
+    /// the heap when it initializes or is assigned to a local variable with
+    /// automatic storage.
+    uint64_t CanAvoidCopyToHeap : 1;
   };
 
   /// Number of non-inherited bits in BlockDeclBitfields.
Index: include/clang/AST/Decl.h
===================================================================
--- include/clang/AST/Decl.h
+++ include/clang/AST/Decl.h
@@ -4008,6 +4008,13 @@
   bool doesNotEscape() const { return BlockDeclBits.DoesNotEscape; }
   void setDoesNotEscape(bool B = true) { BlockDeclBits.DoesNotEscape = B; }
 
+  bool canAvoidCopyToHeap() const {
+    return BlockDeclBits.CanAvoidCopyToHeap;
+  }
+  void setCanAvoidCopyToHeap(bool B = true) {
+    BlockDeclBits.CanAvoidCopyToHeap = B;
+  }
+
   bool capturesVariable(const VarDecl *var) const;
 
   void setCaptures(ASTContext &Context, ArrayRef<Capture> Captures,
_______________________________________________
cfe-commits mailing list
cfe-commits@lists.llvm.org
https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits

Reply via email to