Drain

接下来,让我们来实现 Drain。 Drain 与 IntoIter 大体上相同,只是它不是消耗 Vec,而是借用 Vec,并且不会修改到其分配。现在我们只实现“基本”的全范围版本。

use std::marker::PhantomData; struct Drain<'a, T: 'a> { // 这里需要限制生命周期, 因此我们使用了 `&'a mut Vec<T>`, // 也就是我们语义上包含的内容, // 我们只会调用 `pop()` 和 `remove(0)` 两个方法 vec: PhantomData<&'a mut Vec<T>>, start: *const T, end: *const T, } impl<'a, T> Iterator for Drain<'a, T> { type Item = T; fn next(&mut self) -> Option<T> { if self.start == self.end { None

——等等,这看着好像很眼熟?IntoIter 和 Drain 有完全相同的结构,让我们再做一些抽象:

struct RawValIter<T> { start: *const T, end: *const T, } impl<T> RawValIter<T> { // 构建 RawValIter 是不安全的,因为它没有关联的生命周期, // 将 RawValIter 存储在与它实际分配相同的结构体中是非常有必要的, // 但这里是具体的实现细节,不用对外公开 unsafe fn new(slice: &[T]) -> Self { RawValIter { start: slice.as_ptr(), end: if slice.len() == 0 { // 如果 `len = 0`, 说明没有分配内存,需要避免使用 offset, // 因为那样会给 LLVM 的 GEP 传递错误的信息 slice.as_ptr() } else { slice.as_ptr().add(slice.len()) } } } } // Iterator 和 DoubleEndedIterator 和 IntoIter 实现起来很类似

IntoIter 我们可以改成这样:

pub struct IntoIter<T> { _buf: RawVec<T>, iter: RawValIter<T>, } impl<T> Iterator for IntoIter<T> { type Item = T; fn next(&mut self) -> Option<T> { self.iter.next() } fn size_hint(&self) -> (usize, Option<usize>) { self.iter.size_hint() } } impl<T> DoubleEndedIterator for IntoIter<T> { fn next_back(&mut self) -> Option<T> { self.iter.next_back() } } impl<T> Drop for IntoIter<T> { fn drop(&mut self) { for _ in &mut *self {} } } impl<T> IntoIterator for Vec<T> { type Item = T; type IntoIter = IntoIter<T>; fn into_iter(self) -> IntoIter<T> { unsafe { let iter = RawValIter::new(&self); let buf = ptr::read(&self.buf); mem::forget(self); IntoIter { iter, _buf: buf, } } } }

请注意,我在这个设计中留下了一些奇怪之处,以使升级 Drain 来处理任意的子范围更容易一些。特别是我们可以让 RawValIter 在 drop 时 drain 自身,但这对更复杂的 Drain 来说是不合适的。我们还使用了一个 slice 来简化 Drain 的初始化。

好了,现在实现 Drain 真的很容易了:

use std::marker::PhantomData; pub struct Drain<'a, T: 'a> { vec: PhantomData<&'a mut Vec<T>>, iter: RawValIter<T>, } impl<'a, T> Iterator for Drain<'a, T> { type Item = T; fn next(&mut self) -> Option<T> { self.iter.next() } fn size_hint(&self) -> (usize, Option<usize>) { self.iter.size_hint() } } impl<'a, T> DoubleEndedIterator for Drain<'a, T> { fn next_back(&mut self) -> Option<T> { self.iter.next_back() } } impl<'a, T> Drop for Drain<'a, T> { fn drop(&mut self) { for _ in &mut *self {} } } impl<T> Vec<T> { pub fn drain(&mut self) -> Drain<T> { let iter = unsafe { RawValIter::new(&self) }; // 这里事关 mem::forget 的安全。 // 如果 Drain 被 forget,我们就会泄露整个 Vec 的内存, // 既然我们始终要做这一步,为何不在这里完成呢? self.len = 0; Drain { iter: iter, vec: PhantomData, } } }

关于mem::forget问题的更多细节,参见关于泄漏的章节