本文将深入探讨 Rust 语言中的 vec::Drain
及其 Drop
实现,作为所有权如何防止内存及其他微妙错误的一个例子。
目标读者:能够阅读 Rust 代码,并对其所有权语义及 Drop
特性有基本(且仅是基本)理解的开发者。
引言
在读《The Rust Programming Language》书籍时,偶然发现了 Vec::drain
方法。
常见的编程语言中,没有见过以 drain
命名的方法,这引起了我的好奇心。
什么是 Vec::drain
?
如果你不熟悉 Vec::drain
,可以像下面这样使用它从 Vec
中抽取元素(类似的方法还存在于 String
、HashMap
及其他多种集合类型中):
let mut values = 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::drain
的内部实现
一种合理的实现方式
一种完全合理的实现方法是:接收 Vec
,复制出所有要移除的元素并放入新的 Vec
,更新原始 Vec
以移除这些元素,并返回由新分配的 Vec
支持的迭代器。
然而,这种方法对于计算机来说可能需要大量的前期工作。
如果 Vec
很大(包含数千或数百万个元素)且操作的是中间的一部分,
这将导致大量额外的内存分配和复制操作,甚至在确定是否使用这些值之前。
Rust 的独特实现
因此,Rust 在这里采取了完全不同的方法:它保留了对原始 Vec
的可变引用,并仅从原始存储中读取和更新。
这得益于 Rust 的所有权规则:
- 只要
Vec::drain
返回的迭代器存在,其他任何地方都无法对原始Vec
进行读写访问, - 因此无法通过使迭代器或其支持的存储失效来使其处于错误状态(例如,改变
Vec
中的值,改变其长度等)。
Rust 通过创建一个新的数据结构 Drain
来实现这一点,该结构合理地命名,并持有对原始 Vec
的可变引用以及一个用于访问 Vec
值的迭代器。
使用 Drain
上的迭代器方法时,它会转发到对切片的迭代器。这意味着它不需要自己实现迭代,而是可以使用与对切片迭代相同的(经过优化的)实现。
唯一的区别(也是关键 的区别)是 drain
通过 unsafe std::ptr::read
调用立即从切片中返回值。
内存安全性保障
如果在 Drain
迭代器访问期间或之后有人能访问 Vec
中的值,或者 Vec
记住所有 Drain
访问的元素,这将是不安全的。
然而,如前所述,只要 Drain
迭代器存在,就无法访问 Vec
,因为它通过可变引用持有自身。
因此,对于熟悉 Rust 的开发者来说,这部分内容应该相当直接,std::ptr::read
是唯一不寻常的部分。
Drain
的 Drop 实现
Drop 实现的重要性
当迭代器被丢弃——无论是因为在 for
循环中遍历完毕,还是因为在迭代部分元素后手动丢弃——Drain
类型的 Drop
实现会接管。
这意味着 impl Drop for Drain
负责确保 Drain
的不安全清理保持合理,同时避免内存泄漏,
即 Drain
创建时 Vec
遗忘的内存(稍后会详细讨论)。
常见模式与安全性
这是 Rust 中一种常见且值得理解的模式,同时也非常巧妙。 我们将一步步详细讲解——包括所有实现中的代码(截至 Rust 1.85 nightly 版本)。 你可能需要将这些代码与本文并排查看,以便更好地理解上下文!
Drop Trait 的实现
impl<T> Drop for Drain<'_, T> {
fn drop(&mut self) {
// ...
}
}
需要注意的是,drop
接受 &mut self
。这意味着我们无法做任何需要拥有 self
所有权的操作,这进一步引出了下一个内容。
DropGuard 结构体
/// 将未被 `Drain` 移除的元素移动回来,以恢复原始的 `Vec`。
struct DropGuard<'r, 'a, T>(&'r mut Drain<'a, T>);
这是一个内部数据结构,仅在此函数体内可用。其目的是确保在正确的位置将所有内容移回原始 Vec
。
impl<'r, 'a, T> Drop for DropGuard<'r, 'a, T> {
fn drop(&mut self) {
// 实现细节(稍后详细讨论)
}
}
内存操作与安全性
DropGuard
的 Drop
实现确保即使在发生恐慌(panic)的情况下,也能正确地将元素移回原始位置,保持内存安全。
处理 Drain
的构造
在 Drain
被构造时,Vec::drain
方法通过以下代码行确保即使 Drain
被泄漏,也不会破坏安全性:
// set self.vec length's to start, to be safe in case Drain is leaked
self.set_len(start);
这实际上是截断 Vec
,确保它不包含被 drain
移除的范围内的元素。
即使有人调用 std::mem::forget
来泄漏 Drain
,也不会破坏安全性,只会导致内存泄漏,这是可控的。
DropGuard 的工作机制
当迭代器被丢弃时,DropGuard
确保所有未被消耗的元素被正确地丢弃,同时恢复原始 Vec
的长度。
这通过以下代码实现:
let iter = mem::take(&mut self.iter);
let drop_len = iter.len();
if drop_len == 0 {
return;
}
let _guard = DropGuard(self);
DropGuard
的存在确保即使在 drop_in_place
调用时发生恐慌,也能保持内存和 Vec
的有效性。
性能与安全性的平衡
Rust 的实现通过仅在必要时移动元素,最大限度地减少了计算和内存开销。
这对于处理大型 Vec
特别重要,因为它避免了不必要的内存分配和复制操作。
结论
通过深入了解 Vec::drain
及其 Drop
实现,我们可以看到 Rust 如何利用所有权语义来提供高性能和内存安全的保障。
具体来说:
- 原始的
Vec
在使用drain
期间及之后永远不会处于无效状态。 - 该
Vec
的迭代器永远不会被无效化。 - 即使在面对不良行为的实现(只要没有不安全代码),上述两点依然成立。
- 最坏的情况是内存泄漏,但这依然是安全的。
此外,DropGuard
的使用展示了 Rust 如何通过所有权和析构函数自动管理资源,而无需依赖特殊的语言结构。