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
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123

// SPDX-License-Identifier: AGPL-3.0-or-later
// Copyright (C) 2021-2022  Philipp Emanuel Weidmann <pew@worldwidemann.com>

use darling::FromMeta;
use proc_macro::TokenStream;
use quote::{format_ident, quote};
use syn::{
    parse_macro_input, punctuated::Punctuated, token::Comma, AttributeArgs, ExprArray, FnArg,
    ItemFn, Path, Type,
};

#[derive(FromMeta)]
struct Arguments {
    name: String,
    description: String,
    examples: ExprArray,
    categories: ExprArray,
}

/// Generates code required for the marked function to be usable in a function expression.
/// Function metadata is generated from the provided attribute arguments.
#[proc_macro_attribute]
pub fn function(attr: TokenStream, item: TokenStream) -> TokenStream {
    let arguments = match Arguments::from_list(&parse_macro_input!(attr as AttributeArgs)) {
        Ok(arguments) => arguments,
        Err(error) => return TokenStream::from(error.write_errors()),
    };

    let name_argument = arguments.name;
    let description_argument = arguments.description;
    let examples_argument = arguments.examples;
    let categories_argument = arguments.categories;

    let item_fn = parse_macro_input!(item as ItemFn);

    let name = &item_fn.sig.ident;
    let metadata_name = format_ident!("{}_METADATA", name.to_string().to_uppercase());
    let proxy_name = format_ident!("{}_proxy", name);

    let parameters = item_fn.sig.inputs.iter().map(|fn_arg| {
        if let FnArg::Typed(pat_type) = fn_arg {
            if let Type::Path(type_path) = &*pat_type.ty {
                match type_path.path.get_ident().unwrap().to_string().as_str() {
                    "Expression" => quote! { crate::functions::Parameter::Expression },
                    "Integer" => quote! { crate::functions::Parameter::Integer },
                    "NonNegativeInteger" => {
                        quote! { crate::functions::Parameter::NonNegativeInteger }
                    }
                    "PositiveInteger" => {
                        quote! { crate::functions::Parameter::PositiveInteger }
                    }
                    "Rational" => quote! { crate::functions::Parameter::Rational },
                    "Complex" => quote! { crate::functions::Parameter::Complex },
                    "Vector" => quote! { crate::functions::Parameter::Vector },
                    "Matrix" => quote! { crate::functions::Parameter::Matrix },
                    "SquareMatrix" => quote! { crate::functions::Parameter::SquareMatrix },
                    "bool" => quote! { crate::functions::Parameter::Boolean },
                    _ => unimplemented!(),
                }
            } else {
                unreachable!();
            }
        } else {
            unreachable!();
        }
    });

    let arguments =
        (0..item_fn.sig.inputs.len()).map(|i| quote! { arguments[#i].clone().try_into()? });

    let tokens = quote! {
        #item_fn

        pub(crate) const #metadata_name: crate::functions::Metadata = crate::functions::Metadata {
            name: #name_argument,
            description: #description_argument,
            parameters: &[#(#parameters),*],
            examples: &#examples_argument,
            categories: &#categories_argument,
        };

        pub(crate) fn #proxy_name(arguments: &[crate::expression::Expression]) ->
            ::std::result::Result<crate::expression::Expression, crate::expression::Expression> {
            ::std::result::Result::Ok(#name(#(#arguments),*).into())
        }
    };

    tokens.into()
}

/// Returns a vector of function definitions generated from the base function paths provided as arguments.
#[proc_macro]
pub fn functions(input: TokenStream) -> TokenStream {
    let mut statements = Vec::new();

    for path in parse_macro_input!(input with Punctuated::<Path, Comma>::parse_terminated) {
        let name = &path.segments.last().unwrap().ident;

        let mut metadata_path = path.clone();
        metadata_path.segments.last_mut().unwrap().ident =
            format_ident!("{}_METADATA", name.to_string().to_uppercase());

        let mut proxy_path = path.clone();
        proxy_path.segments.last_mut().unwrap().ident = format_ident!("{}_proxy", name);

        statements.push(quote! {
            functions.push(Function {
                metadata: #metadata_path,
                implementation: wrap_proxy(#metadata_path.parameters, #proxy_path),
            });
        });
    }

    let tokens = quote! {{
        let mut functions = Vec::new();

        #(#statements)*

        functions
    }};

    tokens.into()
}