Auto merge of rust-lang#136575 - scottmcm:nsuw-math, r=nikic

Set both `nuw` and `nsw` in slice size calculation

There's an old note in the code to do this, and now that [LLVM-C has an API for it](https://github.com/llvm/llvm-project/blob/f0b8ff12519270adcfef93410abff76ab073476a/llvm/include/llvm-c/Core.h#L4403-L4408), we might as well.  And it's been there since what looks like LLVM 17 llvm/llvm-project@de9b6aa so doesn't even need to be conditional.

(There's other places, like `RawVecInner` or `Layout`, that might want to do things like this too, but I'll leave those for a future PR.)

compiler/rustc_codegen_gcc/src/builder.rs

@@ -665,6 +665,7 @@ impl<'a, 'gcc, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'gcc, 'tcx> {
         a + b
     }
 
+    // TODO(antoyo): should we also override the `unchecked_` versions?
     fn sub(&mut self, a: RValue<'gcc>, b: RValue<'gcc>) -> RValue<'gcc> {
         self.gcc_sub(a, b)
     }
@@ -832,31 +833,6 @@ impl<'a, 'gcc, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'gcc, 'tcx> {
         set_rvalue_location(self, self.gcc_not(a))
     }
 
-    fn unchecked_sadd(&mut self, a: RValue<'gcc>, b: RValue<'gcc>) -> RValue<'gcc> {
-        set_rvalue_location(self, self.gcc_add(a, b))
-    }
-
-    fn unchecked_uadd(&mut self, a: RValue<'gcc>, b: RValue<'gcc>) -> RValue<'gcc> {
-        set_rvalue_location(self, self.gcc_add(a, b))
-    }
-
-    fn unchecked_ssub(&mut self, a: RValue<'gcc>, b: RValue<'gcc>) -> RValue<'gcc> {
-        set_rvalue_location(self, self.gcc_sub(a, b))
-    }
-
-    fn unchecked_usub(&mut self, a: RValue<'gcc>, b: RValue<'gcc>) -> RValue<'gcc> {
-        // TODO(antoyo): should generate poison value?
-        set_rvalue_location(self, self.gcc_sub(a, b))
-    }
-
-    fn unchecked_smul(&mut self, a: RValue<'gcc>, b: RValue<'gcc>) -> RValue<'gcc> {
-        set_rvalue_location(self, self.gcc_mul(a, b))
-    }
-
-    fn unchecked_umul(&mut self, a: RValue<'gcc>, b: RValue<'gcc>) -> RValue<'gcc> {
-        set_rvalue_location(self, self.gcc_mul(a, b))
-    }
-
     fn fadd_fast(&mut self, lhs: RValue<'gcc>, rhs: RValue<'gcc>) -> RValue<'gcc> {
         // NOTE: it seems like we cannot enable fast-mode for a single operation in GCC.
         set_rvalue_location(self, lhs + rhs)

compiler/rustc_codegen_llvm/src/builder.rs

@@ -421,6 +421,37 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         unchecked_umul(x, y) => LLVMBuildNUWMul,
     }
 
