Home

在 Rust 中实现 defer

lhz07
| 2026/07/21 00:38:06
50 min read

介绍 DropGuard 与 macro_rules,并实现 defer! 宏

RAII 和 defer 是两种常见的资源管理方式,理论上来说,它们是完全可以互相替代的,但实际上它们更像是互补的,各有各的优势。当对象拥有资源时,比如一个 File 对象,这时候适合用 RAII,在离开作用域时自动关闭文件;而需要在作用域结束时执行特定的清理逻辑,这就更适合用 defer
比如在写包管理器的时候,使用数据库来管理安装的文件,数据库通过事务操作,可以确保数据具有原子性,要么提交,要么回滚,不存在中间状态。但是一般的文件系统不具备这样的功能,这时候可以用 defer 来实现,如果出错了,提前返回,通过 defer 自动回滚文件;如果成功了,就取消 defer,相当于是成功提交了事务。
下面这段代码就通过 defer 实现了这样的逻辑:

// runnable
use std::{
    mem::ManuallyDrop,
    ops::{Deref, DerefMut},
};

pub struct DropGuard<F: FnOnce(T), T> {
    f: ManuallyDrop<F>,
    inner: ManuallyDrop<T>,
}

impl<F: FnOnce(T), T> DropGuard<F, T> {
    pub fn new(inner: T, f: F) -> Self {
        Self {
            f: ManuallyDrop::new(f),
            inner: ManuallyDrop::new(inner),
        }
    }

    /// Consumes the `DropGuard` without invoking the drop function
    pub fn into_inner(self) -> T {
        let mut new_guard = ManuallyDrop::new(self);
        let value = unsafe { ManuallyDrop::take(&mut new_guard.inner) };
        unsafe { ManuallyDrop::drop(&mut new_guard.f) };
        value
    }
}

impl<F: FnOnce(T), T> Drop for DropGuard<F, T> {
    fn drop(&mut self) {
        let value = unsafe { ManuallyDrop::take(&mut self.inner) };
        let f = unsafe { ManuallyDrop::take(&mut self.f) };
        f(value);
    }
}

impl<F: FnOnce(T), T> Deref for DropGuard<F, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl<F: FnOnce(T), T> DerefMut for DropGuard<F, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner
    }
}

#[macro_export]
macro_rules! defer {
    ($($move:ident)? |$($tokens:tt)*) => {
        $crate::defer_impl!(__guard $($move)? [] [] [] [,$($tokens)*]);
    };

    ($guard:ident, $($move:ident)? |$($tokens:tt)*) => {
        $crate::defer_impl!($guard $($move)? [] [] [] [,$($tokens)*] );
    };

    ($guard:ident, $body:expr) => {
        let $guard = DropGuard::new(
            (),
            |_| $body(),
        );
    };

    ($body:expr) => {
        let mut __defer_guard = DropGuard::new(
            (),
            |_| $body(),
        );
    };

}

#[macro_export]
macro_rules! defer_impl {
    ($guard:ident $($move:ident)? [$($caps:ident)*] [$($pats:tt)*]
     [$($comma:tt)?]
     [,mut $var:ident $($rest:tt)*]
    ) => {
        $crate::defer_impl!(
            $guard
            $($move)?
            [$($caps)* $var]
            [$($pats)* $($comma)? mut $var]
            [,]
            [$($rest)*]
        );
    };

    // ident, rest
    ($guard:ident $($move:ident)? [$($caps:ident)*] [$($pats:tt)*]
     [$($comma:tt)?]
     [, $bind:ident: &$var:ident $($rest:tt)*]
    ) => {
        let $bind = &mut $var;
        $crate::defer_impl!(
            $guard
            $($move)?
            [$($caps)* $bind]
            [$($pats)* $($comma)? $bind]
            [,]
            [$($rest)*]
        );
        let $bind = &mut **$bind;
    };

    ($guard:ident $($move:ident)? [$($caps:ident)*] [$($pats:tt)*]
     [$($comma:tt)?]
     [, &$var:ident $($rest:tt)*]
    ) => {
        let $var = &mut $var;
        $crate::defer_impl!(
            $guard
            $($move)?
            [$($caps)* $var]
            [$($pats)* $($comma)? $var]
            [,]
            [$($rest)*]
        );
        let $var = &mut **$var;
    };

    // ident, rest
    ($guard:ident $($move:ident)? [$($caps:ident)*] [$($pats:tt)*]
     [$($comma:tt)?]
     [, $var:ident $($rest:tt)*]
    ) => {
        $crate::defer_impl!(
            $guard
            $($move)?
            [$($caps)* $var]
            [$($pats)* $($comma)? $var]
            [,]
            [$($rest)*]
        );
    };

    // END
    ($guard:ident $($move:ident)? [$($caps:ident)*] [$($pats:tt)*] [,] [$(,)? | $body:block]) => {
        #[allow(unused_parens)]
        let mut $guard = DropGuard::new(
            ($($caps),*),
            $($move)? |($($pats)*)| $body,
        );

        #[allow(unused_parens)]
        let ($($caps),*) = &mut *$guard;
    };

}

// ANCHOR
use std::fs;
use std::io;
use std::path::{Path, PathBuf};

fn copy_files() -> Result<(), io::Error> {
    let mut installed_files: Vec<PathBuf> = Vec::new();
    defer!(guard, |&installed_files| {
        for p in installed_files {
            let res = fs::remove_file(&p);
            println!("delete the file: {}, {res:?}", p.display());
        }
    });

    let files = fs::read_dir(".")?;
    for entry in files {
        let entry = entry?;
        let file = entry.path();
        let copy_to = Path::new("target").join(entry.file_name());
        installed_files.push(copy_to);
        println!(
            "copy file: {}, type: {:?}",
            file.display(),
            entry.file_type()?
        );
        fs::copy(file, installed_files.last().unwrap())?;
    }

    guard.into_inner();
    Ok(())
}

fn main() {
    let res = copy_files();
    println!("result: {res:?}");
}
// ANCHOR_END

这里先用 defer 定义了出错时需要执行哪些逻辑,再开始复制文件,出错时会通过问号提前返回,defer 在退出作用域时自动执行,整个流程非常自然,不妨点击右上角的运行按钮看看结果如何。
关于 Rust 中的 defer,社区中也有不少讨论1,可见需求不小,但官方依旧不太想做这个功能,社区里比较有名的实现是 scopeguard ,但是它提供的 defer 宏并不好用,由于借用检查器的限制,完全无法实现上面那样的代码,所以我尝试自己实现了一个比较好用的 defer 宏,效果如上,用起来很灵活,不会因为借用检查器而受到很大的限制。下面介绍一下如何实现这个宏。

实现 DropGuard

尝试简单实现

defer 宏的作用是让代码写起来更简单,但要实现 defer,首先需要实现一个 DropGuard,以上面的例子来说,除去取消功能,它还至少需要两个功能:

  1. 可以存储一个实现了 FnOnce 的闭包2,并在 drop 时调用。
  2. 可以单独存储数据,虽然闭包也能捕获数据,但这样只能在 drop 的时候访问,而我们上面的代码,是需要在定义 defer 后,往installed_files 里添加文件路径的,所以数据必须单独存储,便于外部访问。

那么最简单的实现大概长这样

// runnable
// ANCHOR
struct DropGuard<T, F: FnOnce(T)> {
    inner: T,
    f: F,
}

impl<T, F: FnOnce(T)> Drop for DropGuard<T, F> {
    fn drop(&mut self) {
        (self.f)(self.inner);
    }
}
// ANCHOR_END

fn main() {}

但很可惜,这段代码没法编译,原因是 FnOnce 被调用时会被消耗,存储的值 T 也需要传给 FnOnce 使用,但 drop 方法的参数是 &mut self,这意味着只能更改 Self 里的值,但不能消耗掉。

为什么这么设计?

乍一看很奇怪,结构体被 drop 之后就不能再使用了,为什么不能消耗所有权?
一方面,drop 的本质是在作用域结束时,会自动帮你调用 drop 方法,像这样:

fn main() {
    {
        let s1 = "hello".to_string();
        let s2 = "hello".to_string();
        // 作用域结束,自动调用 drop 方法
        s2.drop();
        s1.drop();
    }
}

而对于方法,如果方法参数是 self,自然也需要调用 drop 方法。那对于 drop 方法本身来说就会发生无限的递归了。

fn drop(self) {
    self.drop();
}

想解决这个问题,必须要对 drop 方法的实现做一些特殊的调整,但这样就会使语言的行为更加不一致了。

另一方面,对于 unsized type,它们可能也需要 drop 方法,如果方法参数是引用,引用的大小是可以在编译时确定的,但如果是 self ,就没法确定大小了,无法编译。

// runnable
// ANCHOR
struct UnsizedType<T: ?Sized>(T);

impl<T: ?Sized> UnsizedType<T> {
    fn drop(self) {}
}
// ANCHOR_END

fn main() {}

进一步探索实现

