From 7663305dff3b6634de90a9f3bc726e976357dc48 Mon Sep 17 00:00:00 2001
From: Maybe Waffle <waffle.lapkin@gmail.com>
Date: Thu, 27 Oct 2022 22:23:58 +0400
Subject: [PATCH] Implement `ptr_mask` intrinsic

---
 src/tools/miri/src/shims/intrinsics/mod.rs  | 13 ++++++++++++-
 src/tools/miri/tests/pass/shims/ptr_mask.rs | 18 ++++++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)
 create mode 100644 src/tools/miri/tests/pass/shims/ptr_mask.rs

diff --git a/src/tools/miri/src/shims/intrinsics/mod.rs b/src/tools/miri/src/shims/intrinsics/mod.rs
index e0985ace5be..6004e2078ad 100644
--- a/src/tools/miri/src/shims/intrinsics/mod.rs
+++ b/src/tools/miri/src/shims/intrinsics/mod.rs
@@ -11,7 +11,7 @@
     mir,
     ty::{self, FloatTy, Ty},
 };
-use rustc_target::abi::Integer;
+use rustc_target::abi::{Integer, Size};
 
 use crate::*;
 use atomic::EvalContextExt as _;
@@ -120,6 +120,17 @@ fn emulate_intrinsic_by_name(
                 this.write_bytes_ptr(ptr, iter::repeat(val_byte).take(byte_count.bytes_usize()))?;
             }
 
+            "ptr_mask" => {
+                let [ptr, mask] = check_arg_count(args)?;
+
+                let ptr = this.read_pointer(ptr)?;
+                let mask = this.read_scalar(mask)?.to_machine_usize(this)?;
+
+                let masked_addr = Size::from_bytes(ptr.addr().bytes() & mask);
+
+                this.write_pointer(Pointer::new(ptr.provenance, masked_addr), dest)?;
+            }
+
             // Floating-point operations
             "fabsf32" => {
                 let [f] = check_arg_count(args)?;
diff --git a/src/tools/miri/tests/pass/shims/ptr_mask.rs b/src/tools/miri/tests/pass/shims/ptr_mask.rs
new file mode 100644
index 00000000000..fb8bb6b13db
--- /dev/null
+++ b/src/tools/miri/tests/pass/shims/ptr_mask.rs
@@ -0,0 +1,18 @@
+#![feature(ptr_mask)]
+#![feature(strict_provenance)]
+
+fn main() {
+    let v: u32 = 0xABCDABCD;
+    let ptr: *const u32 = &v;
+
+    // u32 is 4 aligned,
+    // so the lower `log2(4) = 2` bits of the address are always 0
+    assert_eq!(ptr.addr() & 0b11, 0);
+
+    let tagged_ptr = ptr.map_addr(|a| a | 0b11);
+    let tag = tagged_ptr.addr() & 0b11;
+    let masked_ptr = tagged_ptr.mask(!0b11);
+
+    assert_eq!(tag, 0b11);
+    assert_eq!(unsafe { *masked_ptr }, 0xABCDABCD);
+}