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

use crate::dispatcher::feature_fn_name;
use crate::target::Target;
use syn::{
    spanned::Spanned,
    visit_mut::{self, VisitMut},
    Error, Expr, ItemFn, Result,
};

struct StaticDispatchVisitor<'a> {
    target: Option<&'a Target>,
    status: Result<()>,
}

impl<'a> StaticDispatchVisitor<'a> {
    pub fn new(target: Option<&'a Target>) -> Self {
        Self {
            target,
            status: Ok(()),
        }
    }

    pub fn status(self) -> Result<()> {
        self.status
    }
}

fn dispatch_impl(expr: &mut Expr, target: Option<&Target>) -> Result<()> {
    if let Expr::Macro(macro_expr) = expr {
        if let Some(path_ident) = macro_expr.mac.path.get_ident() {
            if path_ident.to_string().as_str() == "dispatch" {
                let mut call = macro_expr.mac.parse_body::<Expr>()?;
                let ident = match &mut call {
                    Expr::Call(ref mut call) => {
                        if let Expr::Path(ref mut function) = *call.func {
                            Ok(&mut function.path.segments.last_mut().unwrap().ident)
                        } else {
                            Err(Error::new(
                                call.func.span(),
                                "dispatching a function requires a direct function call",
                            ))
                        }
                    }
                    Expr::MethodCall(call) => Ok(&mut call.method),
                    Expr::Path(ref mut path) => {
                        Ok(&mut path.path.segments.last_mut().unwrap().ident)
                    }
                    _ => Err(Error::new(
                        call.span(),
                        "expected a function or method call",
                    )),
                }?;
                *ident = feature_fn_name(ident, target).1;
                *expr = call;
            }
        }
    }
    Ok(())
}

impl VisitMut for StaticDispatchVisitor<'_> {
    fn visit_expr_mut(&mut self, i: &mut Expr) {
        if self.status.as_ref().ok().is_some() {
            if let Err(error) = dispatch_impl(i, self.target) {
                self.status = Err(error);
            }
            visit_mut::visit_expr_mut(self, i);
        }
    }
}

pub(crate) fn process_static_dispatch(item: &mut ItemFn, target: Option<&Target>) -> Result<()> {
    let mut visitor = StaticDispatchVisitor::new(target);
    visitor.visit_item_fn_mut(item);
    visitor.status()
}