既然 drop 方法的签名没法改,就只能想办法从可变引用中取出值了,最简单的想法是用 Option,可以取出里面的值,留下 None 在里面,但是它有两个问题:

  1. 每次都要判断里面是 Some 还是 None ,虽然开销很小,但没法忽略
  2. Option 的本质是 tagged union,需要额外占用 1 字节的空间存储当前的状态,而且由于结构体需要内存对齐,为了存储状态而多占用的空间可能不止 1 字节。
// runnable
struct DropGuard<T, F: FnOnce(T)> {
    inner: Option<T>,
    f: Option<F>,
}

impl<T, F: FnOnce(T)> Drop for DropGuard<T, F> {
    fn drop(&mut self) {
        unsafe {
            let f = self.f.take().unwrap();
            let inner = self.inner.take().unwrap();
            (f)(inner);
        }
    }
}

fn main() {
    println!("u64 size: {}", size_of::<u64>());
    println!("u64 with Option size: {}", size_of::<Option<u64>>());
}

其实我们根本不需要那个多余的 tag,直接用 union 不就行了,像这样

// runnable
// ANCHOR
use std::ptr;

union Opt<T> {
    val: T,
    _nothing: (),
}

struct DropGuard<T, F: FnOnce(T)> {
    inner: Opt<T>,
    f: Opt<F>,
}

impl<T, F: FnOnce(T)> Drop for DropGuard<T, F> {
    fn drop(&mut self) {
        unsafe {
            let f = ptr::read(&self.f.val);
            let inner = ptr::read(&self.inner.val);
            (f)(inner);
        }
    }
}
// ANCHOR_END

fn main() {}

Oops,被编译器阻止了,为什么之前用 Option 没问题,但现在用 union 就不行了呢?
主要问题是,对于 Option,它被 drop 的时候,因为有 tag 记录了当前状态,可以判断当前是否存储了值,如果有值就 drop 掉,如果没有值就什么都不做。
但是对于 union,无法得知它的状态,保守的做法应该是什么也不干,但是这样就内存泄漏了,所以这里编译器要求 union 里存储的类型要么是 Copy 的,无需 drop,要么是 ManuallyDrop<T>,也不需要 drop
顾名思义,被 ManuallyDrop 包裹的值,将不会被 drop,需要手动进行,这似乎也正好是我们需要的。

我们刚刚还没有尝试,如果直接读出值会怎么样,请运行下面的例子看看:

// runnable
use std::ptr;

struct DropGuard<T, F: FnOnce(T)> {
    inner: T,
    f: F,
}

impl<T, F: FnOnce(T)> Drop for DropGuard<T, F> {
    fn drop(&mut self) {
        unsafe {
            let f = ptr::read(&self.f);
            let inner = ptr::read(&self.inner);
            (f)(inner);
        }
    }
}

fn main() {
    let a = Box::new(3);
    let _guard = DropGuard {
        inner: a,
        f: |_| {
            // do nothing
        },
    };
}

我们成功触发了 double free,这是因为 drop 方法只用来处理当前结构体的析构过程,而结构体内部的各个字段,则由它们自己的 drop 方法来处理。
在上面的代码中,当 DropGuard 离开作用域时,先运行了它自己的 drop 方法,读出了 a,并传入闭包,闭包结束时,a 自然被 drop,Box3drop 方法里运行了一次 free,之后则开始析构各个字段,析构 inner 字段的时候,又 drop 了一次 aBoxdrop 方法里又运行了一次 free,于是就触发 double free 了。

显而易见,ManuallyDrop 在这里非常有用,它可以阻止第二次 drop,从而让这段代码正常运行。

最终实现

// runnable
use std::mem::ManuallyDrop;

pub struct DropGuard<F: FnOnce(T), T> {
    f: ManuallyDrop<F>,
    inner: ManuallyDrop<T>,
}

impl<F: FnOnce(T), T> DropGuard<F, T> {
    pub fn new(inner: T, f: F) -> Self {
        Self {
            f: ManuallyDrop::new(f),
            inner: ManuallyDrop::new(inner),
        }
    }
}

impl<F: FnOnce(T), T> Drop for DropGuard<F, T> {
    fn drop(&mut self) {
        let value = unsafe { ManuallyDrop::take(&mut self.inner) };
        let f = unsafe { ManuallyDrop::take(&mut self.f) };
        f(value);
    }
}

fn main() {
    {
        let a = Box::new(3);
        let _guard = DropGuard::new(a, |_| {
            println!("drop `a`");
        });
    }
    println!("dropped `a` safely");
}

这个实现就比较完善了,也没有内存错误,但作为 DropGuard,还需要实现两个功能

Deref Trait

当把值交给 DropGuard 后,可能还需要修改值,每次都调用 guard.inner 非常麻烦,可以通过实现 Deref trait 来改善这一点。
顾名思义,Deref trait 就是定义当解引用这个结构体的引用时,应该返回什么,因为 Rust 区分可变性,所以 DerefDerefMut 分别对应着不可变引用和可变引用两种情况。

// runnable
use std::mem::ManuallyDrop;

pub struct DropGuard<F: FnOnce(T), T> {
    f: ManuallyDrop<F>,
    inner: ManuallyDrop<T>,
}

impl<F: FnOnce(T), T> DropGuard<F, T> {
    pub fn new(inner: T, f: F) -> Self {
        Self {
            f: ManuallyDrop::new(f),
            inner: ManuallyDrop::new(inner),
        }
    }
}

impl<F: FnOnce(T), T> Drop for DropGuard<F, T> {
    fn drop(&mut self) {
        let value = unsafe { ManuallyDrop::take(&mut self.inner) };
        let f = unsafe { ManuallyDrop::take(&mut self.f) };
        f(value);
    }
}

// ANCHOR
use std::ops::{Deref, DerefMut};

impl<F: FnOnce(T), T> Deref for DropGuard<F, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl<F: FnOnce(T), T> DerefMut for DropGuard<F, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner
    }
}

fn main() {
    let a = vec![1, 2];
    let mut guarded_a = DropGuard::new(a, |a| {
        println!("drop `a`");
        println!("{:?}", a);
    });
    println!("modify a");
    // 这里调用 `push` 方法时,因为 `DropGuard` 没有这个方法,
    // 编译器会加上足够多的 `*` 号进行解引用,直到找到对应的方法为止。
    guarded_a.push(3);
    // 显式地写出来就是这样
    let v: &mut Vec<i32> = &mut *guarded_a;
    v.push(4);
}
// ANCHOR_END

取消 guard

文章开头,我举了一个包管理器的例子,安装文件时,如果出错了,就通过 DropGuard 自动回滚,如果成功了,则需要取消它,如何实现呢?
显然,我们没法直接取消 drop 方法,但我们可以把 DropGuard 本身也塞进 ManuallyDrop 里,这样就可以避免它的 drop 方法被执行了。

// runnable
use std::mem::ManuallyDrop;
use std::ops::{Deref, DerefMut};

impl<F: FnOnce(T), T> Drop for DropGuard<F, T> {
    fn drop(&mut self) {
        let value = unsafe { ManuallyDrop::take(&mut self.inner) };
        let f = unsafe { ManuallyDrop::take(&mut self.f) };
        f(value);
    }
}

impl<F: FnOnce(T), T> Deref for DropGuard<F, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl<F: FnOnce(T), T> DerefMut for DropGuard<F, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner
    }
}

// ANCHOR
pub struct DropGuard<F: FnOnce(T), T> {
    f: ManuallyDrop<F>,
    inner: ManuallyDrop<T>,
}

impl<F: FnOnce(T), T> DropGuard<F, T> {
    // ANCHOR_END
    pub fn new(inner: T, f: F) -> Self {
        Self {
            f: ManuallyDrop::new(f),
            inner: ManuallyDrop::new(inner),
        }
    }
    // ANCHOR
    pub fn into_inner(self) -> T {
        // 塞进 ManuallyDrop
        let mut new_guard = ManuallyDrop::new(self);
        // 取出存储的值
        let value = unsafe { ManuallyDrop::take(&mut new_guard.inner) };
        // 不需要运行闭包了,直接 drop 掉
        unsafe { ManuallyDrop::drop(&mut new_guard.f) };
        // 返回值
        value
    }
}

fn main() {
    let a = vec![1, 2];
    let mut guarded_a = DropGuard::new(a, |a| {
        println!("drop `a`");
        println!("{:?}", a);
    });
    println!("modify a");
    guarded_a.push(3);
    // 取消 guard
    let a = guarded_a.into_inner();
}
// ANCHOR_END

灵活使用

这个 DropGuard 看起来简单,但是因为使用了泛型,用起来十分灵活。
比如同时处理多个值,可以把值打包成元组塞进去,重新借用后再解构元组
也可以只传引用进去,不把整个值 move 进去,这样之后还能继续用。

// runnable
use std::mem::ManuallyDrop;
use std::ops::{Deref, DerefMut};

impl<F: FnOnce(T), T> Drop for DropGuard<F, T> {
    fn drop(&mut self) {
        let value = unsafe { ManuallyDrop::take(&mut self.inner) };
        let f = unsafe { ManuallyDrop::take(&mut self.f) };
        f(value);
    }
}

