From 7e64a29e58382b2616096d64a39c38075e51ebc3 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov@redhat.com>
Date: Mon, 13 Jun 2022 12:32:22 +0200
Subject: [PATCH] [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 |  4 ++--
 llvm/lib/IR/Constants.cpp                 | 28 +++++++++++++----------
 llvm/test/Transforms/InstSimplify/phi.ll  |  3 ++-
 3 files changed, 20 insertions(+), 15 deletions(-)

diff --git a/llvm/lib/Analysis/InstructionSimplify.cpp b/llvm/lib/Analysis/InstructionSimplify.cpp
index da1a67c7338882..b810f32643b65d 100644
--- a/llvm/lib/Analysis/InstructionSimplify.cpp
+++ b/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
diff --git a/llvm/lib/IR/Constants.cpp b/llvm/lib/IR/Constants.cpp
index 965012e7f5ed06..139ff532589909 100644
--- a/llvm/lib/IR/Constants.cpp
+++ b/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);
 }
 
diff --git a/llvm/test/Transforms/InstSimplify/phi.ll b/llvm/test/Transforms/InstSimplify/phi.ll
index 8565883b0492d6..822023edd52512 100644
--- a/llvm/test/Transforms/InstSimplify/phi.ll
+++ b/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