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