nordic-dev.net/rust
nordic-dev.net/rust
2
0
Fork 0
mirror of https://github.com/rust-lang/rust.git synced 2025-05-14 02:49:40 +00:00
Code Issues Packages Projects Releases Wiki Activity

BTree: lazily locate leaves in rangeless iterators

Browse Source
This commit is contained in:
Stein Somers 2021-06-01 16:51:56 +02:00
parent aa65b08b1d
commit 35d02e2c6a
2 changed files with 174 additions and 71 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

57
library/alloc/src/collections/btree/map.rs
View File

@ -9,7 +9,7 @@ use core::ops::{Index, RangeBounds};
use core::ptr;
use super::borrow::DormantMutRef;
use super::navigate::LeafRange;
use super::navigate::{LazyLeafRange, LeafRange};
use super::node::{self, marker, ForceResult::*, Handle, NodeRef, Root};
use super::search::SearchResult::*;
@ -278,7 +278,7 @@ where
/// [`iter`]: BTreeMap::iter
#[stable(feature = "rust1", since = "1.0.0")]
pub struct Iter<'a, K: 'a, V: 'a> {
range: Range<'a, K, V>,
range: LazyLeafRange<marker::Immut<'a>, K, V>,
length: usize,
}
@ -296,10 +296,20 @@ impl<K: fmt::Debug, V: fmt::Debug> fmt::Debug for Iter<'_, K, V> {
///
/// [`iter_mut`]: BTreeMap::iter_mut
#[stable(feature = "rust1", since = "1.0.0")]
#[derive(Debug)]
pub struct IterMut<'a, K: 'a, V: 'a> {
range: RangeMut<'a, K, V>,
range: LazyLeafRange<marker::ValMut<'a>, K, V>,
length: usize,
// Be invariant in `K` and `V`
_marker: PhantomData<&'a mut (K, V)>,
}
#[stable(feature = "collection_debug", since = "1.17.0")]
impl<K: fmt::Debug, V: fmt::Debug> fmt::Debug for IterMut<'_, K, V> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let range = Iter { range: self.range.reborrow(), length: self.length };
f.debug_list().entries(range).finish()
}
}
/// An owning iterator over the entries of a `BTreeMap`.
@ -310,7 +320,7 @@ pub struct IterMut<'a, K: 'a, V: 'a> {
/// [`into_iter`]: IntoIterator::into_iter
#[stable(feature = "rust1", since = "1.0.0")]
pub struct IntoIter<K, V> {
range: LeafRange<marker::Dying, K, V>,
range: LazyLeafRange<marker::Dying, K, V>,
length: usize,
}
@ -318,8 +328,7 @@ impl<K, V> IntoIter<K, V> {
/// Returns an iterator of references over the remaining items.
#[inline]
pub(super) fn iter(&self) -> Iter<'_, K, V> {
let range = Range { inner: self.range.reborrow() };
Iter { range: range, length: self.length }
Iter { range: self.range.reborrow(), length: self.length }
}
}
@ -1298,7 +1307,7 @@ impl<'a, K: 'a, V: 'a> Iterator for Iter<'a, K, V> {
None
} else {
self.length -= 1;
Some(unsafe { self.range.inner.next_unchecked() })
Some(unsafe { self.range.next_unchecked() })
}
}
@ -1329,7 +1338,7 @@ impl<'a, K: 'a, V: 'a> DoubleEndedIterator for Iter<'a, K, V> {
None
} else {
self.length -= 1;
Some(unsafe { self.range.inner.next_back_unchecked() })
Some(unsafe { self.range.next_back_unchecked() })
}
}
}
@ -1367,7 +1376,7 @@ impl<'a, K: 'a, V: 'a> Iterator for IterMut<'a, K, V> {
None
} else {
self.length -= 1;
Some(unsafe { self.range.inner.next_unchecked() })
Some(unsafe { self.range.next_unchecked() })
}
}
@ -1395,7 +1404,7 @@ impl<'a, K: 'a, V: 'a> DoubleEndedIterator for IterMut<'a, K, V> {
None
} else {
self.length -= 1;
Some(unsafe { self.range.inner.next_back_unchecked() })
Some(unsafe { self.range.next_back_unchecked() })
}
}
}
@ -1414,7 +1423,7 @@ impl<'a, K, V> IterMut<'a, K, V> {
/// Returns an iterator of references over the remaining items.
#[inline]
pub(super) fn iter(&self) -> Iter<'_, K, V> {
Iter { range: self.range.iter(), length: self.length }
Iter { range: self.range.reborrow(), length: self.length }
}
}
@ -1430,7 +1439,7 @@ impl<K, V> IntoIterator for BTreeMap<K, V> {
IntoIter { range: full_range, length: me.length }
} else {
IntoIter { range: LeafRange::none(), length: 0 }
IntoIter { range: LazyLeafRange::none(), length: 0 }
}
}
}
@ -1888,14 +1897,6 @@ impl<'a, K, V> Iterator for RangeMut<'a, K, V> {
}
}
impl<'a, K, V> RangeMut<'a, K, V> {
/// Returns an iterator of references over the remaining items.
#[inline]
pub(super) fn iter(&self) -> Range<'_, K, V> {
Range { inner: self.inner.reborrow() }
}
}
#[stable(feature = "btree_range", since = "1.17.0")]
impl<'a, K, V> DoubleEndedIterator for RangeMut<'a, K, V> {
fn next_back(&mut self) -> Option<(&'a K, &'a mut V)> {
@ -2038,9 +2039,9 @@ impl<K, V> BTreeMap<K, V> {
if let Some(root) = &self.root {
let full_range = root.reborrow().full_range();
Iter { range: Range { inner: full_range }, length: self.length }
Iter { range: full_range, length: self.length }
} else {
Iter { range: Range { inner: LeafRange::none() }, length: 0 }
Iter { range: LazyLeafRange::none(), length: 0 }
}
}
@ -2070,15 +2071,9 @@ impl<K, V> BTreeMap<K, V> {
if let Some(root) = &mut self.root {
let full_range = root.borrow_valmut().full_range();
IterMut {
range: RangeMut { inner: full_range, _marker: PhantomData },
length: self.length,
}
IterMut { range: full_range, length: self.length, _marker: PhantomData }
} else {
IterMut {
range: RangeMut { inner: LeafRange::none(), _marker: PhantomData },
length: 0,
}
IterMut { range: LazyLeafRange::none(), length: 0, _marker: PhantomData }
}
}

188
library/alloc/src/collections/btree/navigate.rs
View File

@ -1,4 +1,5 @@
use core::borrow::Borrow;
use core::hint;
use core::ops::RangeBounds;
use core::ptr;
@ -37,53 +38,143 @@ impl<BorrowType, K, V> LeafRange<BorrowType, K, V> {
impl<'a, K, V> LeafRange<marker::Immut<'a>, K, V> {
#[inline]
pub fn next_checked(&mut self) -> Option<(&'a K, &'a V)> {
if self.is_empty() { None } else { Some(unsafe { self.next_unchecked() }) }
self.perform_next_checked(|kv| kv.into_kv())
}
#[inline]
pub fn next_back_checked(&mut self) -> Option<(&'a K, &'a V)> {
if self.is_empty() { None } else { Some(unsafe { self.next_back_unchecked() }) }
}
#[inline]
pub unsafe fn next_unchecked(&mut self) -> (&'a K, &'a V) {
unsafe { self.front.as_mut().unwrap().next_unchecked() }
}
#[inline]
pub unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a V) {
unsafe { self.back.as_mut().unwrap().next_back_unchecked() }
self.perform_next_back_checked(|kv| kv.into_kv())
}
}
impl<'a, K, V> LeafRange<marker::ValMut<'a>, K, V> {
#[inline]
pub fn next_checked(&mut self) -> Option<(&'a K, &'a mut V)> {
if self.is_empty() { None } else { Some(unsafe { self.next_unchecked() }) }
self.perform_next_checked(|kv| unsafe { ptr::read(kv) }.into_kv_valmut())
}
#[inline]
pub fn next_back_checked(&mut self) -> Option<(&'a K, &'a mut V)> {
if self.is_empty() { None } else { Some(unsafe { self.next_back_unchecked() }) }
self.perform_next_back_checked(|kv| unsafe { ptr::read(kv) }.into_kv_valmut())
}
}
impl<BorrowType: marker::BorrowType, K, V> LeafRange<BorrowType, K, V> {
/// If possible, extract some result from the following KV and move to the edge beyond it.
fn perform_next_checked<F, R>(&mut self, f: F) -> Option<R>
where
F: Fn(&Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV>) -> R,
{
if self.is_empty() {
None
} else {
super::mem::replace(self.front.as_mut().unwrap(), |front| {
let kv = front.next_kv().ok().unwrap();
let result = f(&kv);
(kv.next_leaf_edge(), Some(result))
})
}
}
/// If possible, extract some result from the preceding KV and move to the edge beyond it.
fn perform_next_back_checked<F, R>(&mut self, f: F) -> Option<R>
where
F: Fn(&Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV>) -> R,
{
if self.is_empty() {
None
} else {
super::mem::replace(self.back.as_mut().unwrap(), |back| {
let kv = back.next_back_kv().ok().unwrap();
let result = f(&kv);
(kv.next_back_leaf_edge(), Some(result))
})
}
}
}
enum LazyLeafHandle<BorrowType, K, V> {
Root(NodeRef<BorrowType, K, V, marker::LeafOrInternal>), // not yet descended
Edge(Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>),
}
impl<'a, K: 'a, V: 'a> Clone for LazyLeafHandle<marker::Immut<'a>, K, V> {
fn clone(&self) -> Self {
match self {
LazyLeafHandle::Root(root) => LazyLeafHandle::Root(*root),
LazyLeafHandle::Edge(edge) => LazyLeafHandle::Edge(*edge),
}
}
}
impl<BorrowType, K, V> LazyLeafHandle<BorrowType, K, V> {
fn reborrow(&self) -> LazyLeafHandle<marker::Immut<'_>, K, V> {
match self {
LazyLeafHandle::Root(root) => LazyLeafHandle::Root(root.reborrow()),
LazyLeafHandle::Edge(edge) => LazyLeafHandle::Edge(edge.reborrow()),
}
}
}
// `front` and `back` are always both `None` or both `Some`.
pub struct LazyLeafRange<BorrowType, K, V> {
front: Option<LazyLeafHandle<BorrowType, K, V>>,
back: Option<LazyLeafHandle<BorrowType, K, V>>,
}
impl<'a, K: 'a, V: 'a> Clone for LazyLeafRange<marker::Immut<'a>, K, V> {
fn clone(&self) -> Self {
LazyLeafRange { front: self.front.clone(), back: self.back.clone() }
}
}
impl<BorrowType, K, V> LazyLeafRange<BorrowType, K, V> {
pub fn none() -> Self {
LazyLeafRange { front: None, back: None }
}
/// Temporarily takes out another, immutable equivalent of the same range.
pub fn reborrow(&self) -> LazyLeafRange<marker::Immut<'_>, K, V> {
LazyLeafRange {
front: self.front.as_ref().map(|f| f.reborrow()),
back: self.back.as_ref().map(|b| b.reborrow()),
}
}
}
impl<'a, K, V> LazyLeafRange<marker::Immut<'a>, K, V> {
#[inline]
pub unsafe fn next_unchecked(&mut self) -> (&'a K, &'a V) {
unsafe { self.init_front().unwrap().next_unchecked() }
}
#[inline]
pub unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a V) {
unsafe { self.init_back().unwrap().next_back_unchecked() }
}
}
impl<'a, K, V> LazyLeafRange<marker::ValMut<'a>, K, V> {
#[inline]
pub unsafe fn next_unchecked(&mut self) -> (&'a K, &'a mut V) {
unsafe { self.front.as_mut().unwrap().next_unchecked() }
unsafe { self.init_front().unwrap().next_unchecked() }
}
#[inline]
pub unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a mut V) {
unsafe { self.back.as_mut().unwrap().next_back_unchecked() }
unsafe { self.init_back().unwrap().next_back_unchecked() }
}
}
impl<K, V> LeafRange<marker::Dying, K, V> {
impl<K, V> LazyLeafRange<marker::Dying, K, V> {
#[inline]
pub fn take_front(
&mut self,
) -> Option<Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge>> {
self.front.take()
match self.front.take()? {
LazyLeafHandle::Root(root) => Some(root.first_leaf_edge()),
LazyLeafHandle::Edge(edge) => Some(edge),
}
}
#[inline]
@ -91,7 +182,7 @@ impl<K, V> LeafRange<marker::Dying, K, V> {
&mut self,
) -> Handle<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, marker::KV> {
debug_assert!(self.front.is_some());
let front = self.front.as_mut().unwrap();
let front = self.init_front().unwrap();
unsafe { front.deallocating_next_unchecked() }
}
@ -100,11 +191,41 @@ impl<K, V> LeafRange<marker::Dying, K, V> {
&mut self,
) -> Handle<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, marker::KV> {
debug_assert!(self.back.is_some());
let back = self.back.as_mut().unwrap();
let back = self.init_back().unwrap();
unsafe { back.deallocating_next_back_unchecked() }
}
}
impl<BorrowType: marker::BorrowType, K, V> LazyLeafRange<BorrowType, K, V> {
fn init_front(
&mut self,
) -> Option<&mut Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>> {
if let Some(LazyLeafHandle::Root(root)) = &self.front {
self.front = Some(LazyLeafHandle::Edge(unsafe { ptr::read(root) }.first_leaf_edge()));
}
match &mut self.front {
None => None,
Some(LazyLeafHandle::Edge(edge)) => Some(edge),
// SAFETY: the code above would have replaced it.
Some(LazyLeafHandle::Root(_)) => unsafe { hint::unreachable_unchecked() },
}
}
fn init_back(
&mut self,
) -> Option<&mut Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>> {
if let Some(LazyLeafHandle::Root(root)) = &self.back {
self.back = Some(LazyLeafHandle::Edge(unsafe { ptr::read(root) }.last_leaf_edge()));
}
match &mut self.back {
None => None,
Some(LazyLeafHandle::Edge(edge)) => Some(edge),
// SAFETY: the code above would have replaced it.
Some(LazyLeafHandle::Root(_)) => unsafe { hint::unreachable_unchecked() },
}
}
}
impl<BorrowType: marker::BorrowType, K, V> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
/// Finds the distinct leaf edges delimiting a specified range in a tree.
///
@ -156,26 +277,13 @@ impl<BorrowType: marker::BorrowType, K, V> NodeRef<BorrowType, K, V, marker::Lea
}
}
/// Equivalent to `(root1.first_leaf_edge(), root2.last_leaf_edge())` but more efficient.
fn full_range<BorrowType: marker::BorrowType, K, V>(
root1: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
root2: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
) -> LeafRange<BorrowType, K, V> {
let mut min_node = root1;
let mut max_node = root2;
loop {
let front = min_node.first_edge();
let back = max_node.last_edge();
match (front.force(), back.force()) {
(Leaf(f), Leaf(b)) => {
return LeafRange { front: Some(f), back: Some(b) };
}
(Internal(min_int), Internal(max_int)) => {
min_node = min_int.descend();
max_node = max_int.descend();
}
_ => unreachable!("BTreeMap has different depths"),
};
) -> LazyLeafRange<BorrowType, K, V> {
LazyLeafRange {
front: Some(LazyLeafHandle::Root(root1)),
back: Some(LazyLeafHandle::Root(root2)),
}
}
@ -195,7 +303,7 @@ impl<'a, K: 'a, V: 'a> NodeRef<marker::Immut<'a>, K, V, marker::LeafOrInternal>
}
/// Finds the pair of leaf edges delimiting an entire tree.
pub fn full_range(self) -> LeafRange<marker::Immut<'a>, K, V> {
pub fn full_range(self) -> LazyLeafRange<marker::Immut<'a>, K, V> {
full_range(self, self)
}
}
@ -222,7 +330,7 @@ impl<'a, K: 'a, V: 'a> NodeRef<marker::ValMut<'a>, K, V, marker::LeafOrInternal>
/// Splits a unique reference into a pair of leaf edges delimiting the full range of the tree.
/// The results are non-unique references allowing mutation (of values only), so must be used
/// with care.
pub fn full_range(self) -> LeafRange<marker::ValMut<'a>, K, V> {
pub fn full_range(self) -> LazyLeafRange<marker::ValMut<'a>, K, V> {
// We duplicate the root NodeRef here -- we will never visit the same KV
// twice, and never end up with overlapping value references.
let self2 = unsafe { ptr::read(&self) };
@ -234,7 +342,7 @@ impl<K, V> NodeRef<marker::Dying, K, V, marker::LeafOrInternal> {
/// Splits a unique reference into a pair of leaf edges delimiting the full range of the tree.
/// The results are non-unique references allowing massively destructive mutation, so must be
/// used with the utmost care.
pub fn full_range(self) -> LeafRange<marker::Dying, K, V> {
pub fn full_range(self) -> LazyLeafRange<marker::Dying, K, V> {
// We duplicate the root NodeRef here -- we will never access it in a way
// that overlaps references obtained from the root.
let self2 = unsafe { ptr::read(&self) };