高级 The Rust Project Developers 著,Pure White 中译 2026-09-13 14:59:14 · 2 阅读

第10章 实现 Vec

为了将所有的东西整合起来,我们将从头开始编写std::Vec。我们将限制自己使用稳定的 Rust。特别是我们不会使用任何可以让我们的代码变得更漂亮或更高效的内建指令,因为内建指令是永远不稳定的。尽管许多内建指令确实在其他地方变得稳定了(std::ptrstd::mem由许多内建指令组成)。

最终,这意味着我们的实现可能不会利用所有可能的优化,但它也绝不是简陋的。我们肯定会在细枝末节的细节上钻牛角尖,即使问题并不真的值得这样做。

你想要高级的。我们要的就是高级。

布局

首先,我们需要想出结构布局。一个 Vec 有三个部分:一个指向分配的指针,分配的大小,以及已经初始化的元素数量。

直观来说,这意味着我们只需要这样的设计:

pub struct Vec<T> {
    ptr: *mut T,
    cap: usize,
    len: usize,
}

这确实可以编译成功。但是不幸的是,这有些过于严格了。编译器会给我们太严格的可变性(variance)。比如一个&Vec<&'static str>不能用在预期&Vec<&'a str>的地方。参见所有权和生命周期一章中关于可变和 drop checker 的所有细节。

正如我们在所有权一章中看到的,当标准库拥有一个分配对象的原始指针时,它使用Unique<T>来代替*mut T。Unique 是不稳定的,所以如果可能的话,我们希望不要使用它。

简而言之,Unique 是一个原始指针的包装,并声明以下内容:

  • 我们对T是协变的
  • 我们可以拥有一个T类型的值(这和我们在这的例子无关,但是可以参考PhantonData那章来看看为什么真正的std::vec::Vec<T>需要这个)
  • 如果TSend/Sync,我们就是Send/Sync
  • 我们的指针从不为空(所以Option<Vec<T>>是空指针优化的)

我们可以在稳定的 Rust 中实现上述所有的要求。为此,我们不使用Unique<T>,而是使用NonNull<T>,这是对原始指针的另一种包装,它为我们提供了上述的两个属性,即它在T上是协变的,并且被声明为永不为空。通过在TSend/Sync的情况下实现Send/Sync,我们得到与使用Unique<T>相同的结果:

use std::ptr::NonNull;

pub struct Vec<T> {
    ptr: NonNull<T>,
    cap: usize,
    len: usize,
}

unsafe impl<T: Send> Send for Vec<T> {}
unsafe impl<T: Sync> Sync for Vec<T> {}
fn main() {}

分配内存

使用NonNull会给 Vec(甚至是所有的 std collections)的一个重要特性带来麻烦:创建一个空的 Vec 实际上根本就没有分配。这与分配一个零大小的内存块不同,因为全局分配器不允许这样做(会导致未定义行为!)。所以,如果我们不能分配,但也不能在ptr里放一个空指针,我们在Vec::new里做什么?好吧,我们只是在里面放一些其他的垃圾值。

这并不会有问题,因为我们已经有了cap == 0作为尚未分配的哨兵。我们甚至不需要在任何代码中特别处理它,因为我们通常需要检查cap > lenlen > 0。在这里,Rust 推荐使用的值是mem::align_of::<T>()NonNull为此提供了一个便利。NonNull::dangling()。有相当多的地方我们会想使用dangling,因为没有真正的分配可言,但null会让编译器做坏事。

所以,代码如下:

use std::mem;

impl<T> Vec<T> {
    pub fn new() -> Self {
        assert!(mem::size_of::<T>() != 0, "We're not ready to handle ZSTs");
        Vec {
            ptr: NonNull::dangling(),
            len: 0,
            cap: 0,
        }
    }
}
fn main() {}

我在这里使用了断言,是因为零大小的类型需要在我们的代码中进行一些特殊的处理,我想把这个问题暂时延后。如果没有这个断言,我们早期的一些实现会导致一些非常糟糕的事情。

接下来,我们需要弄清楚,当我们确实想要分配内存时,究竟该怎么做。为此,我们使用全局分配函数allocreallocdealloc,这些函数在稳定的 Rust 中可以使用std::alloc。在std::alloc::Global类型稳定后,这些函数将被废弃。

我们还需要一种方法来处理内存不足(OOM)的情况。标准库提供了一个函数alloc::handle_alloc_error,它将以特定平台的方式中止程序。我们选择中止而不是 panic 的原因是,unwinding 会导致分配的发生,而当你的分配器刚刚回来说“嘿,我没有更多的内存了”时,这似乎是一件坏事。

当然,这看起来有点蠢,因为大多数平台实际上不会以传统方式耗尽内存。如果你顺理成章地用完了所有的内存,你的操作系统可能会通过其他方式杀死这个应用程序。我们最有可能触发 OOM 的方式是一次性要求大量的内存(例如,理论地址空间的一半)。因此,panic 可能是没问题的,不会发生什么坏事。不过,我们还是想尽可能地像标准库一样,所以我们就把整个程序杀掉。

好了,现在我们可以写 grow 的代码了,简单来说,逻辑应该是这样的:

if cap == 0:
    allocate()
    cap = 1
else:
    reallocate()
    cap *= 2

但是 Rust 唯一支持的分配器 API 太低级了,我们需要做相当多的额外工作。我们还需要防范一些特殊情况,这些情况可能发生在真正的大分配或空分配中。

