1 | use crate::intrinsics::{unchecked_add, unchecked_sub}; |
2 | use crate::iter::{FusedIterator, TrustedLen}; |
3 | use crate::num::NonZero; |
4 | use 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)] |
13 | pub(crate) struct IndexRange { |
14 | start: usize, |
15 | end: usize, |
16 | } |
17 | |
18 | impl 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 | |
115 | impl 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 | |
141 | impl 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 | |
159 | impl 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. |
167 | unsafe impl TrustedLen for IndexRange {} |
168 | |
169 | impl FusedIterator for IndexRange {} |
170 | |