impl<F: FnOnce(T), T> Deref for DropGuard<F, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl<F: FnOnce(T), T> DerefMut for DropGuard<F, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner
    }
}

pub struct DropGuard<F: FnOnce(T), T> {
    f: ManuallyDrop<F>,
    inner: ManuallyDrop<T>,
}

impl<F: FnOnce(T), T> DropGuard<F, T> {
    pub fn new(inner: T, f: F) -> Self {
        Self {
            f: ManuallyDrop::new(f),
            inner: ManuallyDrop::new(inner),
        }
    }

    pub fn into_inner(self) -> T {
        // 塞进 ManuallyDrop
        let mut new_guard = ManuallyDrop::new(self);
        // 取出存储的值
        let value = unsafe { ManuallyDrop::take(&mut new_guard.inner) };
        // 不需要运行闭包了,直接 drop 掉
        unsafe { ManuallyDrop::drop(&mut new_guard.f) };
        // 返回值
        value
    }
}

// ANCHOR
fn handle_multiple_values() {
    let a = vec![1, 2];
    let b = "hello".to_string();
    // 把 a, b 打包成元组放进去,创建 guard 之后,使用模式匹配解构元组
    let (a, b) = &mut *DropGuard::new((a, b), |(a, b)| {
        println!("drop `a`");
        println!("{:?}", a);
        println!("drop b: {b}");
    });
    println!("modify a");
    a.push(3);
    println!("print b: {b}");
}

fn return_value() -> Vec<i32> {
    let mut a = vec![1, 2];
    let b = "hello".to_string();
    // 为 a 创建可变引用,之后只把引用放进 guard,外面的 a 不受影响
    let a_ref = &mut a;
    let mut guard = DropGuard::new((a_ref, b), |(a_ref, b)| {
        a_ref.push(4);
        println!("drop b: {b}");
    });
    let (a_ref, b) = &mut *guard;
    println!("modify a");
    a_ref.push(3);
    println!("print b: {b}");
    // guard 持有对 a 的引用,必须先 drop 掉 guard,才能返回 a
    drop(guard);
    a
}

fn main() {
    println!("----------------- handle multiple values -----------------");
    handle_multiple_values();
    println!("----------------- return value -----------------");
    let a = return_value();
    println!("returned a: {a:?}");
}
// ANCHOR_END

其实就功能性来说,目前这个 DropGuard 已经完全可以实现开头提到的包管理器的例子了,但显然的问题是,写起来太繁琐了,要处理 a, b 两个变量,需要写整整 3 次 (a, b),如果变量名再长一点,写起来更麻烦;如果希望只放引用进去,还要再多写一行 let a_ref = &mut a;,实在是繁琐得不能再繁琐了。对于这种重复的工作,在 Rust 中可以用宏来解决。

宏与 macro_rules!

Rust 中有两种宏:声明式宏和过程宏,就像 C 里的 define 一样,可以在编译之前,先处理一遍代码。
其中过程宏最强大,它几乎可以处理任意输入,感觉上像是编译器的插件,你可以自由地定义一套语法,只需要在过程宏里写 Rust 代码解析,并输出为合法的语法树就行,比如你可以在里面写 html,然后在过程宏里解析并生成合法的 Rust 语法树。而本质上它是个单独的程序,运行 cargo build 时,会先编译它,之后 rustc 会在需要的地方运行它来处理代码,再进入正常的编译流程。显然,这意味着过程宏会大幅拖慢编译速度,不仅需要多编译一个程序,语法树需要在编译器和过程宏的动态库之间传递,而且这一操作是阻塞的,后续的类型检查、借用检查等都要等过程宏的代码生成结束才能运行。

而声明式宏则是内置在 rustc 里的一套语法规则,和普通代码的不同之处是,它处理的是输入的 tokens,并输出为合法的代码。它的规则相对有限,且只能通过模式匹配和递归来处理,无法像过程宏那样拥有判断和循环功能。但它最大的优点就是快,因为内置在编译器里,对编译速度的影响很小。

简单的 macro_rules

声明式宏本质上就是按照分支顺序挨个进行模式匹配,比如一个简单的宏:

// runnable
macro_rules! string {
    ($s:literal) => {{ String::from($s) }};
}

fn main() {
    let s = string!("hello");
    println!("type of s: {}", std::any::type_name_of_val(&s));
}

这里使用美元符号来声明参数,紧跟着的是参数名,冒号后面是参数类型
本文用到的参数类型4有这几种:

  • literal: 表示字面量
  • block: 表示代码块,也就是用大括号包裹的一系列语句,值得注意的是,block 本身就是一个表达式
  • expr: 表示表达式,它也包括 block
  • ident: 表示 identifier,基本上就是匹配一个单词,但是要遵循变量的命名规则,不能以数字开头,中间的符号也有限制,不过比变量的限制更少,这里可以匹配保留的关键字
  • ty: 表示类型,可以匹配 String, &mut Vec<i32> 之类的各种类型
  • tt: 最通用的匹配,表示一个 token,而 Rust 代码就是由一堆 tokens 组成的

=> 符号后面的内容会直接在原地展开,如果多加一个大括号,则展开的内容都被大括号包裹起来了,整个分支的内容作为表达式返回;接单大括号,则直接展开这个分支的所有内容。

宏也可以直接匹配字面 token,比如这样:

// runnable
macro_rules! string {
    // 这个分支返回表达式
    (@expression $s:literal) => {{ String::from($s) }};
    // 这个分支是语句
    (@statement $var:ident, $s:literal) => {
        // 就像 shell 语法一样,使用美元符号来引用匹配到的 token
        let $var = String::from($s);
    };
}

fn main() {
    let s1 = string!(@expression "hello");
    println!("type of s1: {}", std::any::type_name_of_val(&s1));
    string!(@statement s2, "hello");
    println!("type of s2: {}", std::any::type_name_of_val(&s2));
}

原地展开的好处是,可以传入变量名,并在宏里面定义变量,以及宏里面定义的变量和外部是同一作用域
比如这个例子,可以给外部的变量绑定宏内部的变量引用。

// runnable
macro_rules! string {
    ($var:ident, $s:literal) => {
        let temp = String::from($s);
        let $var = &temp;
    };
}

fn main() {
    string!(s, "hello");
    println!("type of s: {}", std::any::type_name_of_val(&s));
}

使用表达式就做不到这点,因为宏内部是单独的块作用域,返回的时候内部的变量就被销毁了。

// runnable
macro_rules! string {
    ($s:literal) => {{
        let temp = String::from($s);
        &temp
    }};
}

fn main() {
    let s = string!("hello");
    println!("type of s: {}", std::any::type_name_of_val(&s));
}

Rust 的宏具有卫生性5 (hygiene),所以即使是原地展开,和外部共享作用域,也不会污染命名空间。

// runnable
macro_rules! string {
    ($var:ident, $s:literal) => {
        let temp = String::from($s);
        let $var = &temp;
    };
}

fn main() {
    string!(s, "hello");
    println!("type of s: {}", std::any::type_name_of_val(&s));
    println!("{}", temp);
}

重复匹配

可以重复匹配 token,比如我想要打印东西,可以写这样的宏

// runnable
macro_rules! myprint {
    ($s:literal, $($args:expr),*) => {
        print!($s, $($args),*);
    };
}

fn main() {
    let s1 = "hello";
    let s2 = "world".to_string();
    myprint!("{} {}\n", s1, s2);
    myprint!("\n",);
}

$($args:expr),* 可以重复匹配一堆表达式,它们以逗号分隔

  • 小括号里是要匹配的模式,可以是任何合法的其他模式,比如字面量 tokens,使用美元符号捕获参数,甚至是另一个重复匹配。在这里是匹配一个表达式
  • 小括号后面紧跟的是分隔符,可以省略,表示没有分隔符,也就是默认的按空格分隔
  • 最后面的 * 表示匹配 0 次或多次,换成 + 则表示匹配 1 次或多次,换成 ? 则表示匹配 0 次或 1 次

匹配之后使用相同的语法来展开,但不需要写类型
可以写成 $($args),*

  • 小括号里是要重复输出的内容,此处可以输出字面量,也可以使用小括号内部匹配到的变量,也可以使用外部匹配的变量,但不能使用其他小括号匹配到的变量
  • 小括号后面紧跟的是分隔符,可以省略,表示没有分隔符,也就是默认的按空格分隔,这里的分隔符不必跟匹配时用的一样
  • 最后面的 * 表示重复次数,可以跟之前匹配时使用的不相同,但仍然遵循相同的规则,也就是说,如果前面用 * 匹配到了 1 次,后面用 + 来重复是没问题的;但如果前面用 * 匹配到了 0 次,那就会在后面用 + 的时候报错。

我们刚刚的 myprint! 宏有点小问题,就是在没有参数时,没法省略逗号,有多个参数时,尾部也不能多一个逗号。

