[InstCombine] Preserve nuw in OptimizePointerDifference

If both the geps and the subs are nuw the new sub is also nuw.

Proof: https://alive2.llvm.org/ce/z/mM8UvF

Author: nikic

llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp

@@ -2109,11 +2109,13 @@ Value *InstCombinerImpl::OptimizePointerDifference(Value *LHS, Value *RHS,
 
   // If we have a 2nd GEP of the same base pointer, subtract the offsets.
   // If both GEPs are inbounds, then the subtract does not have signed overflow.
+  // If both GEPs are nuw and the original sub is nuw, the new sub is also nuw.
   if (GEP2) {
-    bool GEP2IsInBounds = GEP2->isInBounds();
     Value *Offset = EmitGEPOffset(GEP2, RewriteGEPs);
-    Result = Builder.CreateSub(Result, Offset, "gepdiff", /* NUW */ false,
-                               GEP1IsInBounds && GEP2IsInBounds);
+    Result = Builder.CreateSub(Result, Offset, "gepdiff",
+                               IsNUW && GEP1->hasNoUnsignedWrap() &&
+                                   GEP2->hasNoUnsignedWrap(),
+                               GEP1IsInBounds && GEP2->isInBounds());
   }
 
   // If we have p - gep(p, ...)  then we have to negate the result.

llvm/test/Transforms/InstCombine/sub-gep.ll

@@ -156,6 +156,62 @@ define i64 @test_inbounds_nuw_two_gep(ptr %base, i64 %idx, i64 %idx2) {
   ret i64 %d
 }
 
+define i64 @test_nusw_two_gep(ptr %base, i64 %idx, i64 %idx2) {
+; CHECK-LABEL: @test_nusw_two_gep(
+; CHECK-NEXT:    [[TMP1:%.*]] = sub i64 [[IDX2:%.*]], [[IDX:%.*]]
+; CHECK-NEXT:    [[GEPDIFF:%.*]] = shl i64 [[TMP1]], 2
+; CHECK-NEXT:    ret i64 [[GEPDIFF]]
+;
+  %p1 = getelementptr nusw [0 x i32], ptr %base, i64 0, i64 %idx
+  %p2 = getelementptr nusw [0 x i32], ptr %base, i64 0, i64 %idx2
+  %i1 = ptrtoint ptr %p1 to i64
+  %i2 = ptrtoint ptr %p2 to i64
+  %d = sub i64 %i2, %i1
+  ret i64 %d
+}
+
+define i64 @test_nuw_two_gep(ptr %base, i64 %idx, i64 %idx2) {
+; CHECK-LABEL: @test_nuw_two_gep(
+; CHECK-NEXT:    [[TMP1:%.*]] = sub nuw i64 [[IDX2:%.*]], [[IDX:%.*]]
+; CHECK-NEXT:    [[GEPDIFF:%.*]] = shl nuw i64 [[TMP1]], 2
+; CHECK-NEXT:    ret i64 [[GEPDIFF]]
+;
+  %p1 = getelementptr nuw [0 x i32], ptr %base, i64 0, i64 %idx
+  %p2 = getelementptr nuw [0 x i32], ptr %base, i64 0, i64 %idx2
+  %i1 = ptrtoint ptr %p1 to i64
+  %i2 = ptrtoint ptr %p2 to i64
+  %d = sub nuw i64 %i2, %i1
+  ret i64 %d
+}
+
+define i64 @test_nuw_two_gep_missing_nuw_on_sub(ptr %base, i64 %idx, i64 %idx2) {
+; CHECK-LABEL: @test_nuw_two_gep_missing_nuw_on_sub(
+; CHECK-NEXT:    [[TMP1:%.*]] = sub i64 [[IDX2:%.*]], [[IDX:%.*]]
+; CHECK-NEXT:    [[GEPDIFF:%.*]] = shl i64 [[TMP1]], 2
+; CHECK-NEXT:    ret i64 [[GEPDIFF]]
+;
+  %p1 = getelementptr nuw [0 x i32], ptr %base, i64 0, i64 %idx
+  %p2 = getelementptr nuw [0 x i32], ptr %base, i64 0, i64 %idx2
+  %i1 = ptrtoint ptr %p1 to i64
+  %i2 = ptrtoint ptr %p2 to i64
+  %d = sub i64 %i2, %i1
+  ret i64 %d
+}
+
+define i64 @test_nuw_two_gep_missing_nuw_on_one_gep(ptr %base, i64 %idx, i64 %idx2) {
+; CHECK-LABEL: @test_nuw_two_gep_missing_nuw_on_one_gep(
+; CHECK-NEXT:    [[TMP1:%.*]] = sub i64 [[IDX2:%.*]], [[IDX:%.*]]
+; CHECK-NEXT:    [[GEPDIFF:%.*]] = shl i64 [[TMP1]], 2
+; CHECK-NEXT:    ret i64 [[GEPDIFF]]
+;
+  %p1 = getelementptr nuw [0 x i32], ptr %base, i64 0, i64 %idx
+  %p2 = getelementptr [0 x i32], ptr %base, i64 0, i64 %idx2
+  %i1 = ptrtoint ptr %p1 to i64
+  %i2 = ptrtoint ptr %p2 to i64
+  %d = sub nuw i64 %i2, %i1
+  ret i64 %d
+}
+
 define i64 @test_inbounds_nuw_multi_index(ptr %base, i64 %idx, i64 %idx2) {
 ; CHECK-LABEL: @test_inbounds_nuw_multi_index(
 ; CHECK-NEXT:    [[P2_IDX:%.*]] = shl nsw i64 [[IDX:%.*]], 3