[ObjC] Fix lifetime markers of loop variable in EmitObjCForCollectionStmt

[kubamracek]

Summary

CodeGenFunction::EmitObjCForCollectionStmt currently emits lifetime markers for the loop variable in an inconsistent way: lifetime.start is emitted before the loop is entered, but lifetime.end is emitted inside the loop:

; entry block
  %u = alloca %1*, align 8
  %1 = bitcast %1** %u to i8*
  call void @llvm.lifetime.start(i64 8, i8* %1) #5
  ...

; loop body
  ...
  %14 = bitcast %1** %u to i8*
  call void @llvm.lifetime.end(i64 8, i8* %14) #5
  ...
  br i1 %21, ... ; loop

; loop ended
  ret void

AddressSanitizer uses these markers to track out-of-scope accesses to local variables, and we get false positives in Obj-C foreach loops (in the 2nd iteration of the loop). The markers of the loop variable need to be either both inside the loop (so that we poison and unpoison the variable in each iteration), or both outside. This patch implements the "both inside" approach and makes EmitObjCForCollectionStmt emit:

; entry block
  %u = alloca %1*, align 8
  ...

; loop body
  %12 = bitcast %1** %u to i8*
  call void @llvm.lifetime.start(i64 8, i8* %12) #5
  ...
  %14 = bitcast %1** %u to i8*
  call void @llvm.lifetime.end(i64 8, i8* %14) #5
  br label %15

; loop ended
  ret void

The test fixups are only changing the order of allocas. There's some related discussion at https://reviews.llvm.org/D18618.

[kubamracek]

Note that C++ foreach loops also generate lifetime.start and lifetime.end inside of the loop body.

[rjmccall]

Yes, looks good to me.

[kubamracek]

Reverted because this fails for-in.m by crashing the compiler when compiling:

void t2(NSArray *array) {
  for (NSArray *array in array) { // expected-warning {{collection expression type 'NSArray *' may not respond}}
  }
}

[kubamracek]

Trying a different approach: Keeping the loop variable alive for the whole loop by extending ForScope and registering the cleanup function inside EmitAutoVarAlloca.

[rjmccall]

Those test changes are smaller than I thought they might be; great.

cfe/trunk/lib/CodeGen/CGObjC.cpp

void CodeGenFunction::EmitObjCForCollectionStmt(const ObjCForCollectionStmt &S){
   }
   CGCallee EnumerationMutationFn =
     CGCallee::forDirect(EnumerationMutationFnPtr);
 
   CGDebugInfo *DI = getDebugInfo();
   if (DI)
     DI->EmitLexicalBlockStart(Builder, S.getSourceRange().getBegin());
 
-  // The local variable comes into scope immediately.
-  AutoVarEmission variable = AutoVarEmission::invalid();
-  if (const DeclStmt *SD = dyn_cast<DeclStmt>(S.getElement()))
-    variable = EmitAutoVarAlloca(*cast<VarDecl>(SD->getSingleDecl()));
-
   JumpDest LoopEnd = getJumpDestInCurrentScope("forcoll.end");
 
   // Fast enumeration state.
   QualType StateTy = CGM.getObjCFastEnumerationStateType();
   Address StatePtr = CreateMemTemp(StateTy, "state.ptr");
   EmitNullInitialization(StatePtr, StateTy);
 
   // Number of elements in the items array.
