Auto merge of #137271 - nikic:gep-nuw-2, r=scottmcm

Emit getelementptr inbounds nuw for pointer::add()

Lower pointer::add (via intrinsic::offset with unsigned offset) to getelementptr inbounds nuw on LLVM versions that support it. This lets LLVM make use of the pre-condition that the offset addition does not wrap in an unsigned sense. Together with inbounds, this also implies that the offset is non-negative.

Fixes https://github.com/rust-lang/rust/issues/137217.

10 files changed, 93 insertions, 29 deletions

diff --git a/compiler/rustc_codegen_llvm/src/builder.rs b/compiler/rustc_codegen_llvm/src/builder.rs
index e1609e31c07..677a9cd3e90 100644
--- a/compiler/rustc_codegen_llvm/src/builder.rs
+++ b/compiler/rustc_codegen_llvm/src/builder.rs
@@ -31,7 +31,9 @@ use tracing::{debug, instrument};
 use crate::abi::FnAbiLlvmExt;
 use crate::common::Funclet;
 use crate::context::{CodegenCx, SimpleCx};
-use crate::llvm::{self, AtomicOrdering, AtomicRmwBinOp, BasicBlock, False, Metadata, True};
+use crate::llvm::{
+    self, AtomicOrdering, AtomicRmwBinOp, BasicBlock, False, GEPNoWrapFlags, Metadata, True,
+};
 use crate::type_::Type;
 use crate::type_of::LayoutLlvmExt;
 use crate::value::Value;
@@ -910,13 +912,14 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
 
     fn gep(&mut self, ty: &'ll Type, ptr: &'ll Value, indices: &[&'ll Value]) -> &'ll Value {
         unsafe {
-            llvm::LLVMBuildGEP2(
+            llvm::LLVMBuildGEPWithNoWrapFlags(
                 self.llbuilder,
                 ty,
                 ptr,
                 indices.as_ptr(),
                 indices.len() as c_uint,
                 UNNAMED,
+                GEPNoWrapFlags::default(),
             )
         }
     }
@@ -928,13 +931,33 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         indices: &[&'ll Value],
     ) -> &'ll Value {
         unsafe {
-            llvm::LLVMBuildInBoundsGEP2(
+            llvm::LLVMBuildGEPWithNoWrapFlags(
+                self.llbuilder,
+                ty,
+                ptr,
+                indices.as_ptr(),
+                indices.len() as c_uint,
+                UNNAMED,
+                GEPNoWrapFlags::InBounds,
+            )
+        }
+    }
+
+    fn inbounds_nuw_gep(
+        &mut self,
+        ty: &'ll Type,
+        ptr: &'ll Value,
+        indices: &[&'ll Value],
+    ) -> &'ll Value {
+        unsafe {
+            llvm::LLVMBuildGEPWithNoWrapFlags(
                 self.llbuilder,
                 ty,
                 ptr,
                 indices.as_ptr(),
                 indices.len() as c_uint,
                 UNNAMED,
+                GEPNoWrapFlags::InBounds | GEPNoWrapFlags::NUW,
             )
         }
     }
diff --git a/compiler/rustc_codegen_llvm/src/llvm/ffi.rs b/compiler/rustc_codegen_llvm/src/llvm/ffi.rs
index 3b0187b9d37..da91e6edbcf 100644
--- a/compiler/rustc_codegen_llvm/src/llvm/ffi.rs
+++ b/compiler/rustc_codegen_llvm/src/llvm/ffi.rs
@@ -954,6 +954,17 @@ bitflags! {
     }
 }
 
+// These values **must** match with LLVMGEPNoWrapFlags
+bitflags! {
+    #[repr(transparent)]
+    #[derive(Default)]
+    pub struct GEPNoWrapFlags : c_uint {
+        const InBounds = 1 << 0;
+        const NUSW = 1 << 1;
+        const NUW = 1 << 2;
+    }
+}
+
 unsafe extern "C" {
     pub type ModuleBuffer;
 }
@@ -1454,21 +1465,14 @@ unsafe extern "C" {
 
     pub(crate) fn LLVMBuildStore<'a>(B: &Builder<'a>, Val: &'a Value, Ptr: &'a Value) -> &'a Value;
 
-    pub(crate) fn LLVMBuildGEP2<'a>(
-        B: &Builder<'a>,
-        Ty: &'a Type,
-        Pointer: &'a Value,
-        Indices: *const &'a Value,
-        NumIndices: c_uint,
-        Name: *const c_char,
-    ) -> &'a Value;
-    pub(crate) fn LLVMBuildInBoundsGEP2<'a>(
+    pub(crate) fn LLVMBuildGEPWithNoWrapFlags<'a>(
         B: &Builder<'a>,
         Ty: &'a Type,
         Pointer: &'a Value,
         Indices: *const &'a Value,
         NumIndices: c_uint,
         Name: *const c_char,
+        Flags: GEPNoWrapFlags,
     ) -> &'a Value;
 
     // Casts