// runnable
macro_rules! myprint {
    ($s:literal , $($args:expr),*) => {
        print!($s, $($args),*);
    };
}

fn main() {
    let s1 = "hello";
    let s2 = "world".to_string();
    myprint!("{} {}\n", s1, s2,);
    myprint!("hi\n");
}

使用问号来改进一下

// runnable
macro_rules! myprint {
    ($s:literal $(,$args:expr)* $(,)?) => {
        print!($s, $($args),*);
    };
}

fn main() {
    let s1 = "hello";
    let s2 = "world".to_string();
    myprint!("{} {}\n", s1, s2);
    myprint!("{} {}\n", s1, s2,);
    myprint!("hi\n");
    myprint!("hi\n",);
}

我们刚刚提到了 tt,重复匹配 tt,则相当于是匹配任意的 Rust token tree,这是个非常强力的匹配模式
所以可以把所有传进来的 tokens 直接传给其他的宏

// runnable
macro_rules! myprint {
    ($($token:tt)*) => {
        print!($($token)*);
    };
}

fn main() {
    let s1 = "hello";
    let s2 = "world".to_string();
    myprint!("{} {}\n", s1, s2);
    myprint!("{} {}\n", s1, s2,);
    myprint!("hi\n");
    myprint!("hi\n",);
}

利用展开时的语法,还可以做些有趣的事,比如重复匹配了一堆参数,不知道有多少个,可以将每个参数展开为 +1,然后在最前面补个 0,就组合成了一个表达式

// runnable
// 对于多个分支的宏,需要注意顺序,将条件更精确的写在前面
macro_rules! count {
    (@consume $token:tt) => {
        1
    };
    ($($token:tt)*) => {
        0 $(+ count!(@consume $token))*
    };
}

macro_rules! myprint {
    ($($token:tt)*) => {
        let count = count!($($token)*);
        println!("tokens count: {count}");
        print!($($token)*);
    };
}

fn main() {
    let s1 = "hello";
    let s2 = "world".to_string();
    myprint!("{} {}\n", s1, s2);
}

使用 cargo expand6 可以看到宏被展开成了这样:

fn main() {
    let s1 = "hello";
    let s2 = "world".to_string();
    let count = 0 + 1 + 1 + 1 + 1 + 1;
    {
        ::std::io::_print(format_args!("tokens count: {0}\n", count));
    };
    {
        ::std::io::_print(format_args!("{0} {1}\n", s1, s2));
    };
}

当然也有其他写法,比如把每个 token 展开为 (),即 unit type,然后放到数组里,并统计这个数组的长度。
unit type 是 ZST (zero-sized type) 类型,不占用任何空间,所以这里无论有多少参数都不会导致栈溢出,这一写法和刚刚的加法表达式在编译后的机器码应该是等价的。

// runnable
macro_rules! count {
    (@consume $token:tt) => {
        ()
    };
    ($($token:tt)*) => {
        [$(count!(@consume $token)),*].len()
    };
}

macro_rules! myprint {
    ($($token:tt)*) => {
        let count = count!($($token)*);
        println!("tokens count: {count}");
        print!($($token)*);
    };
}

fn main() {
    let s1 = "hello";
    let s2 = "world".to_string();
    myprint!("{} {}\n", s1, s2);
}

展开后是这样:

fn main() {
    let s1 = "hello";
    let s2 = "world".to_string();
    let count = [(), (), (), (), ()].len();
    {
        ::std::io::_print(format_args!("tokens count: {0}\n", count));
    };
    {
        ::std::io::_print(format_args!("{0} {1}\n", s1, s2));
    };
}

Debug

我们的宏越来越复杂了,很多时候没法一次写对,好在编译器提供了一些实用的调试工具

  • trace_macros!
    在启用后,会输出宏每次的展开过程
    需要启用 trace_macros feature,并使用 nightly 版本的 Rust 编译
    这里的在线运行功能好像在编译成功时就不会输出宏的展开过程了,所以我仍然选择使用 stable 版本进行编译
// runnable
#![feature(trace_macros)]

macro_rules! count {
    (@consume $token:tt) => {
        ()
    };
    ($($token:tt)*) => {
        [$(count!(@consume $token)),*].len()
    };
}

macro_rules! myprint {
    ($($token:tt)*) => {
        let count = count!($($token)*);
        println!("tokens count: {count}");
        print!($($token)*);
    };
}

fn main() {
    let s1 = "hello";
    let s2 = "world".to_string();
    trace_macros!(true);
    myprint!("{} {}\n", s1, s2);
    trace_macros!(false);
}
  • log_syntax!($tt)
    需要启用 log_syntax feature,并使用 nightly 版本的 Rust 编译
    它会把输入的 tokens 全部打印出来,比如把 hello world 程序传进去,它会输出所有 tokens
// runnable nightly
#![feature(log_syntax)]

fn main() {
    log_syntax!(
        fn main() {
            println!("hello world");
        }
    );
}

递归

结合递归与重复匹配,可以写出来的东西就更有趣了,比如著名的 tt muncher(Token 吞噬者),即递归下降解析 token tree

// runnable nightly
#![feature(log_syntax)]

macro_rules! tt_muncher{
    // 没有 token 的时候终止
    () => {};
    // 每次匹配一个 token,剩下的继续传递
    ($token:tt $($rest:tt)*) => {
        log_syntax!($token);
        tt_muncher!($($rest)*);
    }
}

fn main() {
    tt_muncher!(
        fn main() {
            println!("hello world");
        }
    );
}

tt 会将大/中/小括号作为一个整体,所以上面的代码会把 main 函数的函数体作为整个 token 打印出来
我们可以让它匹配得更细一些

// runnable
macro_rules! tt_muncher{
    // 终止分支
    () => {};
    // 匹配大括号里面的内容
    ({$($inner:tt)*} $($rest:tt)*) => {
        println!("{{");
        tt_muncher!($($inner)*);
        println!("}}");
        tt_muncher!($($rest)*);
    };
    // 匹配小括号里面的内容
    (($($inner:tt)*) $($rest:tt)*) => {
        println!("(");
        tt_muncher!($($inner)*);
        println!(")");
        tt_muncher!($($rest)*);
    };
    // 匹配单个 token
    ($token:tt $($rest:tt)*) => {
        // stringify 可以将输入的 token 转换为 &'static str
        println!("{}", stringify!($token));
        tt_muncher!($($rest)*);
    };
}

fn main() {
    tt_muncher!(
        fn main() {
            println!("hello world");
        }
    );
}

tt muncher 每次吞噬一个 token,并传递剩下的 token,到最后一个分支的时候就没有任何 token 了
如果希望能保留一些 token 呢,注意我们处理的是 tokens,并且是在编译时最早期的阶段处理,所以没法用 Vec 之类的存起来,而 macro_rules 也不支持把 tokens 作为返回值传回去,所以如果希望保留一些中间值,则需要作为参数继续传递下去,每次调用都带上之前的参数,并加上新的内容,我们称之为 Push-down Accumulator (下推累加器)
我们之前提到过,只有大/中/小括号在解析时,会连带内部的内容视为一整个 token,这其实意味着,只有它们可以终止 $($token:tt)* 的匹配,而不会出现歧义。所以在使用累加器的时候,如果搭配 tt 来存储内容,则一般都会用括号把中间的内容包起来

举个解析函数签名的例子:
乍一看像是一堆乱码,其实只是在 tt muncher 的基础上加了累加器

// runnable
macro_rules! parse_fn{
    // 入口分支,提取出函数名及小括号内部的参数列表
    (fn $name:ident ($($inner:tt)*);) => {
        // 把函数名继续传下去,`[]` 是存放解析结果的,目前为空,最后把参数列表整个传进去
        // 注意这里在最后加了一个逗号,这是为了方便处理逗号匹配,保证最后一个参数后面也有逗号
        // 这样每次都匹配参数和逗号就行了
        // 如果参数列表尾部已经有逗号了,那多余的逗号会在结果分支被处理
        parse_fn!($name [] $($inner)*,);
    };
    // 传递函数名和解析结果的列表,这里不关心列表内部的内容,直接用 tt 匹配
    // 在列表后面,匹配形如“参数名: 参数类型,” 这样的参数,注意这里会匹配一个逗号
    // 最后传递剩下的 tokens
    ($name:ident [$($tokens:tt)*] $var:ident : $type:ty, $($rest:tt)*) => {
        // 原样展开匹配到的结果列表,并把参数存储为 (, 名字, 类型) 的形式
        parse_fn!($name [$($tokens)* (, $var, $type)] $($rest)*);
    };
    // 这里负责匹配参数名前面带 mut 的情况
    ($name:ident [$($tokens:tt)*] mut $var:ident : $type:ty, $($rest:tt)*) => {
        // 把参数存储为 (mut, 名字, 类型) 的形式
        parse_fn!($name [$($tokens)* (mut, $var, $type)] $($rest)*);
    };
    // 结果分支
    // 接收函数名和结果列表
    // 我们之前在列表里存的要么是(, name, type),要么是(mut, name, type)
    // 所以这里可选匹配第一个逗号前面的内容
    // 最后可选匹配一个逗号,处理参数列表尾部有逗号的情况
    ($name:ident [$(($($mut:ident)?, $v:ident, $t:ty))*] $(,)?) => {
        println!("function name: {}", stringify!($name));
        println!("parameters:");
        $(
            print!("\tname: {}, type: {}", stringify!($v), stringify!($t));
            let mutable = stringify!($($mut)?);
            if !mutable.is_empty(){
                println!(", is_mutable: {}", mutable);
            }else{
                println!();
            }
        )*
    };
}

