core/ops/
index_range.rs

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

core/ops/
index_range.rs
