1use crate::intrinsics::{unchecked_add, unchecked_sub};
2use crate::iter::{FusedIterator, TrustedLen};
3use crate::num::NonZero;
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 pub const unsafe fn new_unchecked(start: usize, end: usize) -> Self {
23 ub_checks::assert_unsafe_precondition!(
24 check_library_ub,
25 "IndexRange::new_unchecked requires `start <= end`",
26 (start: usize = start, end: usize = end) => start <= end,
27 );
28 IndexRange { start, end }
29 }
30
31 #[inline]
32 pub const fn zero_to(end: usize) -> Self {
33 IndexRange { start: 0, end }
34 }
35
36 #[inline]
37 pub const fn start(&self) -> usize {
38 self.start
39 }
40
41 #[inline]
42 pub const fn end(&self) -> usize {
43 self.end
44 }
45
46 #[inline]
47 pub const fn len(&self) -> usize {
48 // SAFETY: By invariant, this cannot wrap
49 unsafe { 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 { unchecked_add(value, 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 { unchecked_sub(self.end, 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 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 unsafe { unchecked_add(self.start, n) }
87 } else {
88 self.end
89 };
90 let prefix = Self { start: self.start, end: mid };
91 self.start = mid;
92 prefix
93 }
94
95 /// Removes the last `n` items from this range, returning them as an `IndexRange`.
96 /// If there are fewer than `n`, then the whole range is returned and
97 /// `self` is left empty.
98 ///
99 /// This is designed to help implement `Iterator::advance_back_by`.
100 #[inline]
101 pub fn take_suffix(&mut self, n: usize) -> Self {
102 let mid = if n <= self.len() {
103 // SAFETY: We just checked that this will be between start and end,
104 // and thus the addition cannot overflow.
105 unsafe { unchecked_sub(self.end, n) }
106 } else {
107 self.start
108 };
109 let suffix = Self { start: mid, end: self.end };
110 self.end = mid;
111 suffix
112 }
113}
114
115impl Iterator for IndexRange {
116 type Item = usize;
117
118 #[inline]
119 fn next(&mut self) -> Option<usize> {
120 if self.len() > 0 {
121 // SAFETY: We just checked that the range is non-empty
122 unsafe { Some(self.next_unchecked()) }
123 } else {
124 None
125 }
126 }
127
128 #[inline]
129 fn size_hint(&self) -> (usize, Option<usize>) {
130 let len = self.len();
131 (len, Some(len))
132 }
133
134 #[inline]
135 fn advance_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
136 let taken = self.take_prefix(n);
137 NonZero::new(n - taken.len()).map_or(Ok(()), Err)
138 }
139}
140
141impl DoubleEndedIterator for IndexRange {
142 #[inline]
143 fn next_back(&mut self) -> Option<usize> {
144 if self.len() > 0 {
145 // SAFETY: We just checked that the range is non-empty
146 unsafe { Some(self.next_back_unchecked()) }
147 } else {
148 None
149 }
150 }
151
152 #[inline]
153 fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
154 let taken: IndexRange = self.take_suffix(n);
155 NonZero::new(n - taken.len()).map_or(default:Ok(()), f:Err)
156 }
157}
158
159impl ExactSizeIterator for IndexRange {
160 #[inline]
161 fn len(&self) -> usize {
162 self.len()
163 }
164}
165
166// SAFETY: Because we only deal in `usize`, our `len` is always perfect.
167unsafe impl TrustedLen for IndexRange {}
168
169impl FusedIterator for IndexRange {}
170