core/ops/
index_range.rs

1use crate::iter::{FusedIterator, TrustedLen};
2use crate::num::NonZero;
3use crate::ops::{NeverShortCircuit, Try};
4use crate::ub_checks;
5
6/// Like a `Range<usize>`, but with a safety invariant that `start <= end`.
7///
8/// This means that `end - start` cannot overflow, allowing some μoptimizations.
9///
10/// (Normal `Range` code needs to handle degenerate ranges like `10..0`,
11///  which takes extra checks compared to only handling the canonical form.)
12#[derive(Clone, Debug, PartialEq, Eq)]
13pub(crate) struct IndexRange {
14    start: usize,
15    end: usize,
16}
17
18impl IndexRange {
19    /// # Safety
20    /// - `start <= end`
21    #[inline]
22    #[track_caller]
23    pub(crate) const unsafe fn new_unchecked(start: usize, end: usize) -> Self {
24        ub_checks::assert_unsafe_precondition!(
25            check_library_ub,
26            "IndexRange::new_unchecked requires `start <= end`",
27            (start: usize = start, end: usize = end) => start <= end,
28        );
29        IndexRange { start, end }
30    }
31
32    #[inline]
33    pub(crate) const fn zero_to(end: usize) -> Self {
34        IndexRange { start: 0, end }
35    }
36
37    #[inline]
38    pub(crate) const fn start(&self) -> usize {
39        self.start
40    }
41
42    #[inline]
43    pub(crate) const fn end(&self) -> usize {
44        self.end
45    }
46
47    #[inline]
48    pub(crate) const fn len(&self) -> usize {
49        // SAFETY: By invariant, this cannot wrap
50        // Using the intrinsic because a UB check here impedes LLVM optimization. (#131563)
51        unsafe { crate::intrinsics::unchecked_sub(self.end, self.start) }
52    }
53
54    /// # Safety
55    /// - Can only be called when `start < end`, aka when `len > 0`.
56    #[inline]
57    unsafe fn next_unchecked(&mut self) -> usize {
58        debug_assert!(self.start < self.end);
59
60        let value = self.start;
61        // SAFETY: The range isn't empty, so this cannot overflow
62        self.start = unsafe { value.unchecked_add(1) };
63        value
64    }
65
66    /// # Safety
67    /// - Can only be called when `start < end`, aka when `len > 0`.
68    #[inline]
69    unsafe fn next_back_unchecked(&mut self) -> usize {
70        debug_assert!(self.start < self.end);
71
72        // SAFETY: The range isn't empty, so this cannot overflow
73        let value = unsafe { self.end.unchecked_sub(1) };
74        self.end = value;
75        value
76    }
77
78    /// Removes the first `n` items from this range, returning them as an `IndexRange`.
79    /// If there are fewer than `n`, then the whole range is returned and
80    /// `self` is left empty.
81    ///
82    /// This is designed to help implement `Iterator::advance_by`.
83    #[inline]
84    pub(crate) fn take_prefix(&mut self, n: usize) -> Self {
85        let mid = if n <= self.len() {
86            // SAFETY: We just checked that this will be between start and end,
87            // and thus the addition cannot overflow.
88            // Using the intrinsic avoids a superfluous UB check.
89            unsafe { crate::intrinsics::unchecked_add(self.start, n) }
90        } else {
91            self.end
92        };
93        let prefix = Self { start: self.start, end: mid };
94        self.start = mid;
95        prefix
96    }
97
98    /// Removes the last `n` items from this range, returning them as an `IndexRange`.
99    /// If there are fewer than `n`, then the whole range is returned and
100    /// `self` is left empty.
101    ///
102    /// This is designed to help implement `Iterator::advance_back_by`.
103    #[inline]
104    pub(crate) fn take_suffix(&mut self, n: usize) -> Self {
105        let mid = if n <= self.len() {
106            // SAFETY: We just checked that this will be between start and end,
107            // and thus the subtraction cannot overflow.
108            // Using the intrinsic avoids a superfluous UB check.
109            unsafe { crate::intrinsics::unchecked_sub(self.end, n) }
110        } else {
111            self.start
112        };
113        let suffix = Self { start: mid, end: self.end };
114        self.end = mid;
115        suffix
116    }
117
118    #[inline]
119    fn assume_range(&self) {
120        // SAFETY: This is the type invariant
121        unsafe { crate::hint::assert_unchecked(self.start <= self.end) }
122    }
123}
124
125impl Iterator for IndexRange {
126    type Item = usize;
127
128    #[inline]
129    fn next(&mut self) -> Option<usize> {
130        if self.len() > 0 {
131            // SAFETY: We just checked that the range is non-empty
132            unsafe { Some(self.next_unchecked()) }
133        } else {
134            None
135        }
136    }
137
138    #[inline]
139    fn size_hint(&self) -> (usize, Option<usize>) {
140        let len = self.len();
141        (len, Some(len))
142    }
143
144    #[inline]
145    fn advance_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
146        let taken = self.take_prefix(n);
147        NonZero::new(n - taken.len()).map_or(Ok(()), Err)
148    }
149
150    #[inline]
151    fn fold<B, F: FnMut(B, usize) -> B>(mut self, init: B, f: F) -> B {
152        self.try_fold(init, NeverShortCircuit::wrap_mut_2(f)).0
153    }
154
155    #[inline]
156    fn try_fold<B, F, R>(&mut self, mut accum: B, mut f: F) -> R
157    where
158        Self: Sized,
159        F: FnMut(B, Self::Item) -> R,
160        R: Try<Output = B>,
161    {
162        // `Range` needs to check `start < end`, but thanks to our type invariant
163        // we can loop on the stricter `start != end`.
164
165        self.assume_range();
166        while self.start != self.end {
167            // SAFETY: We just checked that the range is non-empty
168            let i = unsafe { self.next_unchecked() };
169            accum = f(accum, i)?;
170        }
171        try { accum }
172    }
173}
174
175impl DoubleEndedIterator for IndexRange {
176    #[inline]
177    fn next_back(&mut self) -> Option<usize> {
178        if self.len() > 0 {
179            // SAFETY: We just checked that the range is non-empty
180            unsafe { Some(self.next_back_unchecked()) }
181        } else {
182            None
183        }
184    }
185
186    #[inline]
187    fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
188        let taken = self.take_suffix(n);
189        NonZero::new(n - taken.len()).map_or(Ok(()), Err)
190    }
191
192    #[inline]
193    fn rfold<B, F: FnMut(B, usize) -> B>(mut self, init: B, f: F) -> B {
194        self.try_rfold(init, NeverShortCircuit::wrap_mut_2(f)).0
195    }
196
197    #[inline]
198    fn try_rfold<B, F, R>(&mut self, mut accum: B, mut f: F) -> R
199    where
200        Self: Sized,
201        F: FnMut(B, Self::Item) -> R,
202        R: Try<Output = B>,
203    {
204        // `Range` needs to check `start < end`, but thanks to our type invariant
205        // we can loop on the stricter `start != end`.
206
207        self.assume_range();
208        while self.start != self.end {
209            // SAFETY: We just checked that the range is non-empty
210            let i = unsafe { self.next_back_unchecked() };
211            accum = f(accum, i)?;
212        }
213        try { accum }
214    }
215}
216
217impl ExactSizeIterator for IndexRange {
218    #[inline]
219    fn len(&self) -> usize {
220        self.len()
221    }
222}
223
224// SAFETY: Because we only deal in `usize`, our `len` is always perfect.
225unsafe impl TrustedLen for IndexRange {}
226
227impl FusedIterator for IndexRange {}
core/ops/index_range.rs

core/ops/
index_range.rs
