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问题的更多细节,参见关于泄漏的章节