+    fn unchecked_suadd(&mut self, a: &'ll Value, b: &'ll Value) -> &'ll Value {
+        unsafe {
+            let add = llvm::LLVMBuildAdd(self.llbuilder, a, b, UNNAMED);
+            if llvm::LLVMIsAInstruction(add).is_some() {
+                llvm::LLVMSetNUW(add, True);
+                llvm::LLVMSetNSW(add, True);
+            }
+            add
+        }
+    }
+    fn unchecked_susub(&mut self, a: &'ll Value, b: &'ll Value) -> &'ll Value {
+        unsafe {
+            let sub = llvm::LLVMBuildSub(self.llbuilder, a, b, UNNAMED);
+            if llvm::LLVMIsAInstruction(sub).is_some() {
+                llvm::LLVMSetNUW(sub, True);
+                llvm::LLVMSetNSW(sub, True);
+            }
+            sub
+        }
+    }
+    fn unchecked_sumul(&mut self, a: &'ll Value, b: &'ll Value) -> &'ll Value {
+        unsafe {
+            let mul = llvm::LLVMBuildMul(self.llbuilder, a, b, UNNAMED);
+            if llvm::LLVMIsAInstruction(mul).is_some() {
+                llvm::LLVMSetNUW(mul, True);
+                llvm::LLVMSetNSW(mul, True);
+            }
+            mul
+        }
+    }
+
     fn or_disjoint(&mut self, a: &'ll Value, b: &'ll Value) -> &'ll Value {
         unsafe {
             let or = llvm::LLVMBuildOr(self.llbuilder, a, b, UNNAMED);

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -1430,6 +1430,8 @@ unsafe extern "C" {
 
     // Extra flags on arithmetic
     pub(crate) fn LLVMSetIsDisjoint(Instr: &Value, IsDisjoint: Bool);
+    pub(crate) fn LLVMSetNUW(ArithInst: &Value, HasNUW: Bool);
+    pub(crate) fn LLVMSetNSW(ArithInst: &Value, HasNSW: Bool);
 
     // Memory
     pub(crate) fn LLVMBuildAlloca<'a>(

compiler/rustc_codegen_ssa/src/size_of_val.rs

@@ -45,14 +45,9 @@ pub fn size_and_align_of_dst<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
             // The info in this case is the length of the str, so the size is that
             // times the unit size.
             (
-                // All slice sizes must fit into `isize`, so this multiplication cannot (signed)
-                // wrap.
-                // NOTE: ideally, we want the effects of both `unchecked_smul` and `unchecked_umul`
-                // (resulting in `mul nsw nuw` in LLVM IR), since we know that the multiplication
-                // cannot signed wrap, and that both operands are non-negative. But at the time of
-                // writing, the `LLVM-C` binding can't do this, and it doesn't seem to enable any
-                // further optimizations.
-                bx.unchecked_smul(info.unwrap(), bx.const_usize(unit.size.bytes())),
+                // All slice sizes must fit into `isize`, so this multiplication cannot
+                // wrap -- neither signed nor unsigned.
+                bx.unchecked_sumul(info.unwrap(), bx.const_usize(unit.size.bytes())),
                 bx.const_usize(unit.align.abi.bytes()),
             )
         }

compiler/rustc_codegen_ssa/src/traits/builder.rs

@@ -159,12 +159,35 @@ pub trait BuilderMethods<'a, 'tcx>:
     /// must be interpreted as unsigned and can be assumed to be less than the size of the left
     /// operand.
     fn ashr(&mut self, lhs: Self::Value, rhs: Self::Value) -> Self::Value;
-    fn unchecked_sadd(&mut self, lhs: Self::Value, rhs: Self::Value) -> Self::Value;
-    fn unchecked_uadd(&mut self, lhs: Self::Value, rhs: Self::Value) -> Self::Value;
-    fn unchecked_ssub(&mut self, lhs: Self::Value, rhs: Self::Value) -> Self::Value;
-    fn unchecked_usub(&mut self, lhs: Self::Value, rhs: Self::Value) -> Self::Value;
-    fn unchecked_smul(&mut self, lhs: Self::Value, rhs: Self::Value) -> Self::Value;
-    fn unchecked_umul(&mut self, lhs: Self::Value, rhs: Self::Value) -> Self::Value;
+    fn unchecked_sadd(&mut self, lhs: Self::Value, rhs: Self::Value) -> Self::Value {
+        self.add(lhs, rhs)
+    }
+    fn unchecked_uadd(&mut self, lhs: Self::Value, rhs: Self::Value) -> Self::Value {
+        self.add(lhs, rhs)
+    }
+    fn unchecked_suadd(&mut self, lhs: Self::Value, rhs: Self::Value) -> Self::Value {
+        self.unchecked_sadd(lhs, rhs)
+    }
+    fn unchecked_ssub(&mut self, lhs: Self::Value, rhs: Self::Value) -> Self::Value {
+        self.sub(lhs, rhs)
+    }
+    fn unchecked_usub(&mut self, lhs: Self::Value, rhs: Self::Value) -> Self::Value {
+        self.sub(lhs, rhs)
+    }
+    fn unchecked_susub(&mut self, lhs: Self::Value, rhs: Self::Value) -> Self::Value {
+        self.unchecked_ssub(lhs, rhs)
+    }
+    fn unchecked_smul(&mut self, lhs: Self::Value, rhs: Self::Value) -> Self::Value {
+        self.mul(lhs, rhs)
+    }
+    fn unchecked_umul(&mut self, lhs: Self::Value, rhs: Self::Value) -> Self::Value {
+        self.mul(lhs, rhs)
+    }
+    fn unchecked_sumul(&mut self, lhs: Self::Value, rhs: Self::Value) -> Self::Value {
+        // Which to default to is a fairly arbitrary choice,
+        // but this is what slice layout was using before.
+        self.unchecked_smul(lhs, rhs)
+    }
     fn and(&mut self, lhs: Self::Value, rhs: Self::Value) -> Self::Value;
     fn or(&mut self, lhs: Self::Value, rhs: Self::Value) -> Self::Value;
     /// Defaults to [`Self::or`], but guarantees `(lhs & rhs) == 0` so some backends

tests/codegen/issues/issue-96497-slice-size-nowrap.rs

@@ -11,7 +11,7 @@
 pub fn simple_size_of_nowrap(x: &[u32]) -> usize {
     // Make sure the shift used to compute the size has a nowrap flag.
 
-    // CHECK: [[A:%.*]] = shl nsw {{.*}}, 2
+    // CHECK: [[A:%.*]] = shl nuw nsw {{.*}}, 2
     // CHECK-NEXT: ret {{.*}} [[A]]
     core::mem::size_of_val(x)
 }
@@ -26,3 +26,13 @@ pub fn drop_write(mut x: Box<[u32]>) {
     // CHECK-NOT: store i32 42
     x[1] = 42;
 }
+
+// CHECK-LABEL: @slice_size_plus_2
+#[no_mangle]
+pub fn slice_size_plus_2(x: &[u16]) -> usize {
+    // Before #136575 this didn't get the `nuw` in the add.
+
+    // CHECK: [[BYTES:%.+]] = shl nuw nsw {{i16|i32|i64}} %x.1, 1
+    // CHECK: = add nuw {{i16|i32|i64}} [[BYTES]], 2
+    core::mem::size_of_val(x) + 2
+}

tests/codegen/slice-ref-equality.rs

@@ -47,7 +47,7 @@ pub fn is_zero_array(data: &[u8; 4]) -> bool {
 #[no_mangle]
 fn eq_slice_of_nested_u8(x: &[[u8; 3]], y: &[[u8; 3]]) -> bool {
     // CHECK: icmp eq [[USIZE]] %x.1, %y.1
-    // CHECK: %[[BYTES:.+]] = mul nsw [[USIZE]] {{%x.1|%y.1}}, 3
+    // CHECK: %[[BYTES:.+]] = mul nuw nsw [[USIZE]] {{%x.1|%y.1}}, 3
     // CHECK: tail call{{( noundef)?}} i32 @{{bcmp|memcmp}}(ptr
     // CHECK-SAME: , [[USIZE]]{{( noundef)?}} %[[BYTES]])
     x == y
@@ -59,7 +59,7 @@ fn eq_slice_of_nested_u8(x: &[[u8; 3]], y: &[[u8; 3]]) -> bool {
 #[no_mangle]
 fn eq_slice_of_i32(x: &[i32], y: &[i32]) -> bool {
     // CHECK: icmp eq [[USIZE]] %x.1, %y.1
-    // CHECK: %[[BYTES:.+]] = shl nsw [[USIZE]] {{%x.1|%y.1}}, 2
+    // CHECK: %[[BYTES:.+]] = shl nuw nsw [[USIZE]] {{%x.1|%y.1}}, 2
     // CHECK: tail call{{( noundef)?}} i32 @{{bcmp|memcmp}}(ptr
     // CHECK-SAME: , [[USIZE]]{{( noundef)?}} %[[BYTES]])
     x == y
@@ -71,7 +71,7 @@ fn eq_slice_of_i32(x: &[i32], y: &[i32]) -> bool {
 #[no_mangle]
 fn eq_slice_of_nonzero(x: &[NonZero<i32>], y: &[NonZero<i32>]) -> bool {
     // CHECK: icmp eq [[USIZE]] %x.1, %y.1
-    // CHECK: %[[BYTES:.+]] = shl nsw [[USIZE]] {{%x.1|%y.1}}, 2
+    // CHECK: %[[BYTES:.+]] = shl nuw nsw [[USIZE]] {{%x.1|%y.1}}, 2
     // CHECK: tail call{{( noundef)?}} i32 @{{bcmp|memcmp}}(ptr
     // CHECK-SAME: , [[USIZE]]{{( noundef)?}} %[[BYTES]])
     x == y
@@ -83,7 +83,7 @@ fn eq_slice_of_nonzero(x: &[NonZero<i32>], y: &[NonZero<i32>]) -> bool {
 #[no_mangle]
 fn eq_slice_of_option_of_nonzero(x: &[Option<NonZero<i16>>], y: &[Option<NonZero<i16>>]) -> bool {
     // CHECK: icmp eq [[USIZE]] %x.1, %y.1
-    // CHECK: %[[BYTES:.+]] = shl nsw [[USIZE]] {{%x.1|%y.1}}, 1
+    // CHECK: %[[BYTES:.+]] = shl nuw nsw [[USIZE]] {{%x.1|%y.1}}, 1
     // CHECK: tail call{{( noundef)?}} i32 @{{bcmp|memcmp}}(ptr
     // CHECK-SAME: , [[USIZE]]{{( noundef)?}} %[[BYTES]])
     x == y