特别是,ptr::offset会给我们带来很多麻烦,因为它有 LLVM 的 GEP(译者注:GetElementPtr) inbounds 指令的语义。如果你有幸没有处理过这个指令,这里是 GEP 的大致故事:别名分析、别名分析、别名分析!对于一个优化的编译器来说,能够推理出数据的依赖性和别名是超级重要的。

作为一个简单的例子,考虑下面的代码片段:

*x *= 7;
*y *= 3;

如果编译器能够证明xy指向内存中的不同位置,理论上这两个操作可以并行执行(例如将它们加载到不同的寄存器中,并对它们独立工作)。然而,编译器在一般情况下不能这样做,因为如果 x 和 y 指向内存中的同一位置,操作需要对相同的值进行,而且它们不能在事后被合并。

当你使用 GEP inbounds 时,你就是在明确地告诉 LLVM,你要做的偏移是在一个“已分配”对象的范围内(within the bounds of a single "allocated" entity.)。这达到的效果是,LLVM 可以假设,如果两个指针已知指向两个不相干的对象,那么这些指针的所有偏移量被认为不会导致别名(因为你不会在内存中的某个随机地方结束)。LLVM 对 GEP 的偏移量进行了大量的优化,而界内偏移量(inbounds offsets)是所有偏移量中最好的,所以我们尽可能地使用它们是很重要的。

这是 GEP 的作用,它怎么会给我们带来麻烦呢?

第一个问题是,我们用无符号的整数来索引数组,但是 GEP(以及由此产生的ptr::offset)需要一个有符号的整数。这意味着一半的看似有效的数组索引会溢出 GEP,并且在实际上是走错了方向!因此,我们必须将所有的分配限制在isize::MAX元素。这实际上意味着我们只需要担心字节大小的对象,因为例如> isize::MAX``u16s 将真正耗尽系统的所有内存。然而,为了避免出现微妙的边界情况,即有人将一些< isize::MAX对象的数组重新解释为字节,std 将所有分配限制为isize::MAX字节。

在 Rust 目前支持的所有 64 位平台上,我们被人为地限制在明显少于所有 64 位的地址空间(现代 x64 平台只暴露了 48 位寻址),所以我们可以依靠首先耗尽内存。然而,在 32 位目标上,特别是那些有扩展使用更多地址空间的目标(PAE x86 或 x32),理论上是可以成功分配超过isize::MAX字节的内存的。

然而,由于这是一个教程,我们在这里不会特别优化,只是无条件地检查,而不是使用聪明的平台特定的cfgs。

我们需要担心的另一个情况是空分配。我们需要担心两种空分配的情况。对于任意 T:cap = 0;和对于零大小的类型(zero-sized types)cap > 0

这些情况很棘手,因为它们归结于 LLVM 对“分配”的理解。LLVM 的分配概念要比我们通常使用的方式抽象得多。因为 LLVM 需要与不同语言的语义和自定义分配器一起工作,所以它不能真正深入地理解分配。相反,分配背后的主要想法是“不与其他东西重叠”。也就是说,堆分配、栈分配和 globals 不会随机地重叠在一起。没错,这就是别名分析。因此,Rust 在技术上可以对分配的概念做一些快速和松散的处理,只要它是一致的

回到空分配的情况,有几个地方我们想用 0 来抵消,这是通用代码的结果。那么问题来了:这样做是否一致?对于零大小的类型,我们的结论是,用任意数量的元素进行 GEP 界内偏移确实是一致的。这是一个运行时的无用功,因为每个元素都不占用空间,假装在0x01处有无限的零尺寸类型分配也是可以的。没有分配器会分配这个地址,因为他们不会分配0x00,而且他们一般会分配到高于一个字节的最小对齐。另外,一般来说,整个第一页的内存是被保护的,不会被分配(在许多平台上,是整个 4k)。

然而,对于正值大小的类型怎么办呢?这个问题就有点棘手了。原则上,你可以说 0 的偏移量没有给 LLVM 带来任何信息:要么地址之前有一个元素,要么在它之后,但它不能知道是哪个。然而,我们选择了保守的假设,即它可能会做坏事。因此,我们将明确地防止这种情况。

好了,说了这么多废话,让我们实际分配一些内存吧:

use std::alloc::{self, Layout};

impl<T> Vec<T> {
    fn grow(&mut self) {
        let (new_cap, new_layout) = if self.cap == 0 {
            (1, Layout::array::<T>(1))
        } else {
            // 因为 self.cap <= isize::MAX,所以不会溢出
            let new_cap = 2 * self.cap;
            (new_cap, Layout::array::<T>(new_cap))
        };

        // `Layout::array` 会检查分配的字节数是否在 1..=isize::MAX 之间,否则会返回错误。
        // 因此,不可能出现分配 0 字节的情况。
        let new_layout = new_layout.expect("Allocation too large");

        let new_ptr = if self.cap == 0 {
            unsafe { alloc::alloc(new_layout) }
        } else {
            let old_layout = Layout::array::<T>(self.cap).unwrap();
            let old_ptr = self.ptr.as_ptr() as *mut u8;
            unsafe { alloc::realloc(old_ptr, old_layout, new_layout.size()) }
        };

        // 如果分配失败,`new_ptr` 就会成为空指针,我们需要对应 abort 的操作
        self.ptr = match NonNull::new(new_ptr as *mut T) {
            Some(p) => p,
            None => alloc::handle_alloc_error(new_layout),
        };
        self.cap = new_cap;
    }
}
fn main() {}

