proc-macro-workshop:sorted-5

godme發表於2022-07-06
// Get ready for a challenging step -- this test case is going to be a much
// bigger change than the others so far.
//
// Not only do we want #[sorted] to assert that variants of an enum are written
// in order inside the enum definition, but also inside match-expressions that
// match on that enum.
//
//     #[sorted]
//     match conference {
//         RustBeltRust => "...",
//         RustConf => "...",
//         RustFest => "...",
//         RustLatam => "...",
//         RustRush => "...",
//     }
//
// Currently, though, procedural macro invocations on expressions are not
// allowed by the stable compiler! To work around this limitation until the
// feature stabilizes, we'll be implementing a new #[sorted::check] macro which
// the user will need to place on whatever function contains such a match.
//
//     #[sorted::check]
//     fn f() {
//         let conference = ...;
//
//         #[sorted]
//         match conference {
//             ...
//         }
//     }
//
// The #[sorted::check] macro will expand by looking inside the function to find
// any match-expressions carrying a #[sorted] attribute, checking the order of
// the arms in that match-expression, and then stripping away the inner
// #[sorted] attribute to prevent the stable compiler from refusing to compile
// the code.
//
// Note that unlike what we have seen in the previous test cases, stripping away
// the inner #[sorted] attribute will require the new macro to mutate the input
// syntax tree rather than inserting it unchanged into the output TokenStream as
// before.
//
// Overall, the steps to pass this test will be:
//
//   - Introduce a new procedural attribute macro called `check`.
//
//   - Parse the input as a syn::ItemFn.
//
//   - Traverse the function body looking for match-expressions. This part will
//     be easiest if you can use the VisitMut trait from Syn and write a visitor
//     with a visit_expr_match_mut method.
//
//   - For each match-expression, figure out whether it has #[sorted] as one of
//     its attributes. If so, check that the match arms are sorted and delete
//     the #[sorted] attribute from the list of attributes.
//
// The result should be that we get the expected compile-time error pointing out
// that `Fmt` should come before `Io` in the match-expression.
//
//
// Resources:
//
//   - The VisitMut trait to iterate and mutate a syntax tree:
//     https://docs.rs/syn/1.0/syn/visit_mut/trait.VisitMut.html
//
//   - The ExprMatch struct:
//     https://docs.rs/syn/1.0/syn/struct.ExprMatch.html

use sorted::sorted;

use std::fmt::{self, Display};
use std::io;

#[sorted]
pub enum Error {
    Fmt(fmt::Error),
    Io(io::Error),
}

impl Display for Error {
    #[sorted::check]
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::Error::*;

        #[sorted]
        match self {
            Io(e) => write!(f, "{}", e),
            Fmt(e) => write!(f, "{}", e),
        }
    }
}

fn main() {}

這裡主要有兩個驚喜

  • #[sorted::check]:這個和#[tokio::main]何其相像
  • match並不支援函式式宏,但是fn支援
    從錯誤提示我們可以知道將要實現什麼功能
    error: Fmt should sort before Io
    --> tests/05-match-expr.rs:88:13
     |
    88 |             Fmt(e) => write!(f, "{}", e),
     |             ^^^
    就是檢測match中的分支順序。

因此整體的思路是

  1. 透過檢測#[sorted::check]標記的方法,查詢其中的#[sorted]並進行檢測。
  2. 移除#[sorted]保證程式碼解析正確

我們要查詢的被#[sorted]標記的match,雖然核心是#[sorted],但是match才是主體,且不確定數量,因此,需要開啟visit,同時,還要移除#[sorted],所以真正應該開啟的是visit-mut

