From ef7aed3e112b47641ca2704187d3701765a831ce Mon Sep 17 00:00:00 2001
From: Daniel Bertalan <dani@danielbertalan.dev>
Date: Thu, 7 Jul 2022 17:31:59 +0200
Subject: [PATCH] [InstCombine] Do not fold 'and (sext (ashr X, Shift)), C' if
 Shift < 0

The 'and (sext (ashr X, ShiftC)), C' --> 'lshr (sext X), ShiftC'
transformation would access out of bounds bits in APInt::getLowBitsSet
if the shift count was larger than X's bit width or if it was negative.

Fixes #56424
---
 .../InstCombine/InstCombineAndOrXor.cpp       |  3 ++-
 llvm/test/Transforms/InstCombine/pr56424.ll   | 27 +++++++++++++++++++
 2 files changed, 29 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/Transforms/InstCombine/pr56424.ll

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp b/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp
index d3578fe4a2ba97..a8f2cd79830abc 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp
@@ -1808,7 +1808,8 @@ Instruction *InstCombinerImpl::visitAnd(BinaryOperator &I) {
 
     unsigned Width = Ty->getScalarSizeInBits();
     const APInt *ShiftC;
-    if (match(Op0, m_OneUse(m_SExt(m_AShr(m_Value(X), m_APInt(ShiftC)))))) {
+    if (match(Op0, m_OneUse(m_SExt(m_AShr(m_Value(X), m_APInt(ShiftC))))) &&
+        ShiftC->ult(Width)) {
       if (*C == APInt::getLowBitsSet(Width, Width - ShiftC->getZExtValue())) {
         // We are clearing high bits that were potentially set by sext+ashr:
         // and (sext (ashr X, ShiftC)), C --> lshr (sext X), ShiftC
diff --git a/llvm/test/Transforms/InstCombine/pr56424.ll b/llvm/test/Transforms/InstCombine/pr56424.ll
new file mode 100644
index 00000000000000..dda0b6d7e0be10
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/pr56424.ll
@@ -0,0 +1,27 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s --passes=instcombine -S | FileCheck %s
+
+; This would crash if we didn't check for a negative shift.
+; https://github.com/llvm/llvm-project/issues/56424
+define i64 @PR56424(i1 %cond, i32 %arg) {
+; CHECK-LABEL: @PR56424(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br i1 [[COND:%.*]], label [[IF_THEN:%.*]], label [[IF_END:%.*]]
+; CHECK:       if.then:
+; CHECK-NEXT:    br label [[IF_END]]
+; CHECK:       if.end:
+; CHECK-NEXT:    ret i64 0
+;
+entry:
+  br i1 %cond, label %if.then, label %if.end
+
+if.then:
+  %shr  = ashr i32 %arg, -2
+  %sext = sext i32 %shr to i64
+  br label %if.end
+
+if.end:
+  %val = phi i64 [ %sext, %if.then ], [ 0, %entry ]
+  %and = and i64 -81, %val
+  ret i64 %and
+}