Push 和 Pop

好了,我们现在可以初始化,也可以分配了。让我们实际实现一些功能吧! 让我们从push开始。它所需要做的就是检查我们是否已经满了并 grow,然后无条件地写到下一个索引,接着增加我们的长度。

在写入时,我们必须注意不要对我们想要写入的内存做解引用。最坏的情况是,它是来自分配器的真正未初始化的内存(里面是垃圾值)。最好的情况是,它是我们 pop 出的一些旧值的地址。无论是哪种情况,我们都不能索引到那个地址并解引用,因为这将把该内存认为是一个 T 类型的存活的实例;更糟的是,foo[idx] = x会试图在foo[idx]的旧值上调用drop!

正确的方法是使用ptr::write,它只是盲目地用我们提供的值的位来覆盖目标地址,而不会对该地址做解引用。

对于push,如果旧的 len(在 push 被调用之前)是 0,那么我们正好想写到第 0 个索引,所以我们应该用旧的 len 来作为写入的索引。

pub fn push(&mut self, elem: T) {
    if self.len == self.cap { self.grow(); }

    unsafe {
        ptr::write(self.ptr.as_ptr().add(self.len), elem);
    }

    // 不可能出错,因为出错之前一定会 OOM(out of memory)
    self.len += 1;
}

是不是很简单! 那么pop呢?虽然这次我们要访问的索引被初始化了,但 Rust 不会让我们直接解构内存的位置来把实例移动(move)出去,因为这将使内存未被初始化(译者注:和 push 一样,如果 pop 出的是在 Vec 的内存中的值,那么当这个值被丢弃后,Vec 的这块内存会被 drop,这就出大事了)! 为此我们需要ptr::read,它只是从目标地址复制出 bit,并将其解释为 T 类型的值。这将使这个地址的内存在逻辑上未被初始化,尽管事实上那里有一个完美的 T 的实例。

对于pop,举个例子,如果旧的 len 是 1,那我们正好想从第 0 个索引中读出,所以我们应该用新的 len 来作为读出的索引。

pub fn pop(&mut self) -> Option<T> {
    if self.len == 0 {
        None
    } else {
        self.len -= 1;
        unsafe {
            Some(ptr::read(self.ptr.as_ptr().add(self.len)))
        }
    }
}

释放内存

接下来我们应该实现 Drop,这样我们就不会大规模地泄漏大量的资源。最简单的方法是直接调用pop,直到它产生 None,然后再释放我们的 buffer。注意,如果T: !Drop的话,调用pop是不需要的。理论上,我们可以询问 Rust 是否T need_drop并省略对pop的调用。然而在实践中,LLVM 在做类似这样的简单的无副作用的删除代码方面非常好,所以我就省得麻烦了,除非你注意到它没有被优化掉(在这种情况下它被优化了)。

要注意的是,当self.cap == 0时,我们不能调用alloc::dealloc,因为在这种情况下我们实际上没有分配任何内存。

impl<T> Drop for Vec<T> {
    fn drop(&mut self) {
        if self.cap != 0 {
            while let Some(_) = self.pop() { }
            let layout = Layout::array::<T>(self.cap).unwrap();
            unsafe {
                alloc::dealloc(self.ptr.as_ptr() as *mut u8, layout);
            }
        }
    }
}

Deref

好了!我们已经实现了一个基本像样的栈。我们可以 push 和 pop,我们还可以自己 drop。然而,我们还需要一大堆的功能。特别是, 尽管我们有了一个合适的数组,但还没有切片的功能。这其实很容易解决:我们可以实现Deref<Target=[T]>。这将神奇地使我们的 Vec 在各种条件下强制成为一个切片,并且表现得像一个切片。

我们只需要slice::from_raw_parts。它将为我们正确处理空切片。之后当我们设置了零大小的类型支持,它也会对这些类型进行正确的处理。

use std::ops::Deref;

impl<T> Deref for Vec<T> {
    type Target = [T];
    fn deref(&self) -> &[T] {
        unsafe {
            std::slice::from_raw_parts(self.ptr.as_ptr(), self.len)
        }
    }
}

还有 DerefMut:

use std::ops::DerefMut;

impl<T> DerefMut for Vec<T> {
    fn deref_mut(&mut self) -> &mut [T] {
        unsafe {
            std::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.len)
        }
    }
}

现在我们有了lenfirstlast、索引、切片、排序、iteriter_mut,以及 slice 提供的其他各种功能啦!

插入和删除

slice 提供的东西是insertremove,所以我们接下来做这些。

insert 需要将目标索引的所有元素向右移动一个。要做到这一点,我们需要使用ptr::copy,它是 C 语言memmove的 Rust 版本。它将一些内存块从一个位置复制到另一个位置,正确处理源和目标重叠的情况(这在这里肯定会发生)。

如果我们在索引i处插入,我们要使用旧的 len 将[i ... len]转移到[i+1 ... len+1]

