[WebAssembly] Fix reverse mapping in WasmEHFuncInfo

[aheejin]

Summary

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 https://github.com/emscripten-core/emscripten/issues/13514.
(More precisely, fixes
https://github.com/emscripten-core/emscripten/issues/13514#issuecomment-784708744)

[aheejin]

include "llvm/ADT/SmallPtrSet.h"

[dschuff]

LGTM. since the data structure is a SmallPtrSet now, it probably makes sense to update the description and maybe the var names to reflect that.

llvm/include/llvm/CodeGen/WasmEHFuncInfo.h

45

I guess Vec could be named Set now since it's a set rather than a vector.

llvm/lib/Target/WebAssembly/WebAssemblyCFGSort.cpp

266

what's the significance of the iteration order here?

[aheejin]

llvm/include/llvm/CodeGen/WasmEHFuncInfo.h

45

Yeah I made it as a vector first and then changed it to a set, which resulted in this.. Fixed.

llvm/lib/Target/WebAssembly/WebAssemblyCFGSort.cpp

266

Now we shouldn't do a reverse traversal because Succ has to be added to Deferred list of the outermost SortRegion. For example, if exception A contains exception B, and both A's header and B's header's unwind destination is Succ, Entries will be like [..., exception A, ..., exception B, ...], and Succ has to be deferred until all blocks in the exception A are scheduled.

Before when I incorrectly assumed the unwind source is a single BB, it was correct traverse Entries either way. I did it backwards just because I thought it would more likely to hit the target first.

[aheejin]

Vec -> Set

llvm/include/llvm/CodeGen/WasmEHFuncInfo.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CODEGEN_WASMEHFUNCINFO_H
 #define LLVM_CODEGEN_WASMEHFUNCINFO_H
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/PointerUnion.h"
+#include "llvm/ADT/SmallPtrSet.h"
 
 namespace llvm {
 
 class BasicBlock;
 class Function;
 class MachineBasicBlock;
 
 namespace WebAssembly {
 enum EventTag { CPP_EXCEPTION = 0, C_LONGJMP = 1 };
 }
 
 using BBOrMBB = PointerUnion<const BasicBlock *, MachineBasicBlock *>;
 
 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 &Vec = UnwindDestToSrcs.lookup(BB);
+    SmallPtrSet<const BasicBlock *, 4> Ret;
+    for (const auto &P : Vec)
+      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 &Vec = UnwindDestToSrcs.lookup(MBB);
+    SmallPtrSet<MachineBasicBlock *, 4> Ret;
+    for (const auto &P : Vec)
+      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);
   }
 };
 
 // Analyze the IR in the given function to build WasmEHFuncInfo.
 void calculateWasmEHInfo(const Function *F, WasmEHFuncInfo &EHInfo);
 
 } // namespace llvm
 
 #endif // LLVM_CODEGEN_WASMEHFUNCINFO_H

llvm/lib/CodeGen/SelectionDAG/FunctionLoweringInfo.cpp

if (isFuncletEHPersonality(Personality)) {
     for (ClrEHUnwindMapEntry &CME : EHInfo.ClrEHUnwindMap) {
       const auto *BB = CME.Handler.get<const BasicBlock *>();
       CME.Handler = MBBMap[BB];
     }
   }
 
   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);
   }
 }
 
 /// clear - Clear out all the function-specific state. This returns this
 /// FunctionLoweringInfo to an empty state, ready to be used for a
 /// different function.
 void FunctionLoweringInfo::clear() {
   MBBMap.clear();

llvm/lib/Target/WebAssembly/WebAssemblyCFGSort.cpp

for (MachineBasicBlock *Succ : MBB->successors()) {
         // exceptions because there can be cases in which, for example:
         // EHPad A's unwind destination (where the exception lands when it is
         // not caught by EHPad A) is EHPad B, so EHPad B does not belong to the
         // exception dominated by EHPad A. But EHPad B is dominated by EHPad A,
         // so EHPad B can be sorted within EHPad A's exception. This is
         // 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;
             }
           }
           if (IsDeferred)
             continue;
         }

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

 ;       catch
 ;       end_try
 ;             <- (a)
 ;     delegate
 ;   delegate
 ; 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
 
 if.then:                                          ; preds = %entry
   invoke void @__cxa_throw() #1
           to label %unreachable unwind label %catch.dispatch
 

 ehcleanup:                                        ; preds = %try.cont23, %rethrow19, %catch.dispatch16
   %12 = cleanuppad within none []
   cleanupret from %12 unwind to caller
 
 unreachable:                                      ; preds = %if.then, %invoke.cont, %rethrow19
   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
 
 declare void @foo()
 declare void @bar()
 declare i32 @baz()
 declare i32 @qux(i32)