fn main() {
    parse_fn!(
        fn run(title: String, mut tags: Vec<String>);
    );
}

使用 trace_macros! 可以看到这样的展开过程:

// DONT_FORMAT
= note: expanding `parse_fn! { fn run(title: String, mut tags: Vec<String>); }`
= note: to `parse_fn! (run [] title: String, mut tags: Vec<String>,);`
= note: expanding `parse_fn! { run [] title: String, mut tags: Vec<String>, }`
= note: to `parse_fn! (run [(, title, String)] mut tags: Vec<String>,);`
= note: expanding `parse_fn! { run [(, title, String)] mut tags: Vec<String>, }`
= note: to `parse_fn! (run [(, title, String)(mut, tags, Vec<String>)]);`

通过递归下降和累加器来计算,这让人想起了 Haskell 和函数式编程,值得一提的是,macro_rules! 是图灵完备的,玩法很多,比如已经有人用它写出了 brainfxxk 的编译器

当然也可以用它来计算斐波那契数列,不过算不了多少项,n 大了之后会很慢
毕竟 macro_rules 是用来处理 tokens 的,数字也只能表示成 tokens 的形式,编译期计算还是应该用 const fn

// runnable
macro_rules! fib {
    () => {
        0
    };
    (a) => {
        1
    };
    (a a $($n:ident)*) => {
        fib!(a $($n)*) + fib!($($n)*)
    };
}

fn main() {
    println!("fib 18: {}", fib!(a a a a a a a a a a a a a a a a a a));
}

实现 defer!

刚刚介绍了一大堆 macro_rules! 的内容,让我们回忆一下 defer! 宏要简化的东西:

// DONT_FORMAT
let (a, b) = &mut *DropGuard::new((a, b), |(a, b)| {
    println!("drop `a`");
    println!("{:?}", a);
    println!("drop b: {b}");
});

显然,我们希望通过宏实现少写几次 (a, b)

简单实现

// runnable
use std::mem::ManuallyDrop;
use std::ops::{Deref, DerefMut};

impl<F: FnOnce(T), T> Drop for DropGuard<F, T> {
    fn drop(&mut self) {
        let value = unsafe { ManuallyDrop::take(&mut self.inner) };
        let f = unsafe { ManuallyDrop::take(&mut self.f) };
        f(value);
    }
}

impl<F: FnOnce(T), T> Deref for DropGuard<F, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl<F: FnOnce(T), T> DerefMut for DropGuard<F, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner
    }
}

pub struct DropGuard<F: FnOnce(T), T> {
    f: ManuallyDrop<F>,
    inner: ManuallyDrop<T>,
}

impl<F: FnOnce(T), T> DropGuard<F, T> {
    pub fn new(inner: T, f: F) -> Self {
        Self {
            f: ManuallyDrop::new(f),
            inner: ManuallyDrop::new(inner),
        }
    }
    pub fn into_inner(self) -> T {
        let mut new_guard = ManuallyDrop::new(self);
        let value = unsafe { ManuallyDrop::take(&mut new_guard.inner) };
        unsafe { ManuallyDrop::drop(&mut new_guard.f) };
        value
    }
}

// ANCHOR
macro_rules! defer{
    ([$($vars:tt)*] $body:block) => {
        #[allow(unused_parens)]
        let mut __guard = $crate::DropGuard::new(
            ($($vars)*),
            |($($vars)*)| $body,
        );

        #[allow(unused_parens)]
        let ($($vars)*) = &mut *__guard;
    };
}

fn main() {
    let a = vec![1, 2];
    let b = "hello".to_string();
    defer!([a, b] {
        println!("drop `a`");
        println!("{:?}", a);
        println!("drop b: {b}");
    });
    println!("modify a");
    a.push(3);
    println!("print b: {b}");
}
// ANCHOR_END

看起来不错,但是有个问题,我们设计成了类似 C++ 的 lambda 表达式的语法,而 rustfmt 根本不认,它会直接放弃格式化宏内部的所有内容。
想让 rustfmt 能够格式化,我们必须得写成合法的 Rust 语法,这里最适合的是伪装成闭包

伪装成闭包的宏

// runnable
macro_rules! defer{
    (|$($vars:tt)*| $body:block) => {
        #[allow(unused_parens)]
        let mut __guard = $crate::DropGuard::new(
            ($($vars)*),
            |($($vars)*)| $body,
        );

        #[allow(unused_parens)]
        let ($($vars)*) = &mut *__guard;
    };
}

fn main() {
    let a = vec![1, 2];
    let b = "hello".to_string();
    defer!(|a, b| {
        println!("drop `a`");
        println!("{:?}", a);
        println!("drop b: {b}");
    });
    println!("modify a");
    a.push(3);
    println!("print b: {b}");
}
// ANCHOR_END

竟然没法编译!原因是这里产生了歧义,之前我们说过,$($vars:tt)* 是一个非常强力的匹配模式,它会吞噬掉后面的一切 token,直到遇到括号,这也是为什么之前可以编译,因为之前用的是中括号。现在的竖线则会被吞噬进去,所以编译器无法决定这个竖线应该算成重复匹配 tt,还是匹配字面上的竖线。

这就使我们的宏稍微复杂了一些,需要先匹配左竖线,再匹配右竖线。
不妨分成 deferdefer_impl 两个宏

// runnable
use std::mem::ManuallyDrop;
use std::ops::{Deref, DerefMut};

impl<F: FnOnce(T), T> Drop for DropGuard<F, T> {
    fn drop(&mut self) {
        let value = unsafe { ManuallyDrop::take(&mut self.inner) };
        let f = unsafe { ManuallyDrop::take(&mut self.f) };
        f(value);
    }
}

impl<F: FnOnce(T), T> Deref for DropGuard<F, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl<F: FnOnce(T), T> DerefMut for DropGuard<F, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner
    }
}

pub struct DropGuard<F: FnOnce(T), T> {
    f: ManuallyDrop<F>,
    inner: ManuallyDrop<T>,
}

impl<F: FnOnce(T), T> DropGuard<F, T> {
    pub fn new(inner: T, f: F) -> Self {
        Self {
            f: ManuallyDrop::new(f),
            inner: ManuallyDrop::new(inner),
        }
    }
    pub fn into_inner(self) -> T {
        let mut new_guard = ManuallyDrop::new(self);
        let value = unsafe { ManuallyDrop::take(&mut new_guard.inner) };
        unsafe { ManuallyDrop::drop(&mut new_guard.f) };
        value
    }
}

// ANCHOR
macro_rules! defer_impl{
    ([$($vars:tt)*] $(,)? | $body:block) => {
        #[allow(unused_parens)]
        let mut __guard = $crate::DropGuard::new(
            ($($vars)*),
            |($($vars)*)| $body,
        );

        #[allow(unused_parens)]
        let ($($vars)*) = &mut *__guard;
    };
    ([$($vars:tt)*] ,$v:tt $($rest:tt)*) => {
        defer_impl!([$($vars)* $v,] $($rest)*);
    };
}

macro_rules! defer{
    (|$($vars:tt)*) => {
        defer_impl!([] ,$($vars)*);
    };
}

fn main() {
    let a = vec![1, 2];
    let b = "hello".to_string();
    defer!(|a, b| {
        println!("drop `a`");
        println!("{:?}", a);
        println!("drop b: {b}");
    });
    println!("modify a");
    a.push(3);
    println!("print b: {b}");
}
// ANCHOR_END

值得一提的是这里匹配逗号的方式,每次匹配形如 , tt 这样的内容,这样只需要一开始在最前面加一个逗号,匹配完所有输入后,要么剩下的是一个逗号,要么什么都不剩,结果分支那里不需要额外调整。

如果每次匹配 tt , 这样的内容,为了处理末尾没有逗号的情况,要么在输入内容的最后面加一个逗号,但这里做不到,一开始输入进来的内容包括了右竖线和 block;要么在结果分支处理最后一次匹配,会更麻烦。

支持取消功能

需要给 defer! 宏增加一个分支,并给 defer_impl! 宏的每个分支加上 guard ,以便传递到结果分支

// runnable
use std::mem::ManuallyDrop;
use std::ops::{Deref, DerefMut};

impl<F: FnOnce(T), T> Drop for DropGuard<F, T> {
    fn drop(&mut self) {
        let value = unsafe { ManuallyDrop::take(&mut self.inner) };
        let f = unsafe { ManuallyDrop::take(&mut self.f) };
        f(value);
    }
}

