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

src/scheme/memory.rs

@@ -6,7 +6,7 @@ use rmm::PhysicalAddress;
 use crate::{
     context::{
         file::InternalFlags,
-        memory::{handle_notify_files, AddrSpace, AddrSpaceWrapper, Grant, PageSpan},
+        memory::{handle_notify_files, AddrSpace, AddrSpaceWrapper, Grant, PageSpan, UnmapVec},
     },
     memory::{free_frames, used_frames, Frame, VirtualAddress, PAGE_SIZE},
     sync::CleanLockToken,
@@ -78,7 +78,7 @@ impl MemoryScheme {
             .ok_or(Error::new(EINVAL))?;
         let page_count = NonZeroUsize::new(span.count).ok_or(Error::new(EINVAL))?;
 
-        let mut notify_files = Vec::new();
+        let mut notify_files = UnmapVec::new();
 
         if is_phys_contiguous && map.flags.contains(MapFlags::MAP_SHARED) {
             // TODO: Should this be supported?

src/scheme/proc.rs

@@ -3,7 +3,7 @@ use crate::{
         self,
         context::{HardBlockedReason, LockedFdTbl, SignalState},
         file::InternalFlags,
-        memory::{handle_notify_files, AddrSpace, AddrSpaceWrapper, Grant, PageSpan},
+        memory::{handle_notify_files, AddrSpace, AddrSpaceWrapper, Grant, PageSpan, UnmapVec},
         Context, ContextLock, Status,
     },
     memory::{Page, VirtualAddress, PAGE_SIZE},
@@ -521,7 +521,7 @@ impl KernelScheme for ProcScheme {
 
                 let src_page_count = NonZeroUsize::new(src_span.count).ok_or(Error::new(EINVAL))?;
 
-                let mut notify_files = Vec::new();
+                let mut notify_files = UnmapVec::new();
 
                 // TODO: Validate flags
                 let result_base = if consume {

src/scheme/user.rs

@@ -7,6 +7,7 @@ use core::{
     num::NonZeroUsize,
 };
 use slab::Slab;
+use smallvec::SmallVec;
 use syscall::{
     schemev2::{Cqe, CqeOpcode, Opcode, Sqe, SqeFlags},
     CallFlags, FmoveFdFlags, FobtainFdFlags, MunmapFlags, RecvFdFlags, SchemeSocketCall,
@@ -20,7 +21,7 @@ use crate::{
         file::{FileDescription, FileDescriptor, InternalFlags, LockedFileDescription},
         memory::{
             AddrSpace, AddrSpaceWrapper, BorrowedFmapSource, Grant, GrantFileRef, MmapMode,
-            PageSpan, DANGLING,
+            PageSpan, UnmapResult, UnmapVec, DANGLING,
         },
         BorrowedHtBuf, ContextLock, PreemptGuard, PreemptGuardL1, Status,
     },
@@ -1117,7 +1118,7 @@ impl UserInner {
         };
 
         let page_count_nz = NonZeroUsize::new(page_count).expect("already validated map.size != 0");
-        let mut notify_files = Vec::new();
+        let mut notify_files = UnmapVec::new();
         // TODO: Not a Lock ordering violation
         // we've checked Arc::ptr_eq(&src_address_space, &dst_addr_space) before,
         // but it's difficult to apply src.arquire_rewrite