Rust 编程秘籍：深入 Vec::drain 与内存安全的奥秘

文章导读

如果你不熟悉Vec::drain，你可以这样使用它来“清空”一个Vec 中的元素（对于String、HashMap 和许多其他集合类型也有类似的方法）：

let mutvalues = vec![1, 2, 3, 4, 5];
for val in values.drain(1..3) {
    println!("Removed: {val}");
}
println!("Remaining: {values:?}");

上述代码的输出将是：

Removed: 2
Removed: 3
Remaining: [1, 4, 5]

这是 Rust 1.83 版本文档中对Vec::drain 的描述（我强调了一部分）：

批量移除向量中的指定范围，并返回一个迭代器，包含所有移除的元素。如果迭代器在使用前被丢弃，它将丢弃剩余的移除元素。

返回的迭代器保持对向量的可变借用，以优化其实现。

最后一句话，特别是我加粗的部分，引起了我的注意，促使我深入研究实现。

毕竟，有一种完全合理的方式可以做到这一点：接收Vec，复制所有要移除的元素并将它们放入一个新的Vec，更新原始Vec 以移除所有这些元素，并返回一个由新分配的Vec 支持的迭代器。

虽然这样做可能会让计算机提前做很多工作。如果你有一个大型的Vec（数千或数百万的元素），并且你在处理其中的某个部分，这将需要分配很多额外的内存和执行很多复制操作，甚至在我们确定是否会使用这些值之前。

所以 Rust 在这里做了完全不同的事情：它保持对原始Vec 的可变引用，并且只从原始存储中读取和更新。它之所以能做到这一点，是因为 Rust 的所有权规则：只要由Vec::drain 生成的迭代器存在，就不允许其他任何东西获得对原始Vec 的读写访问，因此没有任何东西可以通过使迭代器或其支持的存储无效（例如，通过变异Vec 中的值，改变其长度等）进入错误状态。

Rust 通过创建一个名为Drain 的新数据结构来实现这一点，它保持对原始Vec 的可变引用，并通过使用Vec 的切片来访问Vec 的值的迭代器。当你在Drain 上使用迭代器方法时，它转发到切片上的迭代器。这意味着它不需要自己实现迭代，而是可以使用与任何其他切片迭代器相同的（经过良好优化的！）实现。唯一的区别，这是一个关键的区别，是drain 通过一个不安全的std::ptr::read 调用立即从切片返回值。

如果有人在Drain 迭代器访问期间或之后能够访问Vec 中的值，或者如果Vec “记住”了Drain 访问的所有元素，那将是不合理的。正如我上面提到的，当Drain 迭代器访问时，没有东西可以访问它，因为它通过可变引用获取self。到目前为止，如果你熟悉 Rust，这可能看起来相当直接——std::ptr::read 部分是唯一的不寻常部分。

那么，在处理完排水迭代器后呢？Rust 如何保证合同的那一部分？

这很有趣。

当迭代器被丢弃时——无论是因为你在它上面遍历了一个for 循环的末尾，还是因为你在遍历了一些元素的子集后丢弃了它——Drain 类型的Drop 特征实现接管了。这意味着impl Drop for Drain 负责确保Drain 为Drain 所做的不安全清理保持合理，同时避免泄露我们在创建Drain 时Vec 忘记的内存（我们稍后将回到这一点）。

这篇文章的早期版本说Drain 的Drop 实现负责保持Drain 的合理性。这不仅是错误的，而且是完全相反的！一个Drop 实现永远不能被依赖来保持 Rust 中类型的实现的合理性。我应该知道这一点；在某个时候我确实知道……但我忘记了。这是自 Rust 1.0 之前就已经众所周知的规则，实际上在 Rust 1.0 之前的准备中是一个主要的修复点：“Leakpocalypse”！

这是一个在 Rust 中很常见的模式，值得理解，它也真的很酷，所以让我们一步一步地走过它——全部！我将省略额外的类型参数Allocator，但除此之外，这篇文章包含了实现中的每一段代码（在这种情况下，是 Rust 1.85 夜间构建的）。即便如此，你可能想要将那段代码与这篇文章并排放置，以便于看到所有的上下文！

我们将从特征实现的样板开始：

