Variable auto-init: fix __block initialization

[jfb]

Summary

Automatic initialization [1] of __block variables was trampling over the block's headers after they'd been initialized, which caused self-init usage to crash, such as here:

typedef struct XYZ { void (^block)(); } *xyz_t;
__attribute__((noinline))
xyz_t create(void (^block)()) {
  xyz_t myself = malloc(sizeof(struct XYZ));
  myself->block = block;
  return myself;
}
int main() {
  __block xyz_t captured = create(^(){ (void)captured; });
}

This type of code shouldn't be broken by variable auto-init, even if it's sketchy.

[1] With -ftrivial-auto-var-init=pattern

rdar://problem/47798396

[rjmccall]

lib/CodeGen/CGDecl.cpp

1615

Why don't you just initialize the payload right here, after we've drilled down to it?

1645

Does this check handle flexible arrays on zero-sized structures properly? I suppose they're always initialized.

1646

Can't you just drill to the byref storage right here and avoid all the other complexity and subtle ordering interactions?

1665

This is a late comment on the main patch, but there's an ASTContext::getAsVariableArrayType method.

[rjmccall]

lib/CodeGen/CGDecl.cpp

1615

I meant to delete this comment before submitting, sorry. I am suggesting the other place.

[jfb]

lib/CodeGen/CGDecl.cpp

1645

You mean something like

void foo(int size) {
  struct Z {};
  Z arr[size];
}

?

1646

We're in the lambda that does the initialization here. The tricky order part is that code that calls the lambda does:

• Block (which was missing the early auto-init)
• Trivial initializer (which has auto-init, then early exit)
• Constant aggregate / constexpr (which might auto-init if it wasn't constant, and then early-exit)
• The other stuff

I don't think here is the right place to do anything... and I'm not sure what you're suggesting.

1665

OK I can do in a separate follow-up.

[rjmccall]

lib/CodeGen/CGDecl.cpp

1646

Escaping __block variables are basically a fixed-layout header followed by storage of the variable's formal type. Anything you do at this point in the function to auto-initialize the header is a waste of time and code size because it is precisely at this point in the function that we perform a bunch of stores to initialize that header. You are mitigating nothing by covering the header. So what I'm saying is that, in this lambda which is meant to initialize the user-exposed storage of a variable, you should just make sure you're pointing at the user-exposed storage of the variable by calling emitBlockByrefAddress(false) (which just does a GEP), and then you can initialize only that.

[jfb]

• Only initialize __block's storage.

[jfb]

lib/CodeGen/CGDecl.cpp

1646

Thanks, I now understand what you meant. It's super simple indeed, sorry I took a while to grok. Updated patch should do this now.

[rjmccall]

lib/CodeGen/CGDecl.cpp

1605

Are these changes still needed?

1646

Thanks, this is what I was looking for.

[jfb]

lib/CodeGen/CGDecl.cpp

1724

Note that we still want this to be pulled out in this way because emitByrefStructureInit emits the call to the initializer (in test_block_self_init the call to create). Were we to leave this as it was before, one of the initializations below would be emitted, but it would be *after* the call.

Similarly, we also want the little dance I added to only initialize once.

[jfb]

lib/CodeGen/CGDecl.cpp

1605

I believe so, see concurrent comment here:
https://reviews.llvm.org/D57797#inline-512702

[rjmccall]

lib/CodeGen/CGDecl.cpp

1724

Oh, that's interesting, I hadn't remembered that. Alright, in that case LGTM, I guess.

...unless you want to refactor this so that emitByrefStructureInit only handles the header initialization and the bit about handling the initializer is handled more cleanly down in the rest of this code.

[jfb]

• Simplify patch greatly.

[jfb]

lib/CodeGen/CGDecl.cpp

1605

Actually I'm wrong, updated patch drops this and the test makes sure the call is after this initialization. Sorry for the confusion. This is much cleaner.

1645

Just to circle back, the above code lowers to:

%struct.Z = type { i8 }

@__const._Z3fooi.arr = private unnamed_addr constant %struct.Z { i8 -86 }, align 1

; Function Attrs: noinline nounwind optnone ssp uwtable
define void @_Z3fooi(i32 %size) #0 {
entry:
  %size.addr = alloca i32, align 4
  %saved_stack = alloca i8*, align 8
  %__vla_expr0 = alloca i64, align 8
  store i32 %size, i32* %size.addr, align 4
  %0 = load i32, i32* %size.addr, align 4
  %1 = zext i32 %0 to i64
  %2 = call i8* @llvm.stacksave()
  store i8* %2, i8** %saved_stack, align 8
  %vla = alloca %struct.Z, i64 %1, align 16
  store i64 %1, i64* %__vla_expr0, align 8
  %vla.iszerosized = icmp eq i64 %1, 0
  br i1 %vla.iszerosized, label %vla-init.cont, label %vla-setup.loop

vla-setup.loop:                                   ; preds = %entry
  %vla.begin = bitcast %struct.Z* %vla to i8*
  %vla.end = getelementptr inbounds i8, i8* %vla.begin, i64 %1
  br label %vla-init.loop

vla-init.loop:                                    ; preds = %vla-init.loop, %vla-setup.loop
  %vla.cur = phi i8* [ %vla.begin, %vla-setup.loop ], [ %vla.next, %vla-init.loop ]
  call void @llvm.memcpy.p0i8.p0i8.i64(i8* align 1 %vla.cur, i8* align 1 getelementptr inbounds (%struct.Z, %struct.Z* @__const._Z3fooi.arr, i32 0, i32 0), i64 1, i1 false)
  %vla.next = getelementptr inbounds i8, i8* %vla.cur, i64 1
  %vla-init.isdone = icmp eq i8* %vla.next, %vla.end
  br i1 %vla-init.isdone, label %vla-init.cont, label %vla-init.loop

vla-init.cont:                                    ; preds = %vla-init.loop, %entry
  %3 = load i8*, i8** %saved_stack, align 8
  call void @llvm.stackrestore(i8* %3)
  ret void
}

