Find present pages in xarray

kernel/libs/xarray/src/cursor.rs

@@ -12,7 +12,7 @@ use ostd::{
 use crate::{
     SLOT_SIZE, XArray, XLockGuard,
     entry::NodeEntryRef,
-    mark::{NoneMark, XMark},
+    mark::{NoneMark, PRESENT_MARK, XMark},
     node::{Height, XNode},
 };
 
@@ -167,6 +167,84 @@ impl<'a, P: NonNullPtr + Send + Sync, M> Cursor<'a, P, M> {
         }
     }
 
+    /// Finds the next marked item and moves the cursor to it.
+    ///
+    /// This method will return the index of the marked item, or `None` if no such item exists.
+    fn find_marked(&mut self, mark: usize) -> Option<u64> {
+        let mut index = self.index.checked_add(1)?;
+        let (mut current_node, mut operation_offset) =
+            if let Some((node, offset)) = core::mem::take(&mut self.state).into_node() {
+                (node, offset + 1)
+            } else if let Some(node) = self.xa.head.read_with(self.guard)
+                && index <= node.height().max_index()
+            {
+                let offset = node.entry_offset(index);
+                (node, offset)
+            } else {
+                self.reset();
+                return None;
+            };
+
+        loop {
+            // If we reach the end of the current node, go to its parent node.
+            if operation_offset == SLOT_SIZE as u8 {
+                let Some(parent_node) = current_node.deref_target().parent(self.guard) else {
+                    self.reset();
+                    return None;
+                };
+
+                operation_offset = current_node.offset_in_parent() + 1;
+                current_node = parent_node;
+                continue;
+            }
+
+            // Otherwise, check whether the remaining children contain marked items.
+            let new_operation_offset = current_node
+                .next_marked(operation_offset, mark)
+                .unwrap_or(SLOT_SIZE as u8);
+            let gap = (new_operation_offset - operation_offset) as u64;
+            if gap != 0 {
+                let index_step = current_node.height().index_step();
+                // `index_step` is a power of two. In this case, we want to clear the lower bits
+                // since we should start from the beginning of the next child.
+                let Some(new_index) = (index & !(index_step - 1)).checked_add(gap * index_step)
+                else {
+                    self.reset();
+                    return None;
+                };
+
+                index = new_index;
+                operation_offset = new_operation_offset;
+                if new_operation_offset == SLOT_SIZE as u8 {
+                    continue;
+                }
+            }
+
+            // Due to race conditions, the child may no longer exist. If so, we will retry.
+            let Some(child_node) = current_node
+                .deref_target()
+                .entry_with(self.guard, operation_offset)
+            else {
+                continue;
+            };
+
+            // If we're not at the leaf, then we continue looking down.
+            if !current_node.is_leaf() {
+                current_node = child_node.left().unwrap();
+                operation_offset = current_node.entry_offset(index);
+                continue;
+            }
+
+            // If we're at the leaf, then we have found an item.
+            self.index = index;
+            self.state = CursorState::AtNode {
+                node: current_node,
+                operation_offset,
+            };
+            return Some(index);
+        }
+    }
+
     /**** Public ****/
 
     /// Loads the item at the target index.
@@ -226,6 +304,22 @@ impl<'a, P: NonNullPtr + Send + Sync, M> Cursor<'a, P, M> {
         self.state.move_to(current_node, self.index);
         self.continue_traverse_to_target();
     }
+
+    /// Moves the cursor to the next present item.
+    ///
+    /// If an item is present after the cursor's current index, the cursor will be
+    /// positioned on the corresponding leaf node and the index of the item will be
+    /// returned.
+    ///
+    /// Otherwise, the cursor will stay where it is and a [`None`] will be returned.
+    ///
+    /// Note that this method cannot provide an atomic guarantee for the following
+    /// operations on [`Cursor`]. For example, [`Self::load`] may fail due to
+    /// concurrent removals. If this is a concern, use [`CursorMut`] to avoid the
+    /// issue.
+    pub fn next_present(&mut self) -> Option<u64> {
+        self.find_marked(PRESENT_MARK)
+    }
 }
 
 impl<P: NonNullPtr + Send + Sync, M: Into<XMark>> Cursor<'_, P, M> {
@@ -239,6 +333,22 @@ impl<P: NonNullPtr + Send + Sync, M: Into<XMark>> Cursor<'_, P, M> {
             .map(|(node, off)| node.is_marked(off, mark.into().index()))
             .unwrap_or(false)
     }
+
+    /// Moves the cursor to the next marked item.
+    ///
+    /// If an item is marked after the cursor's current index, the cursor will be
+    /// positioned on the corresponding leaf node and the index of the item will be
+    /// returned.
+    ///
+    /// Otherwise, the cursor will stay where it is and a [`None`] will be returned.
+    ///
+    /// Note that this method cannot provide an atomic guarantee for the following
+    /// operations on [`Cursor`]. For example, [`Self::load`] may return an item that
+    /// is not marked due to concurrent operations. If this is a concern, use
+    /// [`CursorMut`] to avoid the issue.
+    pub fn next_marked(&mut self, mark: M) -> Option<u64> {
+        self.find_marked(mark.into().index())
+    }
 }
 
 /// A `CursorMut` can traverse in the [`XArray`] by setting or increasing the