impl Drop for Drain<'_, T> {
    fn drop(&mut self) {
        // ...
    }
}

关于这一点需要注意的是drop 接受&mut self。这意味着我们不能做任何需要self 所有权的事情，这反过来又激发了我们将要看到的下一件事：

/// 将未 `Drain` 的元素移回以恢复原始 `Vec`。
struct DropGuard<'r, 'a, T>(&'r mut Drain<'a, T>);

这是一个内部数据结构，一个只在特定函数体中可用的类型。正如文档注释所解释的，它的目的是提供一种方法来保证这个实现总是会将原始Vec 的所有内容移回那个Vec，并且位置正确。它这样做，正如它的名字可能暗示的……

impl<'r, 'a, T> Drop for DropGuard<'r, 'a, T> {
    fn drop(&mut self) {
        // 实现的主体（我们即将看到！）
    }
}

……通过它自己的Drop 实现！这是一个相对简单的实现，因此相对容易检查安全性，但它必须在unsafe 块中完成很多事情，因为这一切都只有在你有&mut self 访问DropGuard 并且因此也有Drain 结构的情况下才有效。

首先，它通过检查tail_len 来检查是否有任何事情要做：

if self.0.tail_len > 0 {
    // ...
}

这里计算长度的“尾部”是在调用Vec::drain 时指定范围之后的那些元素。回到我一开始展示的示例代码：

let mutvalues = vec![1, 2, 3, 4, 5];
for val in values.drain(1..3) {
    println!("Removed: {val}");
}
println!("Remaining: {values:?}");

这里的尾部是值4 和5，它们没有被清空。

tail_len 值只在结构体在Vec::drain 中初始化时设置一次，连同tail_start：

pub fndrain(&mutself, range: R) -> Drain<'_, T>
where
    R: RangeBounds<usize>,
{
// 内存安全 //
//
// 当 Drain 首次创建时，它缩短源向量的长度，以确保如果 Drain 的析构函数从未运行，那么未初始化或已移动的元素根本不可访问。
// //
// Drain 将 ptr::read 出要移除的值。
// 当完成时，向量的剩余尾部被复制回覆盖孔洞，并且向量长度恢复到新长度。
// let len = self.len();
letRange { start, end } = slice::range(range, ..len);

unsafe {
// 将 self.vec 长度设置为 start，以在 Drain 泄漏时保持安全
self.set_len(start);
letrange_slice = slice::from_raw_parts(self.as_ptr().add(start), end - start);
        Drain {
            tail_start: end,
            tail_len: len - end,
            iter: range_slice.iter(),
            vec: NonNull::from(self),
        }
    }
}

在这里你可以看到tail_start 和tail_end 代表你使用drain 拉出的部分之后的任何东西，有一个特殊的处理，以保证在处理原始Vec 的内容时的内存安全。

因此，在我的示例代码中，tail_start 将是3，tail_end 将是4：不包括范围的结束，正如我上面描述的值4 和5。

如果有尾部，DropGuard 使用std::ptr::copy 函数重新定位这些项，这个函数类似于 C 函数memmove。它获得原始Vec 的可变引用，并再次只有在尾部不在原始Vec 的末尾时才复制值。

if self.0.tail_len > 0 {
unsafe {
letsource_vec = self.0.vec.as_mut();
// 将未触及的尾部 memmove 回，更新为新长度
letstart = source_vec.len();
lettail = self.0.tail_start;
if tail != start {
letsrc = source_vec.as_ptr().add(tail);
letdst = source_vec.as_mut_ptr().add(start);
            ptr::copy(src, dst, self.0.tail_len);
        }
        source_vec.set_len(start + self.0.tail_len);
    }
}

最后，这个实现更新了原始Vec 的长度。这是一个不安全的操作，因为它甚至不试图维持关于Vec 的正常不变性，即它不包含未初始化的内存，新长度小于或等于Vec 的总分配容量等。在这里，我们可以通过检查保证它是安全的，因为我们在构造时保证tail_len 受原始向量长度的限制——但我们也可以（并且 Rust 确实）使用 Miri 工具进行大量动态分析，以确保通过广泛测试这是真的。

