Skip to content

Commit 123d519

Browse files
Your NameJupeyy
Your Name
authored and
Jupeyy
committed
Add iter_at_key & iter_at_key_mut
An iterator that was starts at a specific key entry, which allows to iterate starting from this entry, instead of starting at the beginning
1 parent 9cd8c09 commit 123d519

File tree

4 files changed

+428
-0
lines changed

4 files changed

+428
-0
lines changed

‎src/linked_hash_map.rs‎

Lines changed: 248 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -531,6 +531,93 @@ where
531531
}
532532
}
533533
}
534+
535+
/// Returns iterator that was skipped to a specific key entry, or None.
536+
/// It does not implement `ExactSizeIterator` because
537+
/// it is unclear where exactly the iterator is.
538+
///
539+
/// It is useful when iterating over a subset of
540+
/// all items in order, e.g. for starting a queue iteration at a specific key
541+
///
542+
/// # Examples
543+
///
544+
/// ```rs
545+
/// let mut map = LinkedHashMap::new();
546+
///
547+
/// map.insert("a", 10);
548+
/// map.insert("b", 20);
549+
/// map.insert("c", 30);
550+
///
551+
/// assert_eq!(map.iter_at_key(&"e").is_none(), true);
552+
///
553+
/// let mut iter = map.iter_at_key(&"b").unwrap();
554+
/// assert_eq!((&"b", &20), iter.next().unwrap());
555+
/// assert_eq!((&"c", &30), iter.next().unwrap());
556+
/// assert_eq!(None, iter.next());
557+
/// assert_eq!(None, iter.next());
558+
/// ```
559+
///
560+
#[inline]
561+
pub fn iter_at_key(&self, k: &K) -> Option<IterAtKey<'_, K, V>> {
562+
let tail = unsafe { self.values?.as_ref().links.value.prev };
563+
564+
let hash = hash_key(&self.hash_builder, k);
565+
let node = unsafe {
566+
*self
567+
.map
568+
.raw_entry()
569+
.from_hash(hash, move |key| k.eq((*key).as_ref().key_ref()))?
570+
.0
571+
};
572+
Some(IterAtKey {
573+
tail: tail.as_ptr(),
574+
marker: PhantomData,
575+
cur: node.as_ptr(),
576+
})
577+
}
578+
579+
/// Returns a mutable iterator that was skipped to a specific key entry, or None.
580+
/// It does not implement `ExactSizeIterator` because
581+
/// it is unclear where exactly the iterator is.
582+
///
583+
/// It is useful when iterating over a subset of
584+
/// all items in order, e.g. for starting a queue iteration at a specific key
585+
///
586+
/// # Examples
587+
///
588+
/// ```rs
589+
/// let mut map = LinkedHashMap::new();
590+
/// map.insert("a", 10);
591+
/// map.insert("c", 30);
592+
/// map.insert("b", 20);
593+
/// map.insert("d", 40);
594+
///
595+
/// assert_eq!(map.iter_at_key_mut(&"e").is_none(), true);
596+
///
597+
/// let mut iter = map.iter_at_key_mut(&"c").unwrap();
598+
/// let entry = iter.next().unwrap();
599+
/// assert_eq!("c", *entry.0);
600+
/// *entry.1 = 17;
601+
///
602+
/// assert_eq!(format!("{:?}", iter), "[(\"b\", 20), (\"d\", 40)]");
603+
/// assert_eq!(17, map[&"c"]);
604+
/// ```
605+
///
606+
#[inline]
607+
pub fn iter_at_key_mut(&mut self, k: &K) -> Option<IterAtKeyMut<'_, K, V>> {
608+
let tail = unsafe { self.values?.as_ref().links.value.prev };
609+
match self.raw_entry_mut().from_key(k) {
610+
RawEntryMut::Occupied(entry) => {
611+
let cur = entry.entry.key();
612+
Some(IterAtKeyMut {
613+
tail: Some(tail),
614+
marker: PhantomData,
615+
cur: Some(*cur),
616+
})
617+
}
618+
RawEntryMut::Vacant(_) => None,
619+
}
620+
}
534621
}
535622

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

1474+
pub struct IterAtKey<'a, K, V> {
1475+
tail: *const Node<K, V>,
1476+
marker: PhantomData<(&'a K, &'a V)>,
1477+
cur: *const Node<K, V>,
1478+
}
1479+
1480+
pub struct IterAtKeyMut<'a, K, V> {
1481+
tail: Option<NonNull<Node<K, V>>>,
1482+
marker: PhantomData<(&'a K, &'a mut V)>,
1483+
cur: Option<NonNull<Node<K, V>>>,
1484+
}
1485+
13871486
impl<K, V> IterMut<'_, K, V> {
13881487
#[inline]
13891488
pub(crate) fn iter(&self) -> Iter<'_, K, V> {
@@ -1420,6 +1519,17 @@ impl<K, V> Drain<'_, K, V> {
14201519
}
14211520
}
14221521