void CodeGenFunction::EmitObjCForCollectionStmt(const ObjCForCollectionStmt &S){
   // Otherwise, or if the mutation function returns, just continue.
   EmitBlock(WasNotMutatedBB);
 
   // Initialize the element variable.
   RunCleanupsScope elementVariableScope(*this);
   bool elementIsVariable;
   LValue elementLValue;
   QualType elementType;
+  AutoVarEmission variable = AutoVarEmission::invalid();
   if (const DeclStmt *SD = dyn_cast<DeclStmt>(S.getElement())) {
     // Initialize the variable, in case it's a __block variable or something.
+    variable = EmitAutoVarAlloca(*cast<VarDecl>(SD->getSingleDecl()));
     EmitAutoVarInit(variable);
 
     const VarDecl* D = cast<VarDecl>(SD->getSingleDecl());
     DeclRefExpr tempDRE(const_cast<VarDecl*>(D), false, D->getType(),
                         VK_LValue, SourceLocation());
     elementLValue = EmitLValue(&tempDRE);
     elementType = D->getType();
     elementIsVariable = true;

cfe/trunk/test/CodeGenObjC/arc-foreach.m

void test0(NSArray *array) {
   // capture should require a retain.
   for (id x in array) {
     use_block(^{ use(x); });
   }
 }
 
 // CHECK-LP64-LABEL:    define void @test0(
 // CHECK-LP64:      [[ARRAY:%.*]] = alloca [[ARRAY_T:%.*]]*,
-// CHECK-LP64-NEXT: [[X:%.*]] = alloca i8*,
 // CHECK-LP64-NEXT: [[STATE:%.*]] = alloca [[STATE_T:%.*]],
 // CHECK-LP64-NEXT: [[BUFFER:%.*]] = alloca [16 x i8*], align 8
+// CHECK-LP64-NEXT: [[X:%.*]] = alloca i8*,
 // CHECK-LP64-NEXT: [[BLOCK:%.*]] = alloca [[BLOCK_T:<{.*}>]],
 
 // Initialize 'array'.
 // CHECK-LP64-NEXT: store [[ARRAY_T]]* null, [[ARRAY_T]]** [[ARRAY]]
 // CHECK-LP64-NEXT: [[ZERO:%.*]] = bitcast [[ARRAY_T]]** [[ARRAY]] to i8**
 // CHECK-LP64-NEXT: [[ONE:%.*]] = bitcast [[ARRAY_T]]* {{%.*}} to i8*
 // CHECK-LP64-NEXT: call void @objc_storeStrong(i8** [[ZERO]], i8* [[ONE]]) [[NUW:#[0-9]+]]
 

 void test1(NSArray *array) {
   for (__weak id x in array) {
     use_block(^{ use(x); });
   }
 }
 
 // CHECK-LP64-LABEL:    define void @test1(
 // CHECK-LP64:      alloca [[ARRAY_T:%.*]]*,
-// CHECK-LP64-NEXT: [[X:%.*]] = alloca i8*,
 // CHECK-LP64-NEXT: [[STATE:%.*]] = alloca [[STATE_T:%.*]],
 // CHECK-LP64-NEXT: alloca [16 x i8*], align 8
+// CHECK-LP64-NEXT: [[X:%.*]] = alloca i8*,
 // CHECK-LP64-NEXT: [[BLOCK:%.*]] = alloca [[BLOCK_T:<{.*}>]],
 
 // CHECK-LP64:      [[T0:%.*]] = getelementptr inbounds [[STATE_T]], [[STATE_T]]* [[STATE]], i32 0, i32 1
 // CHECK-LP64-NEXT: [[T1:%.*]] = load i8**, i8*** [[T0]]
 // CHECK-LP64-NEXT: [[T2:%.*]] = getelementptr i8*, i8** [[T1]], i64
 // CHECK-LP64-NEXT: [[T3:%.*]] = load i8*, i8** [[T2]]
 // CHECK-LP64-NEXT: call i8* @objc_initWeak(i8** [[X]], i8* [[T3]])
 

 void test3(NSArray *array) {
   for (id x in array) {
     if (!x) continue;
     use(x);
   }
 
   // CHECK-LP64-LABEL:    define void @test3(
   // CHECK-LP64:      [[ARRAY:%.*]] = alloca [[ARRAY_T]]*, align 8
-  // CHECK-LP64-NEXT: [[X:%.*]] = alloca i8*, align 8
+  // CHECK-LP64:      [[X:%.*]] = alloca i8*, align 8
   // CHECK-LP64:      [[T0:%.*]] = load i8*, i8** [[X]], align 8
   // CHECK-LP64-NEXT: [[T1:%.*]] = icmp ne i8* [[T0]], null
   // CHECK-LP64-NEXT: br i1 [[T1]],
   // CHECK-LP64:      br label [[L:%[^ ]+]]
   // CHECK-LP64:      [[T0:%.*]] = load i8*, i8** [[X]], align 8
   // CHECK-LP64-NEXT: call void @use(i8* [[T0]])
   // CHECK-LP64-NEXT: br label [[L]]
 }

cfe/trunk/test/CodeGenObjC/arc-ternary-op.m

 // need to find another example of something we need to do this for.
 void test2(int cond) {
   extern id test2_producer(void);
   for (id obj in cond ? test2_producer() : (void*) 0) {
   }
 
   // CHECK-LABEL:    define void @test2(
   // CHECK:      [[COND:%.*]] = alloca i32,
-  // CHECK:      alloca i8*
   // CHECK:      [[CLEANUP_SAVE:%.*]] = alloca i8*
   // CHECK:      [[RUN_CLEANUP:%.*]] = alloca i1
+  // CHECK:      alloca i8*
   //   Evaluate condition; cleanup disabled by default.
   // CHECK:      [[T0:%.*]] = load i32, i32* [[COND]],
   // CHECK-NEXT: icmp ne i32 [[T0]], 0
   // CHECK-NEXT: store i1 false, i1* [[RUN_CLEANUP]]
   // CHECK-NEXT: br i1
   //   Within true branch, cleanup enabled.
   // CHECK:      [[T0:%.*]] = call i8* @test2_producer()
   // CHECK-NEXT: [[T1:%.*]] = call i8* @objc_retainAutoreleasedReturnValue(i8* [[T0]])