Use smallvec for memory and notify_files

src/context/memory.rs

@@ -13,6 +13,7 @@ use core::{
     sync::atomic::{AtomicU32, Ordering},
 };
 use rmm::{Arch as _, PageFlush};
+use smallvec::SmallVec;
 use syscall::{error::*, flag::MapFlags, GrantFlags, MunmapFlags};
 
 use crate::{
@@ -58,6 +59,8 @@ pub struct UnmapResult {
     pub size: usize,
     pub flags: MunmapFlags,
 }
+pub type UnmapVec = SmallVec<[UnmapResult; 16]>;
+
 impl UnmapResult {
     pub fn unmap(mut self, token: &mut CleanLockToken) -> Result<()> {
         let Some(GrantFileRef {
@@ -316,7 +319,7 @@ impl AddrSpaceWrapper {
         requested_span: PageSpan,
         unpin: bool,
         token: &mut CleanLockToken,
-    ) -> Result<Vec<UnmapResult>> {
+    ) -> Result<UnmapVec> {
         let mut guard = self.acquire_write(token.downgrade());
         let guard = &mut *guard;
 
@@ -336,7 +339,7 @@ impl AddrSpaceWrapper {
         requested_dst_base: Option<Page>,
         new_page_count: usize,
         new_flags: MapFlags,
-        mut notify_files_out: Option<&mut Vec<UnmapResult>>,
+        mut notify_files_out: Option<&mut UnmapVec>,
         token: LockToken<L5>,
     ) -> Result<Page> {
         let dst_lock = self;
@@ -413,7 +416,7 @@ impl AddrSpaceWrapper {
 
         if new_page_count < src_span.count {
             let unpin = false;
-            let notify_files: Vec<UnmapResult> = AddrSpace::munmap_inner(
+            let notify_files = AddrSpace::munmap_inner(
                 src_grants,
                 src_mapper,
                 src_flusher,
@@ -593,8 +596,8 @@ impl AddrSpace {
         this_flusher: &mut Flusher,
         mut requested_span: PageSpan,
         unpin: bool,
-    ) -> Result<Vec<UnmapResult>> {
-        let mut notify_files = Vec::new();
+    ) -> Result<UnmapVec> {
+        let mut notify_files = UnmapVec::new();
 
         let next = |grants: &mut UserGrants, span: PageSpan| {
             grants
@@ -692,7 +695,7 @@ impl AddrSpace {
         requested_base_opt: Option<Page>,
         page_count: NonZeroUsize,
         flags: MapFlags,
-        notify_files_out: Option<&mut Vec<UnmapResult>>,
+        notify_files_out: Option<&mut UnmapVec>,
         map: impl FnOnce(Page, PageFlags<RmmA>, &mut PageMapper, &mut Flusher) -> Result<Grant>,
     ) -> Result<Page> {
         assert_eq!(dst_lock.inner.as_mut_ptr(), self as *mut Self);
@@ -2817,7 +2820,7 @@ pub struct BorrowedFmapSource<'a> {
     pub addr_space_guard: RwLockWriteGuard<'a, L5, AddrSpace>,
 }
 
-pub fn handle_notify_files(notify_files: Vec<UnmapResult>, token: &mut CleanLockToken) {
+pub fn handle_notify_files(notify_files: UnmapVec, token: &mut CleanLockToken) {
     for file in notify_files {
         let _ = file.unmap(token);
     }

src/context/switch.rs

@@ -12,12 +12,13 @@ use crate::{
     percpu::PercpuBlock,
     sync::{ArcRwLockWriteGuard, CleanLockToken, L4},
 };
-use alloc::{sync::Arc, vec::Vec};
+use alloc::sync::Arc;
 use core::{
     cell::{Cell, RefCell},
     hint, mem,
     sync::atomic::Ordering,
 };
+use smallvec::SmallVec;
 use syscall::PtraceFlags;
 
 enum UpdateResult {
@@ -320,9 +321,12 @@ pub fn switch(token: &mut CleanLockToken) -> SwitchResult {
     }
 }
 
-fn wakeup_contexts(token: &mut CleanLockToken, switch_time: u128) -> Vec<(usize, WeakContextRef)> {
+fn wakeup_contexts(
+    token: &mut CleanLockToken,
+    switch_time: u128,
+) -> SmallVec<[(usize, WeakContextRef); 16]> {
     // TODO: Optimise this somehow. Perhaps using a separate timer queue?
-    let mut wakeups = Vec::new();
+    let mut wakeups = SmallVec::new();
     let current_context = context::current();
     let Some(idle_contexts) = idle_contexts_try(token.downgrade()) else {
         // other cpus may spawning or killing contexts so let's skip wakeups to avoid contention