refactor: add align_up and align_down functions

Add `align_up` and `align_down` functions to stop reimplementing
them in place all over the codebase.

Signed-off-by: Egor Lazarchuk <[email protected]>

Author: ShadowCurse

src/vmm/src/arch/aarch64/mod.rs

@@ -28,6 +28,7 @@ use crate::arch::{BootProtocol, EntryPoint};
 use crate::cpu_config::aarch64::{CpuConfiguration, CpuConfigurationError};
 use crate::cpu_config::templates::CustomCpuTemplate;
 use crate::initrd::InitrdConfig;
+use crate::utils::{align_up, usize_to_u64};
 use crate::vmm_config::machine_config::MachineConfig;
 use crate::vstate::memory::{Address, Bytes, GuestAddress, GuestMemory, GuestMemoryMmap};
 use crate::vstate::vcpu::KvmVcpuError;
@@ -129,9 +130,11 @@ pub fn get_kernel_start() -> u64 {
 
 /// Returns the memory address where the initrd could be loaded.
 pub fn initrd_load_addr(guest_mem: &GuestMemoryMmap, initrd_size: usize) -> Option<u64> {
-    let round_to_pagesize =
-        |size| (size + (super::GUEST_PAGE_SIZE - 1)) & !(super::GUEST_PAGE_SIZE - 1);
-    match GuestAddress(get_fdt_addr(guest_mem)).checked_sub(round_to_pagesize(initrd_size) as u64) {
+    let rounded_size = align_up(
+        usize_to_u64(initrd_size),
+        usize_to_u64(super::GUEST_PAGE_SIZE),
+    );
+    match GuestAddress(get_fdt_addr(guest_mem)).checked_sub(rounded_size) {
         Some(offset) => {
             if guest_mem.address_in_range(offset) {
                 Some(offset.raw_value())

src/vmm/src/arch/x86_64/mod.rs

@@ -51,7 +51,7 @@ use crate::arch::{BootProtocol, SYSTEM_MEM_SIZE, SYSTEM_MEM_START};
 use crate::cpu_config::templates::{CustomCpuTemplate, GuestConfigError};
 use crate::cpu_config::x86_64::CpuConfiguration;
 use crate::initrd::InitrdConfig;
-use crate::utils::{mib_to_bytes, u64_to_usize};
+use crate::utils::{align_down, mib_to_bytes, u64_to_usize, usize_to_u64};
 use crate::vmm_config::machine_config::MachineConfig;
 use crate::vstate::memory::{
     Address, GuestAddress, GuestMemory, GuestMemoryMmap, GuestMemoryRegion,
@@ -138,8 +138,10 @@ pub fn initrd_load_addr(guest_mem: &GuestMemoryMmap, initrd_size: usize) -> Opti
         return None;
     }
 
-    let align_to_pagesize = |address| address & !(super::GUEST_PAGE_SIZE - 1);
-    Some(align_to_pagesize(lowmem_size - initrd_size) as u64)
+    Some(align_down(
+        usize_to_u64(lowmem_size - initrd_size),
+        usize_to_u64(super::GUEST_PAGE_SIZE),
+    ))
 }
 
 /// Configures the system for booting Linux.

src/vmm/src/devices/virtio/test_utils.rs

@@ -10,7 +10,7 @@ use std::sync::atomic::{AtomicUsize, Ordering};
 
 use crate::devices::virtio::queue::Queue;
 use crate::test_utils::single_region_mem;
-use crate::utils::u64_to_usize;
+use crate::utils::{align_up, u64_to_usize};
 use crate::vstate::memory::{Address, Bytes, GuestAddress, GuestMemoryMmap};
 
 #[macro_export]
@@ -250,7 +250,7 @@ impl<'a> VirtQueue<'a> {
         const USED_ALIGN: u64 = 4;
 
         let mut x = avail.end().0;
-        x = (x + USED_ALIGN - 1) & !(USED_ALIGN - 1);
+        x = align_up(x, USED_ALIGN);
 
         let used = VirtqUsed::new(GuestAddress(x), mem, qsize, u64_to_usize(USED_ALIGN));
 

src/vmm/src/utils/mod.rs

@@ -53,3 +53,15 @@ pub const fn wrap_usize_to_u32(num: usize) -> Wrapping<u32> {
 pub const fn mib_to_bytes(mib: usize) -> usize {
     mib << MIB_TO_BYTES_SHIFT
 }
+
+/// Align address up to the aligment.
+pub const fn align_up(addr: u64, align: u64) -> u64 {
+    debug_assert!(align != 0);
+    (addr + align - 1) & !(align - 1)
+}
+
+/// Align address down to the aligment.
+pub const fn align_down(addr: u64, align: u64) -> u64 {
+    debug_assert!(align != 0);
+    addr & !(align - 1)
+}