From 9e5d497b67c7566001bdcfc8a2767f26a23afc5b Mon Sep 17 00:00:00 2001
From: Lukas Markeffsky <@>
Date: Tue, 15 Nov 2022 15:55:37 +0100
Subject: [PATCH] fix const `align_offset` implementation

---
 .../src/const_eval/machine.rs                 |  9 +-
 library/core/tests/ptr.rs                     | 88 ++++++++++++-------
 2 files changed, 61 insertions(+), 36 deletions(-)

diff --git a/compiler/rustc_const_eval/src/const_eval/machine.rs b/compiler/rustc_const_eval/src/const_eval/machine.rs
index 93afbbd27a4..92ccc94e630 100644
--- a/compiler/rustc_const_eval/src/const_eval/machine.rs
+++ b/compiler/rustc_const_eval/src/const_eval/machine.rs
@@ -278,9 +278,8 @@ fn align_offset_impl(addr: u64, stride: u64, align: u64) -> u64 {
             return u64::MAX;
         }
 
-        let byte_offset = align - addr_mod_align;
-
         if align % stride == 0 {
+            let byte_offset = align - addr_mod_align;
             if byte_offset % stride == 0 {
                 return byte_offset / stride;
             } else {
@@ -296,8 +295,11 @@ fn align_offset_impl(addr: u64, stride: u64, align: u64) -> u64 {
             return u64::MAX;
         }
 
+        // Instead of `(addr + offset * stride) % align == 0`, we solve
+        // `((addr + offset * stride) / gcd) % (align / gcd) == 0`.
+        let addr2 = addr / gcd;
         let align2 = align / gcd;
-        let stride2 = (stride / gcd) % align2;
+        let stride2 = stride / gcd;
 
         let mut stride_inv = 1u64;
         let mut mod_gate = 2u64;
@@ -308,6 +310,7 @@ fn align_offset_impl(addr: u64, stride: u64, align: u64) -> u64 {
             (mod_gate, overflow) = mod_gate.overflowing_mul(mod_gate);
         }
 
+        let byte_offset = align2 - addr2 % align2;
         byte_offset.wrapping_mul(stride_inv) % align2
     }
 
diff --git a/library/core/tests/ptr.rs b/library/core/tests/ptr.rs
index 9f74b0c0410..036acc46dab 100644
--- a/library/core/tests/ptr.rs
+++ b/library/core/tests/ptr.rs
@@ -510,49 +510,53 @@ fn align_offset_various_strides_const() {
         assert!(got == expected);
     }
 
-    // For pointers of stride != 1, we verify the algorithm against the naivest possible
-    // implementation
-    let mut align = 1;
-    let limit = 1024;
-    while align < limit {
-        for ptr in 1usize..4 * align {
-            unsafe {
-                #[repr(packed)]
-                struct A3(u16, u8);
-                test_stride::<A3>(ptr::invalid::<A3>(ptr), ptr, align);
+    const {
+        // For pointers of stride != 1, we verify the algorithm against the naivest possible
+        // implementation
+        let mut align = 1;
+        let limit = 32;
+        while align < limit {
+            let mut ptr = 1;
+            while ptr < 4 * align {
+                unsafe {
+                    #[repr(packed)]
+                    struct A3(u16, u8);
+                    test_stride::<A3>(ptr::invalid::<A3>(ptr), ptr, align);
 
-                struct A4(u32);
-                test_stride::<A4>(ptr::invalid::<A4>(ptr), ptr, align);
+                    struct A4(u32);
+                    test_stride::<A4>(ptr::invalid::<A4>(ptr), ptr, align);
 
-                #[repr(packed)]
-                struct A5(u32, u8);
-                test_stride::<A5>(ptr::invalid::<A5>(ptr), ptr, align);
+                    #[repr(packed)]
+                    struct A5(u32, u8);
+                    test_stride::<A5>(ptr::invalid::<A5>(ptr), ptr, align);
 
-                #[repr(packed)]
-                struct A6(u32, u16);
-                test_stride::<A6>(ptr::invalid::<A6>(ptr), ptr, align);
+                    #[repr(packed)]
+                    struct A6(u32, u16);
+                    test_stride::<A6>(ptr::invalid::<A6>(ptr), ptr, align);
 
-                #[repr(packed)]
-                struct A7(u32, u16, u8);
-                test_stride::<A7>(ptr::invalid::<A7>(ptr), ptr, align);
+                    #[repr(packed)]
+                    struct A7(u32, u16, u8);
+                    test_stride::<A7>(ptr::invalid::<A7>(ptr), ptr, align);
 
-                #[repr(packed)]
-                struct A8(u32, u32);
-                test_stride::<A8>(ptr::invalid::<A8>(ptr), ptr, align);
+                    #[repr(packed)]
+                    struct A8(u32, u32);
+                    test_stride::<A8>(ptr::invalid::<A8>(ptr), ptr, align);
 
-                #[repr(packed)]
-                struct A9(u32, u32, u8);
-                test_stride::<A9>(ptr::invalid::<A9>(ptr), ptr, align);
+                    #[repr(packed)]
+                    struct A9(u32, u32, u8);
+                    test_stride::<A9>(ptr::invalid::<A9>(ptr), ptr, align);
 
-                #[repr(packed)]
-                struct A10(u32, u32, u16);
-                test_stride::<A10>(ptr::invalid::<A10>(ptr), ptr, align);
+                    #[repr(packed)]
+                    struct A10(u32, u32, u16);
+                    test_stride::<A10>(ptr::invalid::<A10>(ptr), ptr, align);
 
-                test_stride::<u32>(ptr::invalid::<u32>(ptr), ptr, align);
-                test_stride::<u128>(ptr::invalid::<u128>(ptr), ptr, align);
+                    test_stride::<u32>(ptr::invalid::<u32>(ptr), ptr, align);
+                    test_stride::<u128>(ptr::invalid::<u128>(ptr), ptr, align);
+                }
+                ptr += 1;
             }
+            align = (align + 1).next_power_of_two();
         }
-        align = (align + 1).next_power_of_two();
     }
 }
 
@@ -632,6 +636,24 @@ fn align_offset_issue_103361() {
     let _ = (SIZE as *const HugeSize).align_offset(SIZE);
 }
 
+#[test]
+#[cfg(not(bootstrap))]
+fn align_offset_issue_103361_const() {
+    #[cfg(target_pointer_width = "64")]
+    const SIZE: usize = 1 << 47;
+    #[cfg(target_pointer_width = "32")]
+    const SIZE: usize = 1 << 30;
+    #[cfg(target_pointer_width = "16")]
+    const SIZE: usize = 1 << 13;
+    struct HugeSize([u8; SIZE - 1]);
+
+    const {
+        assert!(ptr::invalid::<HugeSize>(SIZE - 1).align_offset(SIZE) == SIZE - 1);
+        assert!(ptr::invalid::<HugeSize>(SIZE).align_offset(SIZE) == 0);
+        assert!(ptr::invalid::<HugeSize>(SIZE + 1).align_offset(SIZE) == 1);
+    }
+}
+
 #[test]
 fn is_aligned() {
     let data = 42;