impl<F: FnOnce(T), T> Deref for DropGuard<F, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl<F: FnOnce(T), T> DerefMut for DropGuard<F, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner
    }
}

pub struct DropGuard<F: FnOnce(T), T> {
    f: ManuallyDrop<F>,
    inner: ManuallyDrop<T>,
}

impl<F: FnOnce(T), T> DropGuard<F, T> {
    pub fn new(inner: T, f: F) -> Self {
        Self {
            f: ManuallyDrop::new(f),
            inner: ManuallyDrop::new(inner),
        }
    }
    pub fn into_inner(self) -> T {
        let mut new_guard = ManuallyDrop::new(self);
        let value = unsafe { ManuallyDrop::take(&mut new_guard.inner) };
        unsafe { ManuallyDrop::drop(&mut new_guard.f) };
        value
    }
}

// ANCHOR
macro_rules! defer_impl{
    ($guard:ident [$($vars:tt)*] $(,)? | $body:block) => {
        #[allow(unused_parens)]
        let mut $guard = $crate::DropGuard::new(
            ($($vars)*),
            |($($vars)*)| $body,
        );

        #[allow(unused_parens)]
        let ($($vars)*) = &mut *$guard;
    };
    ($guard:ident [$($vars:tt)*] ,$v:tt $($rest:tt)*) => {
        defer_impl!($guard [$($vars)* $v,] $($rest)*);
    };
}

macro_rules! defer{
    (|$($vars:tt)*) => {
        defer_impl!(__guard [] ,$($vars)*);
    };
    ($guard:ident, |$($vars:tt)*) => {
        defer_impl!($guard [] ,$($vars)*);
    };
}

fn main() {
    let a = vec![1, 2];
    defer!(guard, |a| {
        println!("drop `a`");
        println!("{:?}", a);
    });
    println!("modify a");
    a.push(3);
    guard.into_inner();
}
// ANCHOR_END

支持修饰符

对于上面的例子,无法实现在 defer! 的闭包里修改 a,虽然大部分情况用不上,而且通过重新绑定也能实现,但是既然都伪装成闭包的形状了,干脆支持 mut 关键字吧。
这里宏又变得更复杂了,之前我们只是提取竖线内部的内容,并按原样使用,但现在不一样了,声明闭包的时候需要 mut 关键字,但构造元组和解构元组的时候并不需要,所以我们得存储更多的信息。
这里使用了我们之前解析函数签名时的技巧,把 mut 存成可选的,并用逗号分隔。

顺带一提,之前的实现有个缺点,会把单变量存储为 (var,)(一个单变量元组), 而不是 (var)(单独的变量,带有多余的括号),虽然对 defer! 宏没什么影响,但是调用 guard 的 into_inner 方法时就很反直觉了,明明只存了一个变量,怎么变成了一个单变量元组呢?
而这个缺点在下面的实现里被解决了,因为逗号不再被存储到列表里,而是在展开的时候作为分隔符加入。

这是我写到这里才想到的,之前我设计成了存储成两个列表,一个用于元组,另一个用于闭包,但处理逗号时很麻烦。如果对那个方案感兴趣,可以看看本文最开始的代码块,点击右上角的第一个按钮就会显示出完整的代码。

// runnable
use std::mem::ManuallyDrop;
use std::ops::{Deref, DerefMut};

impl<F: FnOnce(T), T> Drop for DropGuard<F, T> {
    fn drop(&mut self) {
        let value = unsafe { ManuallyDrop::take(&mut self.inner) };
        let f = unsafe { ManuallyDrop::take(&mut self.f) };
        f(value);
    }
}

impl<F: FnOnce(T), T> Deref for DropGuard<F, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl<F: FnOnce(T), T> DerefMut for DropGuard<F, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner
    }
}

pub struct DropGuard<F: FnOnce(T), T> {
    f: ManuallyDrop<F>,
    inner: ManuallyDrop<T>,
}

impl<F: FnOnce(T), T> DropGuard<F, T> {
    pub fn new(inner: T, f: F) -> Self {
        Self {
            f: ManuallyDrop::new(f),
            inner: ManuallyDrop::new(inner),
        }
    }
    pub fn into_inner(self) -> T {
        let mut new_guard = ManuallyDrop::new(self);
        let value = unsafe { ManuallyDrop::take(&mut new_guard.inner) };
        unsafe { ManuallyDrop::drop(&mut new_guard.f) };
        value
    }
}

// ANCHOR
macro_rules! defer_impl{
    ($guard:ident [$($($mut:ident)?, $vars:ident)*] $(,)? | $body:block) => {
        // 在这里展开时,逗号作为分隔符被加入
        // 所以对于单个变量,会变成 `(var)`,这里的括号是多余的
        // 可以使用下面这个属性宏来避免警告
        #[allow(unused_parens)]
        let mut $guard = $crate::DropGuard::new(
            ($($vars),*),
            |($($($mut)? $vars),*)| $body,
        );

        #[allow(unused_parens)]
        let ($($vars),*) = &mut *$guard;
    };
    // 增加一个匹配 mut 的分支,因为逗号都在这些中间的分支被处理了,所以结果完全一样,
    // 仍然只需要一个结果分支
    ($guard:ident [$($($mut:ident)?, $vars:ident)*] , mut $v:ident $($rest:tt)*) => {
        defer_impl!($guard [$($vars)* mut, $v] $($rest)*);
    };
    ($guard:ident [$($($mut:ident)?, $vars:ident)*] ,$v:ident $($rest:tt)*) => {
        defer_impl!($guard [$($vars)* , $v] $($rest)*);
    };
}

macro_rules! defer{
    (|$($vars:tt)*) => {
        defer_impl!(__guard [] ,$($vars)*);
    };
    ($guard:ident, |$($vars:tt)*) => {
        defer_impl!($guard [] ,$($vars)*);
    };
}

fn main() {
    let a = vec![1, 2];
    defer!(|mut a| {
        a.push(4);
        println!("drop `a`");
        println!("{:?}", a);
    });
    println!("modify a");
    a.push(3);
}
// ANCHOR_END

支持捕获引用

之前的实现都是捕获变量本身,把所有权捕获到 DropGuard 结构体里。但有时可能不希望移动所有权,只想捕获引用。
目前想要捕获引用,只能自己手动创建

// runnable
use std::mem::ManuallyDrop;
use std::ops::{Deref, DerefMut};

impl<F: FnOnce(T), T> Drop for DropGuard<F, T> {
    fn drop(&mut self) {
        let value = unsafe { ManuallyDrop::take(&mut self.inner) };
        let f = unsafe { ManuallyDrop::take(&mut self.f) };
        f(value);
    }
}

impl<F: FnOnce(T), T> Deref for DropGuard<F, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl<F: FnOnce(T), T> DerefMut for DropGuard<F, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner
    }
}

pub struct DropGuard<F: FnOnce(T), T> {
    f: ManuallyDrop<F>,
    inner: ManuallyDrop<T>,
}

impl<F: FnOnce(T), T> DropGuard<F, T> {
    pub fn new(inner: T, f: F) -> Self {
        Self {
            f: ManuallyDrop::new(f),
            inner: ManuallyDrop::new(inner),
        }
    }
    pub fn into_inner(self) -> T {
        let mut new_guard = ManuallyDrop::new(self);
        let value = unsafe { ManuallyDrop::take(&mut new_guard.inner) };
        unsafe { ManuallyDrop::drop(&mut new_guard.f) };
        value
    }
}

macro_rules! defer_impl{
    ($guard:ident [$($($mut:ident)?, $vars:ident)*] $(,)? | $body:block) => {
        #[allow(unused_parens)]
        let mut $guard = $crate::DropGuard::new(
            ($($vars),*),
            |($($($mut)? $vars),*)| $body,
        );

        #[allow(unused_parens)]
        let ($($vars),*) = &mut *$guard;
    };
    ($guard:ident [$($($mut:ident)?, $vars:ident)*] , mut $v:ident $($rest:tt)*) => {
        defer_impl!($guard [$($vars)* mut, $v] $($rest)*);
    };
    ($guard:ident [$($($mut:ident)?, $vars:ident)*] ,$v:ident $($rest:tt)*) => {
        defer_impl!($guard [$($vars)* , $v] $($rest)*);
    };
}

macro_rules! defer{
    (|$($vars:tt)*) => {
        defer_impl!(__guard [] ,$($vars)*);
    };
    ($guard:ident, |$($vars:tt)*) => {
        defer_impl!($guard [] ,$($vars)*);
    };
}

// ANCHOR
fn main() {
    let mut a = vec![1, 2];
    let a = &mut a;
    defer!(|a| {
        a.push(4);
        println!("drop `a`");
        println!("{:?}", a);
    });
    println!("modify a");
    a.push(3);
    println!("{}", std::any::type_name_of_val(&a));
}
// ANCHOR_END

这个写法有点麻烦,而且内部重新借用后的 a 的类型变成了 &mut &mut Vec<i32>,这个双重引用完全没必要。
模仿 C++ 的语法,可以给宏加入捕获引用的功能

