Rollup merge of #144192 - RalfJung:atomicrmw-ptr, r=nikic

atomicrmw on pointers: move integer-pointer cast hacks into backend

Conceptually, we want to have atomic operations on pointers of the form `fn atomic_add(ptr: *mut T, offset: usize, ...)`. However, LLVM does not directly support such operations (https://github.com/llvm/llvm-project/issues/120837), so we have to cast the `offset` to a pointer somewhere.

This PR moves that hack into the LLVM backend, so that the standard library, intrinsic, and Miri all work with the conceptual operation we actually want. Hopefully, one day LLVM will gain a way to represent these operations without integer-pointer casts, and then the hack will disappear entirely.

Cc ```@nikic``` -- this is the best we can do right now, right?
Fixes https://github.com/rust-lang/rust/issues/134617

9 files changed, 121 insertions, 48 deletions

diff --git a/compiler/rustc_codegen_cranelift/src/intrinsics/mod.rs b/compiler/rustc_codegen_cranelift/src/intrinsics/mod.rs
index 4ff5773a06c..ed40901ac9b 100644
--- a/compiler/rustc_codegen_cranelift/src/intrinsics/mod.rs
+++ b/compiler/rustc_codegen_cranelift/src/intrinsics/mod.rs
@@ -969,7 +969,7 @@ fn codegen_regular_intrinsic_call<'tcx>(
 
             let layout = amount.layout();
             match layout.ty.kind() {
-                ty::Uint(_) | ty::Int(_) | ty::RawPtr(..) => {}
+                ty::Uint(_) | ty::Int(_) => {}
                 _ => {
                     report_atomic_type_validation_error(fx, intrinsic, source_info.span, layout.ty);
                     return Ok(());
@@ -982,7 +982,7 @@ fn codegen_regular_intrinsic_call<'tcx>(
             let old =
                 fx.bcx.ins().atomic_rmw(ty, MemFlags::trusted(), AtomicRmwOp::Add, ptr, amount);
 
-            let old = CValue::by_val(old, layout);
+            let old = CValue::by_val(old, ret.layout());
             ret.write_cvalue(fx, old);
         }
         sym::atomic_xsub => {
@@ -991,7 +991,7 @@ fn codegen_regular_intrinsic_call<'tcx>(
 
             let layout = amount.layout();
             match layout.ty.kind() {
-                ty::Uint(_) | ty::Int(_) | ty::RawPtr(..) => {}
+                ty::Uint(_) | ty::Int(_) => {}
                 _ => {
                     report_atomic_type_validation_error(fx, intrinsic, source_info.span, layout.ty);
                     return Ok(());
@@ -1004,7 +1004,7 @@ fn codegen_regular_intrinsic_call<'tcx>(
             let old =
                 fx.bcx.ins().atomic_rmw(ty, MemFlags::trusted(), AtomicRmwOp::Sub, ptr, amount);
 
-            let old = CValue::by_val(old, layout);
+            let old = CValue::by_val(old, ret.layout());
             ret.write_cvalue(fx, old);
         }
         sym::atomic_and => {
@@ -1013,7 +1013,7 @@ fn codegen_regular_intrinsic_call<'tcx>(
 
             let layout = src.layout();
             match layout.ty.kind() {
-                ty::Uint(_) | ty::Int(_) | ty::RawPtr(..) => {}
+                ty::Uint(_) | ty::Int(_) => {}
                 _ => {
                     report_atomic_type_validation_error(fx, intrinsic, source_info.span, layout.ty);
                     return Ok(());
@@ -1025,7 +1025,7 @@ fn codegen_regular_intrinsic_call<'tcx>(
 
             let old = fx.bcx.ins().atomic_rmw(ty, MemFlags::trusted(), AtomicRmwOp::And, ptr, src);
 
-            let old = CValue::by_val(old, layout);
+            let old = CValue::by_val(old, ret.layout());
             ret.write_cvalue(fx, old);
         }
         sym::atomic_or => {
@@ -1034,7 +1034,7 @@ fn codegen_regular_intrinsic_call<'tcx>(
 
             let layout = src.layout();
             match layout.ty.kind() {
-                ty::Uint(_) | ty::Int(_) | ty::RawPtr(..) => {}
+                ty::Uint(_) | ty::Int(_) => {}
                 _ => {
                     report_atomic_type_validation_error(fx, intrinsic, source_info.span, layout.ty);
                     return Ok(());
@@ -1046,7 +1046,7 @@ fn codegen_regular_intrinsic_call<'tcx>(
 
             let old = fx.bcx.ins().atomic_rmw(ty, MemFlags::trusted(), AtomicRmwOp::Or, ptr, src);
 
-            let old = CValue::by_val(old, layout);
+            let old = CValue::by_val(old, ret.layout());
             ret.write_cvalue(fx, old);
         }
         sym::atomic_xor => {
@@ -1055,7 +1055,7 @@ fn codegen_regular_intrinsic_call<'tcx>(
 
             let layout = src.layout();
             match layout.ty.kind() {
-                ty::Uint(_) | ty::Int(_) | ty::RawPtr(..) => {}
+                ty::Uint(_) | ty::Int(_) => {}
                 _ => {
                     report_atomic_type_validation_error(fx, intrinsic, source_info.span, layout.ty);
                     return Ok(());
@@ -1067,7 +1067,7 @@ fn codegen_regular_intrinsic_call<'tcx>(
 
             let old = fx.bcx.ins().atomic_rmw(ty, MemFlags::trusted(), AtomicRmwOp::Xor, ptr, src);
 
-            let old = CValue::by_val(old, layout);
+            let old = CValue::by_val(old, ret.layout());
             ret.write_cvalue(fx, old);
         }
         sym::atomic_nand => {
@@ -1076,7 +1076,7 @@ fn codegen_regular_intrinsic_call<'tcx>(
 
             let layout = src.layout();
             match layout.ty.kind() {
-                ty::Uint(_) | ty::Int(_) | ty::RawPtr(..) => {}
+                ty::Uint(_) | ty::Int(_) => {}
                 _ => {
                     report_atomic_type_validation_error(fx, intrinsic, source_info.span, layout.ty);
                     return Ok(());
@@ -1088,7 +1088,7 @@ fn codegen_regular_intrinsic_call<'tcx>(
 
             let old = fx.bcx.ins().atomic_rmw(ty, MemFlags::trusted(), AtomicRmwOp::Nand, ptr, src);
 
-            let old = CValue::by_val(old, layout);
+            let old = CValue::by_val(old, ret.layout());
             ret.write_cvalue(fx, old);
         }
         sym::atomic_max => {
diff --git a/compiler/rustc_codegen_gcc/src/builder.rs b/compiler/rustc_codegen_gcc/src/builder.rs
index 34ade3d025f..f7a7a3f8c7e 100644
--- a/compiler/rustc_codegen_gcc/src/builder.rs
+++ b/compiler/rustc_codegen_gcc/src/builder.rs
@@ -1671,6 +1671,7 @@ impl<'a, 'gcc, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'gcc, 'tcx> {
         dst: RValue<'gcc>,
         src: RValue<'gcc>,
         order: AtomicOrdering,
+        ret_ptr: bool,
     ) -> RValue<'gcc> {
         let size = get_maybe_pointer_size(src);
         let name = match op {
@@ -1698,6 +1699,9 @@ impl<'a, 'gcc, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'gcc, 'tcx> {
         let atomic_function = self.context.get_builtin_function(name);
         let order = self.context.new_rvalue_from_int(self.i32_type, order.to_gcc());
 
+        // FIXME: If `ret_ptr` is true and `src` is an integer, we should really tell GCC
+        // that this is a pointer operation that needs to preserve provenance -- but like LLVM,
+        // GCC does not currently seems to support that.
         let void_ptr_type = self.context.new_type::<*mut ()>();
         let volatile_void_ptr_type = void_ptr_type.make_volatile();
         let dst = self.context.new_cast(self.location, dst, volatile_void_ptr_type);
@@ -1705,7 +1709,8 @@ impl<'a, 'gcc, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'gcc, 'tcx> {
         let new_src_type = atomic_function.get_param(1).to_rvalue().get_type();
         let src = self.context.new_bitcast(self.location, src, new_src_type);
         let res = self.context.new_call(self.location, atomic_function, &[dst, src, order]);
-        self.context.new_cast(self.location, res, src.get_type())
+        let res_type = if ret_ptr { void_ptr_type } else { src.get_type() };
+        self.context.new_cast(self.location, res, res_type)
     }
 
     fn atomic_fence(&mut self, order: AtomicOrdering, scope: SynchronizationScope) {
diff --git a/compiler/rustc_codegen_llvm/src/builder.rs b/compiler/rustc_codegen_llvm/src/builder.rs
index bd175e560c7..32cdef075e7 100644
--- a/compiler/rustc_codegen_llvm/src/builder.rs
+++ b/compiler/rustc_codegen_llvm/src/builder.rs
@@ -1327,15 +1327,13 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         &mut self,
         op: rustc_codegen_ssa::common::AtomicRmwBinOp,
         dst: &'ll Value,
-        mut src: &'ll Value,
+        src: &'ll Value,
         order: rustc_middle::ty::AtomicOrdering,
+        ret_ptr: bool,
     ) -> &'ll Value {
-        // The only RMW operation that LLVM supports on pointers is compare-exchange.
-        let requires_cast_to_int = self.val_ty(src) == self.type_ptr()
-            && op != rustc_codegen_ssa::common::AtomicRmwBinOp::AtomicXchg;
-        if requires_cast_to_int {
-            src = self.ptrtoint(src, self.type_isize());
-        }
+        // FIXME: If `ret_ptr` is true and `src` is not a pointer, we *should* tell LLVM that the
+        // LHS is a pointer and the operation should be provenance-preserving, but LLVM does not
+        // currently support that (https://github.com/llvm/llvm-project/issues/120837).
         let mut res = unsafe {
             llvm::LLVMBuildAtomicRMW(
                 self.llbuilder,
@@ -1346,7 +1344,7 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
                 llvm::False, // SingleThreaded
             )
         };
-        if requires_cast_to_int {
+        if ret_ptr && self.val_ty(res) != self.type_ptr() {
             res = self.inttoptr(res, self.type_ptr());
         }
         res
diff --git a/compiler/rustc_codegen_ssa/messages.ftl b/compiler/rustc_codegen_ssa/messages.ftl
index a70d0011d16..36ad5ede7c2 100644
--- a/compiler/rustc_codegen_ssa/messages.ftl
+++ b/compiler/rustc_codegen_ssa/messages.ftl
@@ -101,6 +101,8 @@ codegen_ssa_invalid_monomorphization_basic_float_type = invalid monomorphization
 
 codegen_ssa_invalid_monomorphization_basic_integer_type = invalid monomorphization of `{$name}` intrinsic: expected basic integer type, found `{$ty}`
 
+codegen_ssa_invalid_monomorphization_basic_integer_or_ptr_type = invalid monomorphization of `{$name}` intrinsic: expected basic integer or pointer type, found `{$ty}`
+
 codegen_ssa_invalid_monomorphization_cannot_return = invalid monomorphization of `{$name}` intrinsic: cannot return `{$ret_ty}`, expected `u{$expected_int_bits}` or `[u8; {$expected_bytes}]`
 
 codegen_ssa_invalid_monomorphization_cast_wide_pointer = invalid monomorphization of `{$name}` intrinsic: cannot cast wide pointer `{$ty}`
diff --git a/compiler/rustc_codegen_ssa/src/errors.rs b/compiler/rustc_codegen_ssa/src/errors.rs
index 3d787d8bdbd..af4adcd1954 100644
--- a/compiler/rustc_codegen_ssa/src/errors.rs
+++ b/compiler/rustc_codegen_ssa/src/errors.rs
@@ -764,6 +764,14 @@ pub enum InvalidMonomorphization<'tcx> {
         ty: Ty<'tcx>,
     },
 
+    #[diag(codegen_ssa_invalid_monomorphization_basic_integer_or_ptr_type, code = E0511)]
+    BasicIntegerOrPtrType {
+        #[primary_span]
+        span: Span,
+        name: Symbol,
+        ty: Ty<'tcx>,
+    },
+
     #[diag(codegen_ssa_invalid_monomorphization_basic_float_type, code = E0511)]
     BasicFloatType {
         #[primary_span]
diff --git a/compiler/rustc_codegen_ssa/src/mir/intrinsic.rs b/compiler/rustc_codegen_ssa/src/mir/intrinsic.rs
index fc95f62b4a4..3c667b8e882 100644
--- a/compiler/rustc_codegen_ssa/src/mir/intrinsic.rs
+++ b/compiler/rustc_codegen_ssa/src/mir/intrinsic.rs
@@ -92,6 +92,13 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         let invalid_monomorphization_int_type = |ty| {
             bx.tcx().dcx().emit_err(InvalidMonomorphization::BasicIntegerType { span, name, ty });
         };
+        let invalid_monomorphization_int_or_ptr_type = |ty| {
+            bx.tcx().dcx().emit_err(InvalidMonomorphization::BasicIntegerOrPtrType {
+                span,
+                name,
+                ty,
+            });
+        };
 
         let parse_atomic_ordering = |ord: ty::Value<'tcx>| {
             let discr = ord.valtree.unwrap_branch()[0].unwrap_leaf();
@@ -351,7 +358,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             sym::atomic_load => {
                 let ty = fn_args.type_at(0);
                 if !(int_type_width_signed(ty, bx.tcx()).is_some() || ty.is_raw_ptr()) {
-                    invalid_monomorphization_int_type(ty);
+                    invalid_monomorphization_int_or_ptr_type(ty);
                     return Ok(());
                 }
                 let ordering = fn_args.const_at(1).to_value();
@@ -367,7 +374,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             sym::atomic_store => {
                 let ty = fn_args.type_at(0);
                 if !(int_type_width_signed(ty, bx.tcx()).is_some() || ty.is_raw_ptr()) {
-                    invalid_monomorphization_int_type(ty);
+                    invalid_monomorphization_int_or_ptr_type(ty);
                     return Ok(());
                 }
                 let ordering = fn_args.const_at(1).to_value();
@@ -377,10 +384,11 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 bx.atomic_store(val, ptr, parse_atomic_ordering(ordering), size);
                 return Ok(());
             }
+            // These are all AtomicRMW ops
             sym::atomic_cxchg | sym::atomic_cxchgweak => {
                 let ty = fn_args.type_at(0);
                 if !(int_type_width_signed(ty, bx.tcx()).is_some() || ty.is_raw_ptr()) {
-                    invalid_monomorphization_int_type(ty);
+                    invalid_monomorphization_int_or_ptr_type(ty);
                     return Ok(());
                 }
                 let succ_ordering = fn_args.const_at(1).to_value();
@@ -407,7 +415,6 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
 
                 return Ok(());
             }
-            // These are all AtomicRMW ops
             sym::atomic_max | sym::atomic_min => {
                 let atom_op = if name == sym::atomic_max {
                     AtomicRmwBinOp::AtomicMax
@@ -420,7 +427,13 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     let ordering = fn_args.const_at(1).to_value();
                     let ptr = args[0].immediate();
                     let val = args[1].immediate();
-                    bx.atomic_rmw(atom_op, ptr, val, parse_atomic_ordering(ordering))
+                    bx.atomic_rmw(
+                        atom_op,
+                        ptr,
+                        val,
+                        parse_atomic_ordering(ordering),
+                        /* ret_ptr */ false,
+                    )
                 } else {
                     invalid_monomorphization_int_type(ty);
                     return Ok(());
@@ -438,21 +451,44 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     let ordering = fn_args.const_at(1).to_value();
                     let ptr = args[0].immediate();
                     let val = args[1].immediate();
-                    bx.atomic_rmw(atom_op, ptr, val, parse_atomic_ordering(ordering))
+                    bx.atomic_rmw(
+                        atom_op,
+                        ptr,
+                        val,
+                        parse_atomic_ordering(ordering),
+                        /* ret_ptr */ false,
+                    )
                 } else {
                     invalid_monomorphization_int_type(ty);
                     return Ok(());
                 }
             }
-            sym::atomic_xchg
-            | sym::atomic_xadd
+            sym::atomic_xchg => {
+                let ty = fn_args.type_at(0);
+                let ordering = fn_args.const_at(1).to_value();
+                if int_type_width_signed(ty, bx.tcx()).is_some() || ty.is_raw_ptr() {
+                    let ptr = args[0].immediate();
+                    let val = args[1].immediate();
+                    let atomic_op = AtomicRmwBinOp::AtomicXchg;
+                    bx.atomic_rmw(
+                        atomic_op,
+                        ptr,
+                        val,
+                        parse_atomic_ordering(ordering),
+                        /* ret_ptr */ ty.is_raw_ptr(),
+                    )
+                } else {
+                    invalid_monomorphization_int_or_ptr_type(ty);
+                    return Ok(());
+                }
+            }
+            sym::atomic_xadd
             | sym::atomic_xsub
             | sym::atomic_and
             | sym::atomic_nand
             | sym::atomic_or
             | sym::atomic_xor => {
                 let atom_op = match name {
-                    sym::atomic_xchg => AtomicRmwBinOp::AtomicXchg,
                     sym::atomic_xadd => AtomicRmwBinOp::AtomicAdd,
                     sym::atomic_xsub => AtomicRmwBinOp::AtomicSub,
                     sym::atomic_and => AtomicRmwBinOp::AtomicAnd,
@@ -462,14 +498,28 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     _ => unreachable!(),
                 };
 
-                let ty = fn_args.type_at(0);
-                if int_type_width_signed(ty, bx.tcx()).is_some() || ty.is_raw_ptr() {
-                    let ordering = fn_args.const_at(1).to_value();
-                    let ptr = args[0].immediate();
-                    let val = args[1].immediate();
-                    bx.atomic_rmw(atom_op, ptr, val, parse_atomic_ordering(ordering))
+                // The type of the in-memory data.
+                let ty_mem = fn_args.type_at(0);
+                // The type of the 2nd operand, given by-value.
+                let ty_op = fn_args.type_at(1);
+
+                let ordering = fn_args.const_at(2).to_value();
+                // We require either both arguments to have the same integer type, or the first to
+                // be a pointer and the second to be `usize`.
+                if (int_type_width_signed(ty_mem, bx.tcx()).is_some() && ty_op == ty_mem)
+                    || (ty_mem.is_raw_ptr() && ty_op == bx.tcx().types.usize)
+                {
+                    let ptr = args[0].immediate(); // of type "pointer to `ty_mem`"
+                    let val = args[1].immediate(); // of type `ty_op`
+                    bx.atomic_rmw(
+                        atom_op,
+                        ptr,
+                        val,
+                        parse_atomic_ordering(ordering),
+                        /* ret_ptr */ ty_mem.is_raw_ptr(),
+                    )
                 } else {
-                    invalid_monomorphization_int_type(ty);
+                    invalid_monomorphization_int_or_ptr_type(ty_mem);
                     return Ok(());
                 }
             }
diff --git a/compiler/rustc_codegen_ssa/src/traits/builder.rs b/compiler/rustc_codegen_ssa/src/traits/builder.rs
index 4b18146863b..f417d1a7bf7 100644
--- a/compiler/rustc_codegen_ssa/src/traits/builder.rs
+++ b/compiler/rustc_codegen_ssa/src/traits/builder.rs
@@ -548,12 +548,15 @@ pub trait BuilderMethods<'a, 'tcx>:
         failure_order: AtomicOrdering,
         weak: bool,
     ) -> (Self::Value, Self::Value);
+    /// `ret_ptr` indicates whether the return type (which is also the type `dst` points to)
+    /// is a pointer or the same type as `src`.
     fn atomic_rmw(
         &mut self,
         op: AtomicRmwBinOp,
         dst: Self::Value,
         src: Self::Value,
         order: AtomicOrdering,
+        ret_ptr: bool,
     ) -> Self::Value;
     fn atomic_fence(&mut self, order: AtomicOrdering, scope: SynchronizationScope);
     fn set_invariant_load(&mut self, load: Self::Value);
diff --git a/compiler/rustc_hir_analysis/src/check/intrinsic.rs b/compiler/rustc_hir_analysis/src/check/intrinsic.rs
index 6e5fe3823ab..4441dd6ebd6 100644
--- a/compiler/rustc_hir_analysis/src/check/intrinsic.rs
+++ b/compiler/rustc_hir_analysis/src/check/intrinsic.rs
@@ -652,16 +652,16 @@ pub(crate) fn check_intrinsic_type(
         sym::atomic_store => (1, 1, vec![Ty::new_mut_ptr(tcx, param(0)), param(0)], tcx.types.unit),
 
         sym::atomic_xchg
-        | sym::atomic_xadd
-        | sym::atomic_xsub
-        | sym::atomic_and
-        | sym::atomic_nand
-        | sym::atomic_or
-        | sym::atomic_xor
         | sym::atomic_max
         | sym::atomic_min
         | sym::atomic_umax
         | sym::atomic_umin => (1, 1, vec![Ty::new_mut_ptr(tcx, param(0)), param(0)], param(0)),
+        sym::atomic_xadd
+        | sym::atomic_xsub
+        | sym::atomic_and
+        | sym::atomic_nand
+        | sym::atomic_or
+        | sym::atomic_xor => (2, 1, vec![Ty::new_mut_ptr(tcx, param(0)), param(1)], param(0)),
         sym::atomic_fence | sym::atomic_singlethreadfence => (0, 1, Vec::new(), tcx.types.unit),
 
         other => {
diff --git a/compiler/rustc_middle/src/ty/generic_args.rs b/compiler/rustc_middle/src/ty/generic_args.rs
index b02abb5ab43..648709e88e6 100644
--- a/compiler/rustc_middle/src/ty/generic_args.rs
+++ b/compiler/rustc_middle/src/ty/generic_args.rs
@@ -527,21 +527,28 @@ impl<'tcx> GenericArgs<'tcx> {
     #[inline]
     #[track_caller]
     pub fn type_at(&self, i: usize) -> Ty<'tcx> {
-        self[i].as_type().unwrap_or_else(|| bug!("expected type for param #{} in {:?}", i, self))
+        self[i].as_type().unwrap_or_else(
+            #[track_caller]
+            || bug!("expected type for param #{} in {:?}", i, self),
+        )
     }
 
     #[inline]
     #[track_caller]
     pub fn region_at(&self, i: usize) -> ty::Region<'tcx> {
-        self[i]
-            .as_region()
-            .unwrap_or_else(|| bug!("expected region for param #{} in {:?}", i, self))
+        self[i].as_region().unwrap_or_else(
+            #[track_caller]
+            || bug!("expected region for param #{} in {:?}", i, self),
+        )
     }
 
     #[inline]
     #[track_caller]
     pub fn const_at(&self, i: usize) -> ty::Const<'tcx> {
-        self[i].as_const().unwrap_or_else(|| bug!("expected const for param #{} in {:?}", i, self))
+        self[i].as_const().unwrap_or_else(
+            #[track_caller]
+            || bug!("expected const for param #{} in {:?}", i, self),
+        )
     }
 
     #[inline]