[WebAssembly] Fix a bug in 'try' placement

[aheejin]

Summary

When searching for local expression tree created by stackified
registers, for block placement, we start the search from the previous
instruction of a BB's terminator. But in try's case, we should start
from the previous instruction of a call that can throw, or a EH_LABEL
that precedes the call, because the return values of the call's previous
instructions can be stackified and consumed by the throwing call.

For example,

i32.call @foo
call @bar         ; may throw
br $label0

In this case, if we start the search from the previous instruction of
the terminator (br here), we end up stopping at call @bar and place
a 'try' between i32.call @foo and call @bar, because call @bar
does not have a return value so it is not a local expression tree of
br.

But in this case, unlike when placing 'block's, we should start the
search from call @bar, because the return value of i32.call @foo is
stackified and used by call @bar.

llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp

if (MI.getOpcode() == WebAssembly::END_BLOCK ||
       BeforeSet.insert(&MI);
 #endif
 
     // Terminators should go after the TRY.
     if (MI.isTerminator())
       AfterSet.insert(&MI);
   }
 
-  // Local expression tree should go after the TRY.
-  for (auto I = Header->getFirstTerminator(), E = Header->begin(); I != E;
-       --I) {
-    if (std::prev(I)->isDebugInstr() || std::prev(I)->isPosition())
-      continue;
-    if (WebAssembly::isChild(*std::prev(I), MFI))
-      AfterSet.insert(&*std::prev(I));
-    else
-      break;
-  }
-
   // If Header unwinds to MBB (= Header contains 'invoke'), the try block should
   // contain the call within it. So the call should go after the TRY. The
   // exception is when the header's terminator is a rethrow instruction, in
   // which case that instruction, not a call instruction before it, is gonna
   // throw.
+  MachineInstr *ThrowingCall = nullptr;
   if (MBB.isPredecessor(Header)) {
     auto TermPos = Header->getFirstTerminator();
     if (TermPos == Header->end() ||
         TermPos->getOpcode() != WebAssembly::RETHROW) {
-      for (const auto &MI : reverse(*Header)) {
+      for (auto &MI : reverse(*Header)) {
         if (MI.isCall()) {
           AfterSet.insert(&MI);
+          ThrowingCall = &MI;
           // Possibly throwing calls are usually wrapped by EH_LABEL
           // instructions. We don't want to split them and the call.
           if (MI.getIterator() != Header->begin() &&
-              std::prev(MI.getIterator())->isEHLabel())
+              std::prev(MI.getIterator())->isEHLabel()) {
             AfterSet.insert(&*std::prev(MI.getIterator()));
+            ThrowingCall = &*std::prev(MI.getIterator());
+          }
           break;
         }
       }
     }
   }
 
+  // Local expression tree should go after the TRY.
+  // For BLOCK placement, we start the search from the previous instruction of a
+  // BB's terminator, but in TRY's case, we should start from the previous
+  // instruction of a call that can throw, or a EH_LABEL that precedes the call,
+  // because the return values of the call's previous instructions can be
+  // stackified and consumed by the throwing call.
+  auto SearchStartPt = ThrowingCall ? MachineBasicBlock::iterator(ThrowingCall)
+                                    : Header->getFirstTerminator();
+  for (auto I = SearchStartPt, E = Header->begin(); I != E; --I) {
+    if (std::prev(I)->isDebugInstr() || std::prev(I)->isPosition())
+      continue;
+    if (WebAssembly::isChild(*std::prev(I), MFI))
+      AfterSet.insert(&*std::prev(I));
+    else
+      break;
+  }
+
   // Add the TRY.
   auto InsertPos = getLatestInsertPos(Header, BeforeSet, AfterSet);
   MachineInstr *Begin =
       BuildMI(*Header, InsertPos, Header->findDebugLoc(InsertPos),
               TII.get(WebAssembly::TRY))
           .addImm(int64_t(WebAssembly::ExprType::Void));
 
   // Decide where in Header to put the END_TRY.

llvm/test/CodeGen/WebAssembly/cfg-stackify-eh.ll

invoke.cont:                                      ; preds = %entry
   ret void
 
 ehcleanup:                                        ; preds = %entry
   %0 = cleanuppad within none []
   %call2 = call %class.Object* @_ZN6ObjectD2Ev(%class.Object* %o) #1 [ "funclet"(token %0) ]
   cleanupret from %0 unwind to caller
 }
 
+; Tests if 'try' marker is placed correctly. In this test, 'try' should be
+; placed before the call to 'nothrow_i32' and not between the call to
+; 'nothrow_i32' and 'fun', because the return value of 'nothrow_i32' is
+; stackified and pushed onto the stack to be consumed by the call to 'fun'.
+
+; CHECK-LABEL: test11
+; CHECK: try
+; CHECK: i32.call  $push{{.*}}=, nothrow_i32
+; CHECK: call      fun, $pop{{.*}}
+define void @test11() personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
+entry:
+  %call = call i32 @nothrow_i32()
+  invoke void @fun(i32 %call)
+          to label %invoke.cont unwind label %terminate
+
+invoke.cont:                                      ; preds = %entry
+  ret void
+
+terminate:                                        ; preds = %entry
+  %0 = cleanuppad within none []
+  %1 = tail call i8* @llvm.wasm.get.exception(token %0)
+  call void @__clang_call_terminate(i8* %1) [ "funclet"(token %0) ]
+  unreachable
+}
+
 ; Check if the unwind destination mismatch stats are correct
 ; NOSORT-STAT: 11 wasm-cfg-stackify    - Number of EH pad unwind mismatches found
 
 declare void @foo()
 declare void @bar()
 declare i32 @baz()
+declare void @fun(i32)
 ; Function Attrs: nounwind
 declare void @nothrow(i32) #0
+declare i32 @nothrow_i32() #0
 ; Function Attrs: nounwind
 declare %class.Object* @_ZN6ObjectD2Ev(%class.Object* returned) #0
 declare i32 @__gxx_wasm_personality_v0(...)
 declare i8* @llvm.wasm.get.exception(token)
 declare i32 @llvm.wasm.get.ehselector(token)
 declare void @llvm.wasm.rethrow.in.catch()
 declare i32 @llvm.eh.typeid.for(i8*)
 declare i8* @__cxa_begin_catch(i8*)