[LoopIdimo] Use tryZExtValue() instead of getZExtValue()

To avoid an assertion for large BECounts.

I also suspect that this code is missing an overflow check.

Fixes llvm#70008.

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

@@ -948,9 +948,13 @@ mayLoopAccessLocation(Value *Ptr, ModRefInfo Access, Loop *L,
   // to be exactly the size of the memset, which is (BECount+1)*StoreSize
   const SCEVConstant *BECst = dyn_cast<SCEVConstant>(BECount);
   const SCEVConstant *ConstSize = dyn_cast<SCEVConstant>(StoreSizeSCEV);
-  if (BECst && ConstSize)
-    AccessSize = LocationSize::precise((BECst->getValue()->getZExtValue() + 1) *
-                                       ConstSize->getValue()->getZExtValue());
+  if (BECst && ConstSize) {
+    std::optional<uint64_t> BEInt = BECst->getAPInt().tryZExtValue();
+    std::optional<uint64_t> SizeInt = ConstSize->getAPInt().tryZExtValue();
+    // FIXME: Should this check for overflow?
+    if (BEInt && SizeInt)
+      AccessSize = LocationSize::precise((*BEInt + 1) * *SizeInt);
+  }
 
   // TODO: For this to be really effective, we have to dive into the pointer
   // operand in the store.  Store to &A[i] of 100 will always return may alias

llvm/test/Transforms/LoopIdiom/pr70008.ll

@@ -0,0 +1,39 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 3
+; RUN: opt -S -passes=loop-idiom < %s | FileCheck %s
+
+; Make sure we don't assert if the BECount is larger than 64 bits.
+
+define void @test(ptr %p) {
+; CHECK-LABEL: define void @test(
+; CHECK-SAME: ptr [[P:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr align 4 [[P]], i8 0, i64 0, i1 false)
+; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
+; CHECK:       for.body:
+; CHECK-NEXT:    [[IV:%.*]] = phi i128 [ 0, [[ENTRY:%.*]] ], [ [[INC:%.*]], [[FOR_BODY]] ]
+; CHECK-NEXT:    [[IV_TRUNC:%.*]] = trunc i128 [[IV]] to i64
+; CHECK-NEXT:    [[GEP1:%.*]] = getelementptr { i64, i64 }, ptr [[P]], i64 [[IV_TRUNC]]
+; CHECK-NEXT:    [[GEP2:%.*]] = getelementptr { i64, i64 }, ptr [[P]], i64 [[IV_TRUNC]], i32 1
+; CHECK-NEXT:    [[INC]] = add i128 [[IV]], 1
+; CHECK-NEXT:    [[TOBOOL_NOT:%.*]] = icmp eq i128 [[INC]], 0
+; CHECK-NEXT:    br i1 [[TOBOOL_NOT]], label [[EXIT:%.*]], label [[FOR_BODY]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %for.body
+
+for.body:
+  %iv = phi i128 [ 0, %entry ], [ %inc, %for.body ]
+  %iv.trunc = trunc i128 %iv to i64
+  %gep1 = getelementptr { i64, i64 }, ptr %p, i64 %iv.trunc
+  %gep2 = getelementptr { i64, i64 }, ptr %p, i64 %iv.trunc, i32 1
+  store i64 0, ptr %gep1
+  store i64 0, ptr %gep2
+  %inc = add i128 %iv, 1
+  %tobool.not = icmp eq i128 %inc, 0
+  br i1 %tobool.not, label %exit, label %for.body
+
+exit:
+  ret void
+}