pub fn insert(&mut self, index: usize, elem: T) {
    // 注意:`<=` 是因为我们可以把值插入到任何索引范围 ([0,length-1]) 内的位置之后
    // 这种情况等同于 push
    assert!(index <= self.len, "index out of bounds");
    if self.len == self.cap { self.grow(); }

    unsafe {
        // ptr::copy(src, dest, len) 的含义: "从 src 复制连续的 len 个元素到 dst "
        ptr::copy(
            self.ptr.as_ptr().add(index),
            self.ptr.as_ptr().add(index + 1),
            self.len - index,
        );
        ptr::write(self.ptr.as_ptr().add(index), elem);
    }

    self.len += 1;
}

remove 的行为方式正好相反。我们需要将所有的元素从[i+1 ... len + 1]转移到[i ... len],使用新的 len。

pub fn remove(&mut self, index: usize) -> T {
    // 注意:使用 `<` 是因为 index 不能删除超出元素下标的范围
    assert!(index < self.len, "index out of bounds");
    unsafe {
        self.len -= 1;
        let result = ptr::read(self.ptr.as_ptr().add(index));
        ptr::copy(
            self.ptr.as_ptr().add(index + 1),
            self.ptr.as_ptr().add(index),
            self.len - index,
        );
        result
    }
}

IntoIter

让我们继续,接下来写迭代器。iteriter_mut已经为我们写好了,感谢 Deref 的魔法。然而,有两个有趣的迭代器是 Vec 提供的,而 slice 不能提供:into_iterdrain

IntoIter 通过消耗掉 Vec 的值(获取 Vec 的所有权),并因此可以产生其元素的值(所有权)。为了实现这个目的,IntoIter 需要控制 Vec 的分配。

IntoIter 也需要是 DoubleEnded,以便能够从两端读取。从后面读取可以通过调用pop来实现,但是从前面读取就比较困难了。我们可以调用remove(0),但这将是非常昂贵的。相反,我们将使用ptr::read来复制 Vec 两端的值,而不改变缓冲区。

为了做到这一点,我们将使用一个非常常见的 C 语言的数组迭代习惯。我们将建立两个指针;一个指向数组的开始,另一个指向数组结束后的一个元素。当我们想从一端获得一个元素时,我们将读出指向那一端的值,并将指针移到另一端。当这两个指针相等时,我们就知道我们已经完成了。

注意,对于nextnext_back来说,读取和偏移的顺序是相反的。对于next_back来说,指针总是在它想读取的元素之后,而对于next来说,指针正好在它想读取的元素上。要想知道为什么会这样,请考虑除一个元素外的所有元素都已经产生的情况。

这个数组看起来像这样:

          S  E
[X, X, X, O, X, X, X]

如果 E 直接指向它想产生的下一个元素,它将与没有更多元素可以产生的情况没有区别。

虽然我们在迭代过程中实际上并不关心它,但我们也需要保留 Vec 的分配信息,以便在 IntoIter 被丢弃后释放它。

所以我们将使用下面的结构。

pub struct IntoIter<T> {
    buf: NonNull<T>,
    cap: usize,
    start: *const T,
    end: *const T,
}

而这就是我们最终的初始化结果:

impl<T> IntoIterator for Vec<T> {
    type Item = T;
    type IntoIter = IntoIter<T>;
    fn into_iter(self) -> IntoIter<T> {
        // 确保 Vec 不会被 drop,因为那样会释放内存
        let vec = ManuallyDrop::new(self);

        // 因为 Vec 实现了 Drop,所以我们不能销毁它
        let ptr = vec.ptr;
        let cap = vec.cap;
        let len = vec.len;

        IntoIter {
            buf: ptr,
            cap,
            start: ptr.as_ptr(),
            end: if cap == 0 {
                // 不能通过这个指针获取偏移,因为没有分配内存
                ptr.as_ptr()
            } else {
                unsafe { ptr.as_ptr().add(len) }
            },
        }
    }
}

向前迭代:

impl<T> Iterator for IntoIter<T> {
    type Item = T;
    fn next(&mut self) -> Option<T> {
        if self.start == self.end {
            None
        } else {
            unsafe {
                let result = ptr::read(self.start);
                self.start = self.start.offset(1);
                Some(result)
            }
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let len = (self.end as usize - self.start as usize)
                  / mem::size_of::<T>();
        (len, Some(len))
    }
}

向后迭代:

impl<T> DoubleEndedIterator for IntoIter<T> {
    fn next_back(&mut self) -> Option<T> {
        if self.start == self.end {
            None
        } else {
            unsafe {
                self.end = self.end.offset(-1);
                Some(ptr::read(self.end))
            }
        }
    }
}

因为 IntoIter 拥有其分配的所有权,它需要实现 Drop 来释放它;并且,它也需要在 Drop 里丢弃它所包含的任何没有被迭代到的元素。

impl<T> Drop for IntoIter<T> {
    fn drop(&mut self) {
        if self.cap != 0 {
            // 将剩下的元素 drop
            for _ in &mut *self {}
            let layout = Layout::array::<T>(self.cap).unwrap();
            unsafe {
                alloc::dealloc(self.buf.as_ptr() as *mut u8, layout);
            }
        }
    }
}

RawVec

我们实际上在这里达到了一个有趣的状态:我们在 Vec 和 IntoIter 中重复了指定缓冲区和释放其内存的逻辑。现在我们已经实现了它,并且确定了实际的逻辑重复,这是一个进行一些逻辑压缩的好时机。

我们将抽象出(ptr, cap)对,并为它们编写分配、增长和释放的逻辑:

struct RawVec<T> {
    ptr: NonNull<T>,
    cap: usize,
}

unsafe impl<T: Send> Send for RawVec<T> {}
unsafe impl<T: Sync> Sync for RawVec<T> {}

impl<T> RawVec<T> {
    fn new() -> Self {
        assert!(mem::size_of::<T>() != 0, "TODO: implement ZST support");
        RawVec {
            ptr: NonNull::dangling(),
            cap: 0,
        }
    }

