Make AddrSpaceWrapper lock lower without solving borrow checker

src/context/context.rs

@@ -22,7 +22,7 @@ use crate::{
     },
     percpu::PercpuBlock,
     scheme::{CallerCtx, FileHandle, SchemeId},
-    sync::{CleanLockToken, LockToken, RwLock, L1, L2, L3, L4, L5},
+    sync::{CleanLockToken, LockToken, RwLock, L1, L4, L5},
     syscall::usercopy::UserSliceRw,
 };
 
@@ -383,7 +383,7 @@ impl Context {
     pub fn set_addr_space(
         &mut self,
         addr_space: Option<Arc<AddrSpaceWrapper>>,
-        token: LockToken<L2>,
+        token: LockToken<L4>,
     ) -> Option<Arc<AddrSpaceWrapper>> {
         if let (Some(old), Some(new)) = (&self.addr_space, &addr_space)
             && Arc::ptr_eq(old, new)
@@ -490,7 +490,7 @@ pub struct BorrowedHtBuf {
     head_and_not_tail: bool,
 }
 impl BorrowedHtBuf {
-    pub fn head_locked(token: LockToken<L3>) -> Result<Self> {
+    pub fn head_locked(token: LockToken<L5>) -> Result<Self> {
         let current = context::current();
         let frame = &mut current.write(token).syscall_head;
         match mem::replace(frame, SyscallFrame::Dummy) {
@@ -506,7 +506,10 @@ impl BorrowedHtBuf {
             SyscallFrame::Used { .. } | SyscallFrame::Dummy => Err(Error::new(EAGAIN)),
         }
     }
-    pub fn tail_locked(token: LockToken<L3>) -> Result<Self> {
+    pub fn tail(token: &mut CleanLockToken) -> Result<Self> {
+        Self::tail_locked(token.downgrade())
+    }
+    pub fn tail_locked(token: LockToken<L5>) -> Result<Self> {
         let current = context::current();
         let frame = &mut current.write(token).syscall_tail;
         match mem::replace(frame, SyscallFrame::Dummy) {

src/context/memory.rs

@@ -23,7 +23,7 @@ use crate::{
     scheme::{self, KernelSchemes},
     sync::{
         CleanLockToken, LockToken, RwLock, RwLockReadGuard, RwLockUpgradableGuard,
-        RwLockWriteGuard, L2, L3,
+        RwLockWriteGuard, L4, L5,
     },
 };
 
@@ -84,7 +84,7 @@ impl UnmapResult {
 
 #[derive(Debug)]
 pub struct AddrSpaceWrapper {
-    pub inner: RwLock<L3, AddrSpace>,
+    pub inner: RwLock<L5, AddrSpace>,
     pub tlb_ack: AtomicU32,
     pub used_by: LogicalCpuSet,
 }
@@ -98,32 +98,32 @@ impl AddrSpaceWrapper {
     }
     pub fn acquire_read<'a>(
         &'a self,
-        lock_token: LockToken<'a, L2>,
-    ) -> RwLockReadGuard<'a, L3, AddrSpace> {
+        lock_token: LockToken<'a, L4>,
+    ) -> RwLockReadGuard<'a, L5, AddrSpace> {
         self.inner.read(lock_token)
     }
     pub fn acquire_upgradeable_read<'a>(
         &'a self,
-        lock_token: LockToken<'a, L2>,
-    ) -> RwLockUpgradableGuard<'a, L3, AddrSpace> {
+        lock_token: LockToken<'a, L4>,
+    ) -> RwLockUpgradableGuard<'a, L5, AddrSpace> {
         self.inner.upgradeable_read(lock_token)
     }
     pub fn acquire_write<'a>(
         &'a self,
-        lock_token: LockToken<'a, L2>,
-    ) -> RwLockWriteGuard<'a, L3, AddrSpace> {
+        lock_token: LockToken<'a, L4>,
+    ) -> RwLockWriteGuard<'a, L5, AddrSpace> {
         self.inner.write(lock_token)
     }
     pub unsafe fn acquire_reupgradeable_read<'a>(
         &'a self,
-        lock_token: LockToken<'a, L3>,
-    ) -> RwLockUpgradableGuard<'a, L3, AddrSpace> {
+        lock_token: LockToken<'a, L5>,
+    ) -> RwLockUpgradableGuard<'a, L5, AddrSpace> {
         unsafe { self.inner.reupgradeable_read(lock_token) }
     }
     pub unsafe fn acquire_rewrite<'a>(
         &'a self,
-        lock_token: LockToken<'a, L3>,
-    ) -> RwLockWriteGuard<'a, L3, AddrSpace> {
+        lock_token: LockToken<'a, L5>,
+    ) -> RwLockWriteGuard<'a, L5, AddrSpace> {
         unsafe { self.inner.rewrite(lock_token) }
     }
     pub fn into_drop(self, token: &mut CleanLockToken) {
@@ -334,7 +334,7 @@ impl AddrSpaceWrapper {
         new_page_count: usize,
         new_flags: MapFlags,
         mut notify_files_out: Option<&mut Vec<UnmapResult>>,
-        token: LockToken<L3>,
+        token: LockToken<L5>,
     ) -> Result<Page> {
         let dst_lock = self;
         // SAFETY: This is moving data between two AddrSpace. Caller ensures the two is a different AddrSpace
@@ -2452,16 +2452,16 @@ pub fn try_correcting_page_tables(
 
 // TODO: maybe refactor the return type into a struct/typedef?
 #[expect(clippy::type_complexity)]
-/// XXX: This require passing L3 addr_space_guard.
+/// XXX: This require passing L5 addr_space_guard.
 /// Caller must ensure there's no other lock being held at this point.
 /// Caller also need to provide clean token for the new AddrSpace.
 fn correct_inner<'l>(
     addr_space_lock: &'l Arc<AddrSpaceWrapper>,
-    mut addr_space: RwLockWriteGuard<'l, L3, AddrSpace>,
+    mut addr_space: RwLockWriteGuard<'l, L5, AddrSpace>,
     faulting_page: Page,
     access: AccessMode,
     recursion_level: u32,
-) -> Result<(Frame, PageFlush<RmmA>, RwLockWriteGuard<'l, L3, AddrSpace>), PfError> {
+) -> Result<(Frame, PageFlush<RmmA>, RwLockWriteGuard<'l, L5, AddrSpace>), PfError> {
     let mut flusher = Flusher::with_cpu_set(&addr_space_lock.used_by, &addr_space_lock.tlb_ack);
 
     let Some((grant_base, grant_info)) = addr_space.grants.contains(faulting_page) else {
@@ -2756,7 +2756,7 @@ pub struct BorrowedFmapSource<'a> {
     pub mode: MmapMode,
     // TODO: There should be a method that obtains the lock from the guard.
     pub addr_space_lock: &'a Arc<AddrSpaceWrapper>,
-    pub addr_space_guard: RwLockWriteGuard<'a, L3, AddrSpace>,
+    pub addr_space_guard: RwLockWriteGuard<'a, L5, AddrSpace>,
 }
 
 pub fn handle_notify_files(notify_files: Vec<UnmapResult>, token: &mut CleanLockToken) {