feat(memory): handle grant merges when mutating

This commit fixes the following panic when running nodejs:
```
KERNEL PANIC: panicked at src/context/memory.rs:267:18:
grant cannot magically disappear while we hold the lock!
  FP ffff80002021f650: PC ffffffff8005e45b
    FFFFFFFF8005E2E0+017B
    kernel::panic::panic_handler_inner
  FP ffff80002021f730: PC ffffffff80057559
  FP ffff80002021f740: PC ffffffff8009a0ff
  FP ffff80002021f770: PC ffffffff8009c214
  FP ffff80002021f7d0: PC ffffffff80010997
    FFFFFFFF8000FDD0+0BC7
    kernel::context::memory::AddrSpaceWrapper::mprotect
  FP ffff80002021fd40: PC ffffffff8008cb0b
    FFFFFFFF8008CA40+00CB
    kernel::syscall::process::mprotect
  FP ffff80002021fd70: PC ffffffff8006ce6c
    FFFFFFFF8006CBE0+028C
    kernel::syscall::syscall
  FP ffff80002021fe90: PC ffffffff8008d3cf
    FFFFFFFF8008D320+00AF
    __inner_syscall_instruction
  FP ffff80002021ff50: PC ffffffff800830c3
    FFFFFFFF80083080+0043
    kernel::arch::x86_64::interrupt::syscall::syscall_instruction
  00007ffffffffaf0: GUARD PAGE
CPU #1, CID 0xffffff7f8015b910
NAME: /usr/bin/node, DEBUG ID: 74
SYSCALL: mprotect(0x203C0000, 262144, Some(MapFlags(0x0)))
HALT
```

The grant did not magically disappear. When going through the
`grant_span_res` regions, the function adds (and removes) grants to the
`self.grants` tree. The insertion and deletion functions also merge
adjacent grants together when possible. This is an issue since we can no
longer use the keys we established before we started iterating. This
commit modifies the places where `remove` is used in this fashion to use
`remove_containing` instead. The `remove_containing` function will
remove the grant that *contains* the page.

Should it be done this way (requires unstable feature `btree_cursors`)?

Signed-off-by: Anhad Singh <andypython@protonmail.com>

src/context/memory.rs

@@ -4,6 +4,7 @@ use core::{
     cmp,
     fmt::Debug,
     num::NonZeroUsize,
+    ops::Bound,
     sync::atomic::{AtomicU32, Ordering},
 };
 use rmm::{Arch as _, PageFlush};
@@ -263,7 +264,7 @@ impl AddrSpaceWrapper {
 
             let grant = guard
                 .grants
-                .remove(grant_span.base)
+                .remove_containing(grant_span.base)
                 .expect("grant cannot magically disappear while we hold the lock!");
             //info!("Mprotecting {:#?} to {:#?} in {:#?}", grant, flags, grant_span);
             let intersection = grant_span.intersection(requested_span);
@@ -439,7 +440,7 @@ impl AddrSpaceWrapper {
                 .as_mut()
                 .map_or(&mut dst.grants, |(g, _, _)| &mut *g);
             let grant = src_grants
-                .remove(grant_base)
+                .remove_containing(grant_base)
                 .expect("grant cannot disappear");
             let grant_span = PageSpan::new(grant.base, grant.info.page_count());
             let (before, middle, after) = grant
@@ -590,7 +591,7 @@ impl AddrSpace {
             let conflicting_span = conflicting_span_res?;
 
             let mut grant = this_grants
-                .remove(conflicting_span.base)
+                .remove_containing(conflicting_span.base)
                 .expect("conflicting region didn't exist");
             if unpin {
                 grant.info.unpin();
@@ -837,6 +838,7 @@ impl UserGrants {
                 .collect::<BTreeMap<_, _>>(),
         }
     }
+
     /// Returns the grant, if any, which occupies the specified page
     pub fn contains(&self, page: Page) -> Option<(Page, &GrantInfo)> {
         self.inner
@@ -845,6 +847,7 @@ impl UserGrants {
             .filter(|(base, info)| (**base..base.next_by(info.page_count)).contains(&page))
             .map(|(base, info)| (*base, info))
     }
+
     /// Returns an iterator over all grants that occupy some part of the
     /// requested region
     pub fn conflicts(&self, span: PageSpan) -> impl Iterator<Item = (Page, &'_ GrantInfo)> + '_ {
@@ -1018,11 +1021,23 @@ impl UserGrants {
 
         self.inner.insert(grant.base, grant.info);
     }
-    pub fn remove(&mut self, base: Page) -> Option<Grant> {
-        let info = self.inner.remove(&base)?;
-        Self::unreserve(&mut self.holes, base, info.page_count);
-        Some(Grant { base, info })
+
+    pub fn remove_containing(&mut self, page: Page) -> Option<Grant> {
+        // Points to the gap *after* the greatest grant smaller than or equal to `page`.
+        let mut cursor = self.inner.upper_bound_mut(Bound::Included(&page));
+        let Some((&base, info)) = cursor.peek_prev() else {
+            return None;
+        };
+
+        if (base..base.next_by(info.page_count())).contains(&page) {
+            let (base, info) = cursor.remove_prev().unwrap();
+            Self::unreserve(&mut self.holes, base, info.page_count());
+            Some(Grant { base, info })
+        } else {
+            None
+        }
     }
+
     pub fn iter(&self) -> impl Iterator<Item = (Page, &GrantInfo)> + '_ {
         self.inner.iter().map(|(base, info)| (*base, info))
     }