// runnable
use std::mem::ManuallyDrop;
use std::ops::{Deref, DerefMut};

impl<F: FnOnce(T), T> Drop for DropGuard<F, T> {
    fn drop(&mut self) {
        let value = unsafe { ManuallyDrop::take(&mut self.inner) };
        let f = unsafe { ManuallyDrop::take(&mut self.f) };
        f(value);
    }
}

impl<F: FnOnce(T), T> Deref for DropGuard<F, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl<F: FnOnce(T), T> DerefMut for DropGuard<F, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner
    }
}

pub struct DropGuard<F: FnOnce(T), T> {
    f: ManuallyDrop<F>,
    inner: ManuallyDrop<T>,
}

impl<F: FnOnce(T), T> DropGuard<F, T> {
    pub fn new(inner: T, f: F) -> Self {
        Self {
            f: ManuallyDrop::new(f),
            inner: ManuallyDrop::new(inner),
        }
    }
    pub fn into_inner(self) -> T {
        let mut new_guard = ManuallyDrop::new(self);
        let value = unsafe { ManuallyDrop::take(&mut new_guard.inner) };
        unsafe { ManuallyDrop::drop(&mut new_guard.f) };
        value
    }
}

// ANCHOR
macro_rules! defer_impl{
    ($guard:ident [$($($mut:ident)?, $vars:ident)*] $(,)? | $body:block) => {
        #[allow(unused_parens)]
        let mut $guard = $crate::DropGuard::new(
            ($($vars),*),
            |($($($mut)? $vars),*)| $body,
        );

        #[allow(unused_parens)]
        let ($($vars),*) = &mut *$guard;
    };
    ($guard:ident [$($($mut:ident)?, $vars:ident)*] , mut $v:ident $($rest:tt)*) => {
        defer_impl!($guard [$($vars)* mut, $v] $($rest)*);
    };
    // 只需要多加一个分支,匹配带 `&` 符号的变量
    ($guard:ident [$($($mut:ident)?, $vars:ident)*] , &$v:ident $($rest:tt)*) => {
        // 先创建可变引用
        let $v = &mut $v;
        // 把可变引用的所有权交给 guard
        defer_impl!($guard [$($vars)* , $v] $($rest)*);
        // 之后重新借用回来,并且多解引用一次,避免双重引用
        let $v = &mut **$v;
    };
    ($guard:ident [$($($mut:ident)?, $vars:ident)*] ,$v:ident $($rest:tt)*) => {
        defer_impl!($guard [$($vars)* , $v] $($rest)*);
    };
}

macro_rules! defer{
    (|$($vars:tt)*) => {
        defer_impl!(__guard [] ,$($vars)*);
    };
    ($guard:ident, |$($vars:tt)*) => {
        defer_impl!($guard [] ,$($vars)*);
    };
}

fn main() {
    let mut a = vec![1, 2];
    defer!(|&a| {
        a.push(4);
        println!("drop `a`");
        println!("{:?}", a);
    });
    println!("modify a");
    a.push(3);
    println!("{}", std::any::type_name_of_val(&a));
}
// ANCHOR_END

但是这样一来少了给引用命名的功能,多数情况下,捕获引用而非所有权,是为了后续能在 drop 之后返回变量
所以再加一个分支

// runnable
use std::mem::ManuallyDrop;
use std::ops::{Deref, DerefMut};

impl<F: FnOnce(T), T> Drop for DropGuard<F, T> {
    fn drop(&mut self) {
        let value = unsafe { ManuallyDrop::take(&mut self.inner) };
        let f = unsafe { ManuallyDrop::take(&mut self.f) };
        f(value);
    }
}

impl<F: FnOnce(T), T> Deref for DropGuard<F, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl<F: FnOnce(T), T> DerefMut for DropGuard<F, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner
    }
}

pub struct DropGuard<F: FnOnce(T), T> {
    f: ManuallyDrop<F>,
    inner: ManuallyDrop<T>,
}

impl<F: FnOnce(T), T> DropGuard<F, T> {
    pub fn new(inner: T, f: F) -> Self {
        Self {
            f: ManuallyDrop::new(f),
            inner: ManuallyDrop::new(inner),
        }
    }
    pub fn into_inner(self) -> T {
        let mut new_guard = ManuallyDrop::new(self);
        let value = unsafe { ManuallyDrop::take(&mut new_guard.inner) };
        unsafe { ManuallyDrop::drop(&mut new_guard.f) };
        value
    }
}

// ANCHOR
macro_rules! defer_impl{
    ($guard:ident [$($($mut:ident)?, $vars:ident)*] $(,)? | $body:block) => {
        #[allow(unused_parens)]
        let mut $guard = $crate::DropGuard::new(
            ($($vars),*),
            |($($($mut)? $vars),*)| $body,
        );

        #[allow(unused_parens)]
        let ($($vars),*) = &mut *$guard;
    };
    ($guard:ident [$($($mut:ident)?, $vars:ident)*] , mut $v:ident $($rest:tt)*) => {
        defer_impl!($guard [$($vars)* mut, $v] $($rest)*);
    };
    // 使用 `bind: &var`,类似构建结构体时的语法
    ($guard:ident [$($($mut:ident)?, $vars:ident)*] , $bind:ident: &$var:ident $($rest:tt)*) => {
        let $bind = &mut $var;
        defer_impl!($guard [$($vars)* , $bind] $($rest)*);
        let $bind = &mut **$bind;
    };
    ($guard:ident [$($($mut:ident)?, $vars:ident)*] , &$v:ident $($rest:tt)*) => {
        let $v = &mut $v;
        defer_impl!($guard [$($vars)* , $v] $($rest)*);
        let $v = &mut **$v;
    };
    ($guard:ident [$($($mut:ident)?, $vars:ident)*] ,$v:ident $($rest:tt)*) => {
        defer_impl!($guard [$($vars)* , $v] $($rest)*);
    };
}

macro_rules! defer{
    (|$($vars:tt)*) => {
        defer_impl!(__guard [] ,$($vars)*);
    };
    ($guard:ident, |$($vars:tt)*) => {
        defer_impl!($guard [] ,$($vars)*);
    };
}

fn return_owned() -> Vec<i32> {
    let mut a = vec![1, 2];
    defer!(guard, |a_ref: &a| {
        a_ref.push(4);
        println!("drop `a_ref`");
        println!("{:?}", a_ref);
    });
    println!("modify a");
    a_ref.push(3);
    println!("a_ref: {}", std::any::type_name_of_val(&a_ref));
    drop(guard);
    a
}

fn main() {
    let a = return_owned();
    println!("return a: {}", std::any::type_name_of_val(&a));
}
// ANCHOR_END

再加点小功能

闭包支持使用 move 来表示内部使用的变量都需要被移动进去,虽然我觉得在这里不太用得到,但是这个功能很好加。
另外就是加上不指定变量的两种情况。

// runnable
use std::mem::ManuallyDrop;
use std::ops::{Deref, DerefMut};

impl<F: FnOnce(T), T> Drop for DropGuard<F, T> {
    fn drop(&mut self) {
        let value = unsafe { ManuallyDrop::take(&mut self.inner) };
        let f = unsafe { ManuallyDrop::take(&mut self.f) };
        f(value);
    }
}

impl<F: FnOnce(T), T> Deref for DropGuard<F, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl<F: FnOnce(T), T> DerefMut for DropGuard<F, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner
    }
}

pub struct DropGuard<F: FnOnce(T), T> {
    f: ManuallyDrop<F>,
    inner: ManuallyDrop<T>,
}

impl<F: FnOnce(T), T> DropGuard<F, T> {
    pub fn new(inner: T, f: F) -> Self {
        Self {
            f: ManuallyDrop::new(f),
            inner: ManuallyDrop::new(inner),
        }
    }
    pub fn into_inner(self) -> T {
        let mut new_guard = ManuallyDrop::new(self);
        let value = unsafe { ManuallyDrop::take(&mut new_guard.inner) };
        unsafe { ManuallyDrop::drop(&mut new_guard.f) };
        value
    }
}

// ANCHOR
macro_rules! defer_impl{
    // 只需将 move 作为可选参数捕获并传递即可,最后在闭包前面展开
    ($guard:ident $($move:ident)? [$($($mut:ident)?, $vars:ident)*] $(,)? | $body:block) => {
        #[allow(unused_parens)]
        let mut $guard = $crate::DropGuard::new(
            ($($vars),*),
            $($move)? |($($($mut)? $vars),*)| $body,
        );

        #[allow(unused_parens)]
        let ($($vars),*) = &mut *$guard;
    };
    ($guard:ident $($move:ident)? [$($($mut:ident)?, $vars:ident)*] , mut $v:ident $($rest:tt)*) => {
        defer_impl!($guard [$($vars)* mut, $v] $($rest)*);
    };
    ($guard:ident $($move:ident)? [$($($mut:ident)?, $vars:ident)*] , $bind:ident: &$var:ident $($rest:tt)*) => {
        let $bind = &mut $var;
        defer_impl!($guard [$($vars)* , $bind] $($rest)*);
        let $bind = &mut **$bind;
    };
    ($guard:ident $($move:ident)? [$($($mut:ident)?, $vars:ident)*] , &$v:ident $($rest:tt)*) => {
        let $v = &mut $v;
        defer_impl!($guard [$($vars)* , $v] $($rest)*);
        let $v = &mut **$v;
    };
    ($guard:ident $($move:ident)? [$($($mut:ident)?, $vars:ident)*] ,$v:ident $($rest:tt)*) => {
        defer_impl!($guard [$($vars)* , $v] $($rest)*);
    };
}

