A Wand Makes Your Ideas Come True

0%

Rust过程宏系列教程(5)--实现proc_macro_workshop项目之sorted题目(属性式过程宏)

发表于 更新于

本系列的上一篇文章中，我们实战了proc_macro_workshop项目的seq题目。前面三篇文章已经给大家介绍过了rust三种过程宏里面的两种，也就是派生式过程宏和函数式过程宏，这一次我们要介绍的是最后一种，也就是属性样式的过程宏。

好了，不废话了，准备好一台电脑，开始我们的第四个挑战任务sorted

先是视频版本的教程，文字版本的在下面~

  • Rust过程宏开发实战系列(5-1)【属性式过程宏】sorted题目1-4关

  • Rust过程宏开发实战系列(5-2)【属性式过程宏】sorted题目5-8关

文字版本开始~

首先打开proc_macro_workshop项目的readme.md文件，看一下sorted这个项目要实现什么样的功能。根据其中的描述，这个题目的最终目标是实现一个可以检测枚举类型中各个成员的排列顺序是否严格按照字母表的顺序递增排列，如果发现有的成员名字与相邻的成员不是递增关系，则抛出异常终止编译，可以说，这个宏是强迫症患者的福音了。

这里提前说明一下，属性式的过程宏和派生式的过程宏写起来很像，所以大家在学习的时候，可以对照着前面的builder题目以及debug题目来学习。在接下来的文章中，我们会先介绍一下属性式过程宏和派生式过程宏的区别，然后与之前不同的是，由于大量知识是前面已经给大家介绍过的，所以这篇文章就不再一关一关的进行了，而是会把相似的几个关卡合成一个主题来讲解，这样会让文章紧凑一些，但是也意味着，如果你还没有对派生式的过程宏有一个了解，那么强烈建议你先去看前面的文章。你可以访问我的个人博客http://blog.ideawand.com 来阅读之前的文章，或者关注我的微信公众号【极客幼稚园】阅读之前的文章。好了，我们先开始对比一下派生式过程宏和属性式过程宏的异同点。

派生式 属性式
调用方法 #[derive(XXX)]
必须以derive开头		 #[XXX(YYY,ZZZ.....)]
#[]内部可以任何的合法标识符开头作为宏的名字，且宏的名字后面可以写任意合法的TokenStream
目标代码块 只能是struct、enum、union 语法树中Item节点对应的代码块，
除了struct、enum、union以外，还可以包含函数、impl块、mod块、trait等等
全部列表可以参见https://docs.rs/syn/1.0.73/syn/enum.Item.html
代码修改能力 只能添加代码，不能修改被修饰的原始代码，即不能修改被装饰的struct、enum、union内部的代码 可以替换、删除、新增代码，对被修饰的代码有完全的掌控权限
入口函数签名 #[proc_macro_derive(Builder)]
pub fn derive(input: TokenStream) -> TokenStream {}

有一个入参，对应被修饰的struct/enum/union代码块		 #[proc_macro_attribute]
pub fn sorted(args: TokenStream, input: TokenStream) -> TokenStream {}

有两个入参，第一个入参对应宏调用本身（也就是#[xxxxx]这一行），第二个参数代表被修饰的代码块

第一关 ~ 第四关

为了加快速度，增加干货，减少废话，我们把前四关的代码放在一起整体实现。毕竟这里面用到的知识点我们在前面已经见过很多次了。

  • 第一关告诉我们将输入的TokenStream解析成syn::Item语法树节点，并且提示我们为了使用syn::Item我们需要开启syn的full特性。
  • 第二关的测试用例要求我们检查过程宏是否作用在枚举体上，同时告诉我们如何设计函数的返回值，从而利用syn::Result和其to_compile_error方法的组合实现统一的编译错误提示。关于syn::Result和to_compile_error的组合模式，我们在之前的《Rust过程宏系列教程(2)–实现proc-macro-workshop项目之Builder题目》和《Rust过程宏系列教程(3)–实现proc_macro_workshop项目之debug题目》两篇文章中都有提及，大家不熟悉的话可以回看一下。
    • 第二关会使用到proc_macro2::Span::call_site()来获取过程宏调用的位置
  • 第三关和第四关都是要检查枚举体的成员是否按字母顺序排列，是本道题目的核心功能，第四关相比第三关，只是支持的枚举类型更加复杂一点，要注意输出结果span的调整。

