Auto merge of #135760 - scottmcm:disjoint-bitor, r=WaffleLapkin

Add `unchecked_disjoint_bitor` per ACP373

Following the names from libs-api in rust-lang/libs-team#373 (comment)

Includes a fallback implementation so this doesn't have to update cg_clif or cg_gcc, and overrides it in cg_llvm to use `or disjoint`, which [is available in LLVM 18](https://releases.llvm.org/18.1.0/docs/LangRef.html#or-instruction) so hopefully we don't need any version checks.

Author: bors

Diff for: compiler/rustc_codegen_llvm/src/builder.rs

@@ -421,6 +421,20 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         unchecked_umul(x, y) => LLVMBuildNUWMul,
     }
 
+    fn or_disjoint(&mut self, a: &'ll Value, b: &'ll Value) -> &'ll Value {
+        unsafe {
+            let or = llvm::LLVMBuildOr(self.llbuilder, a, b, UNNAMED);
+
+            // If a and b are both values, then `or` is a value, rather than
+            // an instruction, so we need to check before setting the flag.
+            // (See also `LLVMBuildNUWNeg` which also needs a check.)
+            if llvm::LLVMIsAInstruction(or).is_some() {
+                llvm::LLVMSetIsDisjoint(or, True);
+            }
+            or
+        }
+    }
+
     set_math_builder_methods! {
         fadd_fast(x, y) => (LLVMBuildFAdd, LLVMRustSetFastMath),
         fsub_fast(x, y) => (LLVMBuildFSub, LLVMRustSetFastMath),

Diff for: compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -1380,6 +1380,9 @@ unsafe extern "C" {
     pub fn LLVMBuildFNeg<'a>(B: &Builder<'a>, V: &'a Value, Name: *const c_char) -> &'a Value;
     pub fn LLVMBuildNot<'a>(B: &Builder<'a>, V: &'a Value, Name: *const c_char) -> &'a Value;
 
+    // Extra flags on arithmetic
+    pub fn LLVMSetIsDisjoint(Instr: &Value, IsDisjoint: Bool);
+
     // Memory
     pub fn LLVMBuildAlloca<'a>(B: &Builder<'a>, Ty: &'a Type, Name: *const c_char) -> &'a Value;
     pub fn LLVMBuildArrayAlloca<'a>(

Diff for: compiler/rustc_codegen_ssa/src/mir/intrinsic.rs

@@ -225,6 +225,11 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 args[1].val.unaligned_volatile_store(bx, dst);
                 return Ok(());
             }
+            sym::disjoint_bitor => {
+                let a = args[0].immediate();
+                let b = args[1].immediate();
+                bx.or_disjoint(a, b)
+            }
             sym::exact_div => {
                 let ty = arg_tys[0];
                 match int_type_width_signed(ty, bx.tcx()) {

Diff for: compiler/rustc_codegen_ssa/src/traits/builder.rs

@@ -167,6 +167,11 @@ pub trait BuilderMethods<'a, 'tcx>:
     fn unchecked_umul(&mut self, lhs: Self::Value, rhs: Self::Value) -> Self::Value;
     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
+    /// can emit something more helpful for optimizations.
+    fn or_disjoint(&mut self, lhs: Self::Value, rhs: Self::Value) -> Self::Value {
+        self.or(lhs, rhs)
+    }
     fn xor(&mut self, lhs: Self::Value, rhs: Self::Value) -> Self::Value;
     fn neg(&mut self, v: Self::Value) -> Self::Value;
     fn fneg(&mut self, v: Self::Value) -> Self::Value;

Diff for: compiler/rustc_hir_analysis/src/check/intrinsic.rs

@@ -472,7 +472,7 @@ pub fn check_intrinsic_type(
                 vec![Ty::new_imm_ptr(tcx, param(0)), Ty::new_imm_ptr(tcx, param(0))],
                 tcx.types.usize,
             ),
-            sym::unchecked_div | sym::unchecked_rem | sym::exact_div => {
+            sym::unchecked_div | sym::unchecked_rem | sym::exact_div | sym::disjoint_bitor => {
                 (1, 0, vec![param(0), param(0)], param(0))
             }
             sym::unchecked_shl | sym::unchecked_shr => (2, 0, vec![param(0), param(1)], param(0)),

Diff for: compiler/rustc_span/src/symbol.rs

@@ -777,6 +777,7 @@ symbols! {
         discriminant_kind,
         discriminant_type,
         discriminant_value,
+        disjoint_bitor,
         dispatch_from_dyn,
         div,
         div_assign,

Diff for: library/core/src/intrinsics/fallback.rs

@@ -110,3 +110,41 @@ impl const CarryingMulAdd for i128 {
         (low, high)
     }
 }
+
+#[const_trait]
+#[rustc_const_unstable(feature = "core_intrinsics_fallbacks", issue = "none")]
+pub trait DisjointBitOr: Copy + 'static {
+    /// See [`super::disjoint_bitor`]; we just need the trait indirection to handle
+    /// different types since calling intrinsics with generics doesn't work.
+    unsafe fn disjoint_bitor(self, other: Self) -> Self;
+}
+macro_rules! zero {
+    (bool) => {
+        false
+    };
+    ($t:ident) => {
+        0
+    };
+}
+macro_rules! impl_disjoint_bitor {
+    ($($t:ident,)+) => {$(
+        #[rustc_const_unstable(feature = "core_intrinsics_fallbacks", issue = "none")]
+        impl const DisjointBitOr for $t {
+            #[cfg_attr(miri, track_caller)]
+            #[inline]
+            unsafe fn disjoint_bitor(self, other: Self) -> Self {
+                // Note that the assume here is required for UB detection in Miri!
+
+                // SAFETY: our precondition is that there are no bits in common,
+                // so this is just telling that to the backend.
+                unsafe { super::assume((self & other) == zero!($t)) };
+                self | other
+            }
+        }
+    )+};
+}
+impl_disjoint_bitor! {
+    bool,
+    u8, u16, u32, u64, u128, usize,
+    i8, i16, i32, i64, i128, isize,
+}

Diff for: library/core/src/intrinsics/mod.rs

@@ -3248,6 +3248,26 @@ pub const fn three_way_compare<T: Copy>(_lhs: T, _rhss: T) -> crate::cmp::Orderi
     unimplemented!()
 }
 
+/// Combine two values which have no bits in common.
+///
+/// This allows the backend to implement it as `a + b` *or* `a | b`,
+/// depending which is easier to implement on a specific target.
+///
+/// # Safety
+///
+/// Requires that `(a & b) == 0`, or equivalently that `(a | b) == (a + b)`.
+///
+/// Otherwise it's immediate UB.
+#[rustc_const_unstable(feature = "disjoint_bitor", issue = "135758")]
+#[rustc_nounwind]
+#[cfg_attr(not(bootstrap), rustc_intrinsic)]
+#[cfg_attr(miri, track_caller)] // even without panics, this helps for Miri backtraces
+#[miri::intrinsic_fallback_is_spec] // the fallbacks all `assume` to tell Miri
+pub const unsafe fn disjoint_bitor<T: ~const fallback::DisjointBitOr>(a: T, b: T) -> T {
+    // SAFETY: same preconditions as this function.
+    unsafe { fallback::DisjointBitOr::disjoint_bitor(a, b) }
+}
+
 /// Performs checked integer addition.
 ///
 /// Note that, unlike most intrinsics, this is safe to call;

Diff for: library/core/src/lib.rs

@@ -117,6 +117,7 @@
 #![feature(const_eval_select)]
 #![feature(core_intrinsics)]
 #![feature(coverage_attribute)]
+#![feature(disjoint_bitor)]
 #![feature(internal_impls_macro)]
 #![feature(ip)]
 #![feature(is_ascii_octdigit)]

Diff for: library/core/src/num/uint_macros.rs

@@ -1187,6 +1187,50 @@ macro_rules! uint_impl {
             self % rhs
         }
 
+        /// Same value as `self | other`, but UB if any bit position is set in both inputs.
+        ///
+        /// This is a situational micro-optimization for places where you'd rather
+        /// use addition on some platforms and bitwise or on other platforms, based
+        /// on exactly which instructions combine better with whatever else you're
+        /// doing.  Note that there's no reason to bother using this for places
+        /// where it's clear from the operations involved that they can't overlap.
+        /// For example, if you're combining `u16`s into a `u32` with
+        /// `((a as u32) << 16) | (b as u32)`, that's fine, as the backend will
+        /// know those sides of the `|` are disjoint without needing help.
+        ///
+        /// # Examples
+        ///
+        /// ```
+        /// #![feature(disjoint_bitor)]
+        ///
+        /// // SAFETY: `1` and `4` have no bits in common.
+        /// unsafe {
+        #[doc = concat!("    assert_eq!(1_", stringify!($SelfT), ".unchecked_disjoint_bitor(4), 5);")]
+        /// }
+        /// ```
+        ///
+        /// # Safety
+        ///
+        /// Requires that `(self & other) == 0`, otherwise it's immediate UB.
+        ///
+        /// Equivalently, requires that `(self | other) == (self + other)`.
+        #[unstable(feature = "disjoint_bitor", issue = "135758")]
+        #[rustc_const_unstable(feature = "disjoint_bitor", issue = "135758")]
+        #[inline]
+        pub const unsafe fn unchecked_disjoint_bitor(self, other: Self) -> Self {
+            assert_unsafe_precondition!(
+                check_language_ub,
+                concat!(stringify!($SelfT), "::unchecked_disjoint_bitor cannot have overlapping bits"),
+                (
+                    lhs: $SelfT = self,
+                    rhs: $SelfT = other,
+                ) => (lhs & rhs) == 0,
+            );
+
+            // SAFETY: Same precondition
+            unsafe { intrinsics::disjoint_bitor(self, other) }
+        }
+
         /// Returns the logarithm of the number with respect to an arbitrary base,
         /// rounded down.
         ///
@@ -2346,15 +2390,22 @@ macro_rules! uint_impl {
         /// assert_eq!((sum1, sum0), (9, 6));
         /// ```
         #[unstable(feature = "bigint_helper_methods", issue = "85532")]
+        #[rustc_const_unstable(feature = "bigint_helper_methods", issue = "85532")]
         #[must_use = "this returns the result of the operation, \
                       without modifying the original"]
         #[inline]
         pub const fn carrying_add(self, rhs: Self, carry: bool) -> (Self, bool) {
             // note: longer-term this should be done via an intrinsic, but this has been shown
             //   to generate optimal code for now, and LLVM doesn't have an equivalent intrinsic
-            let (a, b) = self.overflowing_add(rhs);
-            let (c, d) = a.overflowing_add(carry as $SelfT);
-            (c, b | d)
+            let (a, c1) = self.overflowing_add(rhs);
+            let (b, c2) = a.overflowing_add(carry as $SelfT);
+            // Ideally LLVM would know this is disjoint without us telling them,
+            // but it doesn't <https://github.com/llvm/llvm-project/issues/118162>
+            // SAFETY: Only one of `c1` and `c2` can be set.
+            // For c1 to be set we need to have overflowed, but if we did then
+            // `a` is at most `MAX-1`, which means that `c2` cannot possibly
+            // overflow because it's adding at most `1` (since it came from `bool`)
+            (b, unsafe { intrinsics::disjoint_bitor(c1, c2) })
         }
 
         /// Calculates `self` + `rhs` with a signed `rhs`.

Diff for: src/tools/miri/tests/fail/intrinsics/disjoint_bitor.rs

@@ -0,0 +1,5 @@
+#![feature(core_intrinsics)]
+fn main() {
+    // one bit in common
+    unsafe { std::intrinsics::disjoint_bitor(0b01101001_u8, 0b10001110) }; //~ ERROR: Undefined Behavior
+}

Diff for: src/tools/miri/tests/fail/intrinsics/disjoint_bitor.stderr

@@ -0,0 +1,15 @@
+error: Undefined Behavior: `assume` called with `false`
+  --> tests/fail/intrinsics/disjoint_bitor.rs:LL:CC
+   |
+LL |     unsafe { std::intrinsics::disjoint_bitor(0b01101001_u8, 0b10001110) };
+   |              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `assume` called with `false`
+   |
+   = help: this indicates a bug in the program: it performed an invalid operation, and caused Undefined Behavior
+   = help: see https://doc.rust-lang.org/nightly/reference/behavior-considered-undefined.html for further information
+   = note: BACKTRACE:
+   = note: inside `main` at tests/fail/intrinsics/disjoint_bitor.rs:LL:CC
+
+note: some details are omitted, run with `MIRIFLAGS=-Zmiri-backtrace=full` for a verbose backtrace
+
+error: aborting due to 1 previous error
+

Diff for: tests/codegen/bigint-helpers.rs

@@ -0,0 +1,13 @@
+//@ compile-flags: -C opt-level=3
+
+#![crate_type = "lib"]
+#![feature(bigint_helper_methods)]
+
+// CHECK-LABEL: @u32_carrying_add
+#[no_mangle]
+pub fn u32_carrying_add(a: u32, b: u32, c: bool) -> (u32, bool) {
+    // CHECK: @llvm.uadd.with.overflow.i32
+    // CHECK: @llvm.uadd.with.overflow.i32
+    // CHECK: or disjoint i1
+    u32::carrying_add(a, b, c)
+}

Diff for: tests/codegen/intrinsics/disjoint_bitor.rs

@@ -0,0 +1,30 @@
+//@ compile-flags: -C no-prepopulate-passes -Z mir-opt-level=0
+
+#![crate_type = "lib"]
+#![feature(core_intrinsics)]
+
+use std::intrinsics::disjoint_bitor;
+
+// CHECK-LABEL: @disjoint_bitor_signed
+#[no_mangle]
+pub unsafe fn disjoint_bitor_signed(x: i32, y: i32) -> i32 {
+    // CHECK: or disjoint i32 %x, %y
+    disjoint_bitor(x, y)
+}
+
+// CHECK-LABEL: @disjoint_bitor_unsigned
+#[no_mangle]
+pub unsafe fn disjoint_bitor_unsigned(x: u64, y: u64) -> u64 {
+    // CHECK: or disjoint i64 %x, %y
+    disjoint_bitor(x, y)
+}
+
+// CHECK-LABEL: @disjoint_bitor_literal
+#[no_mangle]
+pub unsafe fn disjoint_bitor_literal() -> u8 {
+    // This is a separate check because even without any passes,
+    // LLVM will fold so it's not an instruction, which can assert in LLVM.
+
+    // CHECK: store i8 3
+    disjoint_bitor(1, 2)
+}