Fixups after review feedback:

- Make sure we try to derive inbound + nuw before doing the
  ADD+GEP->GEP+GEP rewrites. This to make it possible to rely
  on "nuw" being present when trying to preserve flags.

- Make sure the special ADDNSW+SEXTLIKE+GEP->GEP+GEP case is checking
  for ADDNUW+ZEXTNNEG+GEP->GEP+GEP when preserving flags.
  This way we do not need to check for non-negative operands.

Author: bjope

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

@@ -3086,23 +3086,51 @@ Instruction *InstCombinerImpl::visitGetElementPtrInst(GetElementPtrInst &GEP) {
   if (GEPType->isVectorTy())
     return nullptr;
 
+  if (!GEP.isInBounds()) {
+    unsigned IdxWidth =
+        DL.getIndexSizeInBits(PtrOp->getType()->getPointerAddressSpace());
+    APInt BasePtrOffset(IdxWidth, 0);
+    Value *UnderlyingPtrOp =
+            PtrOp->stripAndAccumulateInBoundsConstantOffsets(DL,
+                                                             BasePtrOffset);
+    bool CanBeNull, CanBeFreed;
+    uint64_t DerefBytes = UnderlyingPtrOp->getPointerDereferenceableBytes(
+        DL, CanBeNull, CanBeFreed);
+    if (!CanBeNull && !CanBeFreed && DerefBytes != 0) {
+      if (GEP.accumulateConstantOffset(DL, BasePtrOffset) &&
+          BasePtrOffset.isNonNegative()) {
+        APInt AllocSize(IdxWidth, DerefBytes);
+        if (BasePtrOffset.ule(AllocSize)) {
+          return GetElementPtrInst::CreateInBounds(
+              GEP.getSourceElementType(), PtrOp, Indices, GEP.getName());
+        }
+      }
+    }
+  }
+
+  // nusw + nneg -> nuw
+  if (GEP.hasNoUnsignedSignedWrap() && !GEP.hasNoUnsignedWrap() &&
+      all_of(GEP.indices(), [&](Value *Idx) {
+        return isKnownNonNegative(Idx, SQ.getWithInstruction(&GEP));
+      })) {
+    GEP.setNoWrapFlags(GEP.getNoWrapFlags() | GEPNoWrapFlags::noUnsignedWrap());
+    return &GEP;
+  }
+
+  // These rewrites is trying to preserve inbounds/nuw attributes. So we want to
+  // do this after having tried to derive "nuw" above.
   if (GEP.getNumIndices() == 1) {
-    auto CanPreserveNoWrapFlags = [&](bool AddIsNSW, bool AddIsNUW, Value *Idx1,
-                                      Value *Idx2) {
-      // Preserve "inbounds nuw" if the original gep is "inbounds nuw",
-      // and the add is "nuw".
-      if (GEP.isInBounds() && GEP.hasNoUnsignedWrap() && AddIsNUW)
-        return GEPNoWrapFlags::inBounds() | GEPNoWrapFlags::noUnsignedWrap();
-      // Preserve "inbounds" if the original gep is "inbounds", the add
-      // is "nsw", and the add operands are non-negative.
-      SimplifyQuery Q = SQ.getWithInstruction(&GEP);
-      if (GEP.isInBounds() && AddIsNSW && isKnownNonNegative(Idx1, Q) &&
-          isKnownNonNegative(Idx2, Q))
-        return GEPNoWrapFlags::inBounds();
-      // Preserve "nuw" if the original gep is "nuw", and the add is "nuw".
-      if (GEP.hasNoUnsignedWrap() && AddIsNUW)
-        return GEPNoWrapFlags::noUnsignedWrap();
-      return GEPNoWrapFlags::none();
+    auto GetPreservedNoWrapFlags = [&](bool AddIsNUW, Value *Idx1, Value *Idx2) {
+      // Preserve "inbounds nuw" if the original gep is "inbounds nuw", and the
+      // add is "nuw". Preserve "nuw" if the original gep is "nuw", and the add
+      // is "nuw".
+      GEPNoWrapFlags Flags = GEPNoWrapFlags::none();
+      if (GEP.hasNoUnsignedWrap() && AddIsNUW) {
+        Flags |= GEPNoWrapFlags::noUnsignedWrap();
+        if (GEP.isInBounds())
+          Flags |= GEPNoWrapFlags::inBounds();
+      }
+      return Flags;
     };
 
     // Try to replace ADD + GEP with GEP + GEP.
@@ -3114,9 +3142,8 @@ Instruction *InstCombinerImpl::visitGetElementPtrInst(GetElementPtrInst &GEP) {
       // as:
       //   %newptr = getelementptr i32, ptr %ptr, i64 %idx1
       //   %newgep = getelementptr i32, ptr %newptr, i64 %idx2
-      bool NSW = match(GEP.getOperand(1), m_NSWAddLike(m_Value(), m_Value()));
       bool NUW = match(GEP.getOperand(1), m_NUWAddLike(m_Value(), m_Value()));
-      GEPNoWrapFlags NWFlags = CanPreserveNoWrapFlags(NSW, NUW, Idx1, Idx2);
+      GEPNoWrapFlags NWFlags = GetPreservedNoWrapFlags(NUW, Idx1, Idx2);
       auto *NewPtr =
           Builder.CreateGEP(GEP.getSourceElementType(), GEP.getPointerOperand(),
                             Idx1, "", NWFlags);
@@ -3133,8 +3160,9 @@ Instruction *InstCombinerImpl::visitGetElementPtrInst(GetElementPtrInst &GEP) {
       // as:
       // %newptr = getelementptr i32, ptr %ptr, i32 %idx1
       // %newgep = getelementptr i32, ptr %newptr, i32 idx2
-      GEPNoWrapFlags NWFlags = CanPreserveNoWrapFlags(
-          /*IsNSW=*/true, /*IsNUW=*/false, Idx1, C);
+      bool NUW = match(GEP.getOperand(1), m_NNegZExt(m_NUWAddLike(m_Value(),
+                                                                  m_Value())));
+      GEPNoWrapFlags NWFlags = GetPreservedNoWrapFlags(NUW, Idx1, C);
       auto *NewPtr = Builder.CreateGEP(
           GEP.getSourceElementType(), GEP.getPointerOperand(),
           Builder.CreateSExt(Idx1, GEP.getOperand(1)->getType()), "", NWFlags);
@@ -3146,37 +3174,6 @@ Instruction *InstCombinerImpl::visitGetElementPtrInst(GetElementPtrInst &GEP) {
     }
   }
 
-  if (!GEP.isInBounds()) {
-    unsigned IdxWidth =
-        DL.getIndexSizeInBits(PtrOp->getType()->getPointerAddressSpace());
-    APInt BasePtrOffset(IdxWidth, 0);
-    Value *UnderlyingPtrOp =
-            PtrOp->stripAndAccumulateInBoundsConstantOffsets(DL,
-                                                             BasePtrOffset);
-    bool CanBeNull, CanBeFreed;
-    uint64_t DerefBytes = UnderlyingPtrOp->getPointerDereferenceableBytes(
-        DL, CanBeNull, CanBeFreed);
-    if (!CanBeNull && !CanBeFreed && DerefBytes != 0) {
-      if (GEP.accumulateConstantOffset(DL, BasePtrOffset) &&
-          BasePtrOffset.isNonNegative()) {
-        APInt AllocSize(IdxWidth, DerefBytes);
-        if (BasePtrOffset.ule(AllocSize)) {
-          return GetElementPtrInst::CreateInBounds(
-              GEP.getSourceElementType(), PtrOp, Indices, GEP.getName());
-        }
-      }
-    }
-  }
-
-  // nusw + nneg -> nuw
-  if (GEP.hasNoUnsignedSignedWrap() && !GEP.hasNoUnsignedWrap() &&
-      all_of(GEP.indices(), [&](Value *Idx) {
-        return isKnownNonNegative(Idx, SQ.getWithInstruction(&GEP));
-      })) {
-    GEP.setNoWrapFlags(GEP.getNoWrapFlags() | GEPNoWrapFlags::noUnsignedWrap());
-    return &GEP;
-  }
-
   if (Instruction *R = foldSelectGEP(GEP, Builder))
     return R;