[WebAssembly] Fix leak in Emscripten SjLj

For SjLj, we allocate a table to record setjmp buffer info in the entry
of each setjmp-calling function by inserting a `malloc` call, and insert
a `free` call to free the buffer before each `ret` instruction.

But this is not sufficient; we have to free the buffer before we throw.
In SjLj handling, normal functions that can possibly throw or longjmp
are wrapped with an invoke and caught within the function so they don't
end up escaping the function. But three functions throw and escape the
function:
- `__resumeException` (Emscripten library function used for Emscripten
  EH)
- `emscripten_longjmp` (Emscripten library function used for Emscripten
  SjLj)
- `__cxa_throw` (libc++abi function called when for C++ `throw` keyword)

The first two functions are used to rethrow the current
exception/longjmp when the caught exception/longjmp is not for the
current function. `__cxa_throw` is used for exception, and because we
consider that a function that cannot longjmp, it escapes the function
right away, before which we should free the buffer.

Currently `lsan.test_longjmp3` and `lsan.test_exceptions_longjmp3` fail
in Emscripten; this CL fixes these.

Reviewed By: dschuff

Differential Revision: https://reviews.llvm.org/D107852

Author: aheejin

llvm/lib/Target/WebAssembly/WebAssemblyLowerEmscriptenEHSjLj.cpp

@@ -1244,21 +1244,33 @@ bool WebAssemblyLowerEmscriptenEHSjLj::runSjLjOnFunction(Function &F) {
   for (Instruction *I : ToErase)
     I->eraseFromParent();
 
-  // Free setjmpTable buffer before each return instruction
+  // Free setjmpTable buffer before each return instruction + function-exiting
+  // call
+  SmallVector<Instruction *, 16> ExitingInsts;
   for (BasicBlock &BB : F) {
     Instruction *TI = BB.getTerminator();
-    if (isa<ReturnInst>(TI)) {
-      DebugLoc DL = getOrCreateDebugLoc(TI, F.getSubprogram());
-      auto *Free = CallInst::CreateFree(SetjmpTable, TI);
-      Free->setDebugLoc(DL);
-      // CallInst::CreateFree may create a bitcast instruction if its argument
-      // types mismatch. We need to set the debug loc for the bitcast too.
-      if (auto *FreeCallI = dyn_cast<CallInst>(Free)) {
-        if (auto *BitCastI = dyn_cast<BitCastInst>(FreeCallI->getArgOperand(0)))
-          BitCastI->setDebugLoc(DL);
+    if (isa<ReturnInst>(TI))
+      ExitingInsts.push_back(TI);
+    for (auto &I : BB) {
+      if (auto *CB = dyn_cast<CallBase>(&I)) {
+        StringRef CalleeName = CB->getCalledOperand()->getName();
+        if (CalleeName == "__resumeException" ||
+            CalleeName == "emscripten_longjmp" || CalleeName == "__cxa_throw")
+          ExitingInsts.push_back(&I);
       }
     }
   }
+  for (auto *I : ExitingInsts) {
+    DebugLoc DL = getOrCreateDebugLoc(I, F.getSubprogram());
+    auto *Free = CallInst::CreateFree(SetjmpTable, I);
+    Free->setDebugLoc(DL);
+    // CallInst::CreateFree may create a bitcast instruction if its argument
+    // types mismatch. We need to set the debug loc for the bitcast too.
+    if (auto *FreeCallI = dyn_cast<CallInst>(Free)) {
+      if (auto *BitCastI = dyn_cast<BitCastInst>(FreeCallI->getArgOperand(0)))
+        BitCastI->setDebugLoc(DL);
+    }
+  }
 
   // Every call to saveSetjmp can change setjmpTable and setjmpTableSize
   // (when buffer reallocation occurs)

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

@@ -89,7 +89,8 @@ try.cont:                                         ; preds = %entry, %lpad
 
 ; This function contains a setjmp call and no invoke, so we only handle longjmp
 ; here. But @foo can also throw an exception, so we check if an exception is
-; thrown and if so rethrow it by calling @__resumeException.
+; thrown and if so rethrow it by calling @__resumeException. Also we have to
+; free the setjmpTable buffer before calling @__resumeException.
 define void @rethrow_exception() {
 ; CHECK-LABEL: @rethrow_exception
 entry:
@@ -112,13 +113,41 @@ if.end:                                           ; preds = %entry
 
 ; CHECK:    eh.rethrow:
 ; CHECK-NEXT: %exn = call i8* @__cxa_find_matching_catch_2()
+; CHECK-NEXT: %[[BUF:.*]] = bitcast i32* %setjmpTable1 to i8*
+; CHECK-NEXT: call void @free(i8* %[[BUF]])
 ; CHECK-NEXT: call void @__resumeException(i8* %exn)
 ; CHECK-NEXT: unreachable
 
 return:                                           ; preds = %entry, %if.end
   ret void
 }
 
+; The same as 'rethrow_exception' but contains a __cxa_throw call. We have to
+; free the setjmpTable buffer before calling __cxa_throw.
+define void @rethrow_exception2() {
+; CHECK-LABEL: @rethrow_exception2
+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
+  %cmp = icmp ne i32 %call, 0
+  br i1 %cmp, label %throw, label %if.end
+
+if.end:                                           ; preds = %entry
+  call void @foo()
+  br label %throw
+
+throw:                                            ; preds = %entry, %if.end
+  call void @__cxa_throw(i8* null, i8* null, i8* null) #1
+  unreachable
+
+; CHECK: throw:
+; CHECK:      %[[BUF:.*]] = bitcast i32* %setjmpTable5 to i8*
+; CHECK-NEXT: call void @free(i8* %[[BUF]])
+; CHECK-NEXT: call void @__cxa_throw(i8* null, i8* null, i8* null)
+; CHECK-NEXT: unreachable
+}
+
 declare void @foo()
 ; Function Attrs: returns_twice
 declare i32 @setjmp(%struct.__jmp_buf_tag*)
@@ -127,6 +156,7 @@ declare void @longjmp(%struct.__jmp_buf_tag*, i32)
 declare i32 @__gxx_personality_v0(...)
 declare i8* @__cxa_begin_catch(i8*)
 declare void @__cxa_end_catch()
+declare void @__cxa_throw(i8*, i8*, i8*)
 
 attributes #0 = { returns_twice }
 attributes #1 = { noreturn }

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

@@ -71,6 +71,8 @@ entry:
 ; CHECK-NEXT: ]
 
 ; CHECK: if.then2:
+; CHECK-NEXT: %[[BUF:.*]] = bitcast i32* %[[SETJMP_TABLE1]] to i8*
+; CHECK-NEXT: call void @free(i8* %[[BUF]])
 ; CHECK-NEXT: call void @emscripten_longjmp([[PTR]] %[[__THREW__VAL]], i32 %[[THREWVALUE_VAL]])
 ; CHECK-NEXT: unreachable