Skip to content

Commit a421e3f

Browse files
committed
ADD: drain() method and Drain iterator
This add the drain functionality similar to std Vec's drain to HeaderVec. The `with_weakfix()` things are not needed for Drain (I was wrong in a earlier commit message) but they will be required for upcoming Splice functionality. Since vec::Drain depends on a few nightly features internally but we want to stay compatible with stable a few things are backported from nightly in `future_slice`. OTOH we can already stabilize Drain::keep_rest(). Most code was taken from std::vec and minimally adapted to work for HeaderVec.
1 parent efc96c0 commit a421e3f

File tree

4 files changed

+400
-1
lines changed

4 files changed

+400
-1
lines changed

src/drain.rs

Lines changed: 242 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,242 @@
1+
#![cfg(feature = "std")]
2+
3+
use core::{
4+
any::type_name,
5+
fmt,
6+
mem::{self},
7+
ptr::{self, NonNull},
8+
};
9+
10+
use std::{iter::FusedIterator, mem::ManuallyDrop, slice};
11+
12+
use crate::HeaderVec;
13+
14+
/// A draining iterator for `HeaderVec<H, T>`.
15+
///
16+
/// This `struct` is created by [`HeaderVec::drain`].
17+
/// See its documentation for more.
18+
///
19+
/// # Feature compatibility
20+
///
21+
/// The `drain()` API and [`Drain`] iterator are only available when the `std` feature is
22+
/// enabled.
23+
///
24+
/// # Example
25+
///
26+
/// ```
27+
/// # use header_vec::HeaderVec;
28+
/// let mut hv: HeaderVec<(), _> = HeaderVec::from([0, 1, 2]);
29+
/// let iter: header_vec::Drain<'_, _, _> = hv.drain(..);
30+
/// ```
31+
pub struct Drain<'a, H, T> {
32+
/// Index of tail to preserve
33+
pub(super) tail_start: usize,
34+
/// Length of tail
35+
pub(super) tail_len: usize,
36+
/// Current remaining range to remove
37+
pub(super) iter: slice::Iter<'a, T>,
38+
pub(super) vec: NonNull<HeaderVec<H, T>>,
39+
}
40+
41+
impl<H: fmt::Debug, T: fmt::Debug> fmt::Debug for Drain<'_, H, T> {
42+
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
43+
f.debug_struct(&format!(
44+
"Drain<{}, {}>",
45+
type_name::<H>(),
46+
type_name::<T>()
47+
))
48+
.field("header", unsafe { self.vec.as_ref() })
49+
.field("iter", &self.iter.as_slice())
50+
.finish()
51+
}
52+
}
53+
54+
impl<H, T> Drain<'_, H, T> {
55+
/// Returns the remaining items of this iterator as a slice.
56+
///
57+
/// # Examples
58+
///
59+
/// ```
60+
/// # use header_vec::HeaderVec;
61+
/// let mut hv: HeaderVec<(), _> = HeaderVec::from(['a', 'b', 'c']);
62+
/// let mut drain = hv.drain(..);
63+
/// assert_eq!(drain.as_slice(), &['a', 'b', 'c']);
64+
/// let _ = drain.next().unwrap();
65+
/// assert_eq!(drain.as_slice(), &['b', 'c']);
66+
/// ```
67+
#[must_use]
68+
pub fn as_slice(&self) -> &[T] {
69+
self.iter.as_slice()
70+
}
71+
72+
/// Keep unyielded elements in the source `HeaderVec`.
73+
///
74+
/// # Examples
75+
///
76+
/// ```
77+
/// # use header_vec::HeaderVec;
78+
/// let mut hv: HeaderVec<(), _> = HeaderVec::from(['a', 'b', 'c']);
79+
/// let mut drain = hv.drain(..);
80+
///
81+
/// assert_eq!(drain.next().unwrap(), 'a');
82+
///
83+
/// // This call keeps 'b' and 'c' in the vec.
84+
/// drain.keep_rest();
85+
///
86+
/// // If we wouldn't call `keep_rest()`,
87+
/// // `hv` would be empty.
88+
/// assert_eq!(hv.as_slice(), ['b', 'c']);
89+
/// ```
90+
pub fn keep_rest(self) {
91+
let mut this = ManuallyDrop::new(self);
92+
93+
unsafe {
94+
let source_vec = this.vec.as_mut();
95+
96+
let start = source_vec.len();
97+
let tail = this.tail_start;
98+
99+
let unyielded_len = this.iter.len();
100+
let unyielded_ptr = this.iter.as_slice().as_ptr();
101+
102+
// ZSTs have no identity, so we don't need to move them around.
103+
if std::mem::size_of::<T>() != 0 {
104+
let start_ptr = source_vec.as_mut_ptr().add(start);
105+
106+
// memmove back unyielded elements
107+
if unyielded_ptr != start_ptr {
108+
let src = unyielded_ptr;
109+
let dst = start_ptr;
110+
111+
ptr::copy(src, dst, unyielded_len);
112+
}
113+
114+
// memmove back untouched tail
115+
if tail != (start + unyielded_len) {
116+
let src = source_vec.as_ptr().add(tail);
117+
let dst = start_ptr.add(unyielded_len);
118+
ptr::copy(src, dst, this.tail_len);
119+
}
120+
}
121+
122+
source_vec.set_len(start + unyielded_len + this.tail_len);
123+
}
124+
}
125+
}
126+
127+
impl<H, T> AsRef<[T]> for Drain<'_, H, T> {
128+
fn as_ref(&self) -> &[T] {
129+
self.as_slice()
130+
}
131+
}
132+
133+
unsafe impl<H: Sync, T: Sync> Sync for Drain<'_, H, T> {}
134+
unsafe impl<H: Send, T: Send> Send for Drain<'_, H, T> {}
135+
136+
impl<H, T> Iterator for Drain<'_, H, T> {
137+
type Item = T;
138+
139+
#[inline]
140+
fn next(&mut self) -> Option<T> {
141+
self.iter
142+
.next()
143+
.map(|elt| unsafe { ptr::read(elt as *const _) })
144+
}
145+
146+
fn size_hint(&self) -> (usize, Option<usize>) {
147+
self.iter.size_hint()
148+
}
149+
}
150+
151+
impl<H, T> DoubleEndedIterator for Drain<'_, H, T> {
152+
#[inline]
153+
fn next_back(&mut self) -> Option<T> {
154+
self.iter
155+
.next_back()
156+
.map(|elt| unsafe { ptr::read(elt as *const _) })
157+
}
158+
}
159+
160+
impl<H, T> Drop for Drain<'_, H, T> {
161+
fn drop(&mut self) {
162+
/// Moves back the un-`Drain`ed elements to restore the original `Vec`.
163+
struct DropGuard<'r, 'a, H, T>(&'r mut Drain<'a, H, T>);
164+
165+
impl<H, T> Drop for DropGuard<'_, '_, H, T> {
166+
fn drop(&mut self) {
167+
if self.0.tail_len > 0 {
168+
unsafe {
169+
let source_vec = self.0.vec.as_mut();
170+
// memmove back untouched tail, update to new length
171+
let start = source_vec.len();
172+
let tail = self.0.tail_start;
173+
if tail != start {
174+
let src = source_vec.as_ptr().add(tail);
175+
let dst = source_vec.as_mut_ptr().add(start);
176+
ptr::copy(src, dst, self.0.tail_len);
177+
}
178+
source_vec.set_len(start + self.0.tail_len);
179+
}
180+
}
181+
}
182+
}
183+
184+
let iter = mem::take(&mut self.iter);
185+
let drop_len = iter.len();
186+
187+
let mut vec = self.vec;
188+
189+
// unstable: if T::IS_ZST { instead we use size_of
190+
if mem::size_of::<T>() == 0 {
191+
// ZSTs have no identity, so we don't need to move them around, we only need to drop the correct amount.
192+
// this can be achieved by manipulating the Vec length instead of moving values out from `iter`.
193+
unsafe {
194+
let vec = vec.as_mut();
195+
let old_len = vec.len();
196+
vec.set_len(old_len + drop_len + self.tail_len);
197+
vec.truncate(old_len + self.tail_len);
198+
}
199+
200+
return;
201+
}
202+
203+
// ensure elements are moved back into their appropriate places, even when drop_in_place panics
204+
let _guard = DropGuard(self);
205+
206+
if drop_len == 0 {
207+
return;
208+
}
209+
210+
// as_slice() must only be called when iter.len() is > 0 because
211+
// it also gets touched by vec::Splice which may turn it into a dangling pointer
212+
// which would make it and the vec pointer point to different allocations which would
213+
// lead to invalid pointer arithmetic below.
214+
let drop_ptr = iter.as_slice().as_ptr();
215+
216+
unsafe {
217+
// drop_ptr comes from a slice::Iter which only gives us a &[T] but for drop_in_place
218+
// a pointer with mutable provenance is necessary. Therefore we must reconstruct
219+
// it from the original vec but also avoid creating a &mut to the front since that could
220+
// invalidate raw pointers to it which some unsafe code might rely on.
221+
let vec_ptr = vec.as_mut().as_mut_ptr();
222+
223+
// PLANNED: let drop_offset = drop_ptr.sub_ptr(vec_ptr); is in nightly
224+
let drop_offset = usize::try_from(drop_ptr.offset_from(vec_ptr)).unwrap_unchecked();
225+
let to_drop = ptr::slice_from_raw_parts_mut(vec_ptr.add(drop_offset), drop_len);
226+
ptr::drop_in_place(to_drop);
227+
}
228+
}
229+
}
230+
231+
impl<H, T> FusedIterator for Drain<'_, H, T> {}
232+
233+
// PLANNED: unstable features
234+
// impl<H, T> ExactSizeIterator for Drain<'_, H, T> {
235+
// fn is_empty(&self) -> bool {
236+
// self.iter.is_empty()
237+
// }
238+
// }
239+
//
240+
// #[unstable(feature = "trusted_len", issue = "37572")]
241+
// unsafe impl<H, T> TrustedLen for Drain<'_, H, T> {}
242+
//

