[InstCombine] Avoid iteration/mutation conflict

[JosephTremoulet]

Summary

When iterating users of a multiply in processUMulZExtIdiom, the
call to setOperand in the truncation case may replace the use
being visited; make sure the iterator has been advanced before
doing that replacement.

[davide]

sigh at iterator invalidation issues. lgtm

lib/Transforms/InstCombine/InstCombineCompares.cpp

static Instruction *processUMulZExtIdiom(ICmpInst &I, Value *MulVal,
   CallInst *Call = Builder.CreateCall(F, {MulA, MulB}, "umul");
   IC.Worklist.Add(MulInstr);
 
   // If there are uses of mul result other than the comparison, we know that
   // they are truncation or binary AND. Change them to use result of
   // mul.with.overflow and adjust properly mask/size.
   if (MulVal->hasNUsesOrMore(2)) {
     Value *Mul = Builder.CreateExtractValue(Call, 0, "umul.value");
-    for (User *U : MulVal->users()) {
+    for (auto UI = MulVal->user_begin(), UE = MulVal->user_end(); UI != UE;) {
+      User *U = *UI++;
       if (U == &I || U == OtherVal)
         continue;
       if (TruncInst *TI = dyn_cast<TruncInst>(U)) {
         if (TI->getType()->getPrimitiveSizeInBits() == MulWidth)
           IC.replaceInstUsesWith(*TI, Mul);
         else
           TI->setOperand(0, Mul);
       } else if (BinaryOperator *BO = dyn_cast<BinaryOperator>(U)) {

test/Transforms/InstCombine/icmp-mul-zext.ll

if.then9:
   %tinky = load i16, i16* @glob
   %conv13 = sext i16 %tinky to i32
   %and = and i32 %mul, %conv13
   %conv14 = trunc i32 %and to i16
   store i16 %conv14, i16* @glob
   ret void
 }
 
+; Repro case for bug involving mutating a list while
+; iterating it.
+
+declare i16 @aux(i8)
+
+define i16 @iter_breaker(i16 %a, i16 %b) {
+; CHECK-LABEL: @iter_breaker(
+; CHECK-NEXT:    [[UMUL:%.*]] = call { i16, i1 } @llvm.umul.with.overflow.i16(i16 [[A:%.*]], i16 [[B:%.*]])
+; CHECK-NEXT:    [[UMUL_VALUE:%.*]] = extractvalue { i16, i1 } [[UMUL]], 0
+; CHECK-NEXT:    [[DID_OVF:%.*]] = extractvalue { i16, i1 } [[UMUL]], 1
+; CHECK-NEXT:    br i1 [[DID_OVF]], label [[RET1:%.*]], label [[RET2:%.*]]
+; CHECK:       ret1:
+; CHECK-NEXT:    [[TRUNC_REMAIN:%.*]] = trunc i16 [[UMUL_VALUE]] to i8
+; CHECK-NEXT:    [[VAL:%.*]] = call i16 @aux(i8 [[TRUNC_REMAIN]])
+; CHECK-NEXT:    ret i16 [[VAL]]
+; CHECK:       ret2:
+; CHECK-NEXT:    ret i16 [[UMUL_VALUE]]
+;
+  %a_wide = zext i16 %a to i32
+  %b_wide = zext i16 %b to i32
+  %mul_wide = mul i32 %a_wide, %b_wide              ; uses of %mul_wide will be iterated
+
+  %trunc_remain = trunc i32 %mul_wide to i8         ; this use will be replaced w/ new value
+                                                    ; when iteration visits it, switching
+                                                    ; iteration to the uses of new value
+
+  %trunc_unnecessary = trunc i32 %mul_wide to i16   ; uses of %trunc_unnecessary will have
+                                                    ; been updated to uses of new value
+  %did_ovf = icmp ugt i32 %mul_wide, 65535
+  br i1 %did_ovf, label %ret1, label %ret2
+
+ret1:
+  %val = call i16 @aux(i8 %trunc_remain)
+  ret i16 %val
+
+ret2:
+  ret i16 %trunc_unnecessary              ; crash visiting this use after corrupting iterator
+}