Add precondition checks to ptr::offset, ptr::add, ptr::sub

rust-lang · Sep 12, 2024 · 557d607 · 557d607
1 parent f827364
commit 557d607
Show file tree

Hide file tree

Showing 2 changed files with 170 additions and 1 deletion.
diff --git a/library/core/src/ptr/const_ptr.rs b/library/core/src/ptr/const_ptr.rs
@@ -394,6 +394,37 @@ impl<T: ?Sized> *const T {
     where
         T: Sized,
     {
+        #[inline]
+        const fn runtime_offset_nowrap(this: *const (), count: isize, size: usize) -> bool {
+            #[inline]
+            fn runtime(this: *const (), count: isize, size: usize) -> bool {
+                let Some(byte_offset) = count.checked_mul(size as isize) else {
+                    return false;
+                };
+                if byte_offset < 0 {
+                    -byte_offset as usize <= this.addr()
+                } else {
+                    this.addr().checked_add(byte_offset as usize).is_some()
+                }
+            }
+
+            const fn comptime(_: *const (), _: isize, _: usize) -> bool {
+                true
+            }
+
+            intrinsics::const_eval_select((this, count, size), comptime, runtime)
+        }
+
+        ub_checks::assert_unsafe_precondition!(
+            check_language_ub,
+            "ptr::offset requires the address calculation to not overflow",
+            (
+                this: *const () = self as *const (),
+                count: isize = count,
+                size: usize = size_of::<T>(),
+            ) => runtime_offset_nowrap(this, count, size)
+        );
+
         // SAFETY: the caller must uphold the safety contract for `offset`.
         unsafe { intrinsics::offset(self, count) }
     }
@@ -728,7 +759,6 @@ impl<T: ?Sized> *const T {
                 true
             }
 
-            #[allow(unused_unsafe)]
             intrinsics::const_eval_select((this, origin), comptime, runtime)
         }
 