当我说我们将最小化重新定位项所做的工作时，这就是我的意思：我们只有在实际使用Drain 后才移动这些项；我们不会预先将原始项移出位置并将这些项移过来。

（如果你想知道，这意味着如果你正在清空一个大型Vec 的一小部分，你可能在完成Drain 后看到一个性能问题。像往常一样，关于性能，你应该在假设这是一个问题之前进行测量！）

一旦内存移动完成，DropGuard 也完成了。我们很快就会看到它是如何被使用的，我将解释为什么以这种方式使用它。回到Drain 的drop 实现的其余部分。

现在，在我们继续之前，我们需要停下来再次回顾Drain 是如何构建的，因为正如我上面提到的，Drop 实现不能负责维护实现的合理性。它可能永远不会被运行！相反，Vec::drain 在构建Drain 时通过这一行保持合理性：

// 将 self.vec 长度设置为 start，以在 Drain 泄漏时保持安全
self.set_len(start);

语义上，这“截断”了Vec（这里的self），使其不包括从清空范围的开始。注意消息：“以在 Drain 泄漏时保持安全”。这意味着如果有人做了像调用std::mem::forget 函数这样的操作，从而“泄漏”了Drain，没有安全不变量会被打破。内存会被泄漏——明确且有意地，在这种情况下！——但它不会受到无效别名或其他不合理的问题的影响。

这是关于 Rust 安全性的一个非常重要的观点：泄漏通常是一个错误，但它不是一个安全问题。退一步说，如果我们考虑其他完全安全的语言，这应该是相当明显的。你可以在浏览器中的 JavaScript 中创建内存泄漏，尽管它是一个垃圾收集语言，基本上没有不安全的逃生舱！我几乎花了 2022 年和 2023 年第一季度的全部时间追逐 JavaScript 中的内存泄漏，而且没有不安全的代码出现！

首先，它从Drain 中提取范围迭代器，并使用它来计算在清理所有内容时需要丢弃多少项——因为，正如文档所注，“如果迭代器在使用前被丢弃，它将丢弃剩余的移除元素。”

let iter = mem::take(&mut self.iter);
let drop_len = iter.len();

mem::take 函数将给定值替换为其Default 值，根据其std::default::Default 特征的实现。对于Range<usize>（就像我们对Vec 索引那样），那是0..0，从0 包含到0 独占。换句话说，它是空范围，这应该是你期望的。这一步将self.iter 设置为一个无用值，并使iter 值可用于进一步操作——并且，至关重要的是，可在Drain 范围外丢弃，但不受其影响，这对于稍后发生的一些指针算术很重要。

接下来，这个drop 实现获得了原始Vec 的可变引用：

let mutvec = self.vec;

乍一看，可能不明显我们在这里获得的是一个引用，但self.vec 的类型是NonNull<Vec<T>>，它总是包装一个引用。在这种情况下，它是通过在Drain 被构建时调用NonNull::from(self) 构建的，其中self 是&mut self 引用我们调用drain 的Vec：

pub fndrain(&mutself, range: R) -> Drain<'_, T>
where
    R: RangeBounds<usize>,
{
// 所有安全设置...
unsafe {
// 所有不安全设置...
        Drain {
// 其他字段，最后...
            vec: NonNull::from(self),
        }
    }
}

所以let mut vec = self.vec 是一个可变引用到NonNull 指针的Vec，我们可以通过其Deref 特征的实现使用所有正常的Vec 方法。这正是我们接下来要做的。

首先，有一个零大小类型的特殊情况。零大小类型是这样的类型——

struct TotallyEmpty;

——即，一个没有任何数据与之关联的类型，编译器将保证它根本不占用任何内存。我们之所以要区分处理这种情况，是因为没有什么可以移动的！

典型的 Rust 使用这些类型的几个原因——它们既不是非常常见，但也不是特别不常见：

作为区分其他类型的“标记”。这可以帮助在没有任何额外运行时成本的情况下提供类型安全性，因为 Rust 将区分两个零大小类型。这在实现类型安全状态机时非常有用。

作为一个有用的地方来实现特征。你不能在什么都没有的情况下实现特征。然而，你可以impl SomeTrait for TotallyEmpty。这在与用作标记类型结合时特别有用！你可能想要这样做的一次是作为一个标记，使给定类型退出被Send 或Sync。

