CodeGen: block capture shouldn't ICE

[jfb]

Summary

When a lambda capture captures a __block in the same statement, the compiler asserts out because isCapturedBy assumes that an Expr can only be a BlockExpr, StmtExpr, or if it's a Stmt then all the statement's children are expressions. That's wrong, we need to visit all sub-statements even if they're not expressions to see if they also capture.

Fix this issue by pulling out the isCapturedBy logic to use RecursiveASTVisitor.

rdar://problem/39926584

[rjmccall]

RecursiveASTVisitor instantiations are huge. Can you just make the function take a Stmt and then do the first few checks if it happens to be an Expr?

[jfb]

In D47096#1105368, @rjmccall wrote:

RecursiveASTVisitor instantiations are huge. Can you just make the function take a Stmt and then do the first few checks if it happens to be an Expr?

I'm not super-familiar with the code, so I might be doing something silly.

I did something like this initially (leave the top of the function as-is, and instead of cast do dyn_cast to Expr and if that fails to CompoundStmt, recursively iterating all the children of the CompoundStmt). My worry was that I wasn't traversing all actual children (just CompountStmt's children), and AFAICT there's no easy way to say "take any Stmt, and visit its children if it has such a method". I could hard-code more Stmt derivatives but that seems brittle, I could use the "detection idiom" but that's silly if there's already a visitor which does The Right Thing through tablegen magic.

What I can do is what I did earlier, and conservatively say it was captured if it's neither an Expr nor a CompoundStmt? Or should I special-case other things as well?

[rjmccall]

In D47096#1105374, @jfb wrote:

In D47096#1105368, @rjmccall wrote:

RecursiveASTVisitor instantiations are huge. Can you just make the function take a Stmt and then do the first few checks if it happens to be an Expr?

I'm not super-familiar with the code, so I might be doing something silly.

I did something like this initially (leave the top of the function as-is, and instead of cast do dyn_cast to Expr and if that fails to CompoundStmt, recursively iterating all the children of the CompoundStmt). My worry was that I wasn't traversing all actual children (just CompountStmt's children), and AFAICT there's no easy way to say "take any Stmt, and visit its children if it has such a method". I could hard-code more Stmt derivatives but that seems brittle, I could use the "detection idiom" but that's silly if there's already a visitor which does The Right Thing through tablegen magic.

What I can do is what I did earlier, and conservatively say it was captured if it's neither an Expr nor a CompoundStmt? Or should I special-case other things as well?

children() is actually defined at the Stmt level, and if you look at how it's implemented on e.g. IfStmt, you can see that it visits all of the child Stmts, including the if-condition. So it should be fine.

[jfb]

• Follow John's suggestion.

[jfb]

In D47096#1105455, @rjmccall wrote:

children() is actually defined at the Stmt level, and if you look at how it's implemented on e.g. IfStmt, you can see that it visits all of the child Stmts, including the if-condition. So it should be fine.

Thanks for pointing this out! I was reading the code but missed that, and visitor seemed like the only reliable way to do what I wanted. Here's an update patch.

[rjmccall]

Test case should go in test/CodeGenCXX. Also, there already is a blocks.cpp there.

[jfb]

• move test

[jfb]

In D47096#1105492, @rjmccall wrote:

Test case should go in test/CodeGenCXX. Also, there already is a blocks.cpp there.

I moved it, but didn't merge with the existing block.cpp because it just checks for not crashing. I'd rather make sure that the checks don't inadvertently pick up the wrong test.

[rjmccall]

LGTM.

[jfb]

r332801

lib/CodeGen/CGDecl.cpp

 #include "CodeGenModule.h"
 #include "ConstantEmitter.h"
 #include "TargetInfo.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/CharUnits.h"
 #include "clang/AST/Decl.h"
 #include "clang/AST/DeclObjC.h"
 #include "clang/AST/DeclOpenMP.h"
+#include "clang/AST/RecursiveASTVisitor.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Basic/TargetInfo.h"
 #include "clang/CodeGen/CGFunctionInfo.h"
 #include "clang/Frontend/CodeGenOptions.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/Type.h"
CodeGenFunction::EmitAutoVarAlloca(const VarDecl &D) {
   if (emission.useLifetimeMarkers())
     EHStack.pushCleanup<CallLifetimeEnd>(NormalEHLifetimeMarker,
                                          emission.getOriginalAllocatedAddress(),
                                          emission.getSizeForLifetimeMarkers());
 
   return emission;
 }
 
