[WebAssembly] Fix reverse mapping in WasmEHFuncInfo

aheejin · aheejin · commit aa097ef8d474 · 2021-02-26T17:12:10.000-08:00
D97247 added the reverse mapping from unwind destination to their source, but it had a critical bug; sources can be multiple, because multiple BBs can have a single BB as their unwind destination. This changes `WasmEHFuncInfo::getUnwindSrc` to `getUnwindSrcs` and makes it return a vector rather than a single BB. It does not return the const reference to the existing vector but creates a new vector because `WasmEHFuncInfo` stores not `BasicBlock*` or `MachineBasicBlock*` but `PointerUnion` of them. Also I hoped to unify those methods for `BasicBlock` and `MachineBasicBlock` into one using templates to reduce duplication, but failed because various usages require `BasicBlock*` to be `const` but it's hard to make it `const` for `MachineBasicBlock` usages. Fixes emscripten-core/emscripten#13514. (More precisely, fixes emscripten-core/emscripten#13514 (comment)) Reviewed By: dschuff, tlively Differential Revision: https://reviews.llvm.org/D97583
diff --git a/llvm/include/llvm/CodeGen/WasmEHFuncInfo.h b/llvm/include/llvm/CodeGen/WasmEHFuncInfo.h
@@ -15,6 +15,7 @@
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/PointerUnion.h"
+#include "llvm/ADT/SmallPtrSet.h"
 
 namespace llvm {
 
@@ -32,41 +33,58 @@ struct WasmEHFuncInfo {
   // When there is an entry <A, B>, if an exception is not caught by A, it
   // should next unwind to the EH pad B.
   DenseMap<BBOrMBB, BBOrMBB> SrcToUnwindDest;
-  DenseMap<BBOrMBB, BBOrMBB> UnwindDestToSrc; // reverse map
+  DenseMap<BBOrMBB, SmallPtrSet<BBOrMBB, 4>> UnwindDestToSrcs; // reverse map
 
   // Helper functions
   const BasicBlock *getUnwindDest(const BasicBlock *BB) const {
+    assert(hasUnwindDest(BB));
     return SrcToUnwindDest.lookup(BB).get<const BasicBlock *>();
   }
-  const BasicBlock *getUnwindSrc(const BasicBlock *BB) const {
-    return UnwindDestToSrc.lookup(BB).get<const BasicBlock *>();
+  SmallPtrSet<const BasicBlock *, 4> getUnwindSrcs(const BasicBlock *BB) const {
+    assert(hasUnwindSrcs(BB));
+    const auto &Set = UnwindDestToSrcs.lookup(BB);
+    SmallPtrSet<const BasicBlock *, 4> Ret;
+    for (const auto &P : Set)
+      Ret.insert(P.get<const BasicBlock *>());
+    return Ret;
   }
   void setUnwindDest(const BasicBlock *BB, const BasicBlock *Dest) {
     SrcToUnwindDest[BB] = Dest;
-    UnwindDestToSrc[Dest] = BB;
+    if (!UnwindDestToSrcs.count(Dest))
+      UnwindDestToSrcs[Dest] = SmallPtrSet<BBOrMBB, 4>();
+    UnwindDestToSrcs[Dest].insert(BB);
   }
   bool hasUnwindDest(const BasicBlock *BB) const {
     return SrcToUnwindDest.count(BB);
   }
-  bool hasUnwindSrc(const BasicBlock *BB) const {
-    return UnwindDestToSrc.count(BB);
+  bool hasUnwindSrcs(const BasicBlock *BB) const {
+    return UnwindDestToSrcs.count(BB);
   }
 
   MachineBasicBlock *getUnwindDest(MachineBasicBlock *MBB) const {
+    assert(hasUnwindDest(MBB));
     return SrcToUnwindDest.lookup(MBB).get<MachineBasicBlock *>();
   }
-  MachineBasicBlock *getUnwindSrc(MachineBasicBlock *MBB) const {
-    return UnwindDestToSrc.lookup(MBB).get<MachineBasicBlock *>();
+  SmallPtrSet<MachineBasicBlock *, 4>
+  getUnwindSrcs(MachineBasicBlock *MBB) const {
+    assert(hasUnwindSrcs(MBB));
+    const auto &Set = UnwindDestToSrcs.lookup(MBB);
+    SmallPtrSet<MachineBasicBlock *, 4> Ret;
+    for (const auto &P : Set)
+      Ret.insert(P.get<MachineBasicBlock *>());
+    return Ret;
   }
   void setUnwindDest(MachineBasicBlock *MBB, MachineBasicBlock *Dest) {
     SrcToUnwindDest[MBB] = Dest;
-    UnwindDestToSrc[Dest] = MBB;
+    if (!UnwindDestToSrcs.count(Dest))
+      UnwindDestToSrcs[Dest] = SmallPtrSet<BBOrMBB, 4>();
+    UnwindDestToSrcs[Dest].insert(MBB);
   }
   bool hasUnwindDest(MachineBasicBlock *MBB) const {
     return SrcToUnwindDest.count(MBB);
   }
-  bool hasUnwindSrc(MachineBasicBlock *MBB) const {
-    return UnwindDestToSrc.count(MBB);
+  bool hasUnwindSrcs(MachineBasicBlock *MBB) const {
+    return UnwindDestToSrcs.count(MBB);
   }
 };
 
diff --git a/llvm/lib/CodeGen/SelectionDAG/FunctionLoweringInfo.cpp b/llvm/lib/CodeGen/SelectionDAG/FunctionLoweringInfo.cpp
@@ -333,21 +333,23 @@ void FunctionLoweringInfo::set(const Function &fn, MachineFunction &mf,
 
   else if (Personality == EHPersonality::Wasm_CXX) {
     WasmEHFuncInfo &EHInfo = *MF->getWasmEHFuncInfo();
-    // Map all BB references in the WinEH data to MBBs.
-    DenseMap<BBOrMBB, BBOrMBB> NewMap;
+    // Map all BB references in the Wasm EH data to MBBs.
+    DenseMap<BBOrMBB, BBOrMBB> SrcToUnwindDest;
     for (auto &KV : EHInfo.SrcToUnwindDest) {
       const auto *Src = KV.first.get<const BasicBlock *>();
-      const auto *Dst = KV.second.get<const BasicBlock *>();
-      NewMap[MBBMap[Src]] = MBBMap[Dst];
+      const auto *Dest = KV.second.get<const BasicBlock *>();
+      SrcToUnwindDest[MBBMap[Src]] = MBBMap[Dest];
     }
-    EHInfo.SrcToUnwindDest = std::move(NewMap);
-    NewMap.clear();
-    for (auto &KV : EHInfo.UnwindDestToSrc) {
-      const auto *Src = KV.first.get<const BasicBlock *>();
-      const auto *Dst = KV.second.get<const BasicBlock *>();
-      NewMap[MBBMap[Src]] = MBBMap[Dst];
+    EHInfo.SrcToUnwindDest = std::move(SrcToUnwindDest);
+    DenseMap<BBOrMBB, SmallPtrSet<BBOrMBB, 4>> UnwindDestToSrcs;
+    for (auto &KV : EHInfo.UnwindDestToSrcs) {
+      const auto *Dest = KV.first.get<const BasicBlock *>();
+      UnwindDestToSrcs[MBBMap[Dest]] = SmallPtrSet<BBOrMBB, 4>();
+      for (const auto &P : KV.second)
+        UnwindDestToSrcs[MBBMap[Dest]].insert(
+            MBBMap[P.get<const BasicBlock *>()]);
     }
-    EHInfo.UnwindDestToSrc = std::move(NewMap);
+    EHInfo.UnwindDestToSrcs = std::move(UnwindDestToSrcs);
   }
 }
 
diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyCFGSort.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyCFGSort.cpp
@@ -259,11 +259,12 @@ static void sortBlocks(MachineFunction &MF, const MachineLoopInfo &MLI,
         // incorrect because we may end up delegating/rethrowing to an inner
         // scope in CFGStackify. So here we make sure those unwind destinations
         // are deferred until their unwind source's exception is sorted.
-        if (EHInfo && EHInfo->hasUnwindSrc(Succ)) {
-          auto *UnwindSrc = EHInfo->getUnwindSrc(Succ);
+        if (EHInfo && EHInfo->hasUnwindSrcs(Succ)) {
+          SmallPtrSet<MachineBasicBlock *, 4> UnwindSrcs =
+              EHInfo->getUnwindSrcs(Succ);
           bool IsDeferred = false;
-          for (Entry &E : reverse(Entries)) {
-            if (E.TheRegion->getHeader() == UnwindSrc) {
+          for (Entry &E : Entries) {
+            if (UnwindSrcs.count(E.TheRegion->getHeader())) {
               E.Deferred.push_back(Succ);
               IsDeferred = true;
               break;
diff --git a/llvm/test/CodeGen/WebAssembly/cfg-stackify-eh.ll b/llvm/test/CodeGen/WebAssembly/cfg-stackify-eh.ll
@@ -1222,7 +1222,7 @@ try.cont:                                         ; preds = %catch.start1, %catc
 ; end_block
 ;           <- (b) The br destination should be remapped to here
 ;
-; The test was reduced by bugpoint and should not crash.
+; The test was reduced by bugpoint and should not crash in CFGStackify.
 define void @test21() personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
 entry:
   br i1 undef, label %if.then, label %if.end12
@@ -1298,6 +1298,65 @@ unreachable:                                      ; preds = %if.then, %invoke.co
   unreachable
 }
 
+; Regression test for WasmEHFuncInfo's reverse mapping bug. 'UnwindDestToSrc'
+; should return a vector and not a single BB, which was incorrect.
+; This was reduced by bugpoint and should not crash in CFGStackify.
+define void @test22() personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
+entry:
+  invoke void @foo()
+          to label %invoke.cont unwind label %catch.dispatch
+
+catch.dispatch:                                   ; preds = %entry
+  %0 = catchswitch within none [label %catch.start] unwind label %ehcleanup22
+
+catch.start:                                      ; preds = %catch.dispatch
+  %1 = catchpad within %0 [i8* bitcast (i8** @_ZTIi to i8*)]
+  %2 = call i8* @llvm.wasm.get.exception(token %1)
+  %3 = call i32 @llvm.wasm.get.ehselector(token %1)
+  invoke void @__cxa_throw() #1 [ "funclet"(token %1) ]
+          to label %unreachable unwind label %catch.dispatch2
+
+catch.dispatch2:                                  ; preds = %catch.start
+  %4 = catchswitch within %1 [label %catch.start3] unwind label %ehcleanup
+
+catch.start3:                                     ; preds = %catch.dispatch2
+  %5 = catchpad within %4 [i8* bitcast (i8** @_ZTIi to i8*)]
+  %6 = call i8* @llvm.wasm.get.exception(token %5)
+  %7 = call i32 @llvm.wasm.get.ehselector(token %5)
+  catchret from %5 to label %try.cont
+
+try.cont:                                         ; preds = %catch.start3
+  invoke void @foo() [ "funclet"(token %1) ]
+          to label %invoke.cont8 unwind label %ehcleanup
+
+invoke.cont8:                                     ; preds = %try.cont
+  invoke void @__cxa_throw() #1 [ "funclet"(token %1) ]
+          to label %unreachable unwind label %catch.dispatch11
+
+catch.dispatch11:                                 ; preds = %invoke.cont8
+  %8 = catchswitch within %1 [label %catch.start12] unwind label %ehcleanup
+
+catch.start12:                                    ; preds = %catch.dispatch11
+  %9 = catchpad within %8 [i8* bitcast (i8** @_ZTIi to i8*)]
+  %10 = call i8* @llvm.wasm.get.exception(token %9)
+  %11 = call i32 @llvm.wasm.get.ehselector(token %9)
+  unreachable
+
+invoke.cont:                                      ; preds = %entry
+  unreachable
+
+ehcleanup:                                        ; preds = %try.cont, %catch.dispatch11, %catch.dispatch2
+  %12 = cleanuppad within %1 []
+  cleanupret from %12 unwind label %ehcleanup22
+
+ehcleanup22:                                      ; preds = %ehcleanup, %catch.dispatch
+  %13 = cleanuppad within none []
+  cleanupret from %13 unwind to caller
+
+unreachable:                                      ; preds = %catch.start, %invoke.cont8
+  unreachable
+}
+
 ; Check if the unwind destination mismatch stats are correct
 ; NOSORT: 20 wasm-cfg-stackify    - Number of call unwind mismatches found
 ; NOSORT:  4 wasm-cfg-stackify    - Number of catch unwind mismatches found
