Skip to content

Commit

Permalink
Add iter_at_key & iter_at_key_mut
Browse files Browse the repository at this point in the history 
An iterator that was starts at a specific key entry,
which allows to iterate starting from this entry, instead of starting at
the beginning
  • Loading branch information
@Jupeyy
Your Name authored and Jupeyy committed Jun 27, 2023
1 parent 9cd8c09 commit 9190c6e
Show file tree
Hide file tree
Showing 4 changed files with 428 additions and 0 deletions.
248 changes: 248 additions & 0 deletions src/linked_hash_map.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -531,6 +531,93 @@ where
}
}
}

/// Returns iterator that was skipped to a specific key entry, or None.
/// It does not implement `ExactSizeIterator` because
/// it is unclear where exactly the iterator is.
///
/// It is useful when iterating over a subset of
/// all items in order, e.g. for a starting queue at a specific key
///
/// # Examples
///
/// ```rs
/// let mut map = LinkedHashMap::new();
///
/// map.insert("a", 10);
/// map.insert("b", 20);
/// map.insert("c", 30);
///
/// assert_eq!(map.iter_at_key(&"e").is_none(), true);
///
/// let mut iter = map.iter_at_key(&"b").unwrap();
/// assert_eq!((&"b", &20), iter.next().unwrap());
/// assert_eq!((&"c", &30), iter.next().unwrap());
/// assert_eq!(None, iter.next());
/// assert_eq!(None, iter.next());
/// ```
///
#[inline]
pub fn iter_at_key(&self, k: &K) -> Option<IterAtKey<'_, K, V>> {
let tail = unsafe { self.values?.as_ref().links.value.prev };

let hash = hash_key(&self.hash_builder, k);
let node = unsafe {
*self
.map
.raw_entry()
.from_hash(hash, move |key| k.eq((*key).as_ref().key_ref()))?
.0
};
Some(IterAtKey {
tail: tail.as_ptr(),
marker: PhantomData,
cur: node.as_ptr(),
})
}

/// Returns a mutable iterator that was skipped to a specific key entry, or None.
/// It does not implement `ExactSizeIterator` because
/// it is unclear where exactly the iterator is.
///
/// It is useful when iterating over a subset of
/// all items in order, e.g. for a queue starting at a specific key
///
/// # Examples
///
/// ```rs
/// let mut map = LinkedHashMap::new();
/// map.insert("a", 10);
/// map.insert("c", 30);
/// map.insert("b", 20);
/// map.insert("d", 40);
///
/// assert_eq!(map.iter_at_key_mut(&"e").is_none(), true);
///
/// let mut iter = map.iter_at_key_mut(&"c").unwrap();
/// let entry = iter.next().unwrap();
/// assert_eq!("c", *entry.0);
/// *entry.1 = 17;
///
/// assert_eq!(format!("{:?}", iter), "[(\"b\", 20), (\"d\", 40)]");
/// assert_eq!(17, map[&"c"]);
/// ```
///
#[inline]
pub fn iter_at_key_mut(&mut self, k: &K) -> Option<IterAtKeyMut<'_, K, V>> {
let tail = unsafe { self.values?.as_ref().links.value.prev };
match self.raw_entry_mut().from_key(k) {
RawEntryMut::Occupied(entry) => {
let cur = entry.entry.key();
Some(IterAtKeyMut {
tail: Some(tail),
marker: PhantomData,
cur: Some(*cur),
})
}
RawEntryMut::Vacant(_) => None,
}
}
}

impl<K, V, S> LinkedHashMap<K, V, S>
Expand Down Expand Up @@ -1384,6 +1471,18 @@ pub struct Drain<'a, K, V> {
marker: PhantomData<(K, V, &'a LinkedHashMap<K, V>)>,
}

pub struct IterAtKey<'a, K, V> {
tail: *const Node<K, V>,
marker: PhantomData<(&'a K, &'a V)>,
cur: *const Node<K, V>,
}

pub struct IterAtKeyMut<'a, K, V> {
tail: Option<NonNull<Node<K, V>>>,
marker: PhantomData<(&'a K, &'a mut V)>,
cur: Option<NonNull<Node<K, V>>>,
}