    fn grow(&mut self) {
        // 保证新申请的内存没有超出 `isize::MAX` 字节
        let new_cap = if self.cap == 0 { 1 } else { 2 * self.cap };

        // `Layout::array` 会检查申请的空间是否小于等于 usize::MAX,
        // 但是因为 old_layout.size() <= isize::MAX,
        // 所以这里的 unwrap 永远不可能失败
        let new_layout = Layout::array::<T>(new_cap).unwrap();

        // 保证新申请的内存没有超出 `isize::MAX` 字节
        assert!(new_layout.size() <= isize::MAX as usize, "Allocation too large");

        let new_ptr = if self.cap == 0 {
            unsafe { alloc::alloc(new_layout) }
        } else {
            let old_layout = Layout::array::<T>(self.cap).unwrap();
            let old_ptr = self.ptr.as_ptr() as *mut u8;
            unsafe { alloc::realloc(old_ptr, old_layout, new_layout.size()) }
        };

        // 如果分配失败,`new_ptr` 就会成为空指针,我们需要对应 abort 的操作
        self.ptr = match NonNull::new(new_ptr as *mut T) {
            Some(p) => p,
            None => alloc::handle_alloc_error(new_layout),
        };
        self.cap = new_cap;
    }
}

impl<T> Drop for RawVec<T> {
    fn drop(&mut self) {
        if self.cap != 0 {
            let layout = Layout::array::<T>(self.cap).unwrap();
            unsafe {
                alloc::dealloc(self.ptr.as_ptr() as *mut u8, layout);
            }
        }
    }
}

随后,把 Vec 改成这样:

pub struct Vec<T> {
    buf: RawVec<T>,
    len: usize,
}

impl<T> Vec<T> {
    fn ptr(&self) -> *mut T {
        self.buf.ptr.as_ptr()
    }

    fn cap(&self) -> usize {
        self.buf.cap
    }

    pub fn new() -> Self {
        Vec {
            buf: RawVec::new(),
            len: 0,
        }
    }

    // push/pop/insert/remove 这些操作做了小小的改变,如下所示:
    // * `self.ptr.as_ptr() -> self.ptr()`
    // * `self.cap -> self.cap()`
    // * `self.grow() -> self.buf.grow()`
}

impl<T> Drop for Vec<T> {
    fn drop(&mut self) {
        while let Some(_) = self.pop() {}
        // RawVec 来负责释放内存
    }
}

最后,我们可以把 IntoIter 改得相当简单:

pub struct IntoIter<T> {
    _buf: RawVec<T>, // 我们实际上并不关心这个,只需要他们保证分配的空间不被释放
    start: *const T,
    end: *const T,
}

// next 和 next_back 保持不变,因为它们并没有用到 buf

impl<T> Drop for IntoIter<T> {
    fn drop(&mut self) {
        // 我们只需要确保 Vec 中所有元素都被读取了,
        // 在这之后这些元素会被自动清理
        for _ in &mut *self {}
    }
}

impl<T> IntoIterator for Vec<T> {
    type Item = T;
    type IntoIter = IntoIter<T>;
    fn into_iter(self) -> IntoIter<T> {
        // 需要使用 ptr::read 非安全地把 buf 移出,因为它没有实现 Copy,
        // 而且 Vec 实现了 Drop Trait (因此我们不能销毁它)
        let buf = unsafe { ptr::read(&self.buf) };
        let len = self.len;
        mem::forget(self);

        IntoIter {
            start: buf.ptr.as_ptr(),
            end: if buf.cap == 0 {
                // 不能通过这个指针获取偏移,除非已经分配了内存
                buf.ptr.as_ptr()
            } else {
                unsafe { buf.ptr.as_ptr().add(len) }
            },
            _buf: buf,
        }
    }
}

是不是好多了!

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

处理零大小的类型

是时候了!我们将与 ZST(零大小类型)这个幽灵作斗争。安全的 Rust 从来不需要关心这个问题,但是 Vec 在原始指针和原始分配上非常密集,这正是需要关心零尺寸类型的两种情况。我们需要注意两件事:

  • 如果你在分配大小上传入 0,原始分配器 API 有未定义的行为。
  • 原始指针偏移量对于零大小的类型来说是无效的(no-ops),这将破坏我们的 C 风格指针迭代器。

幸好我们之前把指针迭代器和分配处理分别抽象为RawValIterRawVec。现在回过头来看,多么的方便啊。

分配零大小的类型

那么,如果分配器 API 不支持零大小的分配,我们到底要把什么作为我们的分配来存储呢?当然是NonNull::dangling()! 几乎所有使用 ZST 的操作都是 no-op,因为 ZST 正好有且仅有一个值,因此在存储或加载它们时不需要考虑状态。这实际上延伸到了ptr::readptr::write:它们实际上根本不会去用指针。因此,我们从来不需要改变指针。

然而,请注意,我们之前对在溢出前耗尽内存的防御,在零大小的类型中不再有效了。我们必须明确地防止零大小类型的容量溢出。

由于我们目前的架构,这意味着要写 3 个边界处理,在RawVec的每个方法中都有一个:

impl<T> RawVec<T> {
    fn new() -> Self {
        // 这一段分支代码在编译期间就可以计算出结果返回的结果,返回给 cap
        let cap = if mem::size_of::<T>() == 0 { usize::MAX } else { 0 };

        // `NonNull::dangling()` 有双重含义:
        // `未分配内存 (unallocated)`, `零尺寸 (zero-sized allocation)`
        RawVec {
            ptr: NonNull::dangling(),
            cap: cap,
        }
    }