好，开始搭架子，首先是整个入口框架：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
use proc_macro::TokenStream;
use syn::{self, spanned::Spanned};
use quote::{ToTokens, quote};

#[proc_macro_attribute]
pub fn sorted(args: TokenStream, input: TokenStream) -> TokenStream {
    // 解析为syn::Item语法树节点（源自测试用例01的提示）
    let st = syn::parse_macro_input!(input as syn::Item); 

    // 宏的入口使用syn::Result，配合syn::Result::to_compile_error使用（源自测试用例02的提示）
    match do_expand(&st) {
        Ok(token_stream) => token_stream.into(),
        Err(e) => e.to_compile_error().into(),
    }
}

fn do_expand(st: &syn::Item) -> syn::Result<proc_macro2::TokenStream> {
    let ret = proc_macro2::TokenStream::new();
    return Ok(ret);
}

下面来写前四关的主体逻辑，需要注意的点都标明在注释中了。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
fn do_expand(st: &syn::Item) -> syn::Result<proc_macro2::TokenStream> { 
    let ret = match st {
        syn::Item::Enum(enum_node) => check_enum_order(enum_node),
        // 注意下面call_site()方法的使用，它可以获得过程宏被调用的位置，这样才可以满足测试用例02的输出要求
        _ => syn::Result::Err(syn::Error::new(proc_macro2::Span::call_site(), "expected enum or match expression"))
    }?;

    return Ok(ret);
}

fn check_enum_order(st: &syn::ItemEnum) -> syn::Result<proc_macro2::TokenStream>{
    // 这里是核心的比较算法，构造两个数组，一个是排序的，一个是没有排序的，同时从头遍历，遇到第一个不一致的位置就是需要提示错误的位置

    let origin_names: Vec<_> = st.variants.iter().map(|item|{(item.ident.to_string(), item)}).collect();
    let mut sorted_origin_names = origin_names.clone();
    sorted_origin_names.sort_by(|a,b|{a.0.cmp(&b.0)});

    for (a,b) in origin_names.iter().zip(sorted_origin_names.iter()) {
        if a.0 != b.0 {
            // 注意下面一行要使用`b.1.ident.span()`，而不是`b.1.span()`，因为测试用例04的输出结果中只有Ident部分被特别标出
            return syn::Result::Err(syn::Error::new(b.1.ident.span(), format!("{} should sort before {}", b.0, a.0)))
        }
    }
    return syn::Result::Ok(st.to_token_stream())
}

接下来跑一下测试用例，会发现有两个问题：

  • 首先因为有3关都是校验标注错误输出的，因为我们代码中有未使用的变量和import的内容，所以要修改一下
  • 第二个是，当我们发现第一个顺序不对的元素就返回时，会导致最后返回给编译器的的TokenStream中只包含错误信息，这就等同于被我们的过程宏所修饰的整个枚举体被从代码中删掉了，这样会导致编译器输出其他的异常信息，从而导致标准输出中的错误和预期不一致，因此，不管是否有错误发生，我们都要返回给编译器原来的完整的枚举体。这个地方和派生式过程宏是有区别的，因为派生式的过程宏只能添加代码，不能删掉已有的代码，所以你报错之后直接返回空的TokenStream没有什么问题，但是在属性式的过程宏中麻烦就大了，返回空的TokenStream等于删掉了被过程宏修饰的代码，所以修改后的入口函数如下,修改点参见注释：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
use proc_macro::TokenStream;
use syn;
use quote::ToTokens;
use proc_macro2;

#[proc_macro_attribute]
// 下面的_args入参因为用不到，所以加一个下划线，避免出现警告信息
pub fn sorted(_args: TokenStream, input: TokenStream) -> TokenStream {

    // 解析为syn::Item语法树节点（源自测试用例01的提示）
    let st = syn::parse_macro_input!(input as syn::Item); 

    // 宏的入口使用syn::Result，配合syn::Result::to_compile_error使用（源自测试用例02的提示）
    match do_expand(&st) {
        Ok(token_stream) => token_stream.into(),
        Err(e) => {
            // 下面这一行拿到的TokenStream是空的，里面只包含了错误信息，没有代码信息
            let mut t = e.to_compile_error();
            // 将原始的用户代码塞进去，这样返回结果中既包含代码信息，也包含错误信息
            t.extend(st.to_token_stream());
            t.into()
        }
    }
}