impl<K, V> IterMut<'_, K, V> {
#[inline]
pub(crate) fn iter(&self) -> Iter<'_, K, V> {
Expand Down Expand Up @@ -1420,6 +1519,17 @@ impl<K, V> Drain<'_, K, V> {
}
}

impl<K, V> IterAtKeyMut<'_, K, V> {
#[inline]
pub(crate) fn iter(&self) -> IterAtKey<'_, K, V> {
IterAtKey {
tail: self.tail.as_ptr(),
marker: PhantomData,
cur: self.cur.as_ptr(),
}
}
}

unsafe impl<'a, K, V> Send for Iter<'a, K, V>
where
K: Send,
Expand Down Expand Up @@ -1448,6 +1558,20 @@ where
{
}

unsafe impl<'a, K, V> Send for IterAtKey<'a, K, V>
where
K: Send,
V: Send,
{
}

unsafe impl<'a, K, V> Send for IterAtKeyMut<'a, K, V>
where
K: Send,
V: Send,
{
}

unsafe impl<'a, K, V> Sync for Iter<'a, K, V>
where
K: Sync,
Expand Down Expand Up @@ -1476,13 +1600,38 @@ where
{
}

unsafe impl<'a, K, V> Sync for IterAtKey<'a, K, V>
where
K: Sync,
V: Sync,
{
}

unsafe impl<'a, K, V> Sync for IterAtKeyMut<'a, K, V>
where
K: Sync,
V: Sync,
{
}

impl<'a, K, V> Clone for Iter<'a, K, V> {
#[inline]
fn clone(&self) -> Self {
Iter { ..*self }
}
}

impl<'a, K, V> Clone for IterAtKey<'a, K, V> {
#[inline]
fn clone(&self) -> Self {
IterAtKey {
tail: self.tail,
cur: self.cur,
marker: PhantomData,
}
}
}

impl<K: fmt::Debug, V: fmt::Debug> fmt::Debug for Iter<'_, K, V> {
#[inline]
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
Expand Down Expand Up @@ -1523,6 +1672,28 @@ where
}
}

impl<K, V> fmt::Debug for IterAtKey<'_, K, V>
where
K: fmt::Debug,
V: fmt::Debug,
{
#[inline]
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_list().entries(self.clone()).finish()
}
}

impl<K, V> fmt::Debug for IterAtKeyMut<'_, K, V>
where
K: fmt::Debug,
V: fmt::Debug,
{
#[inline]
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_list().entries(self.iter()).finish()
}
}

impl<'a, K, V> Iterator for Iter<'a, K, V> {
type Item = (&'a K, &'a V);

Expand Down Expand Up @@ -1617,6 +1788,47 @@ impl<'a, K, V> Iterator for Drain<'a, K, V> {
}
}

impl<'a, K, V> Iterator for IterAtKey<'a, K, V> {
type Item = (&'a K, &'a V);

#[inline]
fn next(&mut self) -> Option<(&'a K, &'a V)> {
if self.cur.is_null() {
return None;
}
unsafe {
let last_iter = self.cur == self.tail;
let (key, value) = (*self.cur).entry_ref();
self.cur = (*self.cur).links.value.next.as_ptr();
if last_iter {
self.cur = std::ptr::null();
}
Some((key, value))
}
}
}

impl<'a, K, V> Iterator for IterAtKeyMut<'a, K, V> {
type Item = (&'a K, &'a mut V);

#[inline]
fn next(&mut self) -> Option<(&'a K, &'a mut V)> {
if self.cur.is_none() {
None
} else {
unsafe {
let last_iter = self.cur == self.tail;
let (key, value) = (*self.cur.as_ptr()).entry_mut();
self.cur = Some((*self.cur.as_ptr()).links.value.next);
if last_iter {
self.cur = None;
}
Some((key, value))
}
}
}
}

