Auto merge of #96946 - WaffleLapkin:ptr_mask, r=scottmcm

Add pointer masking convenience functions

This PR adds the following public API:
```rust
impl<T: ?Sized> *const T {
    fn mask(self, mask: usize) -> *const T;
}

impl<T: ?Sized> *mut T {
    fn mask(self, mask: usize) -> *const T;
}

// mod intrinsics
fn mask<T>(ptr: *const T, mask: usize) -> *const T
```
This is equivalent to `ptr.map_addr(|a| a & mask)` but also uses a cool llvm intrinsic.

Proposed in #95643 (comment)

cc `@Gankra` `@scottmcm` `@RalfJung`

r? rust-lang/libs-api

Author: bors

compiler/rustc_codegen_cranelift/src/intrinsics/mod.rs

@@ -577,6 +577,13 @@ fn codegen_regular_intrinsic_call<'tcx>(
             ret.write_cvalue(fx, CValue::by_val(res, base.layout()));
         }
 
+        sym::ptr_mask => {
+            intrinsic_args!(fx, args => (ptr, mask); intrinsic);
+            let ptr = ptr.load_scalar(fx);
+            let mask = mask.load_scalar(fx);
+            fx.bcx.ins().band(ptr, mask);
+        }
+
         sym::transmute => {
             intrinsic_args!(fx, args => (from); intrinsic);
 

compiler/rustc_codegen_gcc/src/intrinsic/mod.rs

@@ -309,6 +309,18 @@ impl<'a, 'gcc, 'tcx> IntrinsicCallMethods<'tcx> for Builder<'a, 'gcc, 'tcx> {
                     return;
                 }
 
+                sym::ptr_mask => {
+                    let usize_type = self.context.new_type::<usize>();
+                    let void_ptr_type = self.context.new_type::<*const ()>();
+
+                    let ptr = args[0].immediate();
+                    let mask = args[1].immediate();
+
+                    let addr = self.bitcast(ptr, usize_type);
+                    let masked = self.and(addr, mask);
+                    self.bitcast(masked, void_ptr_type)
+                },
+                
                 _ if name_str.starts_with("simd_") => {
                     match generic_simd_intrinsic(self, name, callee_ty, args, ret_ty, llret_ty, span) {
                         Ok(llval) => llval,

compiler/rustc_codegen_llvm/src/context.rs

@@ -886,6 +886,9 @@ impl<'ll> CodegenCx<'ll, '_> {
             ifn!("llvm.dbg.declare", fn(t_metadata, t_metadata) -> void);
             ifn!("llvm.dbg.value", fn(t_metadata, t_i64, t_metadata) -> void);
         }
+
+        ifn!("llvm.ptrmask", fn(i8p, t_isize) -> i8p);
+
         None
     }
 

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -71,6 +71,7 @@ fn get_simple_intrinsic<'ll>(
         sym::nearbyintf64 => "llvm.nearbyint.f64",
         sym::roundf32 => "llvm.round.f32",
         sym::roundf64 => "llvm.round.f64",
+        sym::ptr_mask => "llvm.ptrmask",
         _ => return None,
     };
     Some(cx.get_intrinsic(llvm_name))

compiler/rustc_span/src/symbol.rs

@@ -1114,6 +1114,7 @@ symbols! {
         ptr,
         ptr_guaranteed_eq,
         ptr_guaranteed_ne,
+        ptr_mask,
         ptr_null,
         ptr_null_mut,
         ptr_offset_from,

compiler/rustc_typeck/src/check/intrinsic.rs

@@ -105,7 +105,8 @@ pub fn intrinsic_operation_unsafety(intrinsic: Symbol) -> hir::Unsafety {
         | sym::type_name
         | sym::forget
         | sym::black_box
-        | sym::variant_count => hir::Unsafety::Normal,
+        | sym::variant_count
+        | sym::ptr_mask => hir::Unsafety::Normal,
         _ => hir::Unsafety::Unsafe,
     }
 }
@@ -203,6 +204,15 @@ pub fn check_intrinsic_type(tcx: TyCtxt<'_>, it: &hir::ForeignItem<'_>) {
                 ],
                 tcx.mk_ptr(ty::TypeAndMut { ty: param(0), mutbl: hir::Mutability::Not }),
             ),