我可以更多地谈论零大小类型，但相反：回到drop 实现！

因此，实现首先做了一些有趣的事情：检查T 上的一个值。但T 是一个类型！

if T::IS_ZST {
// ZSTs 没有身份，所以我们不需要移动它们，我们只需要丢弃正确数量的实例
// 这可以通过操作 Vec 长度而不是从 `iter` 中移动值来实现。
unsafe {
letvec = vec.as_mut();
letold_len = vec.len();
        vec.set_len(old_len + drop_len + self.tail_len);
        vec.truncate(old_len + self.tail_len);
    }
return;
}

这实际上是在 Rust 的标准库中实现的——技术上在libcore 中——使用一个目前只设计用于此类内部使用的不稳定特性。3 这意味着你不能在自己的代码中这样写，正如你可以用这个 playground 确认的。4 在底层，它做的是一些相当简单的事情：

pub trait SizedTypeProperties: Sized {
    // 其他这样的编译器-only位...
    const IS_ZST: bool = size_of::<Self>() == 0;
}

impl SizedTypeProperties for T {}

也就是说，在编译时，它确定任何类型T 是否是零大小类型，然后 Rust 的内部可以使用那个关联字段。这是你自己可以为其他特征做的——看看这个 playground 的一个愚蠢的例子；它只是不经常看到！

回到drop 实现，我们可以看到在这种情况下我们所需要做的所有事情就是更新原始Vec 的长度，不需要移动内存。然后我们return，因为在这种情况下没有其他事情要做，Drain 实例可以被安全地清理。

unsafe {
    let vec = vec.as_mut();
    let old_len = vec.len();
    vec.set_len(old_len + drop_len + self.tail_len);
    vec.truncate(old_len + self.tail_len);
}
return;

“正常”的模式是针对确实有大小的类型。在这种情况下，DropGuard 最终出现了：

// 确保即使 drop_in_place 恐慌，元素也被移回适当的位置
let _guard = DropGuard(self);

注释解释了DropGuard 存在的原因：我们需要保证我在上面概述的约束，以及所有这些存在的原因——一旦Vec 被包括在Drain 中，原始Vec 的值就永远无法访问，否则我们将违反 Rust 的内存安全保证。我们很快就会看到drop_in_place 部分指的是什么，但首先还有一件事要做：如果没有额外的东西需要丢弃，立即返回！

记住，我们得到了需要丢弃的项目数，即迭代器中剩余的项目数。如果没有东西需要丢弃，我们就完成了：

let iter = mem::take(&mut self.iter);
let drop_len = iter.len();

// 零大小类型处理和创建 drop guard...
if drop_len == 0 {
    return;
}

通过“完成”，我的意思是drop 方法返回，所以声明的drop guard 我们用let _guard = DropGuard(self) 现在超出范围。这意味着它的Drop 实现——我们开始的地方！——将运行。任何需要在原始Vec 中移动的东西都将在这里，在范围的末尾被移动。

对于这个drop 实现的最后部分，基本上也适用，有一点小调整：

// as_slice
() 只能在 iter.len() > 0 时调用，因为
// 它也被 vec::Splice 触摸，这可能会将其变成一个悬空指针
// 这将使它和 vec 指针指向不同的分配，这将导致下面的无效指针算术。
letdrop_ptr = iter.as_slice().as_ptr();

unsafe {
// drop_ptr 来自 slice::Iter，只给我们一个 &[T] 但是
// 对于 drop_in_place，需要一个具有可变出处的指针。
// 因此，我们必须从原始 vec 重建它，同时
// 避免创建一个 &mut 到前面，因为那可能会使依赖于原始指针的一些不安全代码无效。
letvec_ptr = vec.as_mut().as_mut_ptr();
letdrop_offset = drop_ptr.sub_ptr(vec_ptr);
letto_drop = ptr::slice_from_raw_parts_mut(vec_ptr.add(drop_offset), drop_len);
    ptr::drop_in_place(to_drop);
}

