vortex_buffer/
trusted_len.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4/// Trait for all types which have a known upper-bound.
5///
6/// Functions that receive a `TrustedLen` iterator can assume that it's `size_hint` is exact,
7/// and can pre-allocate memory, unroll loops, or otherwise optimize their implementations
8/// accordingly.
9///
10/// # Safety
11///
12/// The type which implements this trait must provide an exact `Some` upper-bound for its
13/// `size_hint` method. Failure to do so can trigger undefined behavior in users of the trait.
14pub unsafe trait TrustedLen: Iterator {}
15
16/// An adapter that turns any iterator into a `TrustedLen` iterator.
17///
18/// # Safety
19///
20/// The caller must guarantee that the wrapped iterator does indeed have an exact length.
21pub struct TrustedLenAdapter<I> {
22    inner: I,
23    len: usize,
24    #[cfg(debug_assertions)]
25    count: usize,
26}
27
28impl<I: Iterator> Iterator for TrustedLenAdapter<I> {
29    type Item = I::Item;
30
31    #[inline]
32    fn next(&mut self) -> Option<Self::Item> {
33        match self.inner.next() {
34            None => {
35                #[cfg(debug_assertions)]
36                {
37                    assert_eq!(
38                        self.len, self.count,
39                        "TrustedLenAdapter: iterator ended early"
40                    );
41                }
42                None
43            }
44            Some(item) => {
45                #[cfg(debug_assertions)]
46                {
47                    self.count += 1;
48                    assert!(
49                        self.count <= self.len,
50                        "TrustedLenAdapter: iterator yielded more items than promised"
51                    );
52                }
53                Some(item)
54            }
55        }
56    }
57
58    #[inline]
59    fn size_hint(&self) -> (usize, Option<usize>) {
60        (self.len, Some(self.len))
61    }
62}
63
64unsafe impl<I: Iterator> TrustedLen for TrustedLenAdapter<I> {}
65
66/// Extension trait for wrapping any iterator in a [`TrustedLenAdapter`].
67pub trait TrustedLenExt: Iterator + Sized {
68    /// Wraps this iterator in a `TrustedLenAdapter`.
69    ///
70    /// # Safety
71    ///
72    /// The caller must guarantee that the iterator does indeed have an exact length.
73    unsafe fn trusted_len(self) -> TrustedLenAdapter<Self> {
74        let (lower_bound, upper_bound_opt) = self.size_hint();
75        if let Some(upper_bound) = upper_bound_opt {
76            assert_eq!(
77                lower_bound, upper_bound,
78                "TrustedLenExt: iterator size hints must match if upper bound is given"
79            );
80        }
81
82        TrustedLenAdapter {
83            inner: self,
84            len: lower_bound,
85            #[cfg(debug_assertions)]
86            count: 0,
87        }
88    }
89}
90
91impl<I: Iterator> TrustedLenExt for I {}
92
93macro_rules! impl_for_range {
94    ($($typ:ty),*) => {
95        $(
96            unsafe impl TrustedLen for std::ops::Range<$typ> {}
97            unsafe impl TrustedLen for std::ops::RangeInclusive<$typ> {}
98            // StepBy
99            // This is only fine for iterators that are TrustedRandomAccess but instead of adding another trait we just declare step by of ranges as supported
100            unsafe impl TrustedLen for std::iter::StepBy<std::ops::Range<$typ>> {}
101            unsafe impl TrustedLen for std::iter::StepBy<std::ops::RangeInclusive<$typ>> {}
102        )*
103    };
104}
105
106impl_for_range!(u8, u16, u32, u64, i8, i16, i32, i64, usize);
107
108// std::slice related types
109unsafe impl<T> TrustedLen for std::slice::Iter<'_, T> {}
110
111unsafe impl<T> TrustedLen for std::slice::IterMut<'_, T> {}
112
113// Iterator types
114unsafe impl<B, I, F> TrustedLen for std::iter::Map<I, F>
115where
116    I: TrustedLen,
117    F: FnMut(I::Item) -> B,
118{
119}
120
121unsafe impl<I> TrustedLen for std::iter::Skip<I> where I: TrustedLen {}
122
123unsafe impl<'a, I, T: 'a> TrustedLen for std::iter::Copied<I>
124where
125    I: TrustedLen<Item = &'a T>,
126    T: Copy,
127{
128}
129
130unsafe impl<'a, I, T: 'a> TrustedLen for std::iter::Cloned<I>
131where
132    I: TrustedLen<Item = &'a T>,
133    T: Clone,
134{
135}
136
137unsafe impl<T> TrustedLen for std::vec::IntoIter<T> {}
138
139// Arrays
140unsafe impl<T, const N: usize> TrustedLen for std::array::IntoIter<T, N> {}
141
142// Buffer
143unsafe impl<T> TrustedLen for crate::Iter<'_, T> {}
144unsafe impl<T: Copy> TrustedLen for crate::BufferIterator<T> {}
145
146// ProcessResults
147unsafe impl<'a, I, T: 'a, E: 'a> TrustedLen for itertools::ProcessResults<'a, I, E> where
148    I: TrustedLen<Item = Result<T, E>>
149{
150}
151
152// Enumerate
153unsafe impl<I, T> TrustedLen for std::iter::Enumerate<I> where I: TrustedLen<Item = T> {}
154
155// Zip
156unsafe impl<T, U> TrustedLen for std::iter::Zip<T, U>
157where
158    T: TrustedLen,
159    U: Iterator,
160{
161}
162
163unsafe impl<T: Clone> TrustedLen for std::iter::RepeatN<T> {}
164
165// Arrow bit iterators
166unsafe impl<'a> TrustedLen for crate::bit::BitIterator<'a> {}
167unsafe impl<'a> TrustedLen for crate::bit::BitChunkIterator<'a> {}
168unsafe impl<'a> TrustedLen for crate::bit::UnalignedBitChunkIterator<'a> {}
vortex_buffer/trusted_len.rs

vortex_buffer/
trusted_len.rs