第五关 ~ 第八关

这五关其实也不难，其更重要的意义是给我们提供了一个通过过程宏来扩展rust编译器功能的案例，并且可以尝试使用syn提供的VisitMutTrait以访问者的模式来对语法树节点进行修改。

在procmarco_workshop作者设置这道题目的时候，rust的编译器的稳定版本还不支持在match语句块上增加属性标签，也就是下面这段代码中的#[sorted]是不能写在这个地方的（不过有可能你看到这篇文章时，这个特性已经稳定加入rust编译器了）。

1
2
3
4
5
6
#[sorted]
match my_works {
    Blog => "blog.ideawand.com",
    BiliBili => "",
    Wechat => "",
}

由于这个特性还没有稳定，所以我们不能使用编译器的这个特性。为了绕过这个问题，我们引入一个新的过程宏来修饰包含match语句的函数，因为函数本身是可以被过程宏修饰的，我们让这个处理函数的过程宏检测代码中是否有match语句被#[sorted]修饰，如果有，就把match语句块上的#[sorted]去掉，同时检测这个match语句块的各个分支是否按照字母顺序排列。下面给出了一个例子：

1
2
3
4
5
6
7
8
9
#[sorted::check]   // 这里是过程宏的调用
fn foo() {
  #[sorted]     // 这一行会被外面的过程宏去掉，在这里就是一个标记，最终编译器看不到这个标签
  match my_works {
      Blog => "blog.ideawand.com",
      BiliBili => "",
      Wechat => "",
  }
}

关于访问者模式，我们在前面的文章《Rust过程宏系列教程(3)–实现proc_macro_workshop项目之debug题目》中已经介绍过，只不过之前使用的是只读模式的Visitor,而在这一关中我们要使用可变的Visitor,记得要在Cargo.toml中为syn增加visit-mut特性以开启可修改的访问者模式。

第六关和第七关是第五关的强化版，我们要用的一个新的语法树节点syn::Pat,它表示了一个匹配模式，由于match语句中可以使用的匹配模式很多，例如列表也可以作为一个匹配模式，但在我们的过程宏中，如果某个match分支的匹配模式是一个列表，那么比较顺序就失去了意义，所以我们只处理Pat::Path、Pat::TupleStruct和Pat::Struct。

最后的第八关，因为在Ascii编码中，_对应的位置本来就排在所有字母的后面，所以我们的比较逻辑不需要做任何调整，应该可以直接通过。

好了，下面开始上代码，下面这段代码和上面的代码写在一个文件里就行，相当于这个时候我们在一个文件里定义了两个属性式的过程宏，重要信息请依然阅读其中的注释：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
#[proc_macro_attribute]
pub fn check(_args: TokenStream, input: TokenStream) -> TokenStream {

    // 这个过程宏的目标是函数，所以我们一上来就将其直接解析为函数块定义的节点
    let mut st = syn::parse_macro_input!(input as syn::ItemFn); 

    match do_match_expand(&mut st) {
        Ok(token_stream) => token_stream.into(),
        Err(e) => {
            let mut t = e.to_compile_error();
            t.extend(st.to_token_stream());
            t.into()
        }
    }
}

fn do_match_expand(st: &mut syn::ItemFn) -> syn::Result<proc_macro2::TokenStream> {

    // 创建Visitor并通过Visitor模式完成核心工作：从语法树节点找到满足条件的match语句块
    let mut visitor = MatchVisitor{err:None};

    visitor.visit_item_fn_mut(st);

    if visitor.err.is_none() {
        return syn::Result::Ok( st.to_token_stream())
    } else {
        return syn::Result::Err(visitor.err.unwrap());
    }
    
}

// 定义一个工具函数，用于把一个Path类型转换为形如AA::BB::CC形式的字符串，这个主要是为了测试用例06使用
fn get_path_string(p: &syn::Path) -> String {
    let mut buf = Vec::new();
    for i in &p.segments {
        buf.push(i.ident.to_string());
    }
    return buf.join("::")
}


