Add AlignExt

src/kxos-frame/src/prelude.rs

@@ -7,4 +7,5 @@ pub(crate) use alloc::sync::Arc;
 pub(crate) use alloc::vec::Vec;
 pub(crate) use core::any::Any;
 
+pub(crate) use crate::util::AlignExt;
 pub use crate::vm::{Paddr, Vaddr};

src/kxos-frame/src/util/align_ext.rs

@@ -0,0 +1,97 @@
+/// An extension trait for Rust integer types, including `u8`, `u16`, `u32`,
+/// `u64`, and `usize`, to provide methods to make integers aligned to a
+/// power of two.
+pub trait AlignExt {
+    /// Returns to the smallest number that is greater than or equal to
+    /// `self` and is a multiple of the given power of two.
+    ///
+    /// The method panics if `power_of_two` is not a
+    /// power of two or is smaller than 2 or the calculation overflows
+    /// because `self` is too large.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// assert!(align_up(12, 2), 12);
+    /// assert!(align_up(12, 4), 12);
+    /// assert!(align_up(12, 8), 16);
+    /// assert!(align_up(12, 16), 16);
+    /// ```
+    fn align_up(self, power_of_two: Self) -> Self;
+
+    /// Returns to the greatest number that is smaller than or equal to
+    /// `self` and is a multiple of the given power of two.
+    ///
+    /// The method panics if `power_of_two` is not a
+    /// power of two or is smaller than 2 or the calculation overflows
+    /// because `self` is too large. In release mode,
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// assert!(align_down(12, 2), 12);
+    /// assert!(align_down(12, 4), 12);
+    /// assert!(align_down(12, 8), 8);
+    /// assert!(align_down(12, 16), 0);
+    /// ```
+    fn align_down(self, power_of_two: Self) -> Self;
+}
+
+macro_rules! impl_align_ext {
+    ($( $uint_type:ty ),+,) => {
+        $(
+            impl AlignExt for $uint_type {
+                fn align_up(self, align: Self) -> Self {
+                    assert!(align.is_power_of_two() && align >= 2);
+                    self.checked_add(align - 1).unwrap() & !(align - 1)
+                }
+
+                fn align_down(self, align: Self) -> Self {
+                    assert!(align.is_power_of_two() && align >= 2);
+                    self & !(align - 1)
+                }
+            }
+        )*
+    }
+}
+
+impl_align_ext! {
+    u8,
+    u16,
+    u32,
+    u64,
+    usize,
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_align_up() {
+        let input_ns = [0, 1, 2, 9, 15, 21, 32, 47, 50];
+        let input_as = [2, 2, 2, 2, 4, 4, 8, 8, 8];
+        let output_ns = [0, 2, 2, 10, 16, 24, 32, 48, 56];
+
+        for i in 0..input_ns.len() {
+            let n = input_ns[i];
+            let a = input_as[i];
+            let n2 = output_ns[i];
+            assert!(align_up(n, a) == n2);
+        }
+    }
+
+    #[test]
+    fn test_align_down() {
+        let input_ns = [0, 1, 2, 9, 15, 21, 32, 47, 50];
+        let input_as = [2, 2, 2, 2, 4, 4, 8, 8, 8];
+        let output_ns = [0, 0, 2, 8, 12, 20, 32, 40, 48];
+
+        for i in 0..input_ns.len() {
+            let n = input_ns[i];
+            let a = input_as[i];
+            let n2 = output_ns[i];
+            assert!(align_down(n, a) == n2);
+        }
+    }
+}

src/kxos-frame/src/util/mod.rs

@@ -1,2 +1,5 @@
+mod align_ext;
 mod type_map;
+
+pub use self::align_ext::AlignExt;
 pub use self::type_map::TypeMap;