Add FrameCount struct, improve zeroing page performance

src/allocator/frame/buddy.rs

@@ -7,6 +7,7 @@ use crate::{
     Arch,
     BumpAllocator,
     FrameAllocator,
+    FrameCount,
     PhysicalAddress,
     VirtualAddress,
 };
@@ -49,7 +50,7 @@ impl<A: Arch> BuddyAllocator<A> {
 
     pub unsafe fn new(mut bump_allocator: BumpAllocator<A>, clear_frees: bool) -> Option<Self> {
         // Allocate buddy table
-        let table_phys = bump_allocator.allocate(1)?;
+        let table_phys = bump_allocator.allocate_one()?;
         let table_virt = A::phys_to_virt(table_phys);
         for i in 0 .. (A::PAGE_SIZE / mem::size_of::<BuddyEntry>()) {
             let virt = table_virt.add(i * mem::size_of::<BuddyEntry>());
@@ -106,7 +107,7 @@ impl<A: Arch> BuddyAllocator<A> {
                 println!("  map pages: {}", map_pages);
 
                 for _ in 0 .. map_pages {
-                    let map_phys = bump_allocator.allocate(1)?;
+                    let map_phys = bump_allocator.allocate_one()?;
                     let map_virt = A::phys_to_virt(map_phys);
                     for i in 0..Self::MAP_PAGE_BYTES {
                         A::write(map_virt.add(i), 0);
@@ -127,7 +128,7 @@ impl<A: Arch> BuddyAllocator<A> {
             if area_offset < area.size {
                 let area_base = area.base.add(area_offset);
                 let area_size = area.size - area_offset;
-                allocator.free(area_base, area_size);
+                allocator.free(area_base, FrameCount::new(area_size / A::PAGE_SIZE));
                 area_offset = 0;
             } else {
                 area_offset -= area.size;
@@ -139,9 +140,9 @@ impl<A: Arch> BuddyAllocator<A> {
 }
 
 impl<A: Arch> FrameAllocator for BuddyAllocator<A> {
-    unsafe fn allocate(&mut self, size: usize) -> Option<PhysicalAddress> {
+    unsafe fn allocate(&mut self, count: FrameCount) -> Option<PhysicalAddress> {
         //TODO: support other sizes
-        if size != A::PAGE_SIZE {
+        if count.data() != 1 {
             return None;
         }
 
@@ -177,27 +178,21 @@ impl<A: Arch> FrameAllocator for BuddyAllocator<A> {
         None
     }
 
-    unsafe fn free(&mut self, base: PhysicalAddress, size: usize) {
-        if self.clear_frees {
-            // Zero freed pages for security, also ensures all allocs are zerod
-            //TODO: improve performance
-            //TODO: assumes linear physical mapping
-            let mut zero_virt = A::phys_to_virt(base);
-            let zero_end = zero_virt.add(size);
-            while zero_virt < zero_end {
-                A::write(zero_virt, 0usize);
-                zero_virt = zero_virt.add(mem::size_of::<usize>());
-            }
-        }
-
+    unsafe fn free(&mut self, base: PhysicalAddress, count: FrameCount) {
+        let size = count.data() * A::PAGE_SIZE;
         for i in 0 .. Self::BUDDY_ENTRIES {
             let virt = self.table_virt.add(i * mem::size_of::<BuddyEntry>());
             let entry = A::read::<BuddyEntry>(virt);
             if base >= entry.base && base.add(size) <= entry.base.add(entry.size) {
                 //TODO: Correct logic
-                let pages = size / A::PAGE_SIZE;
-                for page in 0 .. pages {
+                for page in 0 .. count.data() {
                     let page_base = base.add(page * A::PAGE_SIZE);
+
+                    if self.clear_frees {
+                        let page_virt = A::phys_to_virt(page_base);
+                        A::write_bytes(page_virt, 0, A::PAGE_SIZE);
+                    }
+
                     let index = (page_base.data() - entry.base.data()) / A::PAGE_SIZE;
                     let mut map_page = index / Self::MAP_PAGE_BITS;
                     let map_bit = index % Self::MAP_PAGE_BITS;

src/allocator/frame/bump.rs

@@ -3,6 +3,7 @@ use core::marker::PhantomData;
 use crate::{
     Arch,
     FrameAllocator,
+    FrameCount,
     MemoryArea,
     PhysicalAddress,
 };
@@ -32,9 +33,9 @@ impl<A: Arch> BumpAllocator<A> {
 }
 
 impl<A: Arch> FrameAllocator for BumpAllocator<A> {
-    unsafe fn allocate(&mut self, count: usize) -> Option<PhysicalAddress> {
+    unsafe fn allocate(&mut self, count: FrameCount) -> Option<PhysicalAddress> {
         //TODO: support allocation of multiple pages
-        if count != 1 {
+        if count.data() != 1 {
             return None;
         }
 
@@ -49,7 +50,7 @@ impl<A: Arch> FrameAllocator for BumpAllocator<A> {
         None
     }
 
-    unsafe fn free(&mut self, _address: PhysicalAddress, _count: usize) {
+    unsafe fn free(&mut self, _address: PhysicalAddress, _count: FrameCount) {
         unimplemented!();
     }
 }

src/allocator/frame/mod.rs

@@ -6,7 +6,30 @@ pub use self::bump::*;
 mod buddy;
 mod bump;
 
-pub trait FrameAllocator {
-    unsafe fn allocate(&mut self, count: usize) -> Option<PhysicalAddress>;
-    unsafe fn free(&mut self, address: PhysicalAddress, count: usize);
+#[derive(Clone, Copy, Debug)]
+#[repr(transparent)]
+pub struct FrameCount(usize);
+
+impl FrameCount {
+    pub fn new(count: usize) -> Self {
+        Self(count)
+    }
+
+    pub fn data(&self) -> usize {
+        self.0
+    }
+}
+
+pub trait FrameAllocator {
+    unsafe fn allocate(&mut self, count: FrameCount) -> Option<PhysicalAddress>;
+
+    unsafe fn free(&mut self, address: PhysicalAddress, count: FrameCount);
+
+    unsafe fn allocate_one(&mut self) -> Option<PhysicalAddress> {
+        self.allocate(FrameCount::new(1))
+    }
+
+    unsafe fn free_one(&mut self, address: PhysicalAddress) {
+        self.free(address, FrameCount::new(1));
+    }
 }

src/arch/emulate.rs

@@ -74,6 +74,11 @@ impl Arch for EmulateArch {
         MACHINE.as_mut().unwrap().write(address, value)
     }
 
+    #[inline(always)]
+    unsafe fn write_bytes(address: VirtualAddress, value: u8, count: usize) {
+        MACHINE.as_mut().unwrap().write_bytes(address, value, count)
+    }
+
     #[inline(always)]
     unsafe fn invalidate(address: VirtualAddress) {
         MACHINE.as_mut().unwrap().invalidate(address);
@@ -149,6 +154,16 @@ impl<A: Arch> Machine<A> {
         }
     }
 
+    fn write_phys_bytes(&mut self, phys: PhysicalAddress, value: u8, count: usize) {
+        if phys.add(count).data() <= self.memory.len() {
+            unsafe {
+                ptr::write_bytes(self.memory.as_mut_ptr().add(phys.data()) as *mut u8, value, count);
+            }
+        } else {
+            panic!("write_phys_bytes: 0x{:X} count 0x{:X} outside of memory", phys.data(), count);
+        }
+    }
+
     fn translate(&self, virt: VirtualAddress) -> Option<(PhysicalAddress, usize)> {
         let virt_data = virt.data();
         let page = virt_data & A::PAGE_ADDRESS_MASK;
@@ -194,6 +209,24 @@ impl<A: Arch> Machine<A> {
         }
     }
 
+    fn write_bytes(&mut self, virt: VirtualAddress, value: u8, count: usize) {
+        //TODO: allow writing past page boundaries
+        let virt_data = virt.data();
+        if (virt_data & A::PAGE_ADDRESS_MASK) != ((virt_data + (count - 1)) & A::PAGE_ADDRESS_MASK) {
+            panic!("write_bytes: 0x{:X} count 0x{:X} passes page boundary", virt_data, count);
+        }
+
+        if let Some((phys, flags)) = self.translate(virt) {
+            if flags & A::ENTRY_FLAG_WRITABLE != 0 {
+                self.write_phys_bytes(phys, value, count);
+            } else {
+                panic!("write_bytes: 0x{:X} count 0x{:X} not writable", virt_data, count);
+            }
+        } else {
+            panic!("write_bytes: 0x{:X} count 0x{:X} not present", virt_data, count);
+        }
+    }
+
     fn invalidate(&mut self, _address: VirtualAddress) {
         unimplemented!();
     }

src/arch/mod.rs

@@ -55,6 +55,11 @@ pub trait Arch {
         ptr::write(address.data() as *mut T, value)
     }
 
+    #[inline(always)]
+    unsafe fn write_bytes(address: VirtualAddress, value: u8, count: usize) {
+        ptr::write_bytes(address.data() as *mut u8, value, count)
+    }
+
     unsafe fn invalidate(address: VirtualAddress);
 
     #[inline(always)]

src/main.rs

@@ -202,11 +202,12 @@ unsafe fn new_tables<A: Arch>(areas: &'static [MemoryArea]) {
 
     let mut allocator = BuddyAllocator::<A>::new(bump_allocator, true).unwrap();
     for i in 0..16 {
-        let phys_opt = allocator.allocate(4 * KILOBYTE);
-        println!("4 KB page {}: {:X?}", i, phys_opt);
+        let phys_opt = allocator.allocate_one();
+        println!("page {}: {:X?}", i, phys_opt);
         if i % 2 == 0 {
             if let Some(phys) = phys_opt {
-                allocator.free(phys, 4 * KILOBYTE);
+                println!("free {}: {:X?}", i, phys_opt);
+                allocator.free_one(phys);
             }
         }
     }

src/page/mapper.rs

@@ -17,7 +17,7 @@ pub struct PageMapper<'f, A, F> {
 
 impl<'f, A: Arch, F: FrameAllocator> PageMapper<'f, A, F> {
     pub unsafe fn new(allocator: &'f mut F) -> Option<Self> {
-        let table_addr = allocator.allocate(1)?;
+        let table_addr = allocator.allocate_one()?;
         Some(Self {
             table_addr,
             allocator,
@@ -46,7 +46,7 @@ impl<'f, A: Arch, F: FrameAllocator> PageMapper<'f, A, F> {
                 let next = match next_opt {
                     Some(some) => some,
                     None => {
-                        let phys = self.allocator.allocate(1)?;
+                        let phys = self.allocator.allocate_one()?;
                         //TODO: correct flags?
                         table.set_entry(i, PageEntry::new(phys.data() | A::ENTRY_FLAG_WRITABLE | A::ENTRY_FLAG_PRESENT));
                         table.next(i)?