[InstCombine] fix infinite loop from shift transform

I'm not sure if there is a better way or another bug
still here, but this is enough to avoid the loop from:
https://llvm.org/PR51657

The test requires multiple blocks and datalayout to
trigger the problem path.

Author: rotateright

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

@@ -696,13 +696,14 @@ Instruction *InstCombinerImpl::FoldShiftByConstant(Value *Op0, Constant *Op1,
   // Fold shift2(trunc(shift1(x,c1)), c2) -> trunc(shift2(shift1(x,c1),c2))
   // If 'shift2' is an ashr, we would have to get the sign bit into a funny
   // place.  Don't try to do this transformation in this case.  Also, we
-  // require that the input operand is a shift-by-constant so that we have
-  // confidence that the shifts will get folded together.  We could do this
+  // require that the input operand is a non-poison shift-by-constant so that we
+  // have confidence that the shifts will get folded together.  We could do this
   // xform in more cases, but it is unlikely to be profitable.
   Instruction *TrOp;
   const APInt *TrShiftAmt;
   if (I.isLogicalShift() && match(Op0, m_Trunc(m_Instruction(TrOp))) &&
-      match(TrOp, m_Shift(m_Value(), m_APInt(TrShiftAmt)))) {
+      match(TrOp, m_Shift(m_Value(), m_APInt(TrShiftAmt))) &&
+      TrShiftAmt->ult(TrOp->getType()->getScalarSizeInBits())) {
     Type *SrcTy = TrOp->getType();
 
     // Okay, we'll do this xform.  Make the shift of shift.

llvm/test/Transforms/InstCombine/shift-amount-reassociation-with-truncation-shl.ll

@@ -1,6 +1,8 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -instcombine -S | FileCheck %s
 
+target datalayout = "n8:16:32"
+
 ; Given pattern:
 ;   (trunc (x << Q) to iDst) << K
 ; we should rewrite it as
@@ -227,3 +229,34 @@ define i16 @shl_tr_shl_constant_shift_amount_uses(i32 %x) {
   %r = shl i16 %tr, 4
   ret i16 %r
 }
+
+; This would infinite loop because we did not process the
+; poison shift before trying to fold shift-of-shift. To reach
+; the problematic transform, it needs a datalayout to specify
+; that the narrow types are legal, but i64 is not.
+
+define i1 @PR51657(i64 %x) {
+; CHECK-LABEL: @PR51657(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br i1 undef, label [[COND_FALSE:%.*]], label [[COND_END:%.*]]
+; CHECK:       cond.false:
+; CHECK-NEXT:    br label [[COND_END]]
+; CHECK:       cond.end:
+; CHECK-NEXT:    ret i1 true
+;
+entry:
+  br i1 undef, label %cond.false, label %cond.end
+
+cond.false:
+  %shl = shl i64 %x, 64
+  %conv26 = trunc i64 %shl to i32
+  %t1 = trunc i64 %shl to i8
+  br label %cond.end
+
+cond.end:
+  %cond = phi i32 [ %conv26, %cond.false ], [ 0, %entry ]
+  %t2 = phi i8 [ %t1, %cond.false ], [ 0, %entry ]
+  %conv27 = sext i8 %t2 to i32
+  %eq = icmp eq i32 %cond, %conv27
+  ret i1 %eq
+}