[llvm] [DropUnnecessaryAssumes] Fix iterator invalidation. (PR #203765)

[Florian Hahn via llvm-commits]

https://github.com/fhahn updated https://github.com/llvm/llvm-project/pull/203765

>From f977dd5ce0fc7fb0a62ea7e6139fd9f2fd614df3 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Fri, 12 Jun 2026 20:43:57 +0100
Subject: [PATCH] [DropUnnecessaryAssumes] Fix iterator invalidation.

registerAssumption() below can append to (and reallocate) the cache's
assumption vector. Use integer index for indexing instead of using the
iterator. Stop at the original count, so we don't reprocess assumes
created during the loop.
---
 .../Scalar/DropUnnecessaryAssumes.cpp         |  4 +-
 .../DropUnnecessaryAssumes/dereferenceable.ll | 64 +++++++++++++++++++
 2 files changed, 66 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/DropUnnecessaryAssumes.cpp b/llvm/lib/Transforms/Scalar/DropUnnecessaryAssumes.cpp
index 4a7144fe6c77a..3c8cf93e85935 100644
--- a/llvm/lib/Transforms/Scalar/DropUnnecessaryAssumes.cpp
+++ b/llvm/lib/Transforms/Scalar/DropUnnecessaryAssumes.cpp
@@ -67,8 +67,8 @@ DropUnnecessaryAssumesPass::run(Function &F, FunctionAnalysisManager &FAM) {
   AssumptionCache &AC = FAM.getResult<AssumptionAnalysis>(F);
   bool Changed = false;
 
-  for (const WeakVH &Elem : AC.assumptions()) {
-    auto *Assume = cast_or_null<AssumeInst>(Elem);
+  for (unsigned I = 0, E = AC.assumptions().size(); I != E; ++I) {
+    auto *Assume = cast_or_null<AssumeInst>(AC.assumptions()[I]);
     if (!Assume)
       continue;
 
diff --git a/llvm/test/Transforms/DropUnnecessaryAssumes/dereferenceable.ll b/llvm/test/Transforms/DropUnnecessaryAssumes/dereferenceable.ll
index 43fa08c070828..079c0403227f3 100644
--- a/llvm/test/Transforms/DropUnnecessaryAssumes/dereferenceable.ll
+++ b/llvm/test/Transforms/DropUnnecessaryAssumes/dereferenceable.ll
@@ -52,3 +52,67 @@ define i8 @test_dereferenceable_with_align_ptr_used(ptr %p, i64 %size) {
   %val = load i8, ptr %p
   ret i8 %val
 }
+
+; Make sure newly created assumes are handled properly.
+define i8 @test_dereferenceable_with_align_cache_realloc(ptr %p, ptr %q, i1 %c) {
+; CHECK-LABEL: define i8 @test_dereferenceable_with_align_cache_realloc(
+; CHECK-SAME: ptr [[P:%.*]], ptr [[Q:%.*]], i1 [[C:%.*]]) {
+; CHECK-NEXT:    [[V:%.*]] = load i8, ptr [[P]], align 1
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(ptr [[Q]], i64 8), "align"(ptr [[P]], i64 8) ]
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(ptr [[Q]], i64 8), "align"(ptr [[P]], i64 8) ]
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(ptr [[Q]], i64 8), "align"(ptr [[P]], i64 8) ]
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(ptr [[Q]], i64 8), "align"(ptr [[P]], i64 8) ]
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(ptr [[Q]], i64 8), "align"(ptr [[P]], i64 8) ]
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(ptr [[Q]], i64 8), "align"(ptr [[P]], i64 8) ]
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(ptr [[Q]], i64 8), "align"(ptr [[P]], i64 8) ]
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(ptr [[Q]], i64 8), "align"(ptr [[P]], i64 8) ]
+; CHECK-NEXT:    ret i8 [[V]]
+;
+; DROP-DEREF-LABEL: define i8 @test_dereferenceable_with_align_cache_realloc(
+; DROP-DEREF-SAME: ptr [[P:%.*]], ptr [[Q:%.*]], i1 [[C:%.*]]) {
+; DROP-DEREF-NEXT:    [[V:%.*]] = load i8, ptr [[P]], align 1
+; DROP-DEREF-NEXT:    call void @llvm.assume(i1 true) [ "align"(ptr [[P]], i64 8) ]
+; DROP-DEREF-NEXT:    call void @llvm.assume(i1 true) [ "align"(ptr [[P]], i64 8) ]
+; DROP-DEREF-NEXT:    call void @llvm.assume(i1 true) [ "align"(ptr [[P]], i64 8) ]
+; DROP-DEREF-NEXT:    call void @llvm.assume(i1 true) [ "align"(ptr [[P]], i64 8) ]
+; DROP-DEREF-NEXT:    call void @llvm.assume(i1 true) [ "align"(ptr [[P]], i64 8) ]
+; DROP-DEREF-NEXT:    call void @llvm.assume(i1 true) [ "align"(ptr [[P]], i64 8) ]
+; DROP-DEREF-NEXT:    call void @llvm.assume(i1 true) [ "align"(ptr [[P]], i64 8) ]
+; DROP-DEREF-NEXT:    call void @llvm.assume(i1 true) [ "align"(ptr [[P]], i64 8) ]
+; DROP-DEREF-NEXT:    ret i8 [[V]]
+;
+  %v = load i8, ptr %p
+  call void @llvm.assume(i1 true) [ "dereferenceable"(ptr %q, i64 8), "align"(ptr %p, i64 8) ]
+  call void @llvm.assume(i1 true) [ "dereferenceable"(ptr %q, i64 8), "align"(ptr %p, i64 8) ]
+  call void @llvm.assume(i1 true) [ "dereferenceable"(ptr %q, i64 8), "align"(ptr %p, i64 8) ]
+  call void @llvm.assume(i1 true) [ "dereferenceable"(ptr %q, i64 8), "align"(ptr %p, i64 8) ]
+  call void @llvm.assume(i1 true) [ "dereferenceable"(ptr %q, i64 8), "align"(ptr %p, i64 8) ]
+  call void @llvm.assume(i1 true) [ "dereferenceable"(ptr %q, i64 8), "align"(ptr %p, i64 8) ]
+  call void @llvm.assume(i1 true) [ "dereferenceable"(ptr %q, i64 8), "align"(ptr %p, i64 8) ]
+  call void @llvm.assume(i1 true) [ "dereferenceable"(ptr %q, i64 8), "align"(ptr %p, i64 8) ]
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  ret i8 %v
+}