[llvm] [InstCombine] Fold `(X & Mask) == 0 ? TC : FC -> TC binop (X & Mask)` (PR #100437)

[via llvm-commits]

================
@@ -155,42 +155,43 @@ static Value *foldSelectICmpAnd(SelectInst &Sel, ICmpInst *Cmp,
   } else {
     return nullptr;
   }
+  if (Pred == ICmpInst::ICMP_NE)
+    std::swap(SelTC, SelFC);
 
   // In general, when both constants are non-zero, we would need an offset to
   // replace the select. This would require more instructions than we started
   // with. But there's one special-case that we handle here because it can
   // simplify/reduce the instructions.
-  APInt TC = *SelTC;
-  APInt FC = *SelFC;
+  const APInt &TC = *SelTC;
+  const APInt &FC = *SelFC;
   if (!TC.isZero() && !FC.isZero()) {
-    // If the select constants differ by exactly one bit and that's the same
-    // bit that is masked and checked by the select condition, the select can
-    // be replaced by bitwise logic to set/clear one bit of the constant result.
-    if (TC.getBitWidth() != AndMask.getBitWidth() || (TC ^ FC) != AndMask)
+    if (TC.getBitWidth() != AndMask.getBitWidth())
       return nullptr;
-    if (CreateAnd) {
-      // If we have to create an 'and', then we must kill the cmp to not
-      // increase the instruction count.
-      if (!Cmp->hasOneUse())
-        return nullptr;
-      V = Builder.CreateAnd(V, ConstantInt::get(SelType, AndMask));
-    }
-    bool ExtraBitInTC = TC.ugt(FC);
-    if (Pred == ICmpInst::ICMP_EQ) {
-      // If the masked bit in V is clear, clear or set the bit in the result:
-      // (V & AndMaskC) == 0 ? TC : FC --> (V & AndMaskC) ^ TC
-      // (V & AndMaskC) == 0 ? TC : FC --> (V & AndMaskC) | TC
-      Constant *C = ConstantInt::get(SelType, TC);
-      return ExtraBitInTC ? Builder.CreateXor(V, C) : Builder.CreateOr(V, C);
-    }
-    if (Pred == ICmpInst::ICMP_NE) {
-      // If the masked bit in V is set, set or clear the bit in the result:
-      // (V & AndMaskC) != 0 ? TC : FC --> (V & AndMaskC) | FC
-      // (V & AndMaskC) != 0 ? TC : FC --> (V & AndMaskC) ^ FC
-      Constant *C = ConstantInt::get(SelType, FC);
-      return ExtraBitInTC ? Builder.CreateOr(V, C) : Builder.CreateXor(V, C);
+    // If we have to create an 'and', then we must kill the cmp to not
+    // increase the instruction count.
+    if (CreateAnd && !Cmp->hasOneUse())
+      return nullptr;
+
+    // (V & AndMaskC) == 0 ? TC : FC --> TC | (V & AndMaskC)
+    // (V & AndMaskC) == 0 ? TC : FC --> TC ^ (V & AndMaskC)
+    // (V & AndMaskC) == 0 ? TC : FC --> TC + (V & AndMaskC)
+    // (V & AndMaskC) == 0 ? TC : FC --> TC - (V & AndMaskC)
+    Constant *TCC = ConstantInt::get(SelType, TC);
+    Constant *FCC = ConstantInt::get(SelType, FC);
+    Constant *MaskC = ConstantInt::get(SelType, AndMask);
+    for (auto Opc : {Instruction::Or, Instruction::Xor, Instruction::Add,
+                     Instruction::Sub}) {
+      if (ConstantFoldBinaryOpOperands(Opc, TCC, MaskC, Sel.getDataLayout()) ==
+          FCC) {
+        if (CreateAnd)
+          V = Builder.CreateAnd(V, MaskC);
+        if (Opc == Instruction::Sub)
+          return Builder.CreateSub(TCC, V);
+        return Builder.CreateBinOp(Opc, V, TCC);
----------------
goldsteinn wrote:

nit: (and person preference) Just do `Builder.CreateBinOp(Opc, TCC, V)` for all and let canonicalization clean things up. Code complexity is never worth it.

https://github.com/llvm/llvm-project/pull/100437