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

core/ops/
index_range.rs
