Add iter_at_key & iter_at_key_mut

An iterator that was skipped to a specific key entry,
which allows to iterate starting from this entry, instead of starting at
the beginning

Author: Your Name

src/linked_hash_map.rs

@@ -531,6 +531,114 @@ 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 head_tail = if let Some(values) = self.values {
+            let ValueLinks { next, prev } = unsafe { values.as_ref().links.value };
+            Some((next, prev))
+        } else {
+            None
+        };
+
+        match head_tail {
+            Some((head, tail)) => match self.raw_entry().node_from_key(k) {
+                Some(entry) => Some(IterAtKey {
+                    iter: Iter {
+                        head: head.as_ptr(),
+                        tail: tail.as_ptr(),
+                        remaining: usize::MAX,
+                        marker: PhantomData,
+                    },
+                    cur: entry.as_ptr(),
+                }),
+                None => None,
+            },
+            None => None,
+        }
+    }
+
+    /// 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<IterMutAtKey<'_, K, V>> {
+        let head_tail = if let Some(values) = self.values {
+            let ValueLinks { next, prev } = unsafe { values.as_ref().links.value };
+            Some((next, prev))
+        } else {
+            None
+        };
+
+        match head_tail {
+            Some((head, tail)) => match self.raw_entry_mut().from_key(k) {
+                RawEntryMut::Occupied(entry) => {
+                    let cur = entry.entry.key();
+                    Some(IterMutAtKey {
+                        iter: IterMut {
+                            head: Some(head),
+                            tail: Some(tail),
+                            remaining: usize::MAX,
+                            marker: PhantomData,
+                        },
+                        cur: Some(*cur),
+                    })
+                }
+                RawEntryMut::Vacant(_) => None,
+            },
+            None => None,
+        }
+    }
 }
 
 impl<K, V, S> LinkedHashMap<K, V, S>
@@ -938,6 +1046,16 @@ where
         self.from_key_hashed_nocheck(hash, k)
     }
 
+    #[inline]
+    fn node_from_key<Q>(self, k: &Q) -> Option<NonNull<Node<K, V>>>
+    where
+        K: Borrow<Q>,
+        Q: Hash + Eq + ?Sized,
+    {
+        let hash = hash_key(self.hash_builder, k);
+        self.node_from_key_hashed_nocheck(hash, k)
+    }
+
     #[inline]
     pub fn from_key_hashed_nocheck<Q>(self, hash: u64, k: &Q) -> Option<(&'a K, &'a V)>
     where
@@ -947,6 +1065,15 @@ where
         self.from_hash(hash, move |o| k.eq(o.borrow()))
     }
 
+    #[inline]
+    fn node_from_key_hashed_nocheck<Q>(self, hash: u64, k: &Q) -> Option<NonNull<Node<K, V>>>
+    where
+        K: Borrow<Q>,
+        Q: Hash + Eq + ?Sized,
+    {
+        self.node_from_hash(hash, move |o| k.eq(o.borrow()))
+    }
+
     #[inline]
     pub fn from_hash(
         self,
@@ -963,6 +1090,22 @@ where
             Some((key, value))
         }
     }
+
+    #[inline]
+    fn node_from_hash(
+        self,
+        hash: u64,
+        mut is_match: impl FnMut(&K) -> bool,
+    ) -> Option<NonNull<Node<K, V>>> {
+        unsafe {
+            let node = *self
+                .entry
+                .from_hash(hash, move |k| is_match((*k).as_ref().key_ref()))?
+                .0;
+
+            Some(node)
+        }
+    }
 }
 
 unsafe impl<'a, K, V, S> Send for RawEntryBuilder<'a, K, V, S>
@@ -1384,6 +1527,16 @@ pub struct Drain<'a, K, V> {
     marker: PhantomData<(K, V, &'a LinkedHashMap<K, V>)>,
 }
 
+pub struct IterAtKey<'a, K, V> {
+    iter: Iter<'a, K, V>,
+    cur: *const Node<K, V>,
+}
+
+pub struct IterMutAtKey<'a, K, V> {
+    iter: IterMut<'a, K, V>,
+    cur: Option<NonNull<Node<K, V>>>,
+}
+
 impl<K, V> IterMut<'_, K, V> {
     #[inline]
     pub(crate) fn iter(&self) -> Iter<'_, K, V> {
@@ -1420,6 +1573,21 @@ impl<K, V> Drain<'_, K, V> {
     }
 }
 
+impl<K, V> IterMutAtKey<'_, K, V> {
+    #[inline]
+    pub(crate) fn iter(&self) -> IterAtKey<'_, K, V> {
+        IterAtKey {
+            iter: Iter {
+                head: self.iter.head.as_ptr(),
+                tail: self.iter.tail.as_ptr(),
+                remaining: self.iter.remaining,
+                marker: PhantomData,
+            },
+            cur: self.cur.as_ptr(),
+        }
+    }
+}
+
 unsafe impl<'a, K, V> Send for Iter<'a, K, V>
 where
     K: Send,
@@ -1448,6 +1616,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 IterMutAtKey<'a, K, V>
+where
+    K: Send,
+    V: Send,
+{
+}
+
 unsafe impl<'a, K, V> Sync for Iter<'a, K, V>
 where
     K: Sync,
@@ -1476,13 +1658,37 @@ where
 {
 }
 
+unsafe impl<'a, K, V> Sync for IterAtKey<'a, K, V>
+where
+    K: Sync,
+    V: Sync,
+{
+}
+
+unsafe impl<'a, K, V> Sync for IterMutAtKey<'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 {
+            iter: self.iter.clone(),
+            cur: self.cur,
+        }
+    }
+}
+
 impl<K: fmt::Debug, V: fmt::Debug> fmt::Debug for Iter<'_, K, V> {
     #[inline]
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
@@ -1523,6 +1729,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 IterMutAtKey<'_, 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);
 
@@ -1617,6 +1845,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.iter.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 IterMutAtKey<'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.iter.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)> {
@@ -1683,6 +1952,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.iter.tail {
+                    self.cur = std::ptr::null();
+                }
+                let (key, value) = (*self.iter.tail).entry_ref();
+                self.iter.tail = (*self.iter.tail).links.value.prev.as_ptr();
+                Some((key, value))
+            }
+        }
+    }
+}
+
+impl<'a, K, V> DoubleEndedIterator for IterMutAtKey<'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.iter.tail {
+                    self.cur = None;
+                }
+                let (key, value) = (*self.iter.tail.as_ptr()).entry_mut();
+                self.iter.tail = Some((*self.iter.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> {}