1 | use crate::intrinsics::{unchecked_add, unchecked_sub}; |
2 | use crate::iter::{FusedIterator, TrustedLen}; |
3 | use 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)] |
12 | pub(crate) struct IndexRange { |
13 | start: usize, |
14 | end: usize, |
15 | } |
16 | |
17 | impl 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 | |
113 | impl 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 | |
139 | impl 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 | |
157 | impl 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. |
165 | unsafe impl TrustedLen for IndexRange {} |
166 | |
167 | impl FusedIterator for IndexRange {} |
168 | |