diff --git a/compiler/rustc_codegen_ssa/src/mir/place.rs b/compiler/rustc_codegen_ssa/src/mir/place.rs
index e32d298869b..edd09b9c3c5 100644
--- a/compiler/rustc_codegen_ssa/src/mir/place.rs
+++ b/compiler/rustc_codegen_ssa/src/mir/place.rs
@@ -423,7 +423,7 @@ impl<'a, 'tcx, V: CodegenObject> PlaceRef<'tcx, V> {
             layout.size
         };
 
-        let llval = bx.inbounds_gep(bx.cx().backend_type(layout), self.val.llval, &[llindex]);
+        let llval = bx.inbounds_nuw_gep(bx.cx().backend_type(layout), self.val.llval, &[llindex]);
         let align = self.val.align.restrict_for_offset(offset);
         PlaceValue::new_sized(llval, align).with_type(layout)
     }
diff --git a/compiler/rustc_codegen_ssa/src/mir/rvalue.rs b/compiler/rustc_codegen_ssa/src/mir/rvalue.rs
index daa4fa90ed7..1eebe04225b 100644
--- a/compiler/rustc_codegen_ssa/src/mir/rvalue.rs
+++ b/compiler/rustc_codegen_ssa/src/mir/rvalue.rs
@@ -666,9 +666,15 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                         lhs.layout.ty,
                     ),
 
-                    (OperandValue::Immediate(lhs_val), OperandValue::Immediate(rhs_val)) => {
-                        self.codegen_scalar_binop(bx, op, lhs_val, rhs_val, lhs.layout.ty)
-                    }
+                    (OperandValue::Immediate(lhs_val), OperandValue::Immediate(rhs_val)) => self
+                        .codegen_scalar_binop(
+                            bx,
+                            op,
+                            lhs_val,
+                            rhs_val,
+                            lhs.layout.ty,
+                            rhs.layout.ty,
+                        ),
 
                     _ => bug!(),
                 };
@@ -889,10 +895,11 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         op: mir::BinOp,
         lhs: Bx::Value,
         rhs: Bx::Value,