    fn grow(&mut self) {
        // 因为当 T 的尺寸为 0 时我们设置了 cap 为 usize::MAX
        // 这一步成立意味着 Vec 溢出了
        assert!(mem::size_of::<T>() != 0, "capacity overflow");

        let (new_cap, new_layout) = if self.cap == 0 {
            (1, Layout::array::<T>(1).unwrap())
        } else {
            // 保证新申请的内存没有超出 `isize::MAX` 字节
            let new_cap = 2 * self.cap;

            // `Layout::array` 会检查申请的空间是否小于等于 usize::MAX,
            // 但是因为 old_layout.size() <= isize::MAX,
            // 所以这里的 unwrap 永远不可能失败
            let new_layout = Layout::array::<T>(new_cap).unwrap();
            (new_cap, new_layout)
        };

        // 保证新申请的内存没有超出 `isize::MAX` 字节
        assert!(new_layout.size() <= isize::MAX as usize, "Allocation too large");

        let new_ptr = if self.cap == 0 {
            unsafe { alloc::alloc(new_layout) }
        } else {
            let old_layout = Layout::array::<T>(self.cap).unwrap();
            let old_ptr = self.ptr.as_ptr() as *mut u8;
            unsafe { alloc::realloc(old_ptr, old_layout, new_layout.size()) }
        };

        // 如果分配失败,`new_ptr` 就会成为空指针,我们需要处理这个意外情况
        self.ptr = match NonNull::new(new_ptr as *mut T) {
            Some(p) => p,
            None => alloc::handle_alloc_error(new_layout),
        };
        self.cap = new_cap;
    }
}

impl<T> Drop for RawVec<T> {
    fn drop(&mut self) {
        println!("RawVec<T> Drop called, deallocating memory");
        if self.cap != 0 && std::mem::size_of::<T>() > 0 {
            let layout = std::alloc::Layout::array::<T>(self.cap).unwrap();
            unsafe {
                std::alloc::dealloc(self.ptr.as_ptr() as *mut _, layout);
            }
        }
    }
}

搞定!我们现在支持 push 和 pop 零大小类型。不过,我们的迭代器(不是由 slice Deref 提供的)仍然是一团浆糊。

迭代 ZST

零大小的偏移量是 no-op。这意味着我们目前的设计总是将startend初始化为相同的值,而我们的迭代器将一无所获。目前的解决方案是将指针转为整数,增加,然后再转回。

impl<T> RawValIter<T> {
    unsafe fn new(slice: &[T]) -> Self {
        RawValIter {
            start: slice.as_ptr(),
            end: if mem::size_of::<T>() == 0 {
                ((slice.as_ptr() as usize) + slice.len()) as *const _
            } else if slice.len() == 0 {
                slice.as_ptr()
            } else {
                slice.as_ptr().add(slice.len())
            },
        }
    }
}

现在,我们有了另一个 bug:我们的迭代器不再是完全不运行,而是现在的迭代器永远都在运行。我们需要在我们的迭代器 impls 中做同样的技巧。另外,我们的 size_hint 计算代码将对 ZST 除以 0。既然我们会把这两个指针当作是指向字节的,所以我们就把大小 0 映射到除以 1,这样的话next的代码如下:

fn next(&mut self) -> Option<T> {
    if self.start == self.end {
        None
    } else {
        unsafe {
            let result = ptr::read(self.start);
            self.start = if mem::size_of::<T>() == 0 {
                (self.start as usize + 1) as *const _
            } else {
                self.start.offset(1)
            };
            Some(result)
        }
    }
}

你找到 bug 了嘛?没人看到!连最初的作者也是几年之后闲逛这个页面的时候,觉得这段代码比较可疑,因为这里直接滥用了迭代器的指针当作了计数器,而这就使得了指针不对齐!在使用 ZST 的时候,我们唯一的工作就是必须保证指针对齐!啊这

原始指针在任何时候都不需要对齐,所以使用指针作为计数器的基本技巧是没问题的,但是当它们被传递给ptr::read时,它们应该是对齐的! 这可能是不必要的迂腐操作,因为ptr::read在处理 ZST 时其实是个 noop,但让我们稍微负责一点,当遇到 ZST 时从NonNull::dangling读取。

(或者你也可以在 ZST 路径上调用read_unaligned。两者都可以。因为无论哪种方式,我们都是在无中生有,而且都最终编译成 noop。)

impl<T> Iterator for RawValIter<T> {
    type Item = T;
    fn next(&mut self) -> Option<T> {
        if self.start == self.end {
            None
        } else {
            unsafe {
                if mem::size_of::<T>() == 0 {
                    self.start = (self.start as usize + 1) as *const _;
                    Some(ptr::read(NonNull::<T>::dangling().as_ptr()))
                } else {
                    let old_ptr = self.start;
                    self.start = self.start.offset(1);
                    Some(ptr::read(old_ptr))
                }
            }
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let elem_size = mem::size_of::<T>();
        let len = (self.end as usize - self.start as usize)
                  / if elem_size == 0 { 1 } else { elem_size };
        (len, Some(len))
    }
}

impl<T> DoubleEndedIterator for RawValIter<T> {
    fn next_back(&mut self) -> Option<T> {
        if self.start == self.end {
            None
        } else {
            unsafe {
                if mem::size_of::<T>() == 0 {
                    self.end = (self.end as usize - 1) as *const _;
                    Some(ptr::read(NonNull::<T>::dangling().as_ptr()))
                } else {
                    self.end = self.end.offset(-1);
                    Some(ptr::read(self.end))
                }
            }
        }
    }
}

好啦,迭代器也搞定啦!

我们需要考虑的最后一件事是,当我们的 vector 被 drop 时,它会释放其存活期间分配的内存。对于 ZST,我们没有分配任何内存;事实上,我们从来没有分配过。所以,现在我们的代码存在不健全性:我们仍然试图释放一个用于模拟 vector 中 ZST 的 NonNull::dangling() 指针。这意味着如果我们试图释放一些我们从未分配过的东西,就会导致未定义行为(显然,这是有充分理由的)。为了解决这个问题,在我们的 RawVecDrop trait 中,我们将对其进行调整,以确保我们只释放有大小的类型。

impl<T> Drop for RawVec<T> {
    fn drop(&mut self) {
        println!("RawVec<T> Drop called, deallocating memory");
        if self.cap != 0 && std::mem::size_of::<T>() > 0 {
            let layout = std::alloc::Layout::array::<T>(self.cap).unwrap();
            unsafe {
                std::alloc::dealloc(self.ptr.as_ptr() as *mut _, layout);
            }
        }
    }
}

最终的代码

use std::alloc::{self, Layout};
use std::marker::PhantomData;
use std::mem;
use std::ops::{Deref, DerefMut};
use std::ptr::{self, NonNull};

struct RawVec<T> {
    ptr: NonNull<T>,
    cap: usize,
}

unsafe impl<T: Send> Send for RawVec<T> {}
unsafe impl<T: Sync> Sync for RawVec<T> {}

impl<T> RawVec<T> {
    fn new() -> Self {
        // !0 等价于 usize::MAX, 这一段分支代码在编译期间就可以计算出结果返回的结果,返回给 cap
        let cap = if mem::size_of::<T>() == 0 { !0 } else { 0 };

        // `NonNull::dangling()` 有双重含义:
        // `未分配内存 (unallocated)`, `零尺寸 (zero-sized allocation)`
        RawVec {
            ptr: NonNull::dangling(),
            cap: cap,
        }
    }

