[ValueTracking] Bail out on non-immediate constant expressions

Author: dtcxzyw

llvm/include/llvm/Analysis/ValueTracking.h

@@ -1024,8 +1024,8 @@ findValuesAffectedByCondition(Value *Cond, bool IsAssume,
 LLVM_ABI Value *stripNullTest(Value *V);
 LLVM_ABI const Value *stripNullTest(const Value *V);
 
-/// Enumerates all possible values of V and inserts them into the set \p
-/// Constants. If \p AllowUndefOrPoison is false, it fails when V may contain
+/// Enumerates all possible immediate values of V and inserts them into the set
+/// \p Constants. If \p AllowUndefOrPoison is false, it fails when V may contain
 /// undef/poison elements. Returns true if the result is complete. Otherwise,
 /// the result is incomplete (more than MaxCount values).
 /// NOTE: The constant values are not distinct.

llvm/lib/Analysis/ValueTracking.cpp

@@ -10485,7 +10485,8 @@ bool llvm::collectPossibleValues(const Value *V,
   SmallPtrSet<const Instruction *, 8> Visited;
   SmallVector<const Instruction *, 8> Worklist;
   auto Push = [&](const Value *V) -> bool {
-    if (auto *C = dyn_cast<Constant>(V)) {
+    Constant *C;
+    if (match(const_cast<Value *>(V), m_ImmConstant(C))) {
       if (!AllowUndefOrPoison && !isGuaranteedNotToBeUndefOrPoison(C))
         return false;
       // Check existence first to avoid unnecessary allocations.

llvm/test/Transforms/SimplifyCFG/switch-on-const.ll

@@ -287,7 +287,11 @@ default:
 define i16 @switch_on_nonimmediate_constant_expr() {
 ; CHECK-LABEL: @switch_on_nonimmediate_constant_expr(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    ret i16 789
+; CHECK-NEXT:    [[SWITCH_SELECTCMP:%.*]] = icmp eq i32 ptrtoint (ptr @G to i32), 2
+; CHECK-NEXT:    [[SWITCH_SELECT:%.*]] = select i1 [[SWITCH_SELECTCMP]], i16 456, i16 789
+; CHECK-NEXT:    [[SWITCH_SELECTCMP1:%.*]] = icmp eq i32 ptrtoint (ptr @G to i32), 1
+; CHECK-NEXT:    [[SWITCH_SELECT2:%.*]] = select i1 [[SWITCH_SELECTCMP1]], i16 123, i16 [[SWITCH_SELECT]]
+; CHECK-NEXT:    ret i16 [[SWITCH_SELECT2]]
 ;
 entry:
   switch i32 ptrtoint (ptr @G to i32), label %sw.default [