[What If] NonNull could just offset directly?

compiler/rustc_codegen_ssa/src/mir/rvalue.rs

@@ -2,6 +2,7 @@ use std::assert_matches::assert_matches;
 
 use arrayvec::ArrayVec;
 use rustc_abi::{self as abi, FIRST_VARIANT, FieldIdx};
+use rustc_hir::LangItem;
 use rustc_middle::ty::adjustment::PointerCoercion;
 use rustc_middle::ty::layout::{HasTyCtxt, LayoutOf, TyAndLayout};
 use rustc_middle::ty::{self, Instance, Ty, TyCtxt};
@@ -918,9 +919,15 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             mir::BinOp::BitAnd => bx.and(lhs, rhs),
             mir::BinOp::BitXor => bx.xor(lhs, rhs),
             mir::BinOp::Offset => {
-                let pointee_type = input_ty
-                    .builtin_deref(true)
-                    .unwrap_or_else(|| bug!("deref of non-pointer {:?}", input_ty));
+                let pointee_type = if let ty::Adt(adt_def, generic_args) = input_ty.kind()
+                    && bx.tcx().is_lang_item(adt_def.did(), LangItem::NonNull)
+                {
+                    generic_args.into_type_list(bx.tcx())[0]
+                } else {
+                    input_ty
+                        .builtin_deref(true)
+                        .unwrap_or_else(|| bug!("deref of non-pointer {:?}", input_ty))
+                };
                 let pointee_layout = bx.cx().layout_of(pointee_type);
                 if pointee_layout.is_zst() {
                     // `Offset` works in terms of the size of pointee,

compiler/rustc_hir/src/lang_items.rs

@@ -341,6 +341,8 @@ language_item_table! {
     OwnedBox,                sym::owned_box,           owned_box,                  Target::Struct,         GenericRequirement::Minimum(1);
     GlobalAlloc,             sym::global_alloc_ty,     global_alloc_ty,            Target::Struct,         GenericRequirement::None;
 
+    NonNull,                 sym::NonNull,             non_null_ty,                Target::Struct,         GenericRequirement::Exact(1);
+
     // Experimental lang item for Miri
     PtrUnique,               sym::ptr_unique,          ptr_unique,                 Target::Struct,         GenericRequirement::Exact(1);
 

compiler/rustc_mir_transform/src/validate.rs

@@ -1064,7 +1064,13 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
 
                 match op {
                     Offset => {
-                        check_kinds!(a, "Cannot offset non-pointer type {:?}", ty::RawPtr(..));
+                        if let ty::Adt(adt_def, _) = a.kind()
+                            && self.tcx.is_lang_item(adt_def.did(), LangItem::NonNull)
+                        {
+                            // ok
+                        } else {
+                            check_kinds!(a, "Cannot offset non-pointer type {:?}", ty::RawPtr(..));
+                        }
                         if b != self.tcx.types.isize && b != self.tcx.types.usize {
                             self.fail(location, format!("Cannot offset by non-isize type {b:?}"));
                         }

library/core/src/ptr/non_null.rs

@@ -68,6 +68,8 @@ use crate::{fmt, hash, intrinsics, ptr};
 #[rustc_layout_scalar_valid_range_start(1)]
 #[rustc_nonnull_optimization_guaranteed]
 #[rustc_diagnostic_item = "NonNull"]
+// No special behaviour, but must be transparent over a pointer to the generic parameter.
+#[cfg_attr(not(bootstrap), lang = "NonNull")]
 pub struct NonNull<T: ?Sized> {
     pointer: *const T,
 }
@@ -480,11 +482,22 @@ impl<T: ?Sized> NonNull<T> {
     where
         T: Sized,
     {
-        // SAFETY: the caller must uphold the safety contract for `offset`.
-        // Additionally safety contract of `offset` guarantees that the resulting pointer is
-        // pointing to an allocation, there can't be an allocation at null, thus it's safe to
-        // construct `NonNull`.
-        unsafe { NonNull { pointer: intrinsics::offset(self.pointer, count) } }
+        #[cfg(bootstrap)]
+        {
+            // SAFETY: the caller must uphold the safety contract for `offset`.
+            // Additionally safety contract of `offset` guarantees that the resulting pointer is
+            // pointing to an allocation, there can't be an allocation at null, thus it's safe to
+            // construct `NonNull`.
+            unsafe { NonNull { pointer: intrinsics::offset(self.pointer, count) } }
+        }
+        #[cfg(not(bootstrap))]
+        {
+            // SAFETY: the caller must uphold the safety contract for `offset`.
+            // Additionally safety contract of `offset` guarantees that the resulting pointer is
+            // pointing to an allocation, there can't be an allocation at null, thus it's safe to
+            // construct `NonNull`.
+            unsafe { intrinsics::offset(self, count) }
+        }
     }
 
     /// Calculates the offset from a pointer in bytes.
@@ -556,11 +569,22 @@ impl<T: ?Sized> NonNull<T> {
     where
         T: Sized,
     {
-        // SAFETY: the caller must uphold the safety contract for `offset`.
-        // Additionally safety contract of `offset` guarantees that the resulting pointer is
-        // pointing to an allocation, there can't be an allocation at null, thus it's safe to
-        // construct `NonNull`.
-        unsafe { NonNull { pointer: intrinsics::offset(self.pointer, count) } }
+        #[cfg(bootstrap)]
+        {
+            // SAFETY: the caller must uphold the safety contract for `offset`.
+            // Additionally safety contract of `offset` guarantees that the resulting pointer is
+            // pointing to an allocation, there can't be an allocation at null, thus it's safe to
+            // construct `NonNull`.
+            unsafe { NonNull { pointer: intrinsics::offset(self.pointer, count) } }
+        }
+        #[cfg(not(bootstrap))]
+        {
+            // SAFETY: the caller must uphold the safety contract for `offset`.
+            // Additionally safety contract of `offset` guarantees that the resulting pointer is
+            // pointing to an allocation, there can't be an allocation at null, thus it's safe to
+            // construct `NonNull`.
+            unsafe { intrinsics::offset(self, count) }
+        }
     }
 
     /// Calculates the offset from a pointer in bytes (convenience for `.byte_offset(count as isize)`).