[ValueTracking] Bail out on non-immediate constant expressions (#168084)

In llvm/llvm-project#165748 constant expressions
were allowed in `collectPossibleValues` because we are still using
insertelement + shufflevector idioms to represent a scalable vector
splat. However, it also accepts some unresolved constants like ptrtoint
of globals or pointer difference between two globals. Absolutely we can
ask the user to check this case with the constant folding API. However,
since we don't observe the real-world usefulness of handling constant
expressions, I decide to be more conservative and only handle immediate
constants in the helper function. With this patch, we don't need to
touch the SimplifyCFG part, as the values can only be either ConstantInt
or undef/poison values (NB: switch on undef condition is UB).

Fix the miscompilation reported by
llvm/llvm-project#165748 (comment)
Signed-off-by: Hafidz Muzakky <[email protected]>

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

@@ -10504,7 +10504,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

@@ -280,6 +280,36 @@ default:
   ret void
 }
 
+@G = constant i16 zeroinitializer, align 1
+
+; Make sure we don't remove edges when the condition is a unresolved constant expression.
+
+define i16 @switch_on_nonimmediate_constant_expr() {
+; CHECK-LABEL: @switch_on_nonimmediate_constant_expr(
+; CHECK-NEXT:  entry:
+; 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 [
+  i32 1, label %sw.bb
+  i32 2, label %sw.bb1
+  ]
+
+sw.bb:
+  ret i16 123
+
+sw.bb1:
+  ret i16 456
+
+sw.default:
+  ret i16 789
+}
+
+
 declare void @llvm.trap() nounwind noreturn
 declare void @bees.a() nounwind
 declare void @bees.b() nounwind