From 09ff88b600713a2bfca7cfdfd1828b519c907247 Mon Sep 17 00:00:00 2001
From: Camelid <camelidcamel@gmail.com>
Date: Fri, 2 Apr 2021 19:56:18 -0700
Subject: [PATCH] Add `FromIterator` and `IntoIterator` impls for `ThinVec`

These should make using `ThinVec` feel much more like using `Vec`.
They will allow users of `Vec` to switch to `ThinVec` while continuing
to use `collect()`, `for` loops, and other parts of the iterator API.

I don't know if there were use cases before for using the iterator API
with `ThinVec`, but I would like to start using `ThinVec` in rustdoc,
and having it conform to the iterator API would make the transition
*a lot* easier.

I added a `FromIterator` impl, an `IntoIterator` impl that yields owned
elements, and `IntoIterator` impls that yield immutable or mutable
references to elements. I also added some unit tests for `ThinVec`.
---
 .../rustc_data_structures/src/thin_vec.rs     | 49 +++++++++++++++++++
 .../src/thin_vec/tests.rs                     | 42 ++++++++++++++++
 2 files changed, 91 insertions(+)
 create mode 100644 compiler/rustc_data_structures/src/thin_vec/tests.rs

diff --git a/compiler/rustc_data_structures/src/thin_vec.rs b/compiler/rustc_data_structures/src/thin_vec.rs
index 4d673fd5cf983..00e304734983f 100644
--- a/compiler/rustc_data_structures/src/thin_vec.rs
+++ b/compiler/rustc_data_structures/src/thin_vec.rs
@@ -1,5 +1,7 @@
 use crate::stable_hasher::{HashStable, StableHasher};
 
+use std::iter::FromIterator;
+
 /// A vector type optimized for cases where this size is usually 0 (cf. `SmallVector`).
 /// The `Option<Box<..>>` wrapping allows us to represent a zero sized vector with `None`,
 /// which uses only a single (null) pointer.
@@ -10,6 +12,14 @@ impl<T> ThinVec<T> {
     pub fn new() -> Self {
         ThinVec(None)
     }
+
+    pub fn iter(&self) -> std::slice::Iter<'_, T> {
+        self.into_iter()
+    }
+
+    pub fn iter_mut(&mut self) -> std::slice::IterMut<'_, T> {
+        self.into_iter()
+    }
 }
 
 impl<T> From<Vec<T>> for ThinVec<T> {
@@ -46,6 +56,42 @@ impl<T> ::std::ops::DerefMut for ThinVec<T> {
     }
 }
 
+impl<T> FromIterator<T> for ThinVec<T> {
+    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
+        // `Vec::from_iter()` should not allocate if the iterator is empty.
+        let vec: Vec<_> = iter.into_iter().collect();
+        if vec.is_empty() { ThinVec(None) } else { ThinVec(Some(Box::new(vec))) }
+    }
+}
+
+impl<T> IntoIterator for ThinVec<T> {
+    type Item = T;
+    type IntoIter = std::vec::IntoIter<T>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        // This is still performant because `Vec::new()` does not allocate.
+        self.0.map_or_else(Vec::new, |ptr| *ptr).into_iter()
+    }
+}
+
+impl<'a, T> IntoIterator for &'a ThinVec<T> {
+    type Item = &'a T;
+    type IntoIter = std::slice::Iter<'a, T>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.as_ref().iter()
+    }
+}
+
+impl<'a, T> IntoIterator for &'a mut ThinVec<T> {
+    type Item = &'a mut T;
+    type IntoIter = std::slice::IterMut<'a, T>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.as_mut().iter_mut()
+    }
+}
+
 impl<T> Extend<T> for ThinVec<T> {
     fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
         match *self {
@@ -80,3 +126,6 @@ impl<T> Default for ThinVec<T> {
         Self(None)
     }
 }
+
+#[cfg(test)]
+mod tests;
diff --git a/compiler/rustc_data_structures/src/thin_vec/tests.rs b/compiler/rustc_data_structures/src/thin_vec/tests.rs
new file mode 100644
index 0000000000000..5abfd939373c7
--- /dev/null
+++ b/compiler/rustc_data_structures/src/thin_vec/tests.rs
@@ -0,0 +1,42 @@
+use super::*;
+
+impl<T> ThinVec<T> {
+    fn into_vec(self) -> Vec<T> {
+        self.into()
+    }
+}
+
+#[test]
+fn test_from_iterator() {
+    assert_eq!(std::iter::empty().collect::<ThinVec<String>>().into_vec(), Vec::<String>::new());
+    assert_eq!(std::iter::once(42).collect::<ThinVec<_>>().into_vec(), vec![42]);
+    assert_eq!(vec![1, 2].into_iter().collect::<ThinVec<_>>().into_vec(), vec![1, 2]);
+    assert_eq!(vec![1, 2, 3].into_iter().collect::<ThinVec<_>>().into_vec(), vec![1, 2, 3]);
+}
+
+#[test]
+fn test_into_iterator_owned() {
+    assert_eq!(ThinVec::new().into_iter().collect::<Vec<String>>(), Vec::<String>::new());
+    assert_eq!(ThinVec::from(vec![1]).into_iter().collect::<Vec<_>>(), vec![1]);
+    assert_eq!(ThinVec::from(vec![1, 2]).into_iter().collect::<Vec<_>>(), vec![1, 2]);
+    assert_eq!(ThinVec::from(vec![1, 2, 3]).into_iter().collect::<Vec<_>>(), vec![1, 2, 3]);
+}
+
+#[test]
+fn test_into_iterator_ref() {
+    assert_eq!(ThinVec::new().iter().collect::<Vec<&String>>(), Vec::<&String>::new());
+    assert_eq!(ThinVec::from(vec![1]).iter().collect::<Vec<_>>(), vec![&1]);
+    assert_eq!(ThinVec::from(vec![1, 2]).iter().collect::<Vec<_>>(), vec![&1, &2]);
+    assert_eq!(ThinVec::from(vec![1, 2, 3]).iter().collect::<Vec<_>>(), vec![&1, &2, &3]);
+}
+
+#[test]
+fn test_into_iterator_ref_mut() {
+    assert_eq!(ThinVec::new().iter_mut().collect::<Vec<&mut String>>(), Vec::<&mut String>::new());
+    assert_eq!(ThinVec::from(vec![1]).iter_mut().collect::<Vec<_>>(), vec![&mut 1]);
+    assert_eq!(ThinVec::from(vec![1, 2]).iter_mut().collect::<Vec<_>>(), vec![&mut 1, &mut 2]);
+    assert_eq!(
+        ThinVec::from(vec![1, 2, 3]).iter_mut().collect::<Vec<_>>(),
+        vec![&mut 1, &mut 2, &mut 3],
+    );
+}