Handle `select` in programUndefinedIfPoison.

[jlebar]

Summary

If both the true and false operands of a select are poison, then the select
is poison.

[nikic]

LG

[jlebar]

Thank you for the review!

llvm/lib/Analysis/ValueTracking.cpp

 }
 
 bool llvm::isGuaranteedNotToBePoison(const Value *V, AssumptionCache *AC,
                                      const Instruction *CtxI,
                                      const DominatorTree *DT, unsigned Depth) {
   return ::isGuaranteedNotToBeUndefOrPoison(V, AC, CtxI, DT, Depth, true);
 }
 
-/// Return true if undefined behavior would provable be executed on the path to
+/// Return true if undefined behavior would provably be executed on the path to
 /// OnPathTo if Root produced a posion result.  Note that this doesn't say
 /// anything about whether OnPathTo is actually executed or whether Root is
 /// actually poison.  This can be used to assess whether a new use of Root can
 /// be added at a location which is control equivalent with OnPathTo (such as
 /// immediately before it) without introducing UB which didn't previously
 /// exist.  Note that a false result conveys no information.
 bool llvm::mustExecuteUBIfPoisonOnPathTo(Instruction *Root,
                                          Instruction *OnPathTo,
for (const auto &I : make_range(Begin, End)) {
 
       // If an operand is poison and propagates it, mark I as yielding poison.
       for (const Use &Op : I.operands()) {
         if (YieldsPoison.count(Op) && propagatesPoison(Op)) {
           YieldsPoison.insert(&I);
           break;
         }
       }
+
+      // Special handling for select, which returns poison if its operand 0 is
+      // poison (handled in the loop above) *or* if both its true/false operands
+      // are poison (handled here).
+      if (I.getOpcode() == Instruction::Select &&
+          YieldsPoison.count(I.getOperand(1)) &&
+          YieldsPoison.count(I.getOperand(2))) {
+        YieldsPoison.insert(&I);
+      }
     }
 
     BB = BB->getSingleSuccessor();
     if (!BB || !Visited.insert(BB).second)
       break;
 
     Begin = BB->getFirstNonPHI()->getIterator();
     End = BB->end();

llvm/unittests/Analysis/ValueTrackingTest.cpp

parseAssembly("declare i32 @any_num()"
                 "  %B = or i32 %A, %mask\n"
                 "  udiv i32 1, %B"
                 "  ret void\n"
                 "}\n");
   // If %A was poison, udiv raises UB regardless of %mask's value
   EXPECT_EQ(programUndefinedIfPoison(A), true);
 }
 
+TEST_F(ValueTrackingTest, programUndefinedIfPoisonSelect) {
+  parseAssembly("declare i32 @any_num()"
+                "define void @test(i1 %Cond) {\n"
+                "  %A = call i32 @any_num()\n"
+                "  %B = add i32 %A, 1\n"
+                "  %C = select i1 %Cond, i32 %A, i32 %B\n"
+                "  udiv i32 1, %C"
+                "  ret void\n"
+                "}\n");
+  // If A is poison, B is also poison, and therefore C is poison regardless of
+  // the value of %Cond.
+  EXPECT_EQ(programUndefinedIfPoison(A), true);
+}
+
 TEST_F(ValueTrackingTest, programUndefinedIfUndefOrPoison) {
   parseAssembly("declare i32 @any_num()"
                 "define void @test(i32 %mask) {\n"
                 "  %A = call i32 @any_num()\n"
                 "  %B = or i32 %A, %mask\n"
                 "  udiv i32 1, %B"
                 "  ret void\n"
                 "}\n");