diff --git a/kernel/libs/xarray/src/cursor.rs b/kernel/libs/xarray/src/cursor.rs index 523304a14..c0a74e8e4 100644 --- a/kernel/libs/xarray/src/cursor.rs +++ b/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 { + 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 { + self.find_marked(PRESENT_MARK) + } } impl> Cursor<'_, P, M> { @@ -239,6 +333,22 @@ impl> 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 { + self.find_marked(mark.into().index()) + } } /// A `CursorMut` can traverse in the [`XArray`] by setting or increasing the diff --git a/kernel/libs/xarray/src/mark.rs b/kernel/libs/xarray/src/mark.rs index 72631b199..c747b33b0 100644 --- a/kernel/libs/xarray/src/mark.rs +++ b/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 { + 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. diff --git a/kernel/libs/xarray/src/node.rs b/kernel/libs/xarray/src/node.rs index d167f005d..e290d0b43 100644 --- a/kernel/libs/xarray/src/node.rs +++ b/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 XNode

{ pub(super) fn is_leaf(&self) -> bool { self.height == 1 } + + pub(super) fn next_marked(&self, offset: u8, mark: usize) -> Option { + self.marks[mark].next_marked(offset) + } } impl XNode

{ @@ -229,6 +238,7 @@ impl XNode

{ } _ => false, }; + let is_new_item = matches!(&entry, Some(Either::Right(_))); self.slots[offset as usize].update(entry); @@ -236,7 +246,11 @@ impl XNode

{ 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); + } } } } diff --git a/kernel/libs/xarray/src/test.rs b/kernel/libs/xarray/src/test.rs index 9c92c52c9..10ec2b713 100644 --- a/kernel/libs/xarray/src/test.rs +++ b/kernel/libs/xarray/src/test.rs @@ -164,16 +164,84 @@ fn cursor_store_sparse() { } #[ktest] -fn set_mark() { - let xarray_arc: XArray, XMark> = XArray::new(); - init_continuous_with_arc(&xarray_arc, n!(100)); - +fn cursor_next_present_single() { + let xarray_arc: XArray> = 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> = 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> = 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, 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, 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, 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, XMark> = XArray::new(); diff --git a/kernel/src/vm/vmo/mod.rs b/kernel/src/vm/vmo/mod.rs index b42ad68e6..731335e67 100644 --- a/kernel/src/vm/vmo/mod.rs +++ b/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);