src/future_slice.rs

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
//! This module re-implements a unstable slice functions, these should be removed once they
2+
//! are stabilized. These is copy-pasted with slight modifications from std::slice for
3+
//! functions that do not need language magic.
4+
5+
use std::ops;
6+
7+
#[track_caller]
8+
#[must_use]
9+
pub(crate) fn range<R>(range: R, bounds: ops::RangeTo<usize>) -> ops::Range<usize>
10+
where
11+
R: ops::RangeBounds<usize>,
12+
{
13+
let len = bounds.end;
14+
15+
let start = match range.start_bound() {
16+
ops::Bound::Included(&start) => start,
17+
ops::Bound::Excluded(start) => start
18+
.checked_add(1)
19+
.unwrap_or_else(|| slice_start_index_overflow_fail()),
20+
ops::Bound::Unbounded => 0,
21+
};
22+
23+
let end = match range.end_bound() {
24+
ops::Bound::Included(end) => end
25+
.checked_add(1)
26+
.unwrap_or_else(|| slice_end_index_overflow_fail()),
27+
ops::Bound::Excluded(&end) => end,
28+
ops::Bound::Unbounded => len,
29+
};
30+
31+
if start > end {
32+
slice_index_order_fail(start, end);
33+
}
34+
if end > len {
35+
slice_end_index_len_fail(end, len);
36+
}
37+
38+
ops::Range { start, end }
39+
}
40+
41+
#[track_caller]
42+
const fn slice_start_index_overflow_fail() -> ! {
43+
panic!("attempted to index slice from after maximum usize");
44+
}
45+
46+
#[track_caller]
47+
const fn slice_end_index_overflow_fail() -> ! {
48+
panic!("attempted to index slice up to maximum usize");
49+
}
50+
51+
#[track_caller]
52+
fn slice_index_order_fail(index: usize, end: usize) -> ! {
53+
panic!("slice index start is larger than end, slice index starts at {index} but ends at {end}")
54+
}
55+
56+
#[track_caller]
57+
fn slice_end_index_len_fail(index: usize, len: usize) -> ! {
58+
panic!(
59+
"slice end index is out of range for slice, range end index {index} out of range for slice of length {len}"
60+
)
61+
}

0 commit comments

Comments
 (0)