macro_rules! defer{
    // 指定捕获变量,使用内部 guard,不可取消
    (|$($vars:tt)*) => {
        defer_impl!(__guard [] ,$($vars)*);
    };
    // 指定捕获变量,使用指定的 guard,可取消
    ($guard:ident, $($move:ident)? |$($vars:tt)*) => {
        defer_impl!($guard $($move:ident)? [] ,$($vars)*);
    };
    // 不指定捕获变量,使用指定的 guard,可取消
    // 这里直接捕获表达式,可以接收闭包,move 等功能都由闭包原生支持
    ($guard:ident, $body:expr) => {
        let $guard = $crate::DropGuard::new(
            (),
            |_| $body(),
        );
    };
    // 不指定捕获变量,使用内部 guard,不可取消
    // scopeguard 的 defer 宏就只支持这一种情况
    ($body:expr) => {
        let __guard = $crate::DropGuard::new(
            (),
            |_| $body(),
        );
    };
}

fn main() {
    let mut a = vec![1, 2];
    defer!(guard, move || {
        a.push(3);
        println!("drop `a`");
        println!("{:?}", a);
        println!("type of a: {}", std::any::type_name_of_val(&a));
    });
}
// ANCHOR_END

关于可变性的设计

Rust 里采取的是默认不可变,需要时使用 mut 声明可变这一策略,这个宏也在尽量遵循这个设计。

对于可变性的需求,可以分为在宏内部和在宏之后
这里的可变性,是指对被捕获值的可变性,而不是变量本身的可变性,如捕获 String 类型为可变引用,宏内部使用的是 s: &mut String ,这里可以通过 s 去修改字符串,就称为可变,但 s 这个变量本身不可变(而且似乎也没有可变的需求)

方式 在宏内部 在宏之后
捕获为可变引用 可变 可变
捕获时移动变量所有权 不可变(默认行为) 可变
捕获时移动变量所有权 可变(通过 mut 来声明) 可变
使用闭包自带的捕获功能捕获为不可变引用 不可变 不可变
不支持 可变 不可变

对于捕获引用的情况来说,宏内部和宏之后使用的变量名和类型是一样的,可变性最好保持一样;
对于捕获所有权的情况来说,捕获所有权之后,内部当然可以选择是否需要可变性,但是宏之后的变量是重新绑定的,两者的类型不一致,所以可变性不必相同
理论上来说可以支持最后一种情况,但是相当奇怪,且语法不好设计,就没实现。如果确实有需求,完全可以通过手动的 let 重绑定来实现。

局限性

我们在 Rust 中实现了 defer! 宏,简化了语法,覆盖了基础功能,但使用 macro_rules! 能做到的功能有限,比如它需要手动指定需要的变量,不能完全自动捕获;它不像 Go 的 defer,可以直接 return 函数返回值;也不像 Zig 中的 defer,可以在根据出错与否自动取消或执行。使用过程宏或许能实现更多功能,但更复杂的功能终究还是需要编译器来实现,那样会更简单也更高效,希望 defer 能早日进入 Rust 的 RFC

defer 中的陷阱

defer 与 RAII 的不同之处在于,它更灵活自由,但仍然遵循着先定义的后释放的原则,所以使用时需要特别小心释放顺序
如果是手动管理 C 里的资源,最好像 RAII 一样,获取后立即声明 defer

// runnable
use std::mem::ManuallyDrop;
use std::ops::{Deref, DerefMut};

impl<F: FnOnce(T), T> Drop for DropGuard<F, T> {
    fn drop(&mut self) {
        let value = unsafe { ManuallyDrop::take(&mut self.inner) };
        let f = unsafe { ManuallyDrop::take(&mut self.f) };
        f(value);
    }
}

impl<F: FnOnce(T), T> Deref for DropGuard<F, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl<F: FnOnce(T), T> DerefMut for DropGuard<F, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner
    }
}

pub struct DropGuard<F: FnOnce(T), T> {
    f: ManuallyDrop<F>,
    inner: ManuallyDrop<T>,
}

impl<F: FnOnce(T), T> DropGuard<F, T> {
    pub fn new(inner: T, f: F) -> Self {
        Self {
            f: ManuallyDrop::new(f),
            inner: ManuallyDrop::new(inner),
        }
    }
    pub fn into_inner(self) -> T {
        let mut new_guard = ManuallyDrop::new(self);
        let value = unsafe { ManuallyDrop::take(&mut new_guard.inner) };
        unsafe { ManuallyDrop::drop(&mut new_guard.f) };
        value
    }
}

macro_rules! defer_impl{
    ($guard:ident $($move:ident)? [$($($mut:ident)?, $vars:ident)*] $(,)? | $body:block) => {
        #[allow(unused_parens)]
        let mut $guard = $crate::DropGuard::new(
            ($($vars),*),
            $($move)? |($($($mut)? $vars),*)| $body,
        );

        #[allow(unused_parens)]
        let ($($vars),*) = &mut *$guard;
    };
    ($guard:ident $($move:ident)? [$($($mut:ident)?, $vars:ident)*] , mut $v:ident $($rest:tt)*) => {
        defer_impl!($guard [$($vars)* mut, $v] $($rest)*);
    };
    ($guard:ident $($move:ident)? [$($($mut:ident)?, $vars:ident)*] , $bind:ident: &$var:ident $($rest:tt)*) => {
        let $bind = &mut $var;
        defer_impl!($guard [$($vars)* , $bind] $($rest)*);
        let $bind = &mut **$bind;
    };
    ($guard:ident $($move:ident)? [$($($mut:ident)?, $vars:ident)*] , &$v:ident $($rest:tt)*) => {
        let $v = &mut $v;
        defer_impl!($guard [$($vars)* , $v] $($rest)*);
        let $v = &mut **$v;
    };
    ($guard:ident $($move:ident)? [$($($mut:ident)?, $vars:ident)*] ,$v:ident $($rest:tt)*) => {
        defer_impl!($guard [$($vars)* , $v] $($rest)*);
    };
}

macro_rules! defer{
    (|$($vars:tt)*) => {
        defer_impl!(__guard [] ,$($vars)*);
    };
    ($guard:ident, $($move:ident)? |$($vars:tt)*) => {
        defer_impl!($guard $($move:ident)? [] ,$($vars)*);
    };
    ($guard:ident, $body:expr) => {
        let $guard = $crate::DropGuard::new(
            (),
            |_| $body(),
        );
    };
    ($body:expr) => {
        let __guard = $crate::DropGuard::new(
            (),
            |_| $body(),
        );
    };
}

// ANCHOR
fn correct() {
    let a: *mut i32 = Box::leak(Box::new(1));
    defer!(|| {
        unsafe {
            println!("drop a: {}", *a);
            drop(Box::from_raw(a));
        }
    });
    let b = Box::leak(Box::new(a));
    defer!(|| {
        unsafe {
            println!("drop b: {}", **b);
            drop(Box::from_raw(b));
        }
    });
}

fn wrong() {
    let a: *mut i32 = Box::leak(Box::new(1));
    let b = Box::leak(Box::new(a));
    defer!(|| {
        unsafe {
            println!("drop b: {}", **b);
            drop(Box::from_raw(b));
        }
    });
    defer!(|| {
        unsafe {
            println!("drop a: {}", *a);
            drop(Box::from_raw(a));
        }
    });
}

fn main() {
    println!("correct:");
    correct();
    println!("wrong:");
    wrong();
    println!("finish");
}
// ANCHOR_END

运行程序会发现 wrong 函数里 drop b 的时候发生了 use after free 的经典问题,原因就是先释放了 a,后释放了 b

最后

本来还想再多说说各种语言里的 defer,但就这样吧,这篇文章已经很长很长了……

  1. Perhaps Rust needs "defer"
    Pre-RFC: defer statement
    A defer discussion
    Pre-pre-pre-RFC: implicit code control and defer statements

  2. Rust 的闭包有 3 种类型,Fn, FnMutFnOnce,这里使用 FnOnce,即只能执行一次,会消耗所有权的闭包,另外两种闭包都可以协变为 FnOnce,所以用在这里通用性最好。

  3. Box 是 Rust 中的智能指针,相当于 C++ 里的 unique_ptr

  4. 更多类型请参考Captures

  5. 可以参考Hygiene这一章中的解释

  6. 可以通过 crates.io 安装 cargo-expand

写一个类似doas的程序(一)

已经到底啦!