[InstSimplify][IR] Handle trapping constant aggregate (PR49839)

Unfortunately, it's not just constant expressions that can trap,
we might also have a trapping constant expression nested inside
a constant aggregate.

Perform the check during phi folding on Constant rather than
ConstantExpr, and extend the Constant::mayTrap() implementation
to also recursive into ConstantAggregates, not just ConstantExprs.

Fixes https://github.com/llvm/llvm-project/issues/49839.

llvm/lib/Analysis/InstructionSimplify.cpp

@@ -4846,8 +4846,8 @@ static Value *simplifyPHINode(PHINode *PN, ArrayRef<Value *> IncomingValues,
   if (HasUndefInput) {
     // We cannot start executing a trapping constant expression on more control
     // flow paths.
-    auto *CE = dyn_cast<ConstantExpr>(CommonValue);
-    if (CE && CE->canTrap())
+    auto *C = dyn_cast<Constant>(CommonValue);
+    if (C && C->canTrap())
       return nullptr;
 
     // If we have a PHI node like phi(X, undef, X), where X is defined by some

llvm/lib/IR/Constants.cpp

@@ -564,21 +564,25 @@ void llvm::deleteConstant(Constant *C) {
 }
 
 static bool canTrapImpl(const Constant *C,
-                        SmallPtrSetImpl<const ConstantExpr *> &NonTrappingOps) {
-  assert(C->getType()->isFirstClassType() && "Cannot evaluate aggregate vals!");
-  // The only thing that could possibly trap are constant exprs.
+                        SmallPtrSetImpl<const Constant *> &NonTrappingOps) {
+  assert(C->getType()->isFirstClassType() &&
+         "Cannot evaluate non-first-class types!");
+  // ConstantExpr or ConstantAggregate trap if any operands can trap.
+  if (isa<ConstantExpr>(C) || isa<ConstantAggregate>(C)) {
+    for (unsigned i = 0, e = C->getNumOperands(); i != e; ++i) {
+      const Constant *Op = cast<Constant>(C->getOperand(i));
+      if (isa<ConstantExpr>(Op) || isa<ConstantAggregate>(Op)) {
+        if (NonTrappingOps.insert(Op).second && canTrapImpl(Op, NonTrappingOps))
+          return true;
+      }
+    }
+  }
+
+  // The only leafs that can trap are constant expressions.
   const ConstantExpr *CE = dyn_cast<ConstantExpr>(C);
   if (!CE)
     return false;
 
-  // ConstantExpr traps if any operands can trap.
-  for (unsigned i = 0, e = C->getNumOperands(); i != e; ++i) {
-    if (ConstantExpr *Op = dyn_cast<ConstantExpr>(CE->getOperand(i))) {
-      if (NonTrappingOps.insert(Op).second && canTrapImpl(Op, NonTrappingOps))
-        return true;
-    }
-  }
-
   // Otherwise, only specific operations can trap.
   switch (CE->getOpcode()) {
   default:
@@ -595,7 +599,7 @@ static bool canTrapImpl(const Constant *C,
 }
 
 bool Constant::canTrap() const {
-  SmallPtrSet<const ConstantExpr *, 4> NonTrappingOps;
+  SmallPtrSet<const Constant *, 4> NonTrappingOps;
   return canTrapImpl(this, NonTrappingOps);
 }
 

llvm/test/Transforms/InstSimplify/phi.ll

@@ -203,7 +203,8 @@ define <1 x i64> @pr49839_vector(i1 %c) {
 ; CHECK:       if:
 ; CHECK-NEXT:    br label [[JOIN]]
 ; CHECK:       join:
-; CHECK-NEXT:    ret <1 x i64> <i64 srem (i64 1, i64 ptrtoint (ptr @g to i64))>
+; CHECK-NEXT:    [[PHI:%.*]] = phi <1 x i64> [ poison, [[IF]] ], [ <i64 srem (i64 1, i64 ptrtoint (ptr @g to i64))>, [[ENTRY:%.*]] ]
+; CHECK-NEXT:    ret <1 x i64> [[PHI]]
 ;
 entry:
   br i1 %c, label %if, label %join