core/ops/
index_range.rs

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

core/ops/
index_range.rs