1522+
impl<K, V> IterAtKeyMut<'_, K, V> {
1523+
#[inline]
1524+
pub(crate) fn iter(&self) -> IterAtKey<'_, K, V> {
1525+
IterAtKey {
1526+
tail: self.tail.as_ptr(),
1527+
marker: PhantomData,
1528+
cur: self.cur.as_ptr(),
1529+
}
1530+
}
1531+
}
1532+
14231533
unsafe impl<'a, K, V> Send for Iter<'a, K, V>
14241534
where
14251535
K: Send,
@@ -1448,6 +1558,20 @@ where
14481558
{
14491559
}
14501560

1561+
unsafe impl<'a, K, V> Send for IterAtKey<'a, K, V>
1562+
where
1563+
K: Send,
1564+
V: Send,
1565+
{
1566+
}
1567+
1568+
unsafe impl<'a, K, V> Send for IterAtKeyMut<'a, K, V>
1569+
where
1570+
K: Send,
1571+
V: Send,
1572+
{
1573+
}
1574+
14511575
unsafe impl<'a, K, V> Sync for Iter<'a, K, V>
14521576
where
14531577
K: Sync,
@@ -1476,13 +1600,38 @@ where
14761600
{
14771601
}
14781602

1603+
unsafe impl<'a, K, V> Sync for IterAtKey<'a, K, V>
1604+
where
1605+
K: Sync,
1606+
V: Sync,
1607+
{
1608+
}
1609+
1610+
unsafe impl<'a, K, V> Sync for IterAtKeyMut<'a, K, V>
1611+
where
1612+
K: Sync,
1613+
V: Sync,
1614+
{
1615+
}
1616+
14791617
impl<'a, K, V> Clone for Iter<'a, K, V> {
14801618
#[inline]
14811619
fn clone(&self) -> Self {
14821620
Iter { ..*self }
14831621
}
14841622
}
14851623

1624+
impl<'a, K, V> Clone for IterAtKey<'a, K, V> {
1625+
#[inline]
1626+
fn clone(&self) -> Self {
1627+
IterAtKey {
1628+
tail: self.tail,
1629+
cur: self.cur,
1630+
marker: PhantomData,
1631+
}
1632+
}
1633+
}
1634+
14861635
impl<K: fmt::Debug, V: fmt::Debug> fmt::Debug for Iter<'_, K, V> {
14871636
#[inline]
14881637
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
@@ -1523,6 +1672,28 @@ where
15231672
}
15241673
}
15251674

1675+
impl<K, V> fmt::Debug for IterAtKey<'_, K, V>
1676+
where
1677+
K: fmt::Debug,
1678+
V: fmt::Debug,
1679+
{
1680+
#[inline]
1681+
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1682+
f.debug_list().entries(self.clone()).finish()
1683+
}
1684+
}
1685+
1686+
impl<K, V> fmt::Debug for IterAtKeyMut<'_, K, V>
1687+
where
1688+
K: fmt::Debug,
1689+
V: fmt::Debug,
1690+
{
1691+
#[inline]
1692+
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1693+
f.debug_list().entries(self.iter()).finish()
1694+
}
1695+
}
1696+
15261697
impl<'a, K, V> Iterator for Iter<'a, K, V> {
15271698
type Item = (&'a K, &'a V);
15281699

@@ -1617,6 +1788,47 @@ impl<'a, K, V> Iterator for Drain<'a, K, V> {
16171788
}
16181789
}
16191790

