[BasicAA] Treat returns_twice functions as clobbering unescaped objects (#117902)

Effectively this models all the accesses that occur between the first
and second return as happening at the point of the call.

Fixes #116668.

Author: nikic

llvm/lib/Analysis/BasicAliasAnalysis.cpp

@@ -947,8 +947,14 @@ ModRefInfo BasicAAResult::getModRefInfo(const CallBase *Call,
   //
   // Make sure the object has not escaped here, and then check that none of the
   // call arguments alias the object below.
+  //
+  // We model calls that can return twice (setjmp) as clobbering non-escaping
+  // objects, to model any accesses that may occur prior to the second return.
+  // As an exception, ignore allocas, as setjmp is not required to preserve
+  // non-volatile stores for them.
   if (!isa<Constant>(Object) && Call != Object &&
-      AAQI.CA->isNotCapturedBefore(Object, Call, /*OrAt*/ false)) {
+      AAQI.CA->isNotCapturedBefore(Object, Call, /*OrAt*/ false) &&
+      (isa<AllocaInst>(Object) || !Call->hasFnAttr(Attribute::ReturnsTwice))) {
 
     // Optimistically assume that call doesn't touch Object and check this
     // assumption in the following loop.

llvm/test/Transforms/GVN/setjmp.ll

@@ -5,7 +5,6 @@ declare i32 @setjmp() returns_twice
 declare void @longjmp()
 declare ptr @malloc(i64)
 
-; FIXME: This is a miscompile.
 define i32 @test() {
 ; CHECK-LABEL: define i32 @test() {
 ; CHECK-NEXT:    [[MALLOC:%.*]] = call noalias ptr @malloc(i64 4)
@@ -18,7 +17,8 @@ define i32 @test() {
 ; CHECK-NEXT:    call void @longjmp()
 ; CHECK-NEXT:    unreachable
 ; CHECK:       [[IF_END]]:
-; CHECK-NEXT:    ret i32 10
+; CHECK-NEXT:    [[RES:%.*]] = load i32, ptr [[MALLOC]], align 4
+; CHECK-NEXT:    ret i32 [[RES]]
 ;
   %malloc = call noalias ptr @malloc(i64 4)
   store i32 10, ptr %malloc, align 4
@@ -35,3 +35,65 @@ if.end:
   %res = load i32, ptr %malloc
   ret i32 %res
 }
+
+; We are still allowed to optimize non-volatile accesses to allocas.
+define i32 @test_alloca() {
+; CHECK-LABEL: define i32 @test_alloca() {
+; CHECK-NEXT:    [[ALLOC:%.*]] = alloca i43, align 8
+; CHECK-NEXT:    store i32 10, ptr [[ALLOC]], align 4
+; CHECK-NEXT:    [[SJ:%.*]] = call i32 @setjmp()
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[SJ]], 0
+; CHECK-NEXT:    br i1 [[CMP]], label %[[IF_THEN:.*]], label %[[IF_END:.*]]
+; CHECK:       [[IF_THEN]]:
+; CHECK-NEXT:    store i32 20, ptr [[ALLOC]], align 4
+; CHECK-NEXT:    call void @longjmp()
+; CHECK-NEXT:    unreachable
+; CHECK:       [[IF_END]]:
+; CHECK-NEXT:    ret i32 10
+;
+  %alloc = alloca i43
+  store i32 10, ptr %alloc, align 4
+  %sj = call i32 @setjmp()
+  %cmp = icmp eq i32 %sj, 0
+  br i1 %cmp, label %if.then, label %if.end
+
+if.then:
+  store i32 20, ptr %alloc
+  call void @longjmp()
+  unreachable
+
+if.end:
+  %res = load i32, ptr %alloc
+  ret i32 %res
+}
+
+define i32 @test_alloca_volatile() {
+; CHECK-LABEL: define i32 @test_alloca_volatile() {
+; CHECK-NEXT:    [[ALLOC:%.*]] = alloca i43, align 8
+; CHECK-NEXT:    store volatile i32 10, ptr [[ALLOC]], align 4
+; CHECK-NEXT:    [[SJ:%.*]] = call i32 @setjmp()
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[SJ]], 0
+; CHECK-NEXT:    br i1 [[CMP]], label %[[IF_THEN:.*]], label %[[IF_END:.*]]
+; CHECK:       [[IF_THEN]]:
+; CHECK-NEXT:    store volatile i32 20, ptr [[ALLOC]], align 4
+; CHECK-NEXT:    call void @longjmp()
+; CHECK-NEXT:    unreachable
+; CHECK:       [[IF_END]]:
+; CHECK-NEXT:    [[RES:%.*]] = load volatile i32, ptr [[ALLOC]], align 4
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+  %alloc = alloca i43
+  store volatile i32 10, ptr %alloc, align 4
+  %sj = call i32 @setjmp()
+  %cmp = icmp eq i32 %sj, 0
+  br i1 %cmp, label %if.then, label %if.end
+
+if.then:
+  store volatile i32 20, ptr %alloc
+  call void @longjmp()
+  unreachable
+
+if.end:
+  %res = load volatile i32, ptr %alloc
+  ret i32 %res
+}