core/ops/
index_range.rs

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

core/ops/
index_range.rs
