fix(todo): implement IntoIterator for ArenaMap<IDX, V>

1 files changed, 58 insertions, 6 deletions

diff --git a/lib/la-arena/src/map.rs b/lib/la-arena/src/map.rs
index 610f7d92d65..750f345b539 100644
--- a/lib/la-arena/src/map.rs
+++ b/lib/la-arena/src/map.rs
@@ -1,3 +1,4 @@
+use std::iter::Enumerate;
 use std::marker::PhantomData;
 
 use crate::Idx;
@@ -94,12 +95,6 @@ impl<T, V> ArenaMap<Idx<T>, V> {
             .filter_map(|(idx, o)| Some((Self::from_idx(idx), o.as_mut()?)))
     }
 
-    /// Returns an iterator over the arena indexes and values in the map.
-    // FIXME: Implement `IntoIterator` trait.
-    pub fn into_iter(self) -> impl Iterator<Item = (Idx<T>, V)> + DoubleEndedIterator {
-        self.v.into_iter().enumerate().filter_map(|(idx, o)| Some((Self::from_idx(idx), o?)))
-    }
-
     /// Gets the given key's corresponding entry in the map for in-place manipulation.
     pub fn entry(&mut self, idx: Idx<T>) -> Entry<'_, Idx<T>, V> {
         let idx = Self::to_idx(idx);
@@ -154,6 +149,63 @@ impl<T, V> FromIterator<(Idx<V>, T)> for ArenaMap<Idx<V>, T> {
     }
 }
 
+pub struct ArenaMapIter<IDX, V> {
+    iter: Enumerate<std::vec::IntoIter<Option<V>>>,
+    _ty: PhantomData<IDX>,
+}
+
+impl<T, V> IntoIterator for ArenaMap<Idx<T>, V> {
+    type Item = (Idx<T>, V);
+
+    type IntoIter = ArenaMapIter<Idx<T>, V>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        let iter = self.v.into_iter().enumerate();
+        Self::IntoIter { iter, _ty: PhantomData }
+    }
+}
+
+impl<T, V> ArenaMapIter<Idx<T>, V> {
+    fn mapper((idx, o): (usize, Option<V>)) -> Option<(Idx<T>, V)> {
+        Some((ArenaMap::<Idx<T>, V>::from_idx(idx), o?))
+    }
+}
+
+impl<T, V> Iterator for ArenaMapIter<Idx<T>, V> {
+    type Item = (Idx<T>, V);
+
+    #[inline]
+    fn next(&mut self) -> Option<Self::Item> {
+        for next in self.iter.by_ref() {
+            match Self::mapper(next) {
+                Some(r) => return Some(r),
+                None => continue,
+            }
+        }
+
+        None
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.iter.size_hint()
+    }
+}
+
+impl<T, V> DoubleEndedIterator for ArenaMapIter<Idx<T>, V> {
+    #[inline]
+    fn next_back(&mut self) -> Option<Self::Item> {
+        while let Some(next_back) = self.iter.next_back() {
+            match Self::mapper(next_back) {
+                Some(r) => return Some(r),
+                None => continue,
+            }
+        }
+
+        None
+    }
+}
+
 /// A view into a single entry in a map, which may either be vacant or occupied.
 ///
 /// This `enum` is constructed from the [`entry`] method on [`ArenaMap`].