@@ -855,6 +885,33 @@ impl<T: ?Sized> *const T {
     where
         T: Sized,
     {
+        #[inline]
+        const fn runtime_add_nowrap(this: *const (), count: usize, size: usize) -> bool {
+            #[inline]
+            fn runtime(this: *const (), count: usize, size: usize) -> bool {
+                let Some(byte_offset) = count.checked_mul(size) else {
+                    return false;
+                };
+                this.addr().checked_add(byte_offset as usize).is_some()
+            }
+
+            const fn comptime(_: *const (), _: usize, _: usize) -> bool {
+                true
+            }
+
+            intrinsics::const_eval_select((this, count, size), comptime, runtime)
+        }
+
+        ub_checks::assert_unsafe_precondition!(
+            check_language_ub,
+            "ptr::add requires that the address calculation does not overflow",
+            (
+                this: *const () = self as *const (),
+                count: usize = count,
+                size: usize = size_of::<T>(),
+            ) => runtime_add_nowrap(this, count, size)
+        );
+
         // SAFETY: the caller must uphold the safety contract for `offset`.
         unsafe { intrinsics::offset(self, count) }
     }
@@ -930,6 +987,33 @@ impl<T: ?Sized> *const T {
     where
         T: Sized,
     {
+        #[inline]
+        const fn runtime_sub_nowrap(this: *const (), count: usize, size: usize) -> bool {
+            #[inline]
+            fn runtime(this: *const (), count: usize, size: usize) -> bool {
+                let Some(byte_offset) = count.checked_mul(size) else {
+                    return false;
+                };
+                this.addr() >= byte_offset
+            }
+
+            const fn comptime(_: *const (), _: usize, _: usize) -> bool {
+                true
+            }
+
+            intrinsics::const_eval_select((this, count, size), comptime, runtime)
+        }
+
+        ub_checks::assert_unsafe_precondition!(
+            check_language_ub,
+            "ptr::sub requires that the address calculation does not overflow",
+            (
+                this: *const () = self as *const (),
+                count: usize = count,
+                size: usize = size_of::<T>(),
+            ) => runtime_sub_nowrap(this, count, size)
+        );
+
         if T::IS_ZST {
             // Pointer arithmetic does nothing when the pointee is a ZST.
             self

diff --git a/library/core/src/ptr/mut_ptr.rs b/library/core/src/ptr/mut_ptr.rs
@@ -392,6 +392,37 @@ impl<T: ?Sized> *mut T {
     where
         T: Sized,
     {
+        #[inline]
+        const fn runtime_offset_nowrap(this: *const (), count: isize, size: usize) -> bool {
+            #[inline]
+            fn runtime(this: *const (), count: isize, size: usize) -> bool {
+                let Some(byte_offset) = count.checked_mul(size as isize) else {
+                    return false;
+                };
+                if byte_offset < 0 {
+                    -byte_offset as usize <= this.addr()
+                } else {
+                    this.addr().checked_add(byte_offset as usize).is_some()
+                }
+            }
+
+            const fn comptime(_: *const (), _: isize, _: usize) -> bool {
+                true
+            }
+
+            intrinsics::const_eval_select((this, count, size), comptime, runtime)
+        }
+
+        ub_checks::assert_unsafe_precondition!(
+            check_language_ub,
+            "ptr::offset requires the address calculation to not overflow",
+            (
+                this: *const () = self as *const (),
+                count: isize = count,
+                size: usize = size_of::<T>(),
+            ) => runtime_offset_nowrap(this, count, size)
+        );
+
         // SAFETY: the caller must uphold the safety contract for `offset`.
         // The obtained pointer is valid for writes since the caller must
         // guarantee that it points to the same allocated object as `self`.
@@ -936,6 +967,33 @@ impl<T: ?Sized> *mut T {
     where
         T: Sized,
     {
+        #[inline]
+        const fn runtime_add_nowrap(this: *const (), count: usize, size: usize) -> bool {
+            #[inline]
+            fn runtime(this: *const (), count: usize, size: usize) -> bool {
+                let Some(byte_offset) = count.checked_mul(size) else {
+                    return false;
+                };
+                this.addr().checked_add(byte_offset as usize).is_some()
+            }
+
+            const fn comptime(_: *const (), _: usize, _: usize) -> bool {
+                true
+            }
+
+            intrinsics::const_eval_select((this, count, size), comptime, runtime)
+        }
+
+        ub_checks::assert_unsafe_precondition!(
+            check_language_ub,
+            "ptr::add requires that the address calculation does not overflow",
+            (
+                this: *const () = self as *const (),
+                count: usize = count,
+                size: usize = size_of::<T>(),
+            ) => runtime_add_nowrap(this, count, size)
+        );
+
         // SAFETY: the caller must uphold the safety contract for `offset`.
         unsafe { intrinsics::offset(self, count) }
     }
@@ -1011,6 +1069,33 @@ impl<T: ?Sized> *mut T {
     where
         T: Sized,
     {
+        #[inline]
+        const fn runtime_sub_nowrap(this: *const (), count: usize, size: usize) -> bool {
+            #[inline]
+            fn runtime(this: *const (), count: usize, size: usize) -> bool {
+                let Some(byte_offset) = count.checked_mul(size) else {
+                    return false;
+                };
+                this.addr() >= byte_offset
+            }
+
+            const fn comptime(_: *const (), _: usize, _: usize) -> bool {
+                true
+            }
+
+            intrinsics::const_eval_select((this, count, size), comptime, runtime)
+        }
+
+        ub_checks::assert_unsafe_precondition!(
+            check_language_ub,
+            "ptr::sub requires that the address calculation does not overflow",
+            (
+                this: *const () = self as *const (),
+                count: usize = count,
+                size: usize = size_of::<T>(),
+            ) => runtime_sub_nowrap(this, count, size)
+        );
+
         if T::IS_ZST {
             // Pointer arithmetic does nothing when the pointee is a ZST.
             self