-/// Determines whether the given __block variable is potentially
-/// captured by the given expression.
-static bool isCapturedBy(const VarDecl &var, const Expr *e) {
-  // Skip the most common kinds of expressions that make
-  // hierarchy-walking expensive.
-  e = e->IgnoreParenCasts();
+static bool isCapturedBy(const VarDecl &, const Expr *);
 
-  if (const BlockExpr *be = dyn_cast<BlockExpr>(e)) {
-    const BlockDecl *block = be->getBlockDecl();
-    for (const auto &I : block->captures()) {
-      if (I.getVariable() == &var)
-        return true;
-    }
+class IsCapturedBy : public RecursiveASTVisitor<IsCapturedBy> {
+  bool IsCaptured = false;
+  const VarDecl &Var;
+  bool setCaptured() {
+    IsCaptured = true;
+    // End visitation early, we know it's captured.
+    return false;
 
+public:
+  IsCapturedBy(const VarDecl &Var) : Var(Var) {}
+  bool isCaptured() const { return IsCaptured; }
+
+  bool TraverseBlockExpr(BlockExpr *BE) {
+    const BlockDecl *Block = BE->getBlockDecl();
+    for (const auto &I : Block->captures())
+      if (I.getVariable() == &Var)
+        return setCaptured();
     // No need to walk into the subexpressions.
     return false;
   }
 
-  if (const StmtExpr *SE = dyn_cast<StmtExpr>(e)) {
+  bool TraverseStmtExpr(StmtExpr *SE) {
     const CompoundStmt *CS = SE->getSubStmt();
     for (const auto *BI : CS->body())
       if (const auto *E = dyn_cast<Expr>(BI)) {
-        if (isCapturedBy(var, E))
-            return true;
-      }
-      else if (const auto *DS = dyn_cast<DeclStmt>(BI)) {
-          // special case declarations
-          for (const auto *I : DS->decls()) {
-              if (const auto *VD = dyn_cast<VarDecl>((I))) {
-                const Expr *Init = VD->getInit();
-                if (Init && isCapturedBy(var, Init))
-                  return true;
-              }
-          }
+        if (isCapturedBy(Var, E))
+          return setCaptured();
+      } else if (const auto *DS = dyn_cast<DeclStmt>(BI)) {
+            if (Init && isCapturedBy(Var, Init))
+              return setCaptured();
       }
       else
-        // FIXME. Make safe assumption assuming arbitrary statements cause capturing.
-        // Later, provide code to poke into statements for capture analysis.
-        return true;
-    return false;
+        // FIXME. Make safe assumption assuming arbitrary statements cause
+        // capturing. Later, provide code to poke into statements for capture
+        // analysis.
+        return setCaptured();
   }
+};
 
-  for (const Stmt *SubStmt : e->children())
-    if (isCapturedBy(var, cast<Expr>(SubStmt)))
-      return true;
+/// Determines whether the given __block variable is potentially
+/// captured by the given expression.
+static bool isCapturedBy(const VarDecl &Var, const Expr *E) {
+  // Skip the most common kinds of expressions that make
+  // hierarchy-walking expensive.
+  E = E->IgnoreParenCasts();
 
-  return false;
+  IsCapturedBy Visitor(Var);
+  Visitor.TraverseStmt(const_cast<Expr *>(E));
+  return Visitor.isCaptured();
 }
 
 /// Determine whether the given initializer is trivial in the sense
 /// that it requires no code to be generated.
 bool CodeGenFunction::isTrivialInitializer(const Expr *Init) {
   if (!Init)
     return true;
 

test/CodeGen/block-capture.cpp

+// RUN: %clang_cc1 -x c++ -std=c++11 -fblocks -emit-llvm %s -o - | FileCheck %s
+
+// CHECK: %struct.__block_byref_baz = type { i8*, %struct.__block_byref_baz*, i32, i32, i32 }
+// CHECK: [[baz:%[0-9a-z_]*]] = alloca %struct.__block_byref_baz
+// CHECK: [[bazref:%[0-9a-z_\.]*]] = getelementptr inbounds %struct.__block_byref_baz, %struct.__block_byref_baz* [[baz]], i32 0, i32 1
+// CHECK: store %struct.__block_byref_baz* [[baz]], %struct.__block_byref_baz** [[bazref]]
+// CHECK: [[disposable:%[0-9a-z_]*]] = bitcast %struct.__block_byref_baz* [[baz]] to i8*
+// CHECK: call void @_Block_object_dispose(i8* [[disposable]]
+
+int main() {
+  __block int baz = [&]() { return 0; }();
+  return 0;
+}