    fn grow(&mut self) {
        // 因为当 T 的尺寸为 0 时,我们设置了 cap 为 usize::MAX,
        // 这一步成立便意味着 Vec 溢出了.
        assert!(mem::size_of::<T>() != 0, "capacity overflow");

        let (new_cap, new_layout) = if self.cap == 0 {
            (1, Layout::array::<T>(1))
        } else {
            // 因为 self.cap <= isize::MAX,所以不会溢出
            let new_cap = 2 * self.cap;
            (new_cap, Layout::array::<T>(new_cap))
        };

        // `Layout::array` 会检查分配的字节数是否在 1..=isize::MAX 之间,否则会返回错误。
        // 因此,不可能出现分配 0 字节的情况。
        let new_layout = new_layout.expect("Allocation too large");

        let new_ptr = if self.cap == 0 {
            unsafe { alloc::alloc(new_layout) }
        } else {
            let old_layout = Layout::array::<T>(self.cap).unwrap();
            let old_ptr = self.ptr.as_ptr() as *mut u8;
            unsafe { alloc::realloc(old_ptr, old_layout, new_layout.size()) }
        };

        // 如果分配失败,`new_ptr` 就会成为空指针,我们需要处理该意外情况
        self.ptr = match NonNull::new(new_ptr as *mut T) {
            Some(p) => p,
            None => alloc::handle_alloc_error(new_layout),
        };
        self.cap = new_cap;
    }
}

impl<T> Drop for RawVec<T> {
    fn drop(&mut self) {
        let elem_size = mem::size_of::<T>();

        if self.cap != 0 && elem_size != 0 {
            unsafe {
                alloc::dealloc(
                    self.ptr.as_ptr() as *mut u8,
                    Layout::array::<T>(self.cap).unwrap(),
                );
            }
        }
    }
}

pub struct Vec<T> {
    buf: RawVec<T>,
    len: usize,
}

impl<T> Vec<T> {
    fn ptr(&self) -> *mut T {
        self.buf.ptr.as_ptr()
    }

    fn cap(&self) -> usize {
        self.buf.cap
    }

    pub fn new() -> Self {
        Vec {
            buf: RawVec::new(),
            len: 0,
        }
    }
    pub fn push(&mut self, elem: T) {
        if self.len == self.cap() {
            self.buf.grow();
        }

        unsafe {
            ptr::write(self.ptr().add(self.len), elem);
        }

        // 不会溢出,会先 OOM
        self.len += 1;
    }

    pub fn pop(&mut self) -> Option<T> {
        if self.len == 0 {
            None
        } else {
            self.len -= 1;
            unsafe { Some(ptr::read(self.ptr().add(self.len))) }
        }
    }

