[llvm] [InstCombine] Simplify `(X == Y) ? (X << 1) : (X + Y)` into `(X + Y)` (PR #76738)

[via llvm-commits]

https://github.com/Z572 created https://github.com/llvm/llvm-project/pull/76738

fix https://github.com/llvm/llvm-project/issues/76318
proof: https://alive2.llvm.org/ce/z/eRmmnp

>From 9f6bd00e9c91cc74b21ff8a8f7b88cffab3d6e0c Mon Sep 17 00:00:00 2001
From: Zheng Junjie <zhengjunjie at iscas.ac.cn>
Date: Wed, 3 Jan 2024 00:56:52 +0800
Subject: [PATCH] [InstCombine] Simplify `(X == Y) ? (X << 1) : (X + Y)` into
 `(X + Y)`

proof: https://alive2.llvm.org/ce/z/eRmmnp
---
 .../InstCombine/InstCombineSelect.cpp         |  20 ++++
 llvm/test/Transforms/InstCombine/select.ll    | 103 ++++++++++++++++++
 2 files changed, 123 insertions(+)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index ab55f235920a7b..609f1fe1aba7a0 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -1779,6 +1779,26 @@ Instruction *InstCombinerImpl::foldSelectInstWithICmp(SelectInst &SI,
     }
   }
 
+  // (X == Y) ? (X << 1) : (X + Y) --> (X + Y)
+  // (X == Y) ? (Y << 1) : (X + Y) --> (X + Y)
+  {
+    Value *X, *Y;
+    if ((Pred == ICmpInst::ICMP_EQ) &&
+        (match(CmpLHS, m_Value(X)) && match(CmpRHS, m_Value(Y))) &&
+        (((match(TrueVal, m_Shl(m_Specific(X), m_One()))) &&
+          (match(FalseVal, m_c_Add(m_Specific(X), m_Specific(Y))))) ||
+         ((match(TrueVal, m_Shl(m_Specific(Y), m_One()))) &&
+          (match(FalseVal, m_c_Add(m_Specific(X), m_Specific(Y))))))) {
+      auto *TI = dyn_cast<Instruction>(TrueVal);
+      auto *FI = dyn_cast<Instruction>(FalseVal);
+      Value *V = Builder.CreateAdd(
+          X, Y, "",
+          /*HasNUW*/ TI->hasNoUnsignedWrap() && FI->hasNoUnsignedWrap(),
+          /*HasNSW*/ TI->hasNoSignedWrap() && FI->hasNoSignedWrap());
+      return replaceInstUsesWith(SI, V);
+    }
+  }
+
   if (Instruction *V =
           foldSelectICmpAndAnd(SI.getType(), ICI, TrueVal, FalseVal, Builder))
     return V;
diff --git a/llvm/test/Transforms/InstCombine/select.ll b/llvm/test/Transforms/InstCombine/select.ll
index d3e959b1eaa0e8..51efed82031d21 100644
--- a/llvm/test/Transforms/InstCombine/select.ll
+++ b/llvm/test/Transforms/InstCombine/select.ll
@@ -3658,3 +3658,106 @@ loop:
 exit:
   ret i32 %rem
 }
