safety_parser/
split_attrs.rs

1use proc_macro2::TokenStream;
2use quote::ToTokens;
3use std::mem::take;
4use syn::*;
5
6/// Since `#[safety]` is applied on unsafe function item and any expression,
7/// we have to fully parse to split attributes from them to insert our doc
8/// string to the tail of all attributes.
9pub fn split_attrs_and_rest(ts: TokenStream) -> Input {
10    if let Ok(stmt) = parse2::<Stmt>(ts.clone()) {
11        match stmt {
12            Stmt::Item(item) => parse_item(item),
13            Stmt::Expr(expr, _) => parse_expr(expr),
14            Stmt::Local(mut local) => Input::new(take(&mut local.attrs), local),
15            Stmt::Macro(mut macro_) => Input::new(take(&mut macro_.attrs), macro_),
16        }
17    } else {
18        Input::new(Vec::new(), ts)
19    }
20}
21
22fn parse_item(mut item: Item) -> Input {
23    let attrs = match &mut item {
24        Item::Fn(fun) => take(&mut fun.attrs),
25        Item::Impl(imp) => take(&mut imp.attrs),
26        Item::Trait(trait_) => take(&mut trait_.attrs),
27        _ => Vec::new(),
28    };
29    Input::new(attrs, item).set_gen_doc()
30}
31
32fn parse_expr(mut expr: Expr) -> Input {
33    let attrs = match &mut expr {
34        Expr::Unsafe(unsafe_) => take(&mut unsafe_.attrs),
35        Expr::Let(let_) => take(&mut let_.attrs),
36        Expr::Assign(assign) => take(&mut assign.attrs),
37        Expr::Block(block) => take(&mut block.attrs),
38        Expr::Call(call) => take(&mut call.attrs),
39        Expr::MethodCall(call) => take(&mut call.attrs),
40        Expr::Match(match_) => take(&mut match_.attrs),
41        Expr::If(if_) => take(&mut if_.attrs),
42        Expr::While(while_) => take(&mut while_.attrs),
43        Expr::Loop(loop_) => take(&mut loop_.attrs),
44        Expr::ForLoop(for_) => take(&mut for_.attrs),
45        Expr::Macro(macro_) => take(&mut macro_.attrs),
46        Expr::Struct(struct_) => take(&mut struct_.attrs),
47        Expr::Array(array) => take(&mut array.attrs),
48        Expr::Closure(closure) => take(&mut closure.attrs),
49        Expr::Path(path) => take(&mut path.attrs),
50        Expr::Async(async_) => take(&mut async_.attrs),
51        Expr::Await(await_) => take(&mut await_.attrs),
52        Expr::Const(const_) => take(&mut const_.attrs),
53        Expr::Break(break_) => take(&mut break_.attrs),
54        Expr::Continue(continue_) => take(&mut continue_.attrs),
55        Expr::Repeat(repeat) => take(&mut repeat.attrs),
56        Expr::Return(return_) => take(&mut return_.attrs),
57        Expr::Try(try_) => take(&mut try_.attrs),
58        Expr::TryBlock(block) => take(&mut block.attrs),
59        Expr::Yield(yield_) => take(&mut yield_.attrs),
60        _ => Vec::new(),
61    };
62    Input::new(attrs, expr)
63}
64
65#[derive(Debug)]
66pub struct Input {
67    pub attrs: TokenStream,
68    pub rest: TokenStream,
69    /// Indicate to generate `#[doc]`. True only for items.
70    pub gen_doc: bool,
71}
72
73impl Input {
74    fn new(attrs: Vec<Attribute>, rest: impl ToTokens) -> Input {
75        let attrs = attrs.into_iter().map(|attr| attr.into_token_stream()).collect();
76        Input { attrs, rest: rest.into_token_stream(), gen_doc: false }
77    }
78
79    fn set_gen_doc(mut self) -> Self {
80        self.gen_doc = true;
81        self
82    }
83}
84
85#[test]
86fn split_attrs_on_item() {
87    let ts = quote::quote! {
88        #[a]
89        unsafe fn f() {}
90    };
91    let input = dbg!(split_attrs_and_rest(ts));
92    assert!(!input.attrs.is_empty());
93}
94
95#[test]
96fn split_attrs_on_expr() {
97    let ts = quote::quote! {
98        #[a]
99        unsafe { call() };
100    };
101    dbg!(split_attrs_and_rest(ts));
102}
safety_parser/split_attrs.rs

safety_parser/
split_attrs.rs