impl<'a, K, V> DoubleEndedIterator for Iter<'a, K, V> {
#[inline]
fn next_back(&mut self) -> Option<(&'a K, &'a V)> {
Expand Down Expand Up @@ -1683,6 +1895,42 @@ impl<'a, K, V> DoubleEndedIterator for Drain<'a, K, V> {
}
}

impl<'a, K, V> DoubleEndedIterator for IterAtKey<'a, K, V> {
#[inline]
fn next_back(&mut self) -> Option<(&'a K, &'a V)> {
if self.cur.is_null() {
None
} else {
unsafe {
if self.cur == self.tail {
self.cur = std::ptr::null();
}
let (key, value) = (*self.tail).entry_ref();
self.tail = (*self.tail).links.value.prev.as_ptr();
Some((key, value))
}
}
}
}

impl<'a, K, V> DoubleEndedIterator for IterAtKeyMut<'a, K, V> {
#[inline]
fn next_back(&mut self) -> Option<(&'a K, &'a mut V)> {
if self.cur.is_none() {
None
} else {
unsafe {
if self.cur == self.tail {
self.cur = None;
}
let (key, value) = (*self.tail.as_ptr()).entry_mut();
self.tail = Some((*self.tail.as_ptr()).links.value.prev);
Some((key, value))
}
}
}
}

impl<'a, K, V> ExactSizeIterator for Iter<'a, K, V> {}

impl<'a, K, V> ExactSizeIterator for IterMut<'a, K, V> {}
Expand Down
68 changes: 68 additions & 0 deletions src/linked_hash_set.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -305,6 +305,38 @@ where
{
self.map.retain_with_order(|k, _| f(k));
}

/// Returns iterator that was skipped to a specific key entry, or None.
/// It does not implement `ExactSizeIterator` because
/// it is unclear where exactly the iterator is.
///
/// It is useful when iterating over a subset of
/// all items in order, e.g. for a starting queue at a specific key
///
/// # Examples
///
/// ```rs
/// let mut map = LinkedHashSet::new();
///
/// map.insert("a");
/// map.insert("b");
/// map.insert("c");
///
/// assert_eq!(map.iter_at_key(&"e").is_none(), true);
///
/// // regular iter
/// let mut iter = map.iter_at_key(&"b").unwrap();
/// assert_eq!(&"b", iter.next().unwrap());
/// assert_eq!(&"c", iter.next().unwrap());
/// assert_eq!(None, iter.next());
/// assert_eq!(None, iter.next());
/// ```
///
#[inline]
pub fn iter_at_key(&self, k: &T) -> Option<IterAtKey<'_, T>> {
let iter = self.map.iter_at_key(k)?;
Some(IterAtKey { iter })
}
}

impl<T: Hash + Eq + Clone, S: BuildHasher + Clone> Clone for LinkedHashSet<T, S> {
Expand Down Expand Up @@ -467,6 +499,10 @@ pub struct Drain<'a, K: 'a> {
iter: linked_hash_map::Drain<'a, K, ()>,
}

pub struct IterAtKey<'a, K> {
iter: linked_hash_map::IterAtKey<'a, K, ()>,
}

pub struct Intersection<'a, T, S> {
iter: Iter<'a, T>,
other: &'a LinkedHashSet<T, S>,
Expand Down Expand Up @@ -601,6 +637,38 @@ impl<'a, T, S> Clone for Intersection<'a, T, S> {
}
}

impl<'a, K> Clone for IterAtKey<'a, K> {
#[inline]
fn clone(&self) -> IterAtKey<'a, K> {
IterAtKey {
iter: self.iter.clone(),
}
}
}

impl<'a, K> Iterator for IterAtKey<'a, K> {
type Item = &'a K;

#[inline]
fn next(&mut self) -> Option<&'a K> {
Some(self.iter.next()?.0)
}
}

impl<'a, K> DoubleEndedIterator for IterAtKey<'a, K> {
#[inline]
fn next_back(&mut self) -> Option<&'a K> {
Some(self.iter.next_back()?.0)
}
}

impl<'a, K: fmt::Debug> fmt::Debug for IterAtKey<'a, K> {
#[inline]
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_list().entries(self.clone()).finish()
}
}

impl<'a, T, S> Iterator for Intersection<'a, T, S>
where
T: Eq + Hash,
Expand Down
Loading

0 comments on commit 9190c6e

Please sign in to comment.