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