Skip to content

Commit

Permalink
BTree: lazily locate leaves in rangeless iterators
Browse files Browse the repository at this point in the history
  • Loading branch information
ssomers committed Jul 8, 2021
1 parent aa65b08 commit 35d02e2
Show file tree
Hide file tree
Showing 2 changed files with 171 additions and 68 deletions.
57 changes: 26 additions & 31 deletions library/alloc/src/collections/btree/map.rs
Expand Up @@ -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::*;

Expand Down Expand Up @@ -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,
}

Expand All @@ -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`.
Expand All @@ -310,16 +320,15 @@ 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,
}

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 }
}
}

Expand Down Expand Up @@ -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() })
}
}

Expand Down Expand Up @@ -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() })
}
}
}
Expand Down Expand Up @@ -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() })
}
}

Expand Down Expand Up @@ -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() })
}
}
}
Expand All @@ -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 }
}
}

Expand All @@ -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 }
}
}
}
Expand Down Expand Up @@ -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)> {
Expand Down Expand Up @@ -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 }
}
}

Expand Down Expand Up @@ -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 }
}
}

Expand Down

0 comments on commit 35d02e2

Please sign in to comment.