+            sym::ptr_mask => (
+                1,
+                vec![
+                    tcx.mk_ptr(ty::TypeAndMut { ty: param(0), mutbl: hir::Mutability::Not }),
+                    tcx.types.usize,
+                ],
+                tcx.mk_ptr(ty::TypeAndMut { ty: param(0), mutbl: hir::Mutability::Not }),
+            ),
+
             sym::copy | sym::copy_nonoverlapping => (
                 1,
                 vec![

library/core/src/intrinsics.rs

@@ -1287,6 +1287,17 @@ extern "rust-intrinsic" {
     #[rustc_const_stable(feature = "const_ptr_offset", since = "1.61.0")]
     pub fn arith_offset<T>(dst: *const T, offset: isize) -> *const T;
 
+    /// Masks out bits of the pointer according to a mask.
+    ///
+    /// Note that, unlike most intrinsics, this is safe to call;
+    /// it does not require an `unsafe` block.
+    /// Therefore, implementations must not require the user to uphold
+    /// any safety invariants.
+    ///
+    /// Consider using [`pointer::mask`] instead.
+    #[cfg(not(bootstrap))]
+    pub fn ptr_mask<T>(ptr: *const T, mask: usize) -> *const T;
+
     /// Equivalent to the appropriate `llvm.memcpy.p0i8.0i8.*` intrinsic, with
     /// a size of `count` * `size_of::<T>()` and an alignment of
     /// `min_align_of::<T>()`

library/core/src/ptr/const_ptr.rs

@@ -559,6 +559,21 @@ impl<T: ?Sized> *const T {
         from_raw_parts::<T>(self.cast::<u8>().wrapping_offset(count).cast::<()>(), metadata(self))
     }
 
+    /// Masks out bits of the pointer according to a mask.
+    ///
+    /// This is convenience for `ptr.map_addr(|a| a & mask)`.
+    ///
+    /// For non-`Sized` pointees this operation changes only the data pointer,
+    /// leaving the metadata untouched.
+    #[cfg(not(bootstrap))]
+    #[unstable(feature = "ptr_mask", issue = "98290")]
+    #[must_use = "returns a new pointer rather than modifying its argument"]
+    #[inline(always)]
+    pub fn mask(self, mask: usize) -> *const T {
+        let this = intrinsics::ptr_mask(self.cast::<()>(), mask);
+        from_raw_parts::<T>(this, metadata(self))
+    }
+
     /// Calculates the distance between two pointers. The returned value is in
     /// units of T: the distance in bytes divided by `mem::size_of::<T>()`.
     ///

library/core/src/ptr/mut_ptr.rs

@@ -575,6 +575,21 @@ impl<T: ?Sized> *mut T {
         )
     }
 
+    /// Masks out bits of the pointer according to a mask.
+    ///
+    /// This is convenience for `ptr.map_addr(|a| a & mask)`.
+    ///
+    /// For non-`Sized` pointees this operation changes only the data pointer,
+    /// leaving the metadata untouched.
+    #[cfg(not(bootstrap))]
+    #[unstable(feature = "ptr_mask", issue = "98290")]
+    #[must_use = "returns a new pointer rather than modifying its argument"]
+    #[inline(always)]
+    pub fn mask(self, mask: usize) -> *mut T {
+        let this = intrinsics::ptr_mask(self.cast::<()>(), mask) as *mut ();
+        from_raw_parts_mut::<T>(this, metadata(self))
+    }
+
     /// Returns `None` if the pointer is null, or else returns a unique reference to
     /// the value wrapped in `Some`. If the value may be uninitialized, [`as_uninit_mut`]
     /// must be used instead.

src/test/codegen/intrinsics/mask.rs

@@ -0,0 +1,11 @@
+#![crate_type = "lib"]
+#![feature(core_intrinsics)]
+
+// CHECK-LABEL: @mask_ptr
+// CHECK-SAME: [[WORD:i[0-9]+]] %mask
+#[no_mangle]
+pub fn mask_ptr(ptr: *const u16, mask: usize) -> *const u16 {
+    // CHECK: call
+    // CHECK-SAME: @llvm.ptrmask.{{p0|p0i8}}.[[WORD]]({{ptr|i8\*}} {{%ptr|%0}}, [[WORD]] %mask)
+    core::intrinsics::ptr_mask(ptr, mask)
+}