Skip to main content

polars_arrow/
trusted_len.rs

1//! Declares [`TrustedLen`].
2use std::iter::Scan;
3use std::slice::{Iter, IterMut};
4
5/// An iterator of known, fixed size.
6///
7/// A trait denoting Rusts' unstable [TrustedLen](https://doc.rust-lang.org/std/iter/trait.TrustedLen.html).
8/// This is re-defined here and implemented for some iterators until `std::iter::TrustedLen`
9/// is stabilized.
10///
11/// # Safety
12/// This trait must only be implemented when the contract is upheld.
13/// Consumers of this trait must inspect Iterator::size_hint()’s upper bound.
14pub unsafe trait TrustedLen: Iterator {}
15
16unsafe impl<T> TrustedLen for Iter<'_, T> {}
17unsafe impl<T> TrustedLen for IterMut<'_, T> {}
18
19unsafe impl<'a, I, T: 'a> TrustedLen for std::iter::Copied<I>
20where
21    I: TrustedLen<Item = &'a T>,
22    T: Copy,
23{
24}
25unsafe impl<'a, I, T: 'a> TrustedLen for std::iter::Cloned<I>
26where
27    I: TrustedLen<Item = &'a T>,
28    T: Clone,
29{
30}
31
32unsafe impl<I> TrustedLen for std::iter::Enumerate<I> where I: TrustedLen {}
33
34unsafe impl<A, B> TrustedLen for std::iter::Zip<A, B>
35where
36    A: TrustedLen,
37    B: TrustedLen,
38{
39}
40
41unsafe impl<T> TrustedLen for std::slice::ChunksExact<'_, T> {}
42
43unsafe impl<T> TrustedLen for std::slice::Windows<'_, T> {}
44
45unsafe impl<A, B> TrustedLen for std::iter::Chain<A, B>
46where
47    A: TrustedLen,
48    B: TrustedLen<Item = A::Item>,
49{
50}
51
52unsafe impl<T> TrustedLen for std::iter::Once<T> {}
53
54unsafe impl<T> TrustedLen for std::vec::IntoIter<T> {}
55
56unsafe impl<A: Clone> TrustedLen for std::iter::Repeat<A> {}
57unsafe impl<A, F: FnMut() -> A> TrustedLen for std::iter::RepeatWith<F> {}
58unsafe impl<A: TrustedLen> TrustedLen for std::iter::Take<A> {}
59
60unsafe impl<A: Clone> TrustedLen for std::iter::RepeatN<A> {}
61
62unsafe impl<T> TrustedLen for &mut dyn TrustedLen<Item = T> {}
63unsafe impl<T> TrustedLen for Box<dyn TrustedLen<Item = T> + '_> {}
64
65unsafe impl<B, I: TrustedLen, T: FnMut(I::Item) -> B> TrustedLen for std::iter::Map<I, T> {}
66
67unsafe impl<I: TrustedLen + DoubleEndedIterator> TrustedLen for std::iter::Rev<I> {}
68
69unsafe impl<I: Iterator<Item = J>, J> TrustedLen for TrustMyLength<I, J> {}
70unsafe impl<T> TrustedLen for std::ops::Range<T> where std::ops::Range<T>: Iterator {}
71unsafe impl<T> TrustedLen for std::ops::RangeInclusive<T> where std::ops::RangeInclusive<T>: Iterator
72{}
73unsafe impl<A: TrustedLen> TrustedLen for std::iter::StepBy<A> {}
74
75unsafe impl<I, St, F, B> TrustedLen for Scan<I, St, F>
76where
77    F: FnMut(&mut St, I::Item) -> Option<B>,
78    I: TrustedLen,
79{
80}
81
82unsafe impl<T: Copy> TrustedLen for polars_buffer::buffer::IntoIter<T> {}
83
84unsafe impl<K, V> TrustedLen for hashbrown::hash_map::IntoIter<K, V> {}
85
86#[derive(Clone)]
87pub struct TrustMyLength<I: Iterator<Item = J>, J> {
88    iter: I,
89    len: usize,
90}
91
92impl<I, J> TrustMyLength<I, J>
93where
94    I: Iterator<Item = J>,
95{
96    /// Create a new `TrustMyLength` iterator
97    ///
98    /// # Safety
99    ///
100    /// This is safe if the iterator always has the exact length given by `len`.
101    #[inline]
102    pub unsafe fn new(iter: I, len: usize) -> Self {
103        Self { iter, len }
104    }
105}
106
107impl<J: Clone> TrustMyLength<std::iter::RepeatN<J>, J> {
108    /// Create a new `TrustMyLength` iterator that repeats `value` `len` times.
109    pub fn new_repeat_n(value: J, len: usize) -> Self {
110        // SAFETY: This is always safe since repeat(..).take(n) always repeats exactly `n` times`.
111        unsafe { Self::new(std::iter::repeat_n(value, len), len) }
112    }
113}
114
115impl<I, J> Iterator for TrustMyLength<I, J>
116where
117    I: Iterator<Item = J>,
118{
119    type Item = J;
120
121    #[inline]
122    fn next(&mut self) -> Option<Self::Item> {
123        self.iter.next()
124    }
125
126    #[inline]
127    fn size_hint(&self) -> (usize, Option<usize>) {
128        (self.len, Some(self.len))
129    }
130}
131
132impl<I, J> ExactSizeIterator for TrustMyLength<I, J> where I: Iterator<Item = J> {}
133
134impl<I, J> DoubleEndedIterator for TrustMyLength<I, J>
135where
136    I: Iterator<Item = J> + DoubleEndedIterator,
137{
138    #[inline]
139    fn next_back(&mut self) -> Option<Self::Item> {
140        self.iter.next_back()
141    }
142}