[llvm] [InstCombine] Resolve TODO: Remove one-time check if other logic operand (Y) is constant (PR #77973)

[via llvm-commits]

https://github.com/AtariDreams updated https://github.com/llvm/llvm-project/pull/77973

>From 3dc72941fc97be10ec7b55851807552355e9d0bf Mon Sep 17 00:00:00 2001
From: Rose <83477269+AtariDreams at users.noreply.github.com>
Date: Sun, 14 Jan 2024 12:53:37 -0500
Subject: [PATCH 1/2] [InstCombine] Add pre-commit tests [NFC]

---
 .../Transforms/InstCombine/shift-logic.ll     | 30 +++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/llvm/test/Transforms/InstCombine/shift-logic.ll b/llvm/test/Transforms/InstCombine/shift-logic.ll
index 544694d398431e3..79f0bd824a5d6f7 100644
--- a/llvm/test/Transforms/InstCombine/shift-logic.ll
+++ b/llvm/test/Transforms/InstCombine/shift-logic.ll
@@ -240,6 +240,36 @@ define i32 @lshr_or_extra_use(i32 %x, i32 %y, ptr %p) {
   ret i32 %sh1
 }
 
+define i32 @lshr_or_extra_use_shift(i32 %x, ptr %p) {
+; CHECK-LABEL: @lshr_or_extra_use_shift(
+; CHECK-NEXT:    [[SH0:%.*]] = lshr i32 [[X:%.*]], 5
+; CHECK-NEXT:    store i32 [[SH0]], ptr [[P:%.*]], align 4
+; CHECK-NEXT:    [[R:%.*]] = lshr i32 [[X]], 12
+; CHECK-NEXT:    [[SH1:%.*]] = or i32 [[R]], 192
+; CHECK-NEXT:    ret i32 [[SH1]]
+;
+  %sh0 = lshr i32 %x, 5
+  %r = or i32 %sh0, 24601
+  store i32 %sh0, ptr %p
+  %sh1 = lshr i32 %r, 7
+  ret i32 %sh1
+}
+
+define i32 @lshr_or_extra_use_var(i32 %x, i32 %y, ptr %p) {
+; CHECK-LABEL: @lshr_or_extra_use_var(
+; CHECK-NEXT:    [[SH0:%.*]] = lshr i32 [[X:%.*]], 5
+; CHECK-NEXT:    [[R:%.*]] = or i32 [[SH0]], [[Y:%.*]]
+; CHECK-NEXT:    store i32 [[SH0]], ptr [[P:%.*]], align 4
+; CHECK-NEXT:    [[SH1:%.*]] = lshr i32 [[R]], 7
+; CHECK-NEXT:    ret i32 [[SH1]]
+;
+  %sh0 = lshr i32 %x, 5
+  %r = or i32 %sh0, %y
+  store i32 %sh0, ptr %p
+  %sh1 = lshr i32 %r, 7
+  ret i32 %sh1
+}
+
 ; Avoid crashing on constant expressions.
 
 @g = external global i32

>From 10b2f4938d7d5d4f26072efaded13404a0bfa10e Mon Sep 17 00:00:00 2001
From: Rose <83477269+AtariDreams at users.noreply.github.com>
Date: Fri, 12 Jan 2024 15:15:27 -0500
Subject: [PATCH 2/2] [Transforms] Remove one-time check if other logic operand
 (Y) is constant

By using match(W, m_ImmConstant()), we do not need to worry about one-time use anymore.
---
 llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp b/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp
index b7958978c450c98..cbad2d6f8be7898 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp
@@ -368,12 +368,11 @@ static Instruction *foldShiftOfShiftedBinOp(BinaryOperator &I,
 
   // Find a matching one-use shift by constant. The fold is not valid if the sum
   // of the shift values equals or exceeds bitwidth.
-  // TODO: Remove the one-use check if the other logic operand (Y) is constant.
   Value *X, *Y;
-  auto matchFirstShift = [&](Value *V) {
+  auto matchFirstShift = [&](Value *V, Value *W) {
     APInt Threshold(Ty->getScalarSizeInBits(), Ty->getScalarSizeInBits());
-    return match(V,
-                 m_OneUse(m_BinOp(ShiftOpcode, m_Value(X), m_Constant(C0)))) &&
+    return match(V, m_BinOp(ShiftOpcode, m_Value(X), m_Constant(C0))) &&
+           (V->hasOneUse() || match(W, m_ImmConstant())) &&
            match(ConstantExpr::getAdd(C0, C1),
                  m_SpecificInt_ICMP(ICmpInst::ICMP_ULT, Threshold));
   };
@@ -382,9 +381,9 @@ static Instruction *foldShiftOfShiftedBinOp(BinaryOperator &I,
   // is not so we cannot reoder if we match operand(1) and need to keep the
   // operands in their original positions.
   bool FirstShiftIsOp1 = false;
-  if (matchFirstShift(BinInst->getOperand(0)))
+  if (matchFirstShift(BinInst->getOperand(0), BinInst->getOperand(1)))
     Y = BinInst->getOperand(1);
-  else if (matchFirstShift(BinInst->getOperand(1))) {
+  else if (matchFirstShift(BinInst->getOperand(1), BinInst->getOperand(0))) {
     Y = BinInst->getOperand(0);
     FirstShiftIsOp1 = BinInst->getOpcode() == Instruction::Sub;
   } else