1791+
impl<'a, K, V> Iterator for IterAtKey<'a, K, V> {
1792+
type Item = (&'a K, &'a V);
1793+
1794+
#[inline]
1795+
fn next(&mut self) -> Option<(&'a K, &'a V)> {
1796+
if self.cur.is_null() {
1797+
return None;
1798+
}
1799+
unsafe {
1800+
let last_iter = self.cur == self.tail;
1801+
let (key, value) = (*self.cur).entry_ref();
1802+
self.cur = (*self.cur).links.value.next.as_ptr();
1803+
if last_iter {
1804+
self.cur = std::ptr::null();
1805+
}
1806+
Some((key, value))
1807+
}
1808+
}
1809+
}
1810+
1811+
impl<'a, K, V> Iterator for IterAtKeyMut<'a, K, V> {
1812+
type Item = (&'a K, &'a mut V);
1813+
1814+
#[inline]
1815+
fn next(&mut self) -> Option<(&'a K, &'a mut V)> {
1816+
if self.cur.is_none() {
1817+
None
1818+
} else {
1819+
unsafe {
1820+
let last_iter = self.cur == self.tail;
1821+
let (key, value) = (*self.cur.as_ptr()).entry_mut();
1822+
self.cur = Some((*self.cur.as_ptr()).links.value.next);
1823+
if last_iter {
1824+
self.cur = None;
1825+
}
1826+
Some((key, value))
1827+
}
1828+
}
1829+
}
1830+
}
1831+
16201832
impl<'a, K, V> DoubleEndedIterator for Iter<'a, K, V> {
16211833
#[inline]
16221834
fn next_back(&mut self) -> Option<(&'a K, &'a V)> {
@@ -1683,6 +1895,42 @@ impl<'a, K, V> DoubleEndedIterator for Drain<'a, K, V> {
16831895
}
16841896
}
16851897

1898+
impl<'a, K, V> DoubleEndedIterator for IterAtKey<'a, K, V> {
1899+
#[inline]
1900+
fn next_back(&mut self) -> Option<(&'a K, &'a V)> {
1901+
if self.cur.is_null() {
1902+
None
1903+
} else {
1904+
unsafe {
1905+
if self.cur == self.tail {
1906+
self.cur = std::ptr::null();
1907+
}
1908+
let (key, value) = (*self.tail).entry_ref();
1909+
self.tail = (*self.tail).links.value.prev.as_ptr();
1910+
Some((key, value))
1911+
}
1912+
}
1913+
}
1914+
}
1915+
1916+
impl<'a, K, V> DoubleEndedIterator for IterAtKeyMut<'a, K, V> {
1917+
#[inline]
1918+
fn next_back(&mut self) -> Option<(&'a K, &'a mut V)> {
1919+
if self.cur.is_none() {
1920+
None
1921+
} else {
1922+
unsafe {
1923+
if self.cur == self.tail {
1924+
self.cur = None;
1925+
}
1926+
let (key, value) = (*self.tail.as_ptr()).entry_mut();
1927+
self.tail = Some((*self.tail.as_ptr()).links.value.prev);
1928+
Some((key, value))
1929+
}
1930+
}
1931+
}
1932+
}
1933+
16861934
impl<'a, K, V> ExactSizeIterator for Iter<'a, K, V> {}
16871935

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

‎src/linked_hash_set.rs‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -305,6 +305,38 @@ where
305305
{
306306
self.map.retain_with_order(|k, _| f(k));
307307
}
308+
309+
/// Returns iterator that was skipped to a specific key entry, or None.
310+
/// It does not implement `ExactSizeIterator` because
311+
/// it is unclear where exactly the iterator is.
312+
///
313+
/// It is useful when iterating over a subset of
314+
/// all items in order, e.g. for starting a queue iteration at a specific key
315+
///
316+
/// # Examples
317+
///
318+
/// ```rs
319+
/// let mut map = LinkedHashSet::new();
320+
///
321+
/// map.insert("a");
322+
/// map.insert("b");
323+
/// map.insert("c");
324+
///
325+
/// assert_eq!(map.iter_at_key(&"e").is_none(), true);
326+
///
327+
/// // regular iter
328+
/// let mut iter = map.iter_at_key(&"b").unwrap();
329+
/// assert_eq!(&"b", iter.next().unwrap());
330+
/// assert_eq!(&"c", iter.next().unwrap());
331+
/// assert_eq!(None, iter.next());
332+
/// assert_eq!(None, iter.next());
333+
/// ```
334+
///
335+
#[inline]
336+
pub fn iter_at_key(&self, k: &T) -> Option<IterAtKey<'_, T>> {
337+
let iter = self.map.iter_at_key(k)?;
338+
Some(IterAtKey { iter })
339+
}
308340
}
309341

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

502+
pub struct IterAtKey<'a, K> {
503+
iter: linked_hash_map::IterAtKey<'a, K, ()>,
504+
}
505+
470506
pub struct Intersection<'a, T, S> {
471507
iter: Iter<'a, T>,
472508
other: &'a LinkedHashSet<T, S>,
@@ -601,6 +637,38 @@ impl<'a, T, S> Clone for Intersection<'a, T, S> {
601637
}
602638
}
603639

640+
impl<'a, K> Clone for IterAtKey<'a, K> {
641+
#[inline]
642+
fn clone(&self) -> IterAtKey<'a, K> {
643+
IterAtKey {
644+
iter: self.iter.clone(),
645+
}
646+
}
647+
}
648+
649+
impl<'a, K> Iterator for IterAtKey<'a, K> {
650+
type Item = &'a K;
651+
652+
#[inline]
653+
fn next(&mut self) -> Option<&'a K> {
654+
Some(self.iter.next()?.0)
655+
}
656+
}
657+
658+
impl<'a, K> DoubleEndedIterator for IterAtKey<'a, K> {
659+
#[inline]
660+
fn next_back(&mut self) -> Option<&'a K> {
661+
Some(self.iter.next_back()?.0)
662+
}
663+
}
664+
665+
impl<'a, K: fmt::Debug> fmt::Debug for IterAtKey<'a, K> {
666+
#[inline]
667+
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
668+
f.debug_list().entries(self.clone()).finish()
669+
}
670+
}
671+
604672
impl<'a, T, S> Iterator for Intersection<'a, T, S>
605673
where
606674
T: Eq + Hash,

0 commit comments

Comments
 (0)