fix test

Author: aarkegz

page_table_multiarch/tests/alloc_tests.rs

@@ -1,19 +1,27 @@
 use std::{
     alloc::{self, Layout},
     cell::RefCell,
-    collections::HashSet,
+    collections::{HashMap, HashSet},
     marker::PhantomData,
+    mem,
 };
 
 use memory_addr::{PhysAddr, VirtAddr};
 use page_table_entry::{GenericPTE, MappingFlags};
 use page_table_multiarch::{PageSize, PageTable64, PagingHandler, PagingMetaData, PagingResult};
 use rand::{Rng, SeedableRng, rngs::SmallRng};
 
-const PAGE_LAYOUT: Layout = unsafe { Layout::from_size_align_unchecked(4096, 4096) };
+/// Creates a layout for allocating `num` pages with alignment of `2^align_pow2`
+/// pages.
+const fn pages_layout(num: usize, align_pow2: usize) -> Layout {
+    unsafe { Layout::from_size_align_unchecked(4096 * num, 4096 * (1 << align_pow2)) }
+}
+
+const PAGE_LAYOUT: Layout = pages_layout(1, 0);
 
 thread_local! {
     static ALLOCATED: RefCell<HashSet<usize>> = RefCell::default();
+    static ALIGN_POW2: RefCell<HashMap<usize, usize>> = RefCell::default();
 }
 
 struct TrackPagingHandler<M: PagingMetaData>(PhantomData<M>);
@@ -29,6 +37,24 @@ impl<M: PagingMetaData> PagingHandler for TrackPagingHandler<M> {
         Some(PhysAddr::from_usize(ptr))
     }
 
+    fn alloc_frames(num: usize, align_pow2: usize) -> Option<PhysAddr> {
+        let layout = pages_layout(num, align_pow2);
+        let ptr = unsafe { alloc::alloc(layout) } as usize;
+        assert!(
+            ptr <= M::PA_MAX_ADDR,
+            "allocated frame address exceeds PA_MAX_ADDR"
+        );
+        ALLOCATED.with_borrow_mut(|it| {
+            for i in 0..num {
+                it.insert(ptr + i * 4096);
+            }
+        });
+        ALIGN_POW2.with_borrow_mut(|it| {
+            it.insert(ptr, align_pow2);
+        });
+        Some(PhysAddr::from_usize(ptr))
+    }
+
     fn dealloc_frame(paddr: PhysAddr) {
         let ptr = paddr.as_usize();
         ALLOCATED.with_borrow_mut(|it| {
@@ -39,6 +65,24 @@ impl<M: PagingMetaData> PagingHandler for TrackPagingHandler<M> {
         }
     }
 
+    fn dealloc_frames(paddr: PhysAddr, num: usize) {
+        let ptr = paddr.as_usize();
+        ALLOCATED.with_borrow_mut(|it| {
+            for i in 0..num {
+                let addr = ptr + i * 4096;
+                assert!(it.remove(&addr), "dealloc a frame that was not allocated");
+            }
+        });
+        let align_pow2 = ALIGN_POW2.with_borrow_mut(|it| {
+            it.remove(&ptr)
+                .expect("dealloc frames that were not allocated")
+        });
+        let layout = pages_layout(num, align_pow2);
+        unsafe {
+            alloc::dealloc(ptr as _, layout);
+        }
+    }
+
     fn phys_to_virt(paddr: PhysAddr) -> VirtAddr {
         assert!(paddr.as_usize() > 0);
         VirtAddr::from_usize(paddr.as_usize())