-        input_ty: Ty<'tcx>,
+        lhs_ty: Ty<'tcx>,
+        rhs_ty: Ty<'tcx>,
     ) -> Bx::Value {
-        let is_float = input_ty.is_floating_point();
-        let is_signed = input_ty.is_signed();
+        let is_float = lhs_ty.is_floating_point();
+        let is_signed = lhs_ty.is_signed();
         match op {
             mir::BinOp::Add => {
                 if is_float {
@@ -958,9 +965,9 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             mir::BinOp::BitAnd => bx.and(lhs, rhs),
             mir::BinOp::BitXor => bx.xor(lhs, rhs),
             mir::BinOp::Offset => {
-                let pointee_type = input_ty
+                let pointee_type = lhs_ty
                     .builtin_deref(true)
-                    .unwrap_or_else(|| bug!("deref of non-pointer {:?}", input_ty));
+                    .unwrap_or_else(|| bug!("deref of non-pointer {:?}", lhs_ty));
                 let pointee_layout = bx.cx().layout_of(pointee_type);
                 if pointee_layout.is_zst() {
                     // `Offset` works in terms of the size of pointee,
@@ -968,7 +975,11 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     lhs
                 } else {
                     let llty = bx.cx().backend_type(pointee_layout);
-                    bx.inbounds_gep(llty, lhs, &[rhs])
+                    if !rhs_ty.is_signed() {
+                        bx.inbounds_nuw_gep(llty, lhs, &[rhs])
+                    } else {
+                        bx.inbounds_gep(llty, lhs, &[rhs])
+                    }
                 }
             }
             mir::BinOp::Shl | mir::BinOp::ShlUnchecked => {
diff --git a/compiler/rustc_codegen_ssa/src/traits/builder.rs b/compiler/rustc_codegen_ssa/src/traits/builder.rs
index 2b00ba01946..99fd6b6510f 100644
--- a/compiler/rustc_codegen_ssa/src/traits/builder.rs
+++ b/compiler/rustc_codegen_ssa/src/traits/builder.rs
@@ -325,6 +325,14 @@ pub trait BuilderMethods<'a, 'tcx>:
         ptr: Self::Value,
         indices: &[Self::Value],
     ) -> Self::Value;
+    fn inbounds_nuw_gep(
+        &mut self,
+        ty: Self::Type,
+        ptr: Self::Value,
+        indices: &[Self::Value],
+    ) -> Self::Value {
+        self.inbounds_gep(ty, ptr, indices)
+    }
     fn ptradd(&mut self, ptr: Self::Value, offset: Self::Value) -> Self::Value {
         self.gep(self.cx().type_i8(), ptr, &[offset])
     }
diff --git a/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp b/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp
index b8cef6a7e25..aea2a8dd097 100644
--- a/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp
+++ b/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp
@@ -1767,6 +1767,24 @@ extern "C" LLVMValueRef LLVMRustBuildMaxNum(LLVMBuilderRef B, LLVMValueRef LHS,
   return wrap(unwrap(B)->CreateMaxNum(unwrap(LHS), unwrap(RHS)));
 }
 
+#if LLVM_VERSION_LT(19, 0)
+enum {
+  LLVMGEPFlagInBounds = (1 << 0),
+  LLVMGEPFlagNUSW = (1 << 1),
+  LLVMGEPFlagNUW = (1 << 2),
+};
+extern "C" LLVMValueRef
+LLVMBuildGEPWithNoWrapFlags(LLVMBuilderRef B, LLVMTypeRef Ty,
+                            LLVMValueRef Pointer, LLVMValueRef *Indices,
+                            unsigned NumIndices, const char *Name,
+                            unsigned NoWrapFlags) {
+  if (NoWrapFlags & LLVMGEPFlagInBounds)
+    return LLVMBuildInBoundsGEP2(B, Ty, Pointer, Indices, NumIndices, Name);
+  else
+    return LLVMBuildGEP2(B, Ty, Pointer, Indices, NumIndices, Name);
+}
+#endif
+
 // Transfers ownership of DiagnosticHandler unique_ptr to the caller.
 extern "C" DiagnosticHandler *
 LLVMRustContextGetDiagnosticHandler(LLVMContextRef C) {
diff --git a/tests/codegen/gep-index.rs b/tests/codegen/gep-index.rs
index 1f5e8855910..bfb2511af87 100644
--- a/tests/codegen/gep-index.rs
+++ b/tests/codegen/gep-index.rs
@@ -11,27 +11,27 @@ struct Foo(i32, i32);
 // CHECK-LABEL: @index_on_struct(
 #[no_mangle]
 fn index_on_struct(a: &[Foo], index: usize) -> &Foo {
-    // CHECK: getelementptr inbounds %Foo, ptr %a.0, {{i64|i32}} %index
+    // CHECK: getelementptr inbounds{{( nuw)?}} %Foo, ptr %a.0, {{i64|i32}} %index
     &a[index]
 }
 
 // CHECK-LABEL: @offset_on_struct(
 #[no_mangle]
 fn offset_on_struct(a: *const Foo, index: usize) -> *const Foo {
-    // CHECK: getelementptr inbounds %Foo, ptr %a, {{i64|i32}} %index
+    // CHECK: getelementptr inbounds{{( nuw)?}} %Foo, ptr %a, {{i64|i32}} %index
     unsafe { a.add(index) }
 }
 
 // CHECK-LABEL: @index_on_i32(
 #[no_mangle]
 fn index_on_i32(a: &[i32], index: usize) -> &i32 {
-    // CHECK: getelementptr inbounds i32, ptr %a.0, {{i64|i32}} %index
+    // CHECK: getelementptr inbounds{{( nuw)?}} i32, ptr %a.0, {{i64|i32}} %index
     &a[index]
 }
 
 // CHECK-LABEL: @offset_on_i32(
 #[no_mangle]
 fn offset_on_i32(a: *const i32, index: usize) -> *const i32 {
-    // CHECK: getelementptr inbounds i32, ptr %a, {{i64|i32}} %index
+    // CHECK: getelementptr inbounds{{( nuw)?}} i32, ptr %a, {{i64|i32}} %index
     unsafe { a.add(index) }
 }
diff --git a/tests/codegen/intrinsics/offset.rs b/tests/codegen/intrinsics/offset.rs
index d76d3e705ab..cf0c7c7ac7d 100644
--- a/tests/codegen/intrinsics/offset.rs
+++ b/tests/codegen/intrinsics/offset.rs
@@ -27,7 +27,7 @@ pub unsafe fn offset_isize(p: *const u32, d: isize) -> *const u32 {
 // CHECK-SAME: (ptr noundef %p, [[SIZE]] noundef %d)
 #[no_mangle]
 pub unsafe fn offset_usize(p: *const u64, d: usize) -> *const u64 {
-    // CHECK: %[[R:.*]] = getelementptr inbounds i64, ptr %p, [[SIZE]] %d
+    // CHECK: %[[R:.*]] = getelementptr inbounds{{( nuw)?}} i64, ptr %p, [[SIZE]] %d
     // CHECK-NEXT: ret ptr %[[R]]
     offset(p, d)
 }
diff --git a/tests/codegen/intrinsics/ptr_metadata.rs b/tests/codegen/intrinsics/ptr_metadata.rs
index 87a32fa3d24..044dbc20486 100644
--- a/tests/codegen/intrinsics/ptr_metadata.rs
+++ b/tests/codegen/intrinsics/ptr_metadata.rs
@@ -28,7 +28,7 @@ pub unsafe fn dyn_byte_offset(
     p: *const dyn std::fmt::Debug,
     n: usize,
 ) -> *const dyn std::fmt::Debug {
-    // CHECK: %[[Q:.+]] = getelementptr inbounds i8, ptr %p.0, i64 %n
+    // CHECK: %[[Q:.+]] = getelementptr inbounds{{( nuw)?}} i8, ptr %p.0, i64 %n
     // CHECK: %[[TEMP1:.+]] = insertvalue { ptr, ptr } poison, ptr %[[Q]], 0
     // CHECK: %[[TEMP2:.+]] = insertvalue { ptr, ptr } %[[TEMP1]], ptr %p.1, 1
     // CHECK: ret { ptr, ptr } %[[TEMP2]]
diff --git a/tests/codegen/ptr-arithmetic.rs b/tests/codegen/ptr-arithmetic.rs
index ecb44b30f5c..fc4441ef448 100644
--- a/tests/codegen/ptr-arithmetic.rs
+++ b/tests/codegen/ptr-arithmetic.rs
@@ -6,7 +6,7 @@
 // CHECK-SAME: [[WORD:i[0-9]+]] noundef %n)
 #[no_mangle]
 pub unsafe fn i32_add(p: *const i32, n: usize) -> *const i32 {
-    // CHECK: %[[TEMP:.+]] = getelementptr inbounds i32, ptr %p, [[WORD]] %n
+    // CHECK: %[[TEMP:.+]] = getelementptr inbounds{{( nuw)?}} i32, ptr %p, [[WORD]] %n
     // CHECK: ret ptr %[[TEMP]]
     p.add(n)
 }