+
+
+define i4 @select_icmp_eq_nsw_shl_nsw_add(i4 %x, i4 %y) {
+; CHECK-LABEL: @select_icmp_eq_nsw_shl_nsw_add(
+; CHECK-NEXT:    [[SELECT:%.*]] = add nsw i4 [[X:%.*]], [[Y:%.*]]
+; CHECK-NEXT:    ret i4 [[SELECT]]
+;
+  %icmp = icmp eq i4 %x, %y
+  %add = add nsw i4 %y, %x
+  %shl = shl nsw i4 %x, 1
+  %select = select i1 %icmp, i4 %shl, i4 %add
+  ret i4 %select
+}
+
+define i4 @select_icmp_eq_shl_nsw_add(i4 %x, i4 %y) {
+; CHECK-LABEL: @select_icmp_eq_shl_nsw_add(
+; CHECK-NEXT:    [[SELECT:%.*]] = add i4 [[X:%.*]], [[Y:%.*]]
+; CHECK-NEXT:    ret i4 [[SELECT]]
+;
+  %icmp = icmp eq i4 %x, %y
+  %add = add nsw i4 %y, %x
+  %shl = shl i4 %x, 1
+  %select = select i1 %icmp, i4 %shl, i4 %add
+  ret i4 %select
+}
+
+define i4 @select_icmp_ne_shl_nsw_add(i4 %x, i4 %y) {
+; CHECK-LABEL: @select_icmp_ne_shl_nsw_add(
+; CHECK-NEXT:    [[ICMP_NOT:%.*]] = icmp eq i4 [[X:%.*]], [[Y:%.*]]
+; CHECK-NEXT:    [[ADD:%.*]] = add nsw i4 [[Y]], [[X]]
+; CHECK-NEXT:    [[SHL:%.*]] = shl i4 [[X]], 1
+; CHECK-NEXT:    [[SELECT:%.*]] = select i1 [[ICMP_NOT]], i4 [[ADD]], i4 [[SHL]]
+; CHECK-NEXT:    ret i4 [[SELECT]]
+;
+  %icmp = icmp ne i4 %x, %y
+  %add = add nsw i4 %y, %x
+  %shl = shl i4 %x, 1
+  %select = select i1 %icmp, i4 %shl, i4 %add
+  ret i4 %select
+}
+
+define i4 @select_icmp_ne_nsw_add_nsw_shl(i4 %x, i4 %y) {
+; CHECK-LABEL: @select_icmp_ne_nsw_add_nsw_shl(
+; CHECK-NEXT:    [[SELECT:%.*]] = add nsw i4 [[X:%.*]], [[Y:%.*]]
+; CHECK-NEXT:    ret i4 [[SELECT]]
+;
+  %icmp = icmp ne i4 %x, %y
+  %shl = shl nsw i4 %x, 1
+  %add = add nsw i4 %y, %x
+  %select = select i1 %icmp, i4 %add, i4 %shl
+  ret i4 %select
+}
+
+define i4 @select_icmp_eq_nsw_add_nsw_shl(i4 %x, i4 %y) {
+; CHECK-LABEL: @select_icmp_eq_nsw_add_nsw_shl(
+; CHECK-NEXT:    [[ICMP:%.*]] = icmp eq i4 [[Y:%.*]], [[X:%.*]]
+; CHECK-NEXT:    [[SHL:%.*]] = shl nsw i4 [[X]], 1
+; CHECK-NEXT:    [[ADD:%.*]] = add nsw i4 [[Y]], [[X]]
+; CHECK-NEXT:    [[SELECT:%.*]] = select i1 [[ICMP]], i4 [[ADD]], i4 [[SHL]]
+; CHECK-NEXT:    ret i4 [[SELECT]]
+;
+  %icmp = icmp eq i4 %y, %x
+  %shl = shl nsw i4 %x, 1
+  %add = add nsw i4 %y, %x
+  %select = select i1 %icmp, i4 %add ,i4 %shl
+  ret i4 %select
+}
+
+define i4 @select_icmp_eq_nuw_shl_nuw_add(i4 %x, i4 %y) {
+; CHECK-LABEL: @select_icmp_eq_nuw_shl_nuw_add(
+; CHECK-NEXT:    [[SELECT:%.*]] = add nuw i4 [[X:%.*]], [[Y:%.*]]
+; CHECK-NEXT:    ret i4 [[SELECT]]
+;
+  %icmp = icmp eq i4 %x, %y
+  %add = add nuw i4 %y, %x
+  %shl = shl nuw i4 %x, 1
+  %select = select i1 %icmp, i4 %shl, i4 %add
+  ret i4 %select
+}
+
+define i4 @select_icmp_eq_nuw_shl_add(i4 %x, i4 %y) {
+; CHECK-LABEL: @select_icmp_eq_nuw_shl_add(
+; CHECK-NEXT:    [[SELECT:%.*]] = add i4 [[X:%.*]], [[Y:%.*]]
+; CHECK-NEXT:    ret i4 [[SELECT]]
+;
+  %icmp = icmp eq i4 %x, %y
+  %add = add i4 %y, %x
+  %shl = shl nuw i4 %x, 1
+  %select = select i1 %icmp, i4 %shl, i4 %add
+  ret i4 %select
+}
+
+define i4 @select_icmp_eq_nsw_shl_nsw_add_swap(i4 %x, i4 %y) {
+; CHECK-LABEL: @select_icmp_eq_nsw_shl_nsw_add_swap(
+; CHECK-NEXT:    [[SELECT:%.*]] = add nsw i4 [[Y:%.*]], [[X:%.*]]
+; CHECK-NEXT:    ret i4 [[SELECT]]
+;
+  %icmp = icmp eq i4 %y, %x
+  %add = add nsw i4 %y, %x
+  %shl = shl nsw i4 %x, 1
+  %select = select i1 %icmp, i4 %shl, i4 %add
+  ret i4 %select
+}