[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.

lib/CodeGen/CGDecl.cpp

if (CGDebugInfo *DI = getDebugInfo()) {
         DI->setLocation(D.getLocation());
         DI->EmitDeclareOfAutoVariable(&D, address.getPointer(), Builder);
       }
     }
 
   if (D.hasAttr<AnnotateAttr>())
     EmitVarAnnotations(&D, address.getPointer());
 
+  // Make sure we call @llvm.lifetime.end.
+  if (emission.useLifetimeMarkers())
+    EHStack.pushCleanup<CallLifetimeEnd>(NormalEHLifetimeMarker,
+                                         emission.getAllocatedAddress(),
+                                         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.
void CodeGenFunction::EmitAutoVarCleanups(const AutoVarEmission &emission) {
   if (emission.wasEmittedAsGlobal()) return;
 
   // If we don't have an insertion point, we're done.  Sema prevents
   // us from jumping into any of these scopes anyway.
   if (!HaveInsertPoint()) return;
 
   const VarDecl &D = *emission.Variable;
 
-  // Make sure we call @llvm.lifetime.end.  This needs to happen
-  // *last*, so the cleanup needs to be pushed *first*.
-  if (emission.useLifetimeMarkers())
-    EHStack.pushCleanup<CallLifetimeEnd>(NormalEHLifetimeMarker,
-                                         emission.getAllocatedAddress(),
-                                         emission.getSizeForLifetimeMarkers());
-
   // Check the type for a cleanup.
   if (QualType::DestructionKind dtorKind = D.getType().isDestructedType())
     emitAutoVarTypeCleanup(emission, dtorKind);
 
   // In GC mode, honor objc_precise_lifetime.
   if (getLangOpts().getGC() != LangOptions::NonGC &&
       D.hasAttr<ObjCPreciseLifetimeAttr>()) {
     EHStack.pushCleanup<ExtendGCLifetime>(NormalCleanup, &D);

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());
 
+  RunCleanupsScope ForScope(*this);
+
   // 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.
Selector FastEnumSel =
     CGM.getContext().Selectors.getSelector(llvm::array_lengthof(II), &II[0]);
 
   QualType ItemsTy =
     getContext().getConstantArrayType(getContext().getObjCIdType(),
                                       llvm::APInt(32, NumItems),
                                       ArrayType::Normal, 0);
   Address ItemsPtr = CreateMemTemp(ItemsTy, "items.ptr");
 
-  RunCleanupsScope ForScope(*this);
-
   // Emit the collection pointer.  In ARC, we do a retain.
   llvm::Value *Collection;
   if (getLangOpts().ObjCAutoRefCount) {
     Collection = EmitARCRetainScalarExpr(S.getCollection());
 
     // Enter a cleanup to do the release.
     EmitObjCConsumeObject(S.getCollection()->getType(), Collection);
   } else {

test/CodeGenObjC/arc-blocks.m

void test13(id x) {
   // CHECK-NEXT: [[T2:%.*]] = call i8* @objc_retainBlock(i8* [[T1]])
   // CHECK-NEXT: [[T3:%.*]] = bitcast i8* [[T2]] to void ()*
   // CHECK-NEXT: store void ()* [[T3]], void ()** [[B]], align 8
   // CHECK-NEXT: [[T0:%.*]] = load void ()*, void ()** [[B]], align 8
   // CHECK-NEXT: call void @test13_use(void ()* [[T0]])
   // CHECK-NEXT: [[T0:%.*]] = load void ()*, void ()** [[B]]
   // CHECK-NEXT: [[T1:%.*]] = bitcast void ()* [[T0]] to i8*
   // CHECK-NEXT: call void @objc_release(i8* [[T1]])
-  // CHECK-NEXT: [[BPTR2:%.*]] = bitcast void ()** [[B]] to i8*
-  // CHECK-NEXT: call void @llvm.lifetime.end.p0i8(i64 8, i8* [[BPTR2]])
 
   // CHECK-NEXT: [[T0:%.*]] = load i1, i1* [[CLEANUP_ACTIVE]]
   // CHECK-NEXT: br i1 [[T0]]
   // CHECK:      [[T0:%.*]] = load i8*, i8** [[CLEANUP_ADDR]]
   // CHECK-NEXT: call void @objc_release(i8* [[T0]])
   // CHECK-NEXT: br label
 
-  // CHECK:      [[T0:%.*]] = load i8*, i8** [[X]]
+  // CHECK:      [[BPTR2:%.*]] = bitcast void ()** [[B]] to i8*
+  // CHECK-NEXT: call void @llvm.lifetime.end.p0i8(i64 8, i8* [[BPTR2]])
+  // CHECK-NEXT:      [[T0:%.*]] = load i8*, i8** [[X]]
   // CHECK-NEXT: call void @objc_release(i8* [[T0]])
   // CHECK-NEXT: ret void
 }
 
 // <rdar://problem/10907510>
 void test14() {
   void (^const x[1])(void) = { ^{} };
 }

test/CodeGenObjCXX/arc-references.mm

void test0() {
   callee();
   // CHECK: call i8* @_Z9getObjectv
   // CHECK-NEXT: call i8* @objc_retainAutoreleasedReturnValue
   // CHECK-NEXT: call i8* @objc_autorelease
   const __autoreleasing id &ref2 = getObject();
   // CHECK: call void @_Z6calleev
   callee();
   // CHECK: call void @objc_release
-  // CHECK-NEXT: ret
+  // CHECK: ret
 }
 
 // No lifetime extension when we're binding a reference to an lvalue.
 // CHECK-LABEL: define void @_Z5test1RU8__strongP11objc_objectRU6__weakS0_
 void test1(__strong id &x, __weak id &y) {
   // CHECK-NOT: release
   const __strong id &ref1 = x;
   const __autoreleasing id &ref2 = x;
   const __weak id &ref3 = y;
   // CHECK: ret void
 }
 
 typedef __strong id strong_id;
 
 //CHECK: define void @_Z5test3v
 void test3() {
   // CHECK: [[REF:%.*]] = alloca i8**, align 8
   // CHECK: call i8* @objc_initWeak
   // CHECK-NEXT: store i8**
   const __weak id &ref = strong_id();
   // CHECK-NEXT: call void @_Z6calleev()
   callee();
+  // CHECK-NEXT: call void @objc_destroyWeak
   // CHECK-NEXT: [[PTR:%.*]] = bitcast i8*** [[REF]] to i8*
   // CHECK-NEXT: call void @llvm.lifetime.end.p0i8(i64 8, i8* [[PTR]])
-  // CHECK-NEXT: call void @objc_destroyWeak
   // CHECK-NEXT: ret void
 }
 
 // CHECK-LABEL: define void @_Z5test4RU8__strongP11objc_object
 void test4(__strong id &x) {
   // CHECK: call i8* @objc_retain
   __strong A* const &ar = x;
   // CHECK: store i32 17, i32*