这里代码的第一行注释告诉我们为什么我们在做任何事情之前处理了drop_len 检查。它也暗示了一个主题，这在这里隐含地存在：Rust 使得隔离安全检查变得更容易，但它们通常不可能完全局部化。在这种情况下，通过Vec::splice 方法创建的Splice 类型使用Drain，并且也使用不安全的指针到原始Vec 在它自己的Drop 实现中，所以Drain 必须小心不要违反Splice 所做的假设。这很难做到正确！这就是我们使用 Miri 的原因，正如我上面提到的——这导致了这段代码被写成现在的样子！

下一个评论块和不安全块的第一行深入到 Rust 正在努力改进的安全模型的一个方面：出处。简而言之，出处是关于跟踪不仅仅是指针的地址，还有它来自哪里，因此我们可以证明关于它的什么。出处是一个我知之甚少的非常深刻和迷人的主题，所以我不会再多说。

这里，关键是要确保我们在丢弃调用drain 时指定的原始范围未使用的任何值时，有一个有效的指针出处。Miri 会（正确地！）否则抱怨。

一旦我们有了具有有效出处的指针，我们就用std::ptr::slice_from_raw_parts_mut 函数得到一个“原始切片”。一个“原始切片”基本上只是一块被解释为给定类型连续序列的内存。它被称为“原始”，因为它是不安全的：它直接从一个指针和大小构建；调用者有责任确保它是有效的。5 我们在这里使用这个函数的_mut 版本，因为我们接下来做的下一件事是调用std::ptr::drop_in_place 函数，它在不移动它们的情况下运行它被调用的任何东西的Drop 实现。在切片上调用，将递归调用切片中每个项目的Drop。

drop_in_place 是不安全的，因为它留下了结果内存完全不变，除了给定项目的自己的Drop 实现可能会做的之外。这让我们回到了我们不得不维护的保证，即清除所有这些内存。这也是为什么我要说调用drop_in_place 是这个函数做的“次要”的事情，尽管：这是函数的结尾，因此是范围的结尾，这意味着_guard 实例的DropGuard 超出范围并运行。这意味着保证得到了维护！而且，正如创建DropGuard 实例时的注释所指示的——

// 确保即使 drop_in_place 恐慌，元素也被移回适当的位置
let _guard = DropGuard(self);

——Rust 甚至会运行DropGuard 实现的Drop，即使drop_in_place 恐慌，这可能会发生如果一些内部类型的Drop 实现表现不佳。这允许这个函数保证即使原始Vec 中的数据出了问题，内存本身是有效的，Vec 本身之后也是有效的。也就是说：可能会有严重的错误或问题，但它不会是一个违反 Rust 安全和合理性保证的错误或问题。

这是很多内容，但它展示了一些关于为不安全代码提供安全抽象的有趣部分，以及 Rust 如何利用其所有权语义在保持这些保证的同时提供出色的性能。特别是：

1. 在使用drain 期间或之后，原始Vec 永远不会处于无效状态。
2. 那个Vec 上的迭代器也永远不会被使无效。
3. 无论其他类型的实现有多糟糕，只要没有不安全的代码在那种不良行为中，(1) 和 (2) 都是真的。
4. 最糟糕的潜在结果——即impl Drop for Drain 根本没有运行——是内存被泄漏。这不好，但也是安全的。

还值得看到，虽然这包括内存安全，但公共 API 中的所有权语义的工作方式也消除了其他类型的错误。你可以在 Java 或 JavaScript 中有迭代器失效的错误，只要你不小心！在 Rust 中，你只能通过明确选择unsafe 来有迭代器失效的错误。这很酷，也是我在其他语言工作时想念 Rust 的一个原因！

我还特别想指出这种使用DropGuard 来维护这些保证。这与你可以在 C# 或 JavaScript 中使用using 构造，或在 Python 中使用with 构造做的类似——但在那些情况下，有特殊的语言支持内置来处理那种范围，以便在需要时你可以部署它。在 Rust 中，它直接来自于所有权和具有析构函数的组合，当一个项目超出范围时自动运行。你不需要任何特殊的语言构造来实现它，除了这两个之外。

关于Drop、安全和泄漏的交互的注释，参见 Rust 的官方不安全代码指南 Rustonomicon 的“泄漏”部分。

ref

https://v5.chriskrycho.com/journal/read-the-code/using-drop-safely-in-rust/^[1]

引用链接