[SjLjEHPrepare] Configure call sites correctly

After 9fe78db4, the pass inserts `store volatile i32 -1, ptr %call_site`
before all invoke instruction except the one in the entry block, which
has the effect of bypassing landing pads on exceptions.

When configuring the call site for a potentially throwing instruction
check that it is not `InvokeInst` -- they are handled by earlier code.

Author: s-barannikov

llvm/lib/CodeGen/SjLjEHPrepare.cpp

@@ -435,6 +435,10 @@ bool SjLjEHPrepareImpl::setupEntryBlockAndCallSites(Function &F) {
   // where to look for it.
   Builder.CreateCall(FuncCtxFn, FuncCtx);
 
+  // Register the function context and make sure it's known to not throw.
+  CallInst *Register = Builder.CreateCall(RegisterFn, FuncCtx, "");
+  Register->setDoesNotThrow();
+
   // At this point, we are all set up, update the invoke instructions to mark
   // their call_site values.
   for (unsigned I = 0, E = Invokes.size(); I != E; ++I) {
@@ -457,15 +461,10 @@ bool SjLjEHPrepareImpl::setupEntryBlockAndCallSites(Function &F) {
     if (&BB == &F.front())
       continue;
     for (Instruction &I : BB)
-      if (I.mayThrow())
+      if (!isa<InvokeInst>(I) && I.mayThrow())
         insertCallSiteStore(&I, -1);
   }
 
-  // Register the function context and make sure it's known to not throw
-  CallInst *Register = CallInst::Create(
-      RegisterFn, FuncCtx, "", EntryBB->getTerminator()->getIterator());
-  Register->setDoesNotThrow();
-
   // Following any allocas not in the entry block, update the saved SP in the
   // jmpbuf to the new value.
   for (BasicBlock &BB : F) {

llvm/test/CodeGen/ARM/sjljehprepare-lower-empty-struct.ll

@@ -16,11 +16,15 @@
 define ptr @foo(i8 %a, {} %c) personality ptr @baz {
 entry:
 ; CHECK: bl __Unwind_SjLj_Register
+; CHECK-NEXT: mov r0, #1
+; CHECK-NEXT: str r0, [sp, #{{[0-9]+}}]
 ; CHECK-NEXT: {{[A-Z][a-zA-Z0-9]*}}:
 ; CHECK-NEXT: bl _bar
 ; CHECK: bl __Unwind_SjLj_Resume
 
 ; CHECK-LINUX: bl _Unwind_SjLj_Register
+; CHECK-LINUX-NEXT: mov r0, #1
+; CHECK-LINUX-NEXT: str r0, [sp, #{{[0-9]+}}]
 ; CHECK-LINUX-NEXT: .{{[A-Z][a-zA-Z0-9]*}}:
 ; CHECK-LINUX-NEXT: bl bar
 ; CHECK-LINUX: bl _Unwind_SjLj_Resume

llvm/test/CodeGen/Generic/sjlj-eh-prepare.ll

@@ -0,0 +1,40 @@
+; RUN: opt -p sjlj-eh-prepare %s -S -o - | FileCheck %s
+
+; Check that callsites are set up correctly:
+; 1. Throwing call in the entry block does not set call_site
+;    (function context hasn't been configured yet).
+; 2. Throwing call not in the entry block sets call_site to -1
+;    (reset to the initial state).
+; 3. Invoke instructions set call_site to the correct call site number.
+; 4. Resume instruction sets call_site to -1 (reset to the initial state).
+
+define void @test_call_sites() personality ptr @__gxx_personality_sj0 {
+entry:
+  ; CHECK-NOT: store volatile
+  ; CHECK:     call void @may_throw()
+  call void @may_throw()
+
+  ; CHECK:      store volatile i32 1
+  ; CHECK-NEXT: call void @llvm.eh.sjlj.callsite(i32 1)
+  ; CHECK-NEXT: invoke void @may_throw()
+  invoke void @may_throw() to label %invoke.cont unwind label %lpad
+
+invoke.cont:
+  ; CHECK:      store volatile i32 2
+  ; CHECK-NEXT: call void @llvm.eh.sjlj.callsite(i32 2)
+  ; CHECK-NEXT: invoke void @may_throw()
+  invoke void @may_throw() to label %try.cont unwind label %lpad
+
+lpad:
+  ; CHECK:      store volatile i32 -1
+  ; CHECK-NEXT: resume
+  %lp = landingpad { ptr, i32 } cleanup
+  resume { ptr, i32 } %lp
+
+try.cont:
+  call void @may_throw()
+  ret void
+}
+
+declare void @may_throw()
+declare i32 @__gxx_personality_sj0(...)

llvm/test/CodeGen/X86/indirect-branch-tracking-eh2.ll

@@ -24,9 +24,9 @@ define dso_local i32 @main() #0 personality ptr @__gxx_personality_sj0 {
 ; NUM-NEXT:    movq %rbp, -104(%rbp)
 ; NUM-NEXT:    movq %rsp, -88(%rbp)
 ; NUM-NEXT:    movq $.LBB0_9, -96(%rbp)
-; NUM-NEXT:    movl $1, -144(%rbp)
 ; NUM-NEXT:    leaq -152(%rbp), %rdi
 ; NUM-NEXT:    callq _Unwind_SjLj_Register@PLT
+; NUM-NEXT:    movl $1, -144(%rbp)
 ; NUM-NEXT:  .Ltmp0:
 ; NUM-NEXT:    callq _Z3foov
 ; NUM-NEXT:  .Ltmp1:
@@ -110,9 +110,9 @@ define dso_local i32 @main() #0 personality ptr @__gxx_personality_sj0 {
 ; SJLJ-NEXT:    movq %rbp, -104(%rbp)
 ; SJLJ-NEXT:    movq %rsp, -88(%rbp)
 ; SJLJ-NEXT:    movq $.LBB0_9, -96(%rbp)
-; SJLJ-NEXT:    movl $1, -144(%rbp)
 ; SJLJ-NEXT:    leaq -152(%rbp), %rdi
 ; SJLJ-NEXT:    callq _Unwind_SjLj_Register@PLT
+; SJLJ-NEXT:    movl $1, -144(%rbp)
 ; SJLJ-NEXT:  .Ltmp0:
 ; SJLJ-NEXT:    callq _Z3foov
 ; SJLJ-NEXT:  .Ltmp1:

llvm/test/CodeGen/X86/sjlj-eh.ll

@@ -48,14 +48,13 @@ try.cont:
 ; CHECK: movl %esp, -24(%ebp)
 ;     UFC.__jbuf[1] = $EIP
 ; CHECK: movl $[[RESUME:LBB[0-9]+_[0-9]+]], -28(%ebp)
-;     UFC.__callsite = 1
-; CHECK: movl $1, -60(%ebp)
 ;     _Unwind_SjLj_Register(&UFC);
 ; CHECK: leal -64(%ebp), %eax
 ; CHECK: pushl %eax
 ; CHECK: calll __Unwind_SjLj_Register
 ; CHECK: addl $4, %esp
 ;     function_that_throws();
+; CHECK: movl $1, -60(%ebp)
 ; CHECK: calll __Z20function_that_throwsv
 ;     _Unwind_SjLj_Unregister(&UFC);
 ; CHECK: leal -64(%ebp), %eax
@@ -99,12 +98,11 @@ try.cont:
 ;     UFC.__jbuf[1] = $RIP
 ; CHECK-X64: leaq .[[RESUME:LBB[0-9]+_[0-9]+]](%rip), %rax
 ; CHECK-X64: movq %rax, -256(%rbp)
-;     UFC.__callsite = 1
-; CHECK-X64: movl $1, -304(%rbp)
 ;     _Unwind_SjLj_Register(&UFC);
 ; CHECK-X64: leaq -312(%rbp), %rcx
 ; CHECK-X64: callq _Unwind_SjLj_Register
 ;     function_that_throws();
+; CHECK-X64: movl $1, -304(%rbp)
 ; CHECK-X64: callq _Z20function_that_throwsv
 ;     _Unwind_SjLj_Unregister(&UFC);
 ; CHECK-X64: leaq -312(%rbp), %rcx