kernel/libs/xarray/src/mark.rs

@@ -28,12 +28,12 @@ impl Mark {
 
     pub(super) fn update(&self, _guard: XLockGuard, offset: u8, set: bool) -> bool {
         let old_val = self.inner.load(Ordering::Acquire);
+
         let new_val = if set {
             old_val | (1 << offset as u64)
         } else {
             old_val & !(1 << offset as u64)
         };
-
         self.inner.store(new_val, Ordering::Release);
 
         old_val != new_val
@@ -46,6 +46,15 @@ impl Mark {
     pub(super) fn is_clear(&self) -> bool {
         self.inner.load(Ordering::Acquire) == 0
     }
+
+    pub(super) fn next_marked(&self, offset: u8) -> Option<u8> {
+        let high_marks = self.inner.load(Ordering::Acquire) >> offset;
+        if high_marks == 0 {
+            None
+        } else {
+            Some(offset + (high_marks.trailing_zeros() as u8))
+        }
+    }
 }
 
 /// The mark type used in the [`XArray`].
@@ -65,7 +74,12 @@ pub enum XMark {
     Mark2,
 }
 
-pub(super) const NUM_MARKS: usize = 3;
+pub(super) const NUM_MARKS: usize = 4;
+
+/// A mark carried by every [`XArray`] item.
+///
+/// This is is for internal use only. `XArray` users cannot set or unset this mark.
+pub(super) const PRESENT_MARK: usize = 3;
 
 impl XMark {
     /// Maps the `XMark` to an index in the range 0 to 2.

kernel/libs/xarray/src/node.rs

@@ -15,7 +15,7 @@ use ostd::{
 use crate::{
     BITS_PER_LAYER, SLOT_MASK, SLOT_SIZE, XLockGuard,
     entry::{NodeEntry, NodeEntryRef, XEntry, XEntryRef},
-    mark::{Mark, NUM_MARKS},
+    mark::{Mark, NUM_MARKS, PRESENT_MARK},
 };
 
 /// The height of an `XNode` within an `XArray`.
@@ -95,6 +95,11 @@ impl Height {
     pub(super) fn max_index(&self) -> u64 {
         ((SLOT_SIZE as u64) << self.height_shift()) - 1
     }
+
+    /// Calculates the index step representing one offset at the current height.
+    pub(super) fn index_step(&self) -> u64 {
+        1 << self.height_shift()
+    }
 }
 
 /// The `XNode` is the intermediate node in the tree-like structure of the `XArray`.
@@ -184,6 +189,10 @@ impl<P: NonNullPtr + Send + Sync> XNode<P> {
     pub(super) fn is_leaf(&self) -> bool {
         self.height == 1
     }
+
+    pub(super) fn next_marked(&self, offset: u8, mark: usize) -> Option<u8> {
+        self.marks[mark].next_marked(offset)
+    }
 }
 
 impl<P: NonNullPtr + Send + Sync> XNode<P> {
@@ -229,6 +238,7 @@ impl<P: NonNullPtr + Send + Sync> XNode<P> {
             }
             _ => false,
         };
+        let is_new_item = matches!(&entry, Some(Either::Right(_)));
 
         self.slots[offset as usize].update(entry);
 
@@ -236,7 +246,11 @@ impl<P: NonNullPtr + Send + Sync> XNode<P> {
             self.update_mark(guard, offset);
         } else {
             for i in 0..NUM_MARKS {
-                self.unset_mark(guard, offset, i);
+                if i == PRESENT_MARK && is_new_item {
+                    self.set_mark(guard, offset, i);
+                } else {
+                    self.unset_mark(guard, offset, i);
+                }
             }
         }
     }

kernel/libs/xarray/src/test.rs

@@ -164,16 +164,84 @@ fn cursor_store_sparse() {
 }
 
 #[ktest]
-fn set_mark() {
-    let xarray_arc: XArray<Arc<u32>, XMark> = XArray::new();
-    init_continuous_with_arc(&xarray_arc, n!(100));
-
+fn cursor_next_present_single() {
+    let xarray_arc: XArray<Arc<u32>> = XArray::new();
     let mut locked_xarray = xarray_arc.lock();
+    locked_xarray.store(2, Arc::new(2));
+
+    let mut cursor = locked_xarray.cursor(0);
+    for i in 0..2 {
+        cursor.reset_to(i);
+        assert_eq!(cursor.next_present(), Some(2));
+        assert_eq!(*cursor.load().unwrap().as_ref(), 2);
+    }
+    for i in 2..n!(100) {
+        cursor.reset_to(i);
+        assert_eq!(cursor.next_present(), None);
+    }
+}
+
+#[ktest]
+fn cursor_next_present_sparse() {
+    let xarray_arc: XArray<Arc<u32>> = XArray::new();
+    let mut locked_xarray = xarray_arc.lock();
+    locked_xarray.store(0, Arc::new(1));
+    locked_xarray.store(n!(10), Arc::new(2));
+    locked_xarray.store(n!(100), Arc::new(3));
+
+    let mut cursor = locked_xarray.cursor(0);
+    for i in 0..n!(10) {
+        cursor.reset_to(i);
+        let _ = cursor.load();
+        assert_eq!(cursor.next_present(), Some(n!(10)));
+        assert_eq!(*cursor.load().unwrap().as_ref(), 2);
+    }
+    for i in n!(10)..n!(100) {
+        cursor.reset_to(i);
+        assert_eq!(cursor.next_present(), Some(n!(100)));
+        assert_eq!(*cursor.load().unwrap().as_ref(), 3);
+    }
+}
+
+#[ktest]
+fn cursor_next_present_continuous() {
+    let xarray_arc: XArray<Arc<u32>> = XArray::new();
+    let mut locked_xarray = xarray_arc.lock();
+
+    let mut cursor = locked_xarray.cursor_mut(0);
+    for i in 0..n!(100) {
+        let value = Arc::new(i);
+        cursor.store(value);
+        cursor.next();
+    }
+
+    cursor.reset_to(0);
+    for i in 0..(n!(100) - 1) {
+        assert_eq!(cursor.next_present(), Some(i as u64 + 1));
+        assert_eq!(*cursor.load().unwrap().as_ref(), i + 1);
+    }
+    assert_eq!(cursor.next_present(), None);
+    assert_eq!(*cursor.load().unwrap().as_ref(), n!(100) - 1);
+}
+
+fn init_continuous_with_marks(xarray: &XArray<Arc<u32>, XMark>) {
+    init_continuous_with_arc(xarray, n!(100));
+
+    let mut locked_xarray = xarray.lock();
     let mut cursor = locked_xarray.cursor_mut(n!(10));
     cursor.set_mark(XMark::Mark0).unwrap();
     cursor.set_mark(XMark::Mark1).unwrap();
     cursor.reset_to(n!(20));
     cursor.set_mark(XMark::Mark1).unwrap();
+}
+
+#[ktest]
+fn set_mark() {
+    let xarray_arc: XArray<Arc<u32>, XMark> = XArray::new();
+    init_continuous_with_marks(&xarray_arc);
+
+    let guard = disable_preempt();
+    let mut cursor = xarray_arc.cursor(&guard, 0);
 
     cursor.reset_to(n!(10));
     let value1_mark0 = cursor.is_marked(XMark::Mark0);
@@ -212,6 +280,32 @@ fn unset_mark() {
     assert!(!value1_mark2);
 }
 
+#[ktest]
+fn next_marked() {
+    let xarray_arc: XArray<Arc<u32>, XMark> = XArray::new();
+    init_continuous_with_marks(&xarray_arc);
+
+    let guard = disable_preempt();
+    let mut cursor = xarray_arc.cursor(&guard, 0);
+
+    assert_eq!(cursor.next_marked(XMark::Mark0), Some(n!(10)));
+    assert_eq!(cursor.next_marked(XMark::Mark0), None);
+
+    cursor.reset_to(1);
+
+    assert_eq!(cursor.next_marked(XMark::Mark1), Some(n!(10)));
+    assert_eq!(cursor.next_marked(XMark::Mark1), Some(n!(20)));
+    assert_eq!(*cursor.load().unwrap().as_ref(), n!(20));
+    assert_eq!(cursor.next_marked(XMark::Mark1), None);
+
+    cursor.reset_to(2);
+
+    assert_eq!(cursor.next_marked(XMark::Mark1), Some(n!(10)));
+    assert_eq!(*cursor.load().unwrap().as_ref(), n!(10));
+    assert_eq!(cursor.next_marked(XMark::Mark1), Some(n!(20)));
+    assert_eq!(cursor.next_marked(XMark::Mark1), None);
+}
+
 #[ktest]
 fn mark_overflow() {
     let xarray_arc: XArray<Arc<u32>, XMark> = XArray::new();

kernel/src/vm/vmo/mod.rs

@@ -416,19 +416,24 @@ impl Vmo {
         let mut cursor = locked_pages.cursor_mut(page_idx_range.start as u64);
 
         let Some(pager) = &self.pager else {
-            for _ in page_idx_range {
+            cursor.remove();
+            while let Some(page_idx) = cursor.next_present()
+                && page_idx < page_idx_range.end as u64
+            {
                 cursor.remove();
-                cursor.next();
             }
             return Ok(());
         };
 
         let mut removed_page_idx = Vec::new();
-        for page_idx in page_idx_range {
-            if cursor.remove().is_some() {
-                removed_page_idx.push(page_idx);
-            }
-            cursor.next();
+        if cursor.remove().is_some() {
+            removed_page_idx.push(page_idx_range.start);
+        }
+        while let Some(page_idx) = cursor.next_present()
+            && page_idx < page_idx_range.end as u64
+        {
+            removed_page_idx.push(page_idx as usize);
+            cursor.remove();
         }
 
         drop(locked_pages);