From 3736e1d1cd5c24b554a13e493c8614f458bdf123 Mon Sep 17 00:00:00 2001
From: Simon Pilgrim <llvm-dev@redking.me.uk>
Date: Fri, 22 Dec 2023 13:53:49 +0000
Subject: [PATCH] [SCEV] Ensure shift amount is in range before calling
 getZExtValue()

Fixes #76234
---
 llvm/lib/Analysis/ScalarEvolution.cpp         |  7 +++---
 llvm/test/Analysis/ScalarEvolution/pr76234.ll | 23 +++++++++++++++++++
 2 files changed, 27 insertions(+), 3 deletions(-)
 create mode 100644 llvm/test/Analysis/ScalarEvolution/pr76234.ll

diff --git a/llvm/lib/Analysis/ScalarEvolution.cpp b/llvm/lib/Analysis/ScalarEvolution.cpp
index 580fe112fcd7bd..623814c038a78f 100644
--- a/llvm/lib/Analysis/ScalarEvolution.cpp
+++ b/llvm/lib/Analysis/ScalarEvolution.cpp
@@ -7914,9 +7914,10 @@ const SCEV *ScalarEvolution::createSCEV(Value *V) {
         //    expression. We already checked that ShlAmt < BitWidth, so
         //    the multiplier, 1 << (ShlAmt - AShrAmt), fits into TruncTy as
         //    ShlAmt - AShrAmt < Amt.
-        uint64_t ShlAmt = ShlAmtCI->getZExtValue();
-        if (ShlAmtCI->getValue().ult(BitWidth) && ShlAmt >= AShrAmt) {
-          APInt Mul = APInt::getOneBitSet(BitWidth - AShrAmt, ShlAmt - AShrAmt);
+        const APInt &ShlAmt = ShlAmtCI->getValue();
+        if (ShlAmt.ult(BitWidth) && ShlAmt.uge(AShrAmt)) {
+          APInt Mul = APInt::getOneBitSet(BitWidth - AShrAmt,
+                                          ShlAmtCI->getZExtValue() - AShrAmt);
           const SCEV *CompositeExpr =
               getMulExpr(AddTruncateExpr, getConstant(Mul));
           if (L->getOpcode() != Instruction::Shl)
diff --git a/llvm/test/Analysis/ScalarEvolution/pr76234.ll b/llvm/test/Analysis/ScalarEvolution/pr76234.ll
new file mode 100644
index 00000000000000..0d82f0ed1a81ce
--- /dev/null
+++ b/llvm/test/Analysis/ScalarEvolution/pr76234.ll
@@ -0,0 +1,23 @@
+; NOTE: Assertions have been autogenerated by utils/update_analyze_test_checks.py UTC_ARGS: --version 4
+; RUN: opt < %s -disable-output "-passes=print<scalar-evolution>" 2>&1 | FileCheck %s
+
+; Reduced from https://bugs.chromium.org/p/oss-fuzz/issues/detail?id=65278
+define i32 @PR76234() {
+; CHECK-LABEL: 'PR76234'
+; CHECK-NEXT:  Classifying expressions for: @PR76234
+; CHECK-NEXT:    %B9 = shl i896 0, -264147265567832623176169892458258303259423663018060761063980354513336951278362429737208627943828593947337197496628564339441173779751342768625269489231469788454193341999502542084365758838213220526512116454105594202074014146375780869419198449383518238244769290448868999168
+; CHECK-NEXT:    --> %B9 U: [0,1) S: [0,1)
+; CHECK-NEXT:    %B39 = ashr i896 %B9, 1
+; CHECK-NEXT:    --> %B39 U: [0,1) S: [0,1) Exits: <<Unknown>> LoopDispositions: { %1: Variant }
+; CHECK-NEXT:  Determining loop execution counts for: @PR76234
+; CHECK-NEXT:  Loop %1: <multiple exits> Unpredictable backedge-taken count.
+; CHECK-NEXT:  Loop %1: Unpredictable constant max backedge-taken count.
+; CHECK-NEXT:  Loop %1: Unpredictable symbolic max backedge-taken count.
+; CHECK-NEXT:  Loop %1: Unpredictable predicated backedge-taken count.
+;
+  %B9 = shl i896 0, -264147265567832623176169892458258303259423663018060761063980354513336951278362429737208627943828593947337197496628564339441173779751342768625269489231469788454193341999502542084365758838213220526512116454105594202074014146375780869419198449383518238244769290448868999168
+  br label %1
+1:
+  %B39 = ashr i896 %B9, 1
+  br label %1
+}