fix(mmu): correct align_down implementation

The previously used `align_down` implementation was incorrect as it would cause previously aligned addresses (where the offset was 0) to become unaligned. This patch fixes both this behaviour and makes both `align_down` and `align_up` more efficient
JonasKruckenberg · JonasKruckenberg · Dec 29, 2024 · Dec 18, 2024 · Dec 18, 2024 · Dec 18, 2024
commit ca9f574954c6242a7ee3b135757c0c86b36bb9e0
diff --git a/kernel/src/vm/aspace.rs b/kernel/src/vm/aspace.rs
@@ -193,7 +193,7 @@ impl AddressSpace {
         if flags.contains(PageFaultFlags::ACCESS) {
             return self.access_fault(virt);
         }
-        let virt = virt.checked_align_down(arch::PAGE_SIZE).unwrap();
+        let virt = virt.align_down(arch::PAGE_SIZE);
 
         // check if the address is within the last fault range
         // if so, we can save ourselves a tree lookup

diff --git a/kernel/src/vm/mod.rs b/kernel/src/vm/mod.rs
@@ -218,7 +218,7 @@ fn reserve_wired_regions(
         );
 
         aspace.reserve(
-            virt.checked_align_down(arch::PAGE_SIZE).unwrap()
+            virt.align_down(arch::PAGE_SIZE)
                 ..virt
                     .checked_add(ph.mem_size() as usize)
                     .unwrap()

diff --git a/libs/mmu/src/lib.rs b/libs/mmu/src/lib.rs
@@ -185,62 +185,48 @@ macro_rules! address_impl {
                 }
 
                 // SAFETY: `align` has been checked to be a power of 2 above
-                let offset = unsafe { align_offset(self.0, align) };
-
-                let Some(out) = self.checked_add(offset) else {
-                    return None;
-                };
-                debug_assert!(out.is_aligned_to(align));
-                Some(out)
-            }
-
-            #[must_use]
-            #[inline]
-            pub const fn checked_align_down(self, align: usize) -> Option<Self> {
-                if !align.is_power_of_two() {
-                    panic!("checked_align_up: align is not a power-of-two");
+                let align_minus_one = unsafe { align.unchecked_sub(1) };
+
+                // addr.wrapping_add(align_minus_one) & 0usize.wrapping_sub(align)
+                if let Some(addr_plus_align) = self.0.checked_add(align_minus_one) {
+                    let aligned = Self(addr_plus_align & 0usize.wrapping_sub(align));
+                    debug_assert!(aligned.is_aligned_to(align));
+                    debug_assert!(aligned.0 >= self.0);
+                    Some(aligned)
+                } else {
+                    None
                 }
-
-                // SAFETY: `align` has been checked to be a power of 2 above
-                let offset = align - unsafe { align_offset(self.0, align) };
-
-                let Some(out) = self.checked_sub(offset) else {
-                    return None;
-                };
-                debug_assert!(out.is_aligned_to(align));
-                Some(out)
             }
 
             // #[must_use]
             // #[inline]
-            // pub const fn saturating_align_up(self, align: usize) -> Self {
-            //     if !align.is_power_of_two() {
-            //         panic!("checked_align_up: align is not a power-of-two");
-            //     }
-            //
-            //     // SAFETY: `align` has been checked to be a power of 2 above
-            //     let offset = unsafe { align_offset(self.0, align) };
-            //
-            //     let out = self.saturating_add(offset);
-            //     debug_assert!(out.is_aligned_to(align));
-            //     out
-            // }
-            //
-            // #[must_use]
-            // #[inline]
-            // pub const fn saturating_align_down(self, align: usize) -> Self {
+            // pub const fn wrapping_align_up(self, align: usize) -> Self {
             //     if !align.is_power_of_two() {
             //         panic!("checked_align_up: align is not a power-of-two");
             //     }
             //
             //     // SAFETY: `align` has been checked to be a power of 2 above
-            //     let offset = align - unsafe { align_offset(self.0, align) };
+            //     let align_minus_one = unsafe { align.unchecked_sub(1) };
             //
-            //     let out = self.saturating_sub(offset);
+            //     // addr.wrapping_add(align_minus_one) & 0usize.wrapping_sub(align)
+            //     let out = addr.wrapping_add(align_minus_one) & 0usize.wrapping_sub(align);
             //     debug_assert!(out.is_aligned_to(align));
             //     out
             // }
 
+            #[must_use]
+            #[inline]
+            pub const fn align_down(self, align: usize) -> Self {
+                if !align.is_power_of_two() {
+                    panic!("checked_align_up: align is not a power-of-two");
+                }
+
+                let aligned = Self(self.0 & 0usize.wrapping_sub(align));
+                debug_assert!(aligned.is_aligned_to(align));
+                debug_assert!(aligned.0 <= self.0);
+                aligned
+            }
+
             #[inline]
             pub const fn as_ptr(self) -> *const u8 {
                 self.0 as *const u8
@@ -302,13 +288,17 @@ macro_rules! address_range_impl {
         where
             Self: Sized,
         {
-            Some(self.start.checked_align_up(align)?..self.end.checked_align_down(align)?)
+            let res = self.start.checked_align_up(align)?..self.end.align_down(align);
+            Some(res)
         }
         fn checked_align_out(self, align: usize) -> Option<Self>
         where
             Self: Sized,
         {
-            Some(self.start.checked_align_down(align)?..self.end.checked_align_up(align)?)
+            let res = self.start.align_down(align)..self.end.checked_align_up(align)?;
+            // aligning outwards can only increase the size
+            debug_assert!(res.start.0 <= res.end.0);
+            Some(res)
         }
         // fn saturating_align_in(self, align: usize) -> Self {
         //     self.start.saturating_align_up(align)..self.end.saturating_align_down(align)
@@ -415,37 +405,10 @@ impl AddressRangeExt for Range<VirtualAddress> {
     }
 }
 
-const unsafe fn align_offset(addr: usize, a: usize) -> usize {
-    // SAFETY: `a` is a power-of-two, therefore non-zero.
-    let a_minus_one = unsafe { a.unchecked_sub(1) };
-
-    // SPECIAL_CASE: In cases where the `a` is divisible by `STRIDE`, byte offset to align a
-    // pointer can be computed more simply through `-p (mod a)`. In the off-chance the byte
-    // offset is not a multiple of `STRIDE`, the input pointer was misaligned and no pointer
-    // offset will be able to produce a `p` aligned to the specified `a`.
-    //
-    // The naive `-p (mod a)` equation inhibits LLVM's ability to select instructions
-    // like `lea`. We compute `(round_up_to_next_alignment(p, a) - p)` instead. This
-    // redistributes operations around the load-bearing, but pessimizing `and` instruction
-    // sufficiently for LLVM to be able to utilize the various optimizations it knows about.
-    //
-    // LLVM handles the branch here particularly nicely. If this branch needs to be evaluated
-    // at runtime, it will produce a mask `if addr_mod_stride == 0 { 0 } else { usize::MAX }`
-    // in a branch-free way and then bitwise-OR it with whatever result the `-p mod a`
-    // computation produces.
-
-    let aligned_address = addr.wrapping_add(a_minus_one) & 0usize.wrapping_sub(a);
-    let byte_offset = aligned_address.wrapping_sub(addr);
-    // FIXME: Remove the assume after <https://github.com/llvm/llvm-project/issues/62502>
-    // SAFETY: Masking by `-a` can only affect the low bits, and thus cannot have reduced
-    // the value by more than `a-1`, so even though the intermediate values might have
-    // wrapped, the byte_offset is always in `[0, a)`.
-    unsafe {
-        core::hint::assert_unchecked(byte_offset < a);
-    }
-
-    byte_offset
-}
+static_assertions::const_assert!(VirtualAddress(0xffffffc000000000).is_aligned_to(4096));
+static_assertions::const_assert_eq!(VirtualAddress(0xffffffc0000156e8).align_down(4096).0, 0xffffffc000015000);
+static_assertions::const_assert_eq!(VirtualAddress(0xffffffc000000000).checked_align_up(4096).unwrap().0, 0xffffffc000000000);
+static_assertions::const_assert_eq!(VirtualAddress(0xffffffc0000156e8).checked_align_up(4096).unwrap().0, 0xffffffc000016000);
 
 #[cfg(test)]
 mod tests {

diff --git a/loader/src/machine_info.rs b/loader/src/machine_info.rs
@@ -161,7 +161,7 @@ impl<'dt> MachineInfo<'dt> {
         // but we can't make use of it either way
         info.memories.iter_mut().for_each(|region| {
             region.start = region.start.checked_align_up(arch::PAGE_SIZE).unwrap();
-            region.end = region.end.checked_align_down(arch::PAGE_SIZE).unwrap();
+            region.end = region.end.align_down(arch::PAGE_SIZE);
         });
 
         // ensure the memory regions are sorted.

diff --git a/loader/src/vm.rs b/loader/src/vm.rs
@@ -309,13 +309,11 @@ fn handle_bss_section(
         let last_page = virt_start
             .checked_add(ph.file_size - 1)
             .unwrap()
-            .checked_align_down(ph.align)
-            .unwrap();
+            .align_down(ph.align);
         let last_frame = phys_base
             .checked_add(ph.offset + ph.file_size - 1)
             .unwrap()
-            .checked_align_down(ph.align)
-            .unwrap();
+            .align_down(ph.align);
 
         let new_frame = frame_alloc
             .allocate_contiguous_zeroed(
@@ -489,8 +487,8 @@ fn handle_relro_segment(
     };
 
     let virt_aligned = {
-        virt.start.checked_align_down(arch::PAGE_SIZE).unwrap()
-            ..virt.end.checked_align_down(arch::PAGE_SIZE).unwrap()
+        virt.start.align_down(arch::PAGE_SIZE)
+            ..virt.end.align_down(arch::PAGE_SIZE)
     };
 
     log::debug!("Marking RELRO segment {virt_aligned:?} as read-only");