yew_macro/hook/
body.rs

1use std::sync::{Arc, Mutex};
2
3use proc_macro_error::emit_error;
4use syn::spanned::Spanned;
5use syn::visit_mut::VisitMut;
6use syn::{
7    parse_quote_spanned, visit_mut, Expr, ExprCall, ExprClosure, ExprForLoop, ExprIf, ExprLoop,
8    ExprMatch, ExprWhile, Ident, Item,
9};
10
11#[derive(Debug)]
12pub struct BodyRewriter {
13    branch_lock: Arc<Mutex<()>>,
14    ctx_ident: Ident,
15}
16
17impl BodyRewriter {
18    pub fn new(ctx_ident: Ident) -> Self {
19        Self {
20            branch_lock: Arc::default(),
21            ctx_ident,
22        }
23    }
24
25    fn is_branched(&self) -> bool {
26        self.branch_lock.try_lock().is_err()
27    }
28
29    fn with_branch<F, O>(&mut self, f: F) -> O
30    where
31        F: FnOnce(&mut BodyRewriter) -> O,
32    {
33        let branch_lock = self.branch_lock.clone();
34        let _branched = branch_lock.try_lock();
35        f(self)
36    }
37}
38
39impl VisitMut for BodyRewriter {
40    fn visit_expr_call_mut(&mut self, i: &mut ExprCall) {
41        let ctx_ident = &self.ctx_ident;
42
43        // Only rewrite hook calls.
44        if let Expr::Path(ref m) = &*i.func {
45            if let Some(m) = m.path.segments.last().as_ref().map(|m| &m.ident) {
46                if m.to_string().starts_with("use_") {
47                    if self.is_branched() {
48                        emit_error!(
49                            m,
50                            "hooks cannot be called at this position.";
51                            help = "move hooks to the top-level of your function.";
52                            note = "see: https://yew.rs/docs/next/concepts/function-components/hooks"
53                        );
54                    } else {
55                        *i = parse_quote_spanned! { i.span() => ::yew::functional::Hook::run(#i, #ctx_ident) };
56                    }
57
58                    return;
59                }
60            }
61        }
62
63        visit_mut::visit_expr_call_mut(self, i);
64    }
65
66    fn visit_expr_mut(&mut self, i: &mut Expr) {
67        let ctx_ident = &self.ctx_ident;
68
69        match &mut *i {
70            Expr::Macro(m) => {
71                if let Some(ident) = m.mac.path.segments.last().as_ref().map(|m| &m.ident) {
72                    if ident.to_string().starts_with("use_") {
73                        if self.is_branched() {
74                            emit_error!(
75                                ident,
76                                "hooks cannot be called at this position.";
77                                help = "move hooks to the top-level of your function.";
78                                note = "see: https://yew.rs/docs/next/concepts/function-components/hooks"
79                            );
80                        } else {
81                            *i = parse_quote_spanned! { i.span() => ::yew::functional::Hook::run(#i, #ctx_ident) };
82                        }
83                    } else {
84                        visit_mut::visit_expr_macro_mut(self, m);
85                    }
86                }
87            }
88            _ => visit_mut::visit_expr_mut(self, i),
89        }
90    }
91
92    fn visit_expr_closure_mut(&mut self, i: &mut ExprClosure) {
93        self.with_branch(move |m| visit_mut::visit_expr_closure_mut(m, i))
94    }
95
96    fn visit_expr_if_mut(&mut self, i: &mut ExprIf) {
97        for it in &mut i.attrs {
98            visit_mut::visit_attribute_mut(self, it);
99        }
100
101        visit_mut::visit_expr_mut(self, &mut i.cond);
102
103        self.with_branch(|m| visit_mut::visit_block_mut(m, &mut i.then_branch));
104
105        if let Some(it) = &mut i.else_branch {
106            self.with_branch(|m| visit_mut::visit_expr_mut(m, &mut (it).1));
107        }
108    }
109
110    fn visit_expr_loop_mut(&mut self, i: &mut ExprLoop) {
111        self.with_branch(|m| visit_mut::visit_expr_loop_mut(m, i));
112    }
113
114    fn visit_expr_for_loop_mut(&mut self, i: &mut ExprForLoop) {
115        for it in &mut i.attrs {
116            visit_mut::visit_attribute_mut(self, it);
117        }
118        if let Some(it) = &mut i.label {
119            visit_mut::visit_label_mut(self, it);
120        }
121        visit_mut::visit_pat_mut(self, &mut i.pat);
122        visit_mut::visit_expr_mut(self, &mut i.expr);
123
124        self.with_branch(|m| visit_mut::visit_block_mut(m, &mut i.body));
125    }
126
127    fn visit_expr_match_mut(&mut self, i: &mut ExprMatch) {
128        for it in &mut i.attrs {
129            visit_mut::visit_attribute_mut(self, it);
130        }
131
132        visit_mut::visit_expr_mut(self, &mut i.expr);
133
134        self.with_branch(|m| {
135            for it in &mut i.arms {
136                visit_mut::visit_arm_mut(m, it);
137            }
138        });
139    }
140
141    fn visit_expr_while_mut(&mut self, i: &mut ExprWhile) {
142        for it in &mut i.attrs {
143            visit_mut::visit_attribute_mut(self, it);
144        }
145        if let Some(it) = &mut i.label {
146            visit_mut::visit_label_mut(self, it);
147        }
148
149        self.with_branch(|m| visit_mut::visit_expr_mut(m, &mut i.cond));
150        self.with_branch(|m| visit_mut::visit_block_mut(m, &mut i.body));
151    }
152
153    fn visit_item_mut(&mut self, _i: &mut Item) {
154        // We don't do anything for items.
155        // for components / hooks in other components / hooks, apply the attribute again.
156    }
157}
yew_macro/hook/body.rs

yew_macro/hook/
body.rs
