Auto merge of #61953 - Centril:shared-from-iter, r=RalfJung

Add `impl<T> FromIterator<T> for Arc/Rc<[T]>`

Add implementations of `FromIterator<T> for Arc/Rc<[T]>` with symmetrical logic.

This also takes advantage of specialization in the case of iterators with known length (`TrustedLen`) to elide the final allocation/copying from a `Vec<T>` into `Rc<[T]>` because we can allocate the space for the `Rc<[T]>` directly when the size is known. This is the primary motivation and why this is to be preferred over `iter.collect::<Vec<_>>().into(): Rc<[T]>`.

Moreover, this PR does some refactoring in some places.

r? @RalfJung for the code
cc @alexcrichton from T-libs

3 files changed, 240 insertions, 0 deletions

diff --git a/src/liballoc/tests/arc.rs b/src/liballoc/tests/arc.rs
index 2759b1b1cac..cf2ad2a8e60 100644
--- a/src/liballoc/tests/arc.rs
+++ b/src/liballoc/tests/arc.rs
@@ -2,6 +2,8 @@ use std::any::Any;
 use std::sync::{Arc, Weak};
 use std::cell::RefCell;
 use std::cmp::PartialEq;
+use std::iter::TrustedLen;
+use std::mem;
 
 #[test]
 fn uninhabited() {
@@ -85,3 +87,122 @@ fn eq() {
     assert!(!(x != x));
     assert_eq!(*x.0.borrow(), 0);
 }
+
+// The test code below is identical to that in `rc.rs`.
+// For better maintainability we therefore define this type alias.
+type Rc<T> = Arc<T>;
+
+const SHARED_ITER_MAX: u16 = 100;
+
+fn assert_trusted_len<I: TrustedLen>(_: &I) {}
+
+#[test]
+fn shared_from_iter_normal() {
+    // Exercise the base implementation for non-`TrustedLen` iterators.
+    {
+        // `Filter` is never `TrustedLen` since we don't
+        // know statically how many elements will be kept:
+        let iter = (0..SHARED_ITER_MAX).filter(|x| x % 2 == 0).map(Box::new);
+
+        // Collecting into a `Vec<T>` or `Rc<[T]>` should make no difference:
+        let vec = iter.clone().collect::<Vec<_>>();
+        let rc = iter.collect::<Rc<[_]>>();
+        assert_eq!(&*vec, &*rc);
+
+        // Clone a bit and let these get dropped.
+        {
+            let _rc_2 = rc.clone();
+            let _rc_3 = rc.clone();
+            let _rc_4 = Rc::downgrade(&_rc_3);
+        }
+    } // Drop what hasn't been here.
+}
+
+#[test]
+fn shared_from_iter_trustedlen_normal() {
+    // Exercise the `TrustedLen` implementation under normal circumstances
+    // where `size_hint()` matches `(_, Some(exact_len))`.
+    {
+        let iter = (0..SHARED_ITER_MAX).map(Box::new);
+        assert_trusted_len(&iter);
+
+        // Collecting into a `Vec<T>` or `Rc<[T]>` should make no difference:
+        let vec = iter.clone().collect::<Vec<_>>();
+        let rc = iter.collect::<Rc<[_]>>();
+        assert_eq!(&*vec, &*rc);
+        assert_eq!(mem::size_of::<Box<u16>>() * SHARED_ITER_MAX as usize, mem::size_of_val(&*rc));
+
+        // Clone a bit and let these get dropped.
+        {
+            let _rc_2 = rc.clone();
+            let _rc_3 = rc.clone();
+            let _rc_4 = Rc::downgrade(&_rc_3);
+        }
+    } // Drop what hasn't been here.
+
+    // Try a ZST to make sure it is handled well.
+    {
+        let iter = (0..SHARED_ITER_MAX).map(|_| ());
+        let vec = iter.clone().collect::<Vec<_>>();
+        let rc = iter.collect::<Rc<[_]>>();
+        assert_eq!(&*vec, &*rc);
+        assert_eq!(0, mem::size_of_val(&*rc));
+        {
+            let _rc_2 = rc.clone();
+            let _rc_3 = rc.clone();
+            let _rc_4 = Rc::downgrade(&_rc_3);
+        }
+    }
+}
+
+#[test]
+#[should_panic = "I've almost got 99 problems."]
+fn shared_from_iter_trustedlen_panic() {
+    // Exercise the `TrustedLen` implementation when `size_hint()` matches
+    // `(_, Some(exact_len))` but where `.next()` drops before the last iteration.
+    let iter = (0..SHARED_ITER_MAX)
+        .map(|val| {
+            match val {
+                98 => panic!("I've almost got 99 problems."),
+                _ => Box::new(val),
+            }
+        });
+    assert_trusted_len(&iter);
+    let _ = iter.collect::<Rc<[_]>>();
+
+    panic!("I am unreachable.");
+}
+
+#[test]
+fn shared_from_iter_trustedlen_no_fuse() {
+    // Exercise the `TrustedLen` implementation when `size_hint()` matches
+    // `(_, Some(exact_len))` but where the iterator does not behave in a fused manner.
+    struct Iter(std::vec::IntoIter<Option<Box<u8>>>);
+
+    unsafe impl TrustedLen for Iter {}
+
+    impl Iterator for Iter {
+        fn size_hint(&self) -> (usize, Option<usize>) {
+            (2, Some(2))
+        }
+
+        type Item = Box<u8>;
+
+        fn next(&mut self) -> Option<Self::Item> {
+            self.0.next().flatten()
+        }
+    }
+
+    let vec = vec![
+        Some(Box::new(42)),
+        Some(Box::new(24)),
+        None,
+        Some(Box::new(12)),
+    ];
+    let iter = Iter(vec.into_iter());
+    assert_trusted_len(&iter);
+    assert_eq!(
+        &[Box::new(42), Box::new(24)],
+        &*iter.collect::<Rc<[_]>>()
+    );
+}
diff --git a/src/liballoc/tests/lib.rs b/src/liballoc/tests/lib.rs
index ddb3120e89d..5a43c8e09a2 100644
--- a/src/liballoc/tests/lib.rs
+++ b/src/liballoc/tests/lib.rs
@@ -2,8 +2,10 @@
 #![feature(box_syntax)]
 #![feature(drain_filter)]
 #![feature(exact_size_is_empty)]
+#![feature(option_flattening)]
 #![feature(pattern)]
 #![feature(repeat_generic_slice)]
+#![feature(trusted_len)]
 #![feature(try_reserve)]
 #![feature(unboxed_closures)]
 #![deny(rust_2018_idioms)]
diff --git a/src/liballoc/tests/rc.rs b/src/liballoc/tests/rc.rs
index 18f82e80410..7854ca0fc16 100644
--- a/src/liballoc/tests/rc.rs
+++ b/src/liballoc/tests/rc.rs
@@ -2,6 +2,8 @@ use std::any::Any;
 use std::rc::{Rc, Weak};
 use std::cell::RefCell;
 use std::cmp::PartialEq;
+use std::mem;
+use std::iter::TrustedLen;
 
 #[test]
 fn uninhabited() {
@@ -85,3 +87,118 @@ fn eq() {
     assert!(!(x != x));
     assert_eq!(*x.0.borrow(), 0);
 }
+
+const SHARED_ITER_MAX: u16 = 100;
+
+fn assert_trusted_len<I: TrustedLen>(_: &I) {}
+
+#[test]
+fn shared_from_iter_normal() {
+    // Exercise the base implementation for non-`TrustedLen` iterators.
+    {
+        // `Filter` is never `TrustedLen` since we don't
+        // know statically how many elements will be kept:
+        let iter = (0..SHARED_ITER_MAX).filter(|x| x % 2 == 0).map(Box::new);
+
+        // Collecting into a `Vec<T>` or `Rc<[T]>` should make no difference:
+        let vec = iter.clone().collect::<Vec<_>>();
+        let rc = iter.collect::<Rc<[_]>>();
+        assert_eq!(&*vec, &*rc);
+
+        // Clone a bit and let these get dropped.
+        {
+            let _rc_2 = rc.clone();
+            let _rc_3 = rc.clone();
+            let _rc_4 = Rc::downgrade(&_rc_3);
+        }
+    } // Drop what hasn't been here.
+}
+
+#[test]
+fn shared_from_iter_trustedlen_normal() {
+    // Exercise the `TrustedLen` implementation under normal circumstances
+    // where `size_hint()` matches `(_, Some(exact_len))`.
+    {
+        let iter = (0..SHARED_ITER_MAX).map(Box::new);
+        assert_trusted_len(&iter);
+
+        // Collecting into a `Vec<T>` or `Rc<[T]>` should make no difference:
+        let vec = iter.clone().collect::<Vec<_>>();
+        let rc = iter.collect::<Rc<[_]>>();
+        assert_eq!(&*vec, &*rc);
+        assert_eq!(mem::size_of::<Box<u16>>() * SHARED_ITER_MAX as usize, mem::size_of_val(&*rc));
+
+        // Clone a bit and let these get dropped.
+        {
+            let _rc_2 = rc.clone();
+            let _rc_3 = rc.clone();
+            let _rc_4 = Rc::downgrade(&_rc_3);
+        }
+    } // Drop what hasn't been here.
+
+    // Try a ZST to make sure it is handled well.
+    {
+        let iter = (0..SHARED_ITER_MAX).map(|_| ());
+        let vec = iter.clone().collect::<Vec<_>>();
+        let rc = iter.collect::<Rc<[_]>>();
+        assert_eq!(&*vec, &*rc);
+        assert_eq!(0, mem::size_of_val(&*rc));
+        {
+            let _rc_2 = rc.clone();
+            let _rc_3 = rc.clone();
+            let _rc_4 = Rc::downgrade(&_rc_3);
+        }
+    }
+}
+
+#[test]
+#[should_panic = "I've almost got 99 problems."]
+fn shared_from_iter_trustedlen_panic() {
+    // Exercise the `TrustedLen` implementation when `size_hint()` matches
+    // `(_, Some(exact_len))` but where `.next()` drops before the last iteration.
+    let iter = (0..SHARED_ITER_MAX)
+        .map(|val| {
+            match val {
+                98 => panic!("I've almost got 99 problems."),
+                _ => Box::new(val),
+            }
+        });
+    assert_trusted_len(&iter);
+    let _ = iter.collect::<Rc<[_]>>();
+
+    panic!("I am unreachable.");
+}
+
+#[test]
+fn shared_from_iter_trustedlen_no_fuse() {
+    // Exercise the `TrustedLen` implementation when `size_hint()` matches
+    // `(_, Some(exact_len))` but where the iterator does not behave in a fused manner.
+    struct Iter(std::vec::IntoIter<Option<Box<u8>>>);
+
+    unsafe impl TrustedLen for Iter {}
+
+    impl Iterator for Iter {
+        fn size_hint(&self) -> (usize, Option<usize>) {
+            (2, Some(2))
+        }
+
+        type Item = Box<u8>;
+
+        fn next(&mut self) -> Option<Self::Item> {
+            self.0.next().flatten()
+        }
+    }
+
+    let vec = vec![
+        Some(Box::new(42)),
+        Some(Box::new(24)),
+        None,
+        Some(Box::new(12)),
+    ];
+    let iter = Iter(vec.into_iter());
+    assert_trusted_len(&iter);
+    assert_eq!(
+        &[Box::new(42), Box::new(24)],
+        &*iter.collect::<Rc<[_]>>()
+    );
+}