    pub fn insert(&mut self, index: usize, elem: T) {
        assert!(index <= self.len, "index out of bounds");
        if self.len == self.cap() {
            self.buf.grow();
        }

        unsafe {
            ptr::copy(
                self.ptr().add(index),
                self.ptr().add(index + 1),
                self.len - index,
            );
            ptr::write(self.ptr().add(index), elem);
        }

        self.len += 1;
    }

    pub fn remove(&mut self, index: usize) -> T {
        assert!(index < self.len, "index out of bounds");

        self.len -= 1;

        unsafe {
            let result = ptr::read(self.ptr().add(index));
            ptr::copy(
                self.ptr().add(index + 1),
                self.ptr().add(index),
                self.len - index,
            );
            result
        }
    }

    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,
        }
    }
}

impl<T> Drop for Vec<T> {
    fn drop(&mut self) {
        while let Some(_) = self.pop() {}
        // RawVec 来负责释放内存
    }
}

impl<T> Deref for Vec<T> {
    type Target = [T];
    fn deref(&self) -> &[T] {
        unsafe { std::slice::from_raw_parts(self.ptr(), self.len) }
    }
}

impl<T> DerefMut for Vec<T> {
    fn deref_mut(&mut self) -> &mut [T] {
        unsafe { std::slice::from_raw_parts_mut(self.ptr(), self.len) }
    }
}

impl<T> IntoIterator for Vec<T> {
    type Item = T;
    type IntoIter = IntoIter<T>;
    fn into_iter(self) -> IntoIter<T> {
        let (iter, buf) = unsafe {
            (RawValIter::new(&self), ptr::read(&self.buf))
        };

        mem::forget(self);

        IntoIter {
            iter: iter,
            _buf: buf,
        }
    }
}

struct RawValIter<T> {
    start: *const T,
    end: *const T,
}

impl<T> RawValIter<T> {
    unsafe fn new(slice: &[T]) -> Self {
        RawValIter {
            start: slice.as_ptr(),
            end: if mem::size_of::<T>() == 0 {
                ((slice.as_ptr() as usize) + slice.len()) as *const _
            } else if slice.len() == 0 {
                slice.as_ptr()
            } else {
                slice.as_ptr().add(slice.len())
            },
        }
    }
}

impl<T> Iterator for RawValIter<T> {
    type Item = T;
    fn next(&mut self) -> Option<T> {
        if self.start == self.end {
            None
        } else {
            unsafe {
                if mem::size_of::<T>() == 0 {
                    self.start = (self.start as usize + 1) as *const _;
                    Some(ptr::read(NonNull::<T>::dangling().as_ptr()))
                } else {
                    let old_ptr = self.start;
                    self.start = self.start.offset(1);
                    Some(ptr::read(old_ptr))
                }
            }
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let elem_size = mem::size_of::<T>();
        let len = (self.end as usize - self.start as usize)
                  / if elem_size == 0 { 1 } else { elem_size };
        (len, Some(len))
    }
}

impl<T> DoubleEndedIterator for RawValIter<T> {
    fn next_back(&mut self) -> Option<T> {
        if self.start == self.end {
            None
        } else {
            unsafe {
                if mem::size_of::<T>() == 0 {
                    self.end = (self.end as usize - 1) as *const _;
                    Some(ptr::read(NonNull::<T>::dangling().as_ptr()))
                } else {
                    self.end = self.end.offset(-1);
                    Some(ptr::read(self.end))
                }
            }
        }
    }
}

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 {}
    }
}

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) {
        // 消耗drain
        for _ in &mut *self {}
    }
}

fn main() {
    tests::create_push_pop();
    tests::iter_test();
    tests::test_drain();
    tests::test_zst();
    println!("All tests finished OK");
}

mod tests {
    use super::*;

    pub fn create_push_pop() {
        let mut v = Vec::new();
        v.push(1);
        assert_eq!(1, v.len());
        assert_eq!(1, v[0]);
        for i in v.iter_mut() {
            *i += 1;
        }
        v.insert(0, 5);
        let x = v.pop();
        assert_eq!(Some(2), x);
        assert_eq!(1, v.len());
        v.push(10);
        let x = v.remove(0);
        assert_eq!(5, x);
        assert_eq!(1, v.len());
    }

    pub fn iter_test() {
        let mut v = Vec::new();
        for i in 0..10 {
            v.push(Box::new(i))
        }
        let mut iter = v.into_iter();
        let first = iter.next().unwrap();
        let last = iter.next_back().unwrap();
        drop(iter);
        assert_eq!(0, *first);
        assert_eq!(9, *last);
    }

    pub fn test_drain() {
        let mut v = Vec::new();
        for i in 0..10 {
            v.push(Box::new(i))
        }
        {
            let mut drain = v.drain();
            let first = drain.next().unwrap();
            let last = drain.next_back().unwrap();
            assert_eq!(0, *first);
            assert_eq!(9, *last);
        }
        assert_eq!(0, v.len());
        v.push(Box::new(1));
        assert_eq!(1, *v.pop().unwrap());
    }

    pub fn test_zst() {
        let mut v = Vec::new();
        for _i in 0..10 {
            v.push(())
        }

        let mut count = 0;

        for _ in v.into_iter() {
            count += 1
        }

        assert_eq!(10, count);
    }
}

评论 (0)