[reg2mem] Add special handling to CatchSwitchInst

When promoting a phi in a catchswitch block to memory, we cannot
insert load/store instruction in that block, and need to insert
them inside all successors instead.

Fixes #59185.

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

llvm/lib/Transforms/Utils/DemoteRegToStack.cpp

@@ -92,8 +92,15 @@ AllocaInst *llvm::DemoteRegToStack(Instruction &I, bool VolatileLoads,
   BasicBlock::iterator InsertPt;
   if (!I.isTerminator()) {
     InsertPt = ++I.getIterator();
+    // Don't insert before PHI nodes or landingpad instrs.
     for (; isa<PHINode>(InsertPt) || InsertPt->isEHPad(); ++InsertPt)
-      /* empty */;   // Don't insert before PHI nodes or landingpad instrs.
+      if (isa<CatchSwitchInst>(InsertPt))
+        break;
+    if (isa<CatchSwitchInst>(InsertPt)) {
+      for (BasicBlock *Handler : successors(&*InsertPt))
+        new StoreInst(&I, Slot, &*Handler->getFirstInsertionPt());
+      return Slot;
+    }
   } else {
     InvokeInst &II = cast<InvokeInst>(I);
     InsertPt = II.getNormalDest()->getFirstInsertionPt();
@@ -138,14 +145,27 @@ AllocaInst *llvm::DemotePHIToStack(PHINode *P, Instruction *AllocaPoint) {
 
   // Insert a load in place of the PHI and replace all uses.
   BasicBlock::iterator InsertPt = P->getIterator();
-
+  // Don't insert before PHI nodes or landingpad instrs.
   for (; isa<PHINode>(InsertPt) || InsertPt->isEHPad(); ++InsertPt)
-    /* empty */;   // Don't insert before PHI nodes or landingpad instrs.
-
-  Value *V =
-      new LoadInst(P->getType(), Slot, P->getName() + ".reload", &*InsertPt);
-  P->replaceAllUsesWith(V);
-
+    if (isa<CatchSwitchInst>(InsertPt))
+      break;
+  if (isa<CatchSwitchInst>(InsertPt)) {
+    // We need a separate load before each actual use of the PHI
+    SmallVector<Instruction *, 4> Users;
+    for (User *U : P->users()) {
+      Instruction *User = cast<Instruction>(U);
+      Users.push_back(User);
+    }
+    for (Instruction *User : Users) {
+      Value *V =
+          new LoadInst(P->getType(), Slot, P->getName() + ".reload", User);
+      User->replaceUsesOfWith(P, V);
+    }
+  } else {
+    Value *V =
+        new LoadInst(P->getType(), Slot, P->getName() + ".reload", &*InsertPt);
+    P->replaceAllUsesWith(V);
+  }
   // Delete PHI.
   P->eraseFromParent();
   return Slot;

llvm/test/Transforms/Reg2Mem/catchswitch-crash2.ll

@@ -0,0 +1,76 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes=reg2mem -S < %s | FileCheck %s
+%opaque = type opaque
+
+declare i32 @__CxxFrameHandler3(...)
+
+define void @testreg2mem(i8* %_Val) personality i8* bitcast (i32 (...)* @__CxxFrameHandler3 to i8*) {
+; CHECK-LABEL: @testreg2mem(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[_STATE_3_REG2MEM:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[_STATE_3_REG2MEM1:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    %"reg2mem alloca point" = bitcast i32 0 to i32
+; CHECK-NEXT:    store i32 0, i32* [[_STATE_3_REG2MEM1]], align 4
+; CHECK-NEXT:    [[CALL_I166167:%.*]] = invoke noundef i64 @extfunc_i64()
+; CHECK-NEXT:    to label [[IF_END56:%.*]] unwind label [[CATCH_DISPATCH:%.*]]
+; CHECK:       if.end56:
+; CHECK-NEXT:    store i32 1, i32* [[_STATE_3_REG2MEM1]], align 4
+; CHECK-NEXT:    invoke void @extfunc()
+; CHECK-NEXT:    to label [[INVOKE_CONT75:%.*]] unwind label [[CATCH_DISPATCH]]
+; CHECK:       invoke.cont75:
+; CHECK-NEXT:    unreachable
+; CHECK:       catch.dispatch:
+; CHECK-NEXT:    [[TMP0:%.*]] = catchswitch within none [label %catch] unwind label [[EHCLEANUP105:%.*]]
+; CHECK:       catch:
+; CHECK-NEXT:    [[TMP1:%.*]] = catchpad within [[TMP0]] [i8* null, i32 64, i8* null]
+; CHECK-NEXT:    [[_STATE_3_RELOAD2:%.*]] = load i32, i32* [[_STATE_3_REG2MEM1]], align 4
+; CHECK-NEXT:    store i32 [[_STATE_3_RELOAD2]], i32* [[_STATE_3_REG2MEM]], align 4
+; CHECK-NEXT:    invoke void @extfunc() [ "funclet"(token [[TMP1]]) ]
+; CHECK-NEXT:    to label [[INVOKE_CONT98:%.*]] unwind label [[EHCLEANUP105]]
+; CHECK:       invoke.cont98:
+; CHECK-NEXT:    catchret from [[TMP1]] to label [[IF_END99:%.*]]
+; CHECK:       if.end99:
+; CHECK-NEXT:    [[_STATE_3_RELOAD:%.*]] = load i32, i32* [[_STATE_3_REG2MEM]], align 4
+; CHECK-NEXT:    [[OR_I:%.*]] = or i32 0, [[_STATE_3_RELOAD]]
+; CHECK-NEXT:    unreachable
+; CHECK:       ehcleanup105:
+; CHECK-NEXT:    [[TMP2:%.*]] = cleanuppad within none []
+; CHECK-NEXT:    [[_STATE_3_RELOAD3:%.*]] = load i32, i32* [[_STATE_3_REG2MEM1]], align 4
+; CHECK-NEXT:    store i32 [[_STATE_3_RELOAD3]], i32* [[_STATE_3_REG2MEM]], align 4
+; CHECK-NEXT:    cleanupret from [[TMP2]] unwind to caller
+;
+entry:
+  %call.i166167 = invoke noundef i64 @"extfunc_i64"()
+  to label %if.end56 unwind label %catch.dispatch
+
+if.end56:                                         ; preds = %entry
+  invoke void @"extfunc"()
+  to label %invoke.cont75 unwind label %catch.dispatch
+
+invoke.cont75:                                    ; preds = %if.end56
+  unreachable
+
+catch.dispatch:                                   ; preds = %if.end56, %entry
+  %_State.3 = phi i32 [ 1, %if.end56 ], [ 0, %entry ]
+  %0 = catchswitch within none [label %catch] unwind label %ehcleanup105
+
+catch:                                            ; preds = %catch.dispatch
+  %1 = catchpad within %0 [i8* null, i32 64, i8* null]
+  invoke void @"extfunc"() [ "funclet"(token %1) ]
+  to label %invoke.cont98 unwind label %ehcleanup105
+
+invoke.cont98:                                    ; preds = %catch
+  catchret from %1 to label %if.end99
+
+if.end99:                                         ; preds = %invoke.cont98
+  %or.i = or i32 0, %_State.3
+  unreachable
+
+ehcleanup105:                                     ; preds = %catch, %catch.dispatch
+  %2 = cleanuppad within none []
+  cleanupret from %2 unwind to caller
+}
+
+declare void @"extfunc"()
+
+declare i64 @"extfunc_i64"()