[DAGCombiner] Fix ReplaceAllUsesOfValueWith mutation bug in visitFREEZE

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

@@ -15808,13 +15808,16 @@ SDValue DAGCombiner::visitFREEZE(SDNode *N) {
     }
   }
 
-  SmallSetVector<SDValue, 8> MaybePoisonOperands;
-  for (SDValue Op : N0->ops()) {
+  SmallSet<SDValue, 8> MaybePoisonOperands;
+  SmallVector<unsigned, 8> MaybePoisonOperandNumbers;
+  for (auto [OpNo, Op] : enumerate(N0->ops())) {
     if (DAG.isGuaranteedNotToBeUndefOrPoison(Op, /*PoisonOnly*/ false,
                                              /*Depth*/ 1))
       continue;
     bool HadMaybePoisonOperands = !MaybePoisonOperands.empty();
-    bool IsNewMaybePoisonOperand = MaybePoisonOperands.insert(Op);
+    bool IsNewMaybePoisonOperand = MaybePoisonOperands.insert(Op).second;
+    if (IsNewMaybePoisonOperand)
+      MaybePoisonOperandNumbers.push_back(OpNo);
     if (!HadMaybePoisonOperands)
       continue;
     if (IsNewMaybePoisonOperand && !AllowMultipleMaybePoisonOperands) {
@@ -15826,7 +15829,18 @@ SDValue DAGCombiner::visitFREEZE(SDNode *N) {
   // it could create undef or poison due to it's poison-generating flags.
   // So not finding any maybe-poison operands is fine.
 
-  for (SDValue MaybePoisonOperand : MaybePoisonOperands) {
+  for (unsigned OpNo : MaybePoisonOperandNumbers) {
+    // N0 can mutate during iteration, so make sure to refech the maybe poison
+    // operands via the operand numbers. The typical scenario is that we have
+    // something like this
+    //   t262: i32 = freeze t181
+    //   t150: i32 = ctlz_zero_undef t262
+    //   t184: i32 = ctlz_zero_undef t181
+    //   t268: i32 = select_cc t181, Constant:i32<0>, t184, t186, setne:ch
+    // When freezing the t181 operand we get t262 back, and then the
+    // ReplaceAllUsesOfValueWith call will not only replace t181 by t262, but
+    // also recursively replace t184 by t150.
+    SDValue MaybePoisonOperand = N->getOperand(0).getOperand(OpNo);
     // Don't replace every single UNDEF everywhere with frozen UNDEF, though.
     if (MaybePoisonOperand.getOpcode() == ISD::UNDEF)
       continue;

llvm/test/CodeGen/AArch64/dag-combine-freeze.ll

@@ -0,0 +1,31 @@
+; RUN: llc -mtriple aarch64 -o /dev/null %s
+
+; This used to fail with:
+;    Assertion `N1.getOpcode() != ISD::DELETED_NODE &&
+;               "Operand is DELETED_NODE!"' failed.
+; Just make sure we do not crash here.
+define void @test_fold_freeze_over_select_cc(i15 %a, ptr %p1, ptr %p2) {
+entry:
+  %a2 = add nsw i15 %a, 1
+  %sext = sext i15 %a2 to i32
+  %ashr = ashr i32 %sext, 31
+  %lshr = lshr i32 %ashr, 7
+  ; Setup an already frozen input to ctlz.
+  %freeze = freeze i32 %lshr
+  %ctlz = call i32 @llvm.ctlz.i32(i32 %freeze, i1 true)
+  store i32 %ctlz, ptr %p1, align 1
+  ; Here is another ctlz, which is used by a frozen select.
+  ; DAGCombiner::visitFREEZE will to try to fold the freeze over a SELECT_CC,
+  ; and when dealing with the condition operand the other SELECT_CC operands
+  ; will be replaced/simplified as well. So the SELECT_CC is mutated while
+  ; freezing the "maybe poison operands". This needs to be handled by
+  ; DAGCombiner::visitFREEZE, as it can't store the list of SDValues that
+  ; should be frozen in a separate data structure that isn't updated when the
+  ; SELECT_CC is mutated.
+  %ctlz1 = call i32 @llvm.ctlz.i32(i32 %lshr, i1 true)
+  %icmp = icmp ne i32 %lshr, 0
+  %select = select i1 %icmp, i32 %ctlz1, i32 0
+  %freeze1 = freeze i32 %select
+  store i32 %freeze1, ptr %p2, align 1
+  ret void
+}