[WebAssembly] Fix leak in Emscripten SjLj

[aheejin]

Summary

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.

llvm/lib/Target/WebAssembly/WebAssemblyLowerEmscriptenEHSjLj.cpp

for (Instruction &I : *BB) {
       BBs.push_back(Tail);
     }
   }
 
   // Erase everything we no longer need in this function
   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);
 
   // Every call to saveSetjmp can change setjmpTable and setjmpTableSize
   // (when buffer reallocation occurs)
   // entry:
   //   setjmpTableSize = 4;
   //   setjmpTable = (int *) malloc(40);
   //   setjmpTable[0] = 0;
   // ...

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

lpad:                                             ; preds = %entry
   br label %try.cont
 
 try.cont:                                         ; preds = %entry, %lpad
   ret void
 }
 
 ; 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:
   %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 %return, label %if.end
 
 if.end:                                           ; preds = %entry
   call void @foo()
   br label %return
 
 ; CHECK:    if.end:
 ; CHECK:      %cmp.eq.one = icmp eq i32 %__THREW__.val, 1
 ; CHECK-NEXT: br i1 %cmp.eq.one, label %eh.rethrow, label %normal
 
 ; CHECK:    normal:
 ; CHECK-NEXT: icmp ne i32 %__THREW__.val, 0
 
 ; 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*)
 ; Function Attrs: noreturn
 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

 ; CHECK: if.end:
 ; CHECK-NEXT: %[[LABEL_PHI:.*]] = phi i32 [ %[[LABEL:.*]], %if.end2 ], [ -1, %if.else1 ]
 ; CHECK-NEXT: %[[LONGJMP_RESULT]] = call i32 @getTempRet0()
 ; CHECK-NEXT: switch i32 %[[LABEL_PHI]], label %entry.split.split [
 ; CHECK-NEXT:   i32 1, label %entry.split
 ; 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
 
 ; CHECK: if.end2:
 ; CHECK-NEXT: call void  @setTempRet0(i32 %[[THREWVALUE_VAL]])
 ; CHECK-NEXT: br label %if.end
 }