vortex_buffer/
trusted_len.rs

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

vortex_buffer/
trusted_len.rs