// common.rs
// 遞迴掃描
impl syn::visit_mut::VisitMut for MatchVisitor {
    fn visit_expr_match_mut(&mut self, i: &mut syn::ExprMatch) {
        // 標記屬性
        let mut target_idx: isize = -1;
        for (idx, attr) in i.attrs.iter().enumerate() {
            // 是否標記了#[sorted]
            if path_to_string(&attr.path) == "sorted" {
                target_idx = idx as isize;
                break;
            }
        }
        // 如果標記了才處理
        if target_idx != -1 {
            // 移除#[sorted]
            i.attrs.remove(target_idx as usize);
            // 收集names
            let mut match_arm_names: Vec<(String, &dyn quote::ToTokens)> = Vec::new();
            for arm in i.arms.iter() {
                // 匹配分支
                match &arm.pat {
                    // 路徑匹配
                    syn::Pat::Path(p) => {
                        match_arm_names.push((path_to_string(&p.path), &p.path));
                    }
                    // 元組匹配
                    syn::Pat::TupleStruct(p) => {
                        match_arm_names.push((path_to_string(&p.path), &p.path));
                    }
                    // 結構匹配
                    syn::Pat::Struct(p) => {
                        match_arm_names.push((path_to_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 = std::option::Option::Some(syn::Error::new_spanned(
                            &arm.pat,
                            "unsupported by #[sorted]",
                        ));
                        return;
                    }
                }
            }
            // 檢查錯誤
            if let Some(e) = check_order(match_arm_names) {
                self.err = std::option::Option::Some(e);
                return;
            }
        }
        // 遞迴查詢
        syn::visit_mut::visit_expr_match_mut(self, i)
    }
}
// 全路徑拼接
fn path_to_string(path: &syn::Path) -> String {
    path.segments
        .iter()
        .map(|s| s.ident.to_string())
        .collect::<Vec<String>>()
        .join("::")
}

這裡主要從match的角度進行思考,匹配的場景有很多,這裡只是列舉出了基本的幾種。
雖然程式碼中基本算是TupleStruct,但是其他的情況可能也會有,算是小小的擴充。

還有一點,就是關於visit,之前我們也使用過visit,但是需要主要它的入口和接續。
我們解析這一個之後,其實限定的上下文是在一個fn裡面,因此visit_expr_match_mut

// solution5.rs
pub(crate) fn solution(fn_item: &mut syn::ItemFn) -> syn::Result<proc_macro2::TokenStream> {
    let mut visitor = crate::common::MatchVisitor {
        err: std::option::Option::None,
    };
    syn::visit_mut::visit_item_fn_mut(&mut visitor, fn_item);
    match visitor.err {
        Some(e) => syn::Result::Err(e),
        None => syn::Result::Ok(crate::common::to_token_stream(fn_item)),
    }
}

這裡可以看出來,visit的入口其實是visit_item_fn_mut
雖然都是遞迴遍歷,不過需要注意入口和接續的遍歷位置。

雖然看起來有兩個#[sorted],但是不得不打斷的是,其中有一個假貨。
if path_to_string(&attr.path) == "sorted",這裡解析的其實是一個標記,並不具備排序功能,這裡真正生效的應該是#[sorted::check],分支下面的#[sorted]就像是派生宏的惰性繫結一樣,歸屬於#[sorted::check]之下。

也就是說,我們實際上還定義了一個宏

mod common;
mod solution1;
mod solution2;
mod solution3;
mod solution5;

#[proc_macro_attribute]
pub fn sorted(
    _args: proc_macro::TokenStream,
    input: proc_macro::TokenStream,
) -> proc_macro::TokenStream {
    let item = syn::parse_macro_input!(input as syn::Item);
    match solution1::solution(&item) {
        syn::Result::Ok(stream) => stream,
        syn::Result::Err(e) => {
            let mut res = e.into_compile_error();
            res.extend(crate::common::to_token_stream(item));
            res
        }
    }
    .into()
}

#[proc_macro_attribute]
pub fn check(
    _args: proc_macro::TokenStream,
    input: proc_macro::TokenStream,
) -> proc_macro::TokenStream {
    let mut fn_item = syn::parse_macro_input!(input as syn::ItemFn);
    match solution5::solution(&mut fn_item) {
        syn::Result::Ok(stream) => stream,
        syn::Result::Err(e) => {
            let mut res = e.into_compile_error();
            res.extend(crate::common::to_token_stream(fn_item));
            res
        }
    }
    .into()
}

這裡的#[sorted::check]還能夠給我們一個提示,這裡的sorted並非是一個macro,而是一個crate,並不存在關聯或者歸屬的屬性宏,固然有惰性繫結的宏,但只不過是一個宏解析時候的標記,繫結的宏並不具備解析入口的功能。

我一直以為#[tokio::main]有什麼黑魔法,只不過是長路徑。

這裡比較容易模糊的是標記列舉#[sorted]和標記match#[sorted]的區別。
對於標記列舉的#[sorted]可以修改為#[sorted::sorted],但是標記match#[sorted]想要修改為#[sorted::sorted],需要修改#[sorted::check]的內部處理方式,因為一個是確定的宏使用不同路勁,而一個是直接更換了解析方式。

開始的時候我還搞混了這兩個,還以為check下的sorted解析不完全…

其實,到這裡,sorted已經結束了,後續娓娓道來。

本作品採用《CC 協議》,轉載必須註明作者和本文連結