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