// 定义Visitor
struct MatchVisitor {
    err: Option<syn::Error>,
}
impl syn::visit_mut::VisitMut for MatchVisitor {
    fn visit_expr_match_mut(&mut self, i: &mut syn::ExprMatch) {
        let mut target_idx:isize = -1;
        // 尝试锁定携带了`sorted`标记的match语句块
        for (idx, attr) in i.attrs.iter().enumerate() {
            if get_path_string(&attr.path) == "sorted" {
                target_idx = idx as isize;
                break;
            }
        }

        if target_idx != -1 {
            // 删除掉编译器不支持的写在match语句块上面的属性标签
            i.attrs.remove(target_idx as usize);
            
            let mut match_arm_names = Vec::new();
            for arm in &i.arms {
                // 按照测试用例06的提示，我们要处理三种匹配模式，测试用例07告诉我们对于不支持的模式需要抛出异常
                match &arm.pat {
                    syn::Pat::Path(p) => {
                        match_arm_names.push((get_path_string(&p.path), &p.path));
                    },
                    syn::Pat::TupleStruct(p) => {
                        match_arm_names.push((get_path_string(&p.path), &p.path));
                    },
                    syn::Pat::Struct(p) => {
                        match_arm_names.push((get_path_string(&p.path), &p.path));
                    },
                    _ => {
                        self.err = Some(syn::Error::new_spanned(&arm.pat, "unsupported by #[sorted]"));
                        return
                    }
                }
            }

            // 这里的算法和前面4关完全一样
            let mut sorted_names = match_arm_names.clone();
            sorted_names.sort_by(|a,b|{a.0.cmp(&b.0)});
            for (a,b) in match_arm_names.iter().zip(sorted_names.iter()) {
                if a.0 != b.0 {
                    self.err = Some(syn::Error::new_spanned(b.1, format!("{} should sort before {}", b.0, a.0)));
                    return
                }
            }
        }

        // 继续迭代深层次的match
        visit_mut::visit_expr_match_mut(self, i)
    }
}

上面的代码前面7关都可以通过，但是最后一关过不了，因为我们只考虑测试用例06提示我们的3种匹配模式还不够，测试用例08又引入了两种match模式。比较烦恼的是这两个新的模式和我们上面已经支持的3种模式，他们的语法树节点结构不太一样,我们获取span信息的方法需要修改。通过阅读syn::Error::new_spanned()的定义，我们知道需要传入一个满足quote::ToTokens Trait的对象才可以，所以这里动用一下trait object来解决在列表中存储不同类型语法树节点的问题，修改后的代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
impl syn::visit_mut::VisitMut for MatchVisitor {
    fn visit_expr_match_mut(&mut self, i: &mut syn::ExprMatch) {

        // 省略部分没有改变代码 。。。。。。

        if target_idx != -1 {
            i.attrs.remove(target_idx as usize);
            // 显示指定trait object的类型
            let mut match_arm_names: Vec<(_, &dyn ToTokens)> = Vec::new();

            for arm in &i.arms {
                match &arm.pat {
                    syn::Pat::Path(p) => {
                        match_arm_names.push((get_path_string(&p.path), &p.path));
                    },
                    syn::Pat::TupleStruct(p) => {
                        match_arm_names.push((get_path_string(&p.path), &p.path));
                    },
                    syn::Pat::Struct(p) => {
                        match_arm_names.push((get_path_string(&p.path), &p.path));
                    },
                    // 下面新增两种匹配模式的处理
                    syn::Pat::Ident(p) => {
                        match_arm_names.push((p.ident.to_string(), &p.ident));
                    },
                    syn::Pat::Wild(p) => {
                        match_arm_names.push(("_".to_string(), &p.underscore_token));
                    },
                    _ => {
                        self.err = Some(syn::Error::new_spanned(&arm.pat, "unsupported by #[sorted]"));
                        return
                    }
                }
            }
            // 省略部分没有改变代码 。。。。。。
        }
       // 省略部分没有改变代码 。。。。。。
}

至此，全部测试用例通过~

微信公众号：极客幼稚园
关注阅读更多优质技术文章