[WebAssembly] Handle indirect uses of longjmp

[aheejin]

Summary

In LowerEmscriptenEHSjLj, longjmp used to be replaced with
emscripten_longjmp_jmpbuf(jmp_buf*, i32), which will eventually be
lowered to emscripten_longjmp(i32, i32). The reason we used two
different names was because they had different signatures in the IR
pass.

D88697 fixed this by only using emscripten_longjmp(i32, i32) and
adding a ptrtoint cast to its first argument, so

longjmp(buf, 0)

becomes

emscripten_longjmp((i32)buf, 0)

But this assumed all uses of longjmp was a direct call to it, which
was not the case. This patch handles indirect uses of longjmp by
replacing

longjmp

with

(i32(*)(jmp_buf*, i32))emscripten_longjmp

[tlively]

Nice!

llvm/lib/Target/WebAssembly/WebAssemblyLowerEmscriptenEHSjLj.cpp

 // arguments of type {i32, i32} and longjmp takes {jmp_buf*, i32}, so we need a
 // ptrtoint instruction here to make the type match. jmp_buf* will eventually be
 // lowered to i32 in the wasm backend.
 static void replaceLongjmpWithEmscriptenLongjmp(Function *LongjmpF,
                                                 Function *EmLongjmpF) {
   SmallVector<CallInst *, 8> ToErase;
   LLVMContext &C = LongjmpF->getParent()->getContext();
   IRBuilder<> IRB(C);
+
+  // For calls to longjmp, replace it with emscripten_longjmp and cast its first
+  // argument (jmp_buf*) to int
   for (User *U : LongjmpF->users()) {
     auto *CI = dyn_cast<CallInst>(U);
-    if (!CI)
-      report_fatal_error("Does not support indirect calls to longjmp");
+    if (CI && CI->getCalledFunction() == LongjmpF) {
       IRB.SetInsertPoint(CI);
       Value *Jmpbuf =
           IRB.CreatePtrToInt(CI->getArgOperand(0), IRB.getInt32Ty(), "jmpbuf");
       IRB.CreateCall(EmLongjmpF, {Jmpbuf, CI->getArgOperand(1)});
       ToErase.push_back(CI);
     }
-
+  }
   for (auto *I : ToErase)
     I->eraseFromParent();
+
+  // If we have any remaining uses of longjmp's function pointer, replace it
+  // with (int(*)(jmp_buf*, int))emscripten_longjmp.
+  if (!LongjmpF->uses().empty()) {
+    Value *EmLongjmp =
+        IRB.CreateBitCast(EmLongjmpF, LongjmpF->getType(), "em_longjmp");
+    LongjmpF->replaceAllUsesWith(EmLongjmp);
+  }
 }
 
 bool WebAssemblyLowerEmscriptenEHSjLj::runOnModule(Module &M) {
   LLVM_DEBUG(dbgs() << "********** Lower Emscripten EH & SjLj **********\n");
 
   LLVMContext &C = M.getContext();
   IRBuilder<> IRB(C);
 

llvm/test/CodeGen/WebAssembly/lower-em-sjlj.ll

 ; RUN: opt < %s -wasm-lower-em-ehsjlj -S | FileCheck %s --check-prefixes=CHECK,NO-TLS
 ; RUN: opt < %s -wasm-lower-em-ehsjlj -S --mattr=+atomics,+bulk-memory | FileCheck %s --check-prefixes=CHECK,TLS
 
 target datalayout = "e-m:e-p:32:32-i64:64-n32:64-S128"
 target triple = "wasm32-unknown-unknown"
 
 %struct.__jmp_buf_tag = type { [6 x i32], i32, [32 x i32] }
 
 @global_var = global i32 0, align 4
 ; NO-TLS-DAG: __THREW__ = external global i32
 ; NO-TLS-DAG: __threwValue = external global i32
 ; TLS-DAG: __THREW__ = external thread_local(localexec) global i32
 ; TLS-DAG: __threwValue = external thread_local(localexec) global i32
+@global_longjmp_ptr = global void (%struct.__jmp_buf_tag*, i32)* @longjmp, align 4
+; CHECK-DAG: @global_longjmp_ptr = global void (%struct.__jmp_buf_tag*, i32)* bitcast (void (i32, i32)* @emscripten_longjmp to void (%struct.__jmp_buf_tag*, i32)*)
 
 ; Test a simple setjmp - longjmp sequence
 define void @setjmp_longjmp() {
 ; CHECK-LABEL: @setjmp_longjmp
 entry:
   %buf = alloca [1 x %struct.__jmp_buf_tag], align 16
   %arraydecay = getelementptr inbounds [1 x %struct.__jmp_buf_tag], [1 x %struct.__jmp_buf_tag]* %buf, i32 0, i32 0
   %call = call i32 @setjmp(%struct.__jmp_buf_tag* %arraydecay) #0

 for.inc:                                          ; preds = %for.cond
   %call5 = call i32 @setjmp(%struct.__jmp_buf_tag* undef) #0
   br label %for.cond
 
 ; CHECK: for.inc.split:
   ; CHECK: %var[[VARNO]] = phi i32 [ %var, %for.inc ]
 }
 
+; Tests cases where longjmp function pointer is used in other ways than direct
+; calls. longjmps should be replaced with
+; (int(*)(jmp_buf*, int))emscripten_longjmp.
+declare void @take_longjmp(void (%struct.__jmp_buf_tag*, i32)* %arg_ptr)
+define void @indirect_longjmp() {
+; CHECK-LABEL: @indirect_longjmp
+entry:
+  %local_longjmp_ptr = alloca void (%struct.__jmp_buf_tag*, i32)*, align 4
+  %buf0 = alloca [1 x %struct.__jmp_buf_tag], align 16
+  %buf1 = alloca [1 x %struct.__jmp_buf_tag], align 16
+
+  ; Store longjmp in a local variable, load it, and call it
+  store void (%struct.__jmp_buf_tag*, i32)* @longjmp, void (%struct.__jmp_buf_tag*, i32)** %local_longjmp_ptr, align 4
+  ; CHECK: store void (%struct.__jmp_buf_tag*, i32)* bitcast (void (i32, i32)* @emscripten_longjmp to void (%struct.__jmp_buf_tag*, i32)*), void (%struct.__jmp_buf_tag*, i32)** %local_longjmp_ptr, align 4
+  %longjmp_from_local_ptr = load void (%struct.__jmp_buf_tag*, i32)*, void (%struct.__jmp_buf_tag*, i32)** %local_longjmp_ptr, align 4
+  %arraydecay = getelementptr inbounds [1 x %struct.__jmp_buf_tag], [1 x %struct.__jmp_buf_tag]* %buf0, i32 0, i32 0
+  call void %longjmp_from_local_ptr(%struct.__jmp_buf_tag* %arraydecay, i32 0)
+
+  ; Load longjmp from a global variable and call it
+  %longjmp_from_global_ptr = load void (%struct.__jmp_buf_tag*, i32)*, void (%struct.__jmp_buf_tag*, i32)** @global_longjmp_ptr, align 4
+  %arraydecay1 = getelementptr inbounds [1 x %struct.__jmp_buf_tag], [1 x %struct.__jmp_buf_tag]* %buf1, i32 0, i32 0
+  call void %longjmp_from_global_ptr(%struct.__jmp_buf_tag* %arraydecay1, i32 0)
+
+  ; Pass longjmp as a function argument. This is a call but longjmp is not a
+  ; callee but an argument.
+  call void @take_longjmp(void (%struct.__jmp_buf_tag*, i32)* @longjmp)
+  ; CHECK: call void @take_longjmp(void (%struct.__jmp_buf_tag*, i32)* bitcast (void (i32, i32)* @emscripten_longjmp to void (%struct.__jmp_buf_tag*, i32)*))
+  ret void
+}
+
 declare void @foo()
 ; Function Attrs: returns_twice
 declare i32 @setjmp(%struct.__jmp_buf_tag*) #0
 ; Function Attrs: noreturn
 declare void @longjmp(%struct.__jmp_buf_tag*, i32) #1
 declare i32 @__gxx_personality_v0(...)
 declare i8* @__cxa_begin_catch(i8*)
 declare void @__cxa_end_catch()