So yes it's handled.

1665

https://reviews.llvm.org/rC353490

1724

As noted above, I was wrong here.

[jfb]

• Remove whitespace changes.

[rjmccall]

SGTM, thanks!

lib/CodeGen/CGDecl.cpp

void CodeGenFunction::EmitAutoVarInit(const AutoVarEmission &emission) {
   if (!HaveInsertPoint()) {
     if (!Init || !ContainsLabel(Init)) return;
     EnsureInsertPoint();
   }
 
   // Initialize the structure of a __block variable.
   if (emission.IsEscapingByRef)
     emitByrefStructureInit(emission);
 
   // Initialize the variable here if it doesn't have a initializer and it is a
   // C struct that is non-trivial to initialize or an array containing such a
   // struct.
   if (!Init &&
       type.isNonTrivialToPrimitiveDefaultInitialize() ==
           QualType::PDIK_Struct) {
     LValue Dst = MakeAddrLValue(emission.getAllocatedAddress(), type);
     if (emission.IsEscapingByRef)
       drillIntoBlockVariable(*this, Dst, &D);
     defaultInitNonTrivialCStructVar(Dst);
     return;
   }
 
   // Check whether this is a byref variable that's potentially
   // captured and moved by its own initializer.  If so, we'll need to
   // emit the initializer first, then copy into the variable.
   bool capturedByInit =
       Init && emission.IsEscapingByRef && isCapturedBy(D, Init);
 
   Address Loc =
       capturedByInit ? emission.Addr : emission.getObjectAddress(*this);
 
   // Note: constexpr already initializes everything correctly.
   LangOptions::TrivialAutoVarInitKind trivialAutoVarInit =
       (D.isConstexpr()
            ? LangOptions::TrivialAutoVarInitKind::Uninitialized
            : (D.getAttr<UninitializedAttr>()
                   ? LangOptions::TrivialAutoVarInitKind::Uninitialized
                   : getContext().getLangOpts().getTrivialAutoVarInit()));
 
-  auto initializeWhatIsTechnicallyUninitialized = [&]() {
+  auto initializeWhatIsTechnicallyUninitialized = [&](Address Loc) {
     if (trivialAutoVarInit ==
         LangOptions::TrivialAutoVarInitKind::Uninitialized)
       return;
 
+    // Only initialize a __block's storage: we always initialize the header.
+    if (emission.IsEscapingByRef)
+      Loc = emitBlockByrefAddress(Loc, &D, /*follow=*/false);
+
     CharUnits Size = getContext().getTypeSizeInChars(type);
     if (!Size.isZero()) {
       switch (trivialAutoVarInit) {
       case LangOptions::TrivialAutoVarInitKind::Uninitialized:
         llvm_unreachable("Uninitialized handled above");
       case LangOptions::TrivialAutoVarInitKind::Zero:
         emitStoresForZeroInit(CGM, D, Loc, isVolatile, Builder);
         break;
       case LangOptions::TrivialAutoVarInitKind::Pattern:
         emitStoresForPatternInit(CGM, D, Loc, isVolatile, Builder);
         break;
       }
       return;
     }
 
     // VLAs look zero-sized to getTypeInfo. We can't emit constant stores to
     // them, so emit a memcpy with the VLA size to initialize each element.
     // Technically zero-sized or negative-sized VLAs are undefined, and UBSan
     // will catch that code, but there exists code which generates zero-sized
     // VLAs. Be nice and initialize whatever they requested.
     const auto *VlaType = getContext().getAsVariableArrayType(type);
     if (!VlaType)
       return;
     auto VlaSize = getVLASize(VlaType);
     auto SizeVal = VlaSize.NumElts;
     CharUnits EltSize = getContext().getTypeSizeInChars(VlaSize.Type);
     switch (trivialAutoVarInit) {
     case LangOptions::TrivialAutoVarInitKind::Uninitialized:
       llvm_unreachable("Uninitialized handled above");
case LangOptions::TrivialAutoVarInitKind::Pattern: {
       Builder.CreateCondBr(Done, ContBB, LoopBB);
       Cur->addIncoming(Next, LoopBB);
       EmitBlock(ContBB);
     } break;
     }
   };
 
   if (isTrivialInitializer(Init)) {
-    initializeWhatIsTechnicallyUninitialized();
+    initializeWhatIsTechnicallyUninitialized(Loc);
     return;
   }
 
   llvm::Constant *constant = nullptr;
   if (emission.IsConstantAggregate || D.isConstexpr()) {
     assert(!capturedByInit && "constant init contains a capturing block?");
     constant = ConstantEmitter(*this).tryEmitAbstractForInitializer(D);
     if (constant && trivialAutoVarInit !=
                         LangOptions::TrivialAutoVarInitKind::Uninitialized)
       constant = replaceUndef(constant);
   }
 
   if (!constant) {
-    initializeWhatIsTechnicallyUninitialized();
+    initializeWhatIsTechnicallyUninitialized(Loc);
     LValue lv = MakeAddrLValue(Loc, type);
     lv.setNonGC(true);
     return EmitExprAsInit(Init, &D, lv, capturedByInit);
   }
 
   if (!emission.IsConstantAggregate) {
     // For simple scalar/complex initialization, store the value directly.
     LValue lv = MakeAddrLValue(Loc, type);

test/CodeGenCXX/trivial-auto-var-init.cpp

 // ZERO: store i32 0, i32* %block
 // PATTERN-LABEL: test_block(
 // PATTERN: store i32 -1431655766, i32* %block
 void test_block() {
   __block int block;
   used(block);
 }
 
+// Using the variable being initialized is typically UB in C, but for blocks we
+// can be nice: they imply extra book-keeping and we can do the auto-init before
+// any of said book-keeping.
+//
+// UNINIT-LABEL:  test_block_self_init(
+// ZERO-LABEL:    test_block_self_init(
+// ZERO:          %block = alloca <{ i8*, i32, i32, i8*, %struct.__block_descriptor*, i8* }>, align 8
+// ZERO:          %captured1 = getelementptr inbounds %struct.__block_byref_captured, %struct.__block_byref_captured* %captured, i32 0, i32 4
+// ZERO-NEXT:     store %struct.XYZ* null, %struct.XYZ** %captured1, align 8
+// ZERO:          %call = call %struct.XYZ* @create(
+// PATTERN-LABEL: test_block_self_init(
+// PATTERN:       %block = alloca <{ i8*, i32, i32, i8*, %struct.__block_descriptor*, i8* }>, align 8
+// PATTERN:       %captured1 = getelementptr inbounds %struct.__block_byref_captured, %struct.__block_byref_captured* %captured, i32 0, i32 4
+// PATTERN-NEXT:  store %struct.XYZ* inttoptr (i64 -6148914691236517206 to %struct.XYZ*), %struct.XYZ** %captured1, align 8
+// PATTERN:       %call = call %struct.XYZ* @create(
+void test_block_self_init() {
+  using Block = void (^)();
+  typedef struct XYZ {
+    Block block;
+  } * xyz_t;
+  extern xyz_t create(Block block);
+  __block xyz_t captured = create(^() {
+    (void)captured;
+  });
+}
+
 // This type of code is currently not handled by zero / pattern initialization.
 // The test will break when that is fixed.
 // UNINIT-LABEL:  test_goto_unreachable_value(
 // ZERO-LABEL:    test_goto_unreachable_value(
 // ZERO-NOT: store {{.*}}%oops
 // PATTERN-LABEL: test_goto_unreachable_value(
 // PATTERN-NOT: store {{.*}}%oops
 void test_goto_unreachable_value() {