1
use proc_macro::TokenStream;
2

            
3
use proc_macro2::Span;
4
use quote::quote;
5
use syn::punctuated::Punctuated;
6
use syn::token::Comma;
7
use syn::{
8
    ExprClosure, Ident, ItemFn, LitInt, LitStr, Result, parenthesized, parse::Parse,
9
    parse::ParseStream, parse_macro_input,
10
};
11

            
12
#[derive(Debug)]
13
struct RuleSetAndPriority {
14
    rule_set: LitStr,
15
    priority: LitInt,
16
}
17

            
18
impl Parse for RuleSetAndPriority {
19
581
    fn parse(input: ParseStream) -> Result<Self> {
20
        let content;
21
581
        parenthesized!(content in input);
22
581
        let rule_set: LitStr = content.parse()?;
23
581
        let _: Comma = content.parse()?;
24
581
        let priority: LitInt = content.parse()?;
25
581
        Ok(RuleSetAndPriority { rule_set, priority })
26
581
    }
27
}
28

            
29
#[derive(Debug)]
30
struct RegisterRuleArgs {
31
    pub rule_sets: Vec<RuleSetAndPriority>,
32
}
33

            
34
impl Parse for RegisterRuleArgs {
35
583
    fn parse(input: ParseStream) -> Result<Self> {
36
583
        let rule_sets = Punctuated::<RuleSetAndPriority, Comma>::parse_terminated(input)?;
37
583
        Ok(RegisterRuleArgs {
38
583
            rule_sets: rule_sets.into_iter().collect(),
39
583
        })
40
583
    }
41
}
42

            
43
/**
44
 * Register a rule with the given rule sets and priorities.
45
 */
46
#[proc_macro_attribute]
47
583
pub fn register_rule(arg_tokens: TokenStream, item: TokenStream) -> TokenStream {
48
583
    let func = parse_macro_input!(item as ItemFn);
49
583
    let rule_ident = &func.sig.ident;
50
583
    let static_name = format!("CONJURE_GEN_RULE_{rule_ident}").to_uppercase();
51
583
    let static_ident = Ident::new(&static_name, rule_ident.span());
52

            
53
583
    let args = parse_macro_input!(arg_tokens as RegisterRuleArgs);
54
583
    let rule_sets = args
55
583
        .rule_sets
56
583
        .iter()
57
583
        .map(|rule_set| {
58
581
            let rule_set_name = &rule_set.rule_set;
59
581
            let priority = &rule_set.priority;
60
581
            quote! {
61
                (#rule_set_name, #priority as u16)
62
            }
63
581
        })
64
583
        .collect::<Vec<_>>();
65

            
66
583
    let expanded = quote! {
67
        #func
68

            
69
        use ::conjure_cp::rule_engine::_dependencies::*; // ToDo idk if we need to explicitly do that?
70

            
71
        #[::conjure_cp::rule_engine::_dependencies::distributed_slice(::conjure_cp::rule_engine::RULES_DISTRIBUTED_SLICE)]
72
        pub static #static_ident: ::conjure_cp::rule_engine::Rule<'static> = ::conjure_cp::rule_engine::Rule {
73
            name: stringify!(#rule_ident),
74
            application: #rule_ident,
75
            rule_sets: &[#(#rule_sets),*],
76
        };
77
    };
78

            
79
583
    TokenStream::from(expanded)
80
583
}
81

            
82
52
fn parse_parenthesized<T: Parse>(input: ParseStream) -> Result<Vec<T>> {
83
    let content;
84
52
    parenthesized!(content in input);
85

            
86
52
    let mut paths = Vec::new();
87
54
    while !content.is_empty() {
88
40
        let path = content.parse()?;
89
40
        paths.push(path);
90
40
        if content.is_empty() {
91
38
            break;
92
2
        }
93
2
        content.parse::<Comma>()?;
94
    }
95

            
96
52
    Ok(paths)
97
52
}
98

            
99
struct RuleSetArgs {
100
    name: LitStr,
101
    dependencies: Vec<LitStr>,
102
    applies_fn: Option<ExprClosure>,
103
}
104

            
105
impl Parse for RuleSetArgs {
106
53
    fn parse(input: ParseStream) -> Result<Self> {
107
53
        let name = input.parse()?;
108

            
109
53
        if input.is_empty() {
110
1
            return Ok(Self {
111
1
                name,
112
1
                dependencies: Vec::new(),
113
1
                applies_fn: None,
114
1
            });
115
52
        }
116

            
117
52
        input.parse::<Comma>()?;
118
52
        let dependencies = parse_parenthesized::<LitStr>(input)?;
119

            
120
52
        if input.is_empty() {
121
24
            return Ok(Self {
122
24
                name,
123
24
                dependencies,
124
24
                applies_fn: None,
125
24
            });
126
28
        }
127

            
128
28
        input.parse::<Comma>()?;
129
28
        let applies_fn = input.parse::<ExprClosure>()?;
130

            
131
28
        Ok(Self {
132
28
            name,
133
28
            dependencies,
134
28
            applies_fn: Some(applies_fn),
135
28
        })
136
53
    }
137
}
138

            
139
/**
140
* Register a rule set with the given name, dependencies, and metadata.
141
*
142
* # Example
143
* ```rust
144
 * use conjure_cp_rule_macros::register_rule_set;
145
 * register_rule_set!("MyRuleSet", ("DependencyRuleSet", "AnotherRuleSet"));
146
* ```
147
 */
148
#[proc_macro]
149
53
pub fn register_rule_set(args: TokenStream) -> TokenStream {
150
    let RuleSetArgs {
151
53
        name,
152
53
        dependencies,
153
53
        applies_fn,
154
53
    } = parse_macro_input!(args as RuleSetArgs);
155

            
156
53
    let static_name = format!("CONJURE_GEN_RULE_SET_{}", name.value()).to_uppercase();
157
53
    let static_ident = Ident::new(&static_name, Span::call_site());
158

            
159
53
    let dependencies = quote! {
160
        #(#dependencies),*
161
    };
162

            
163
53
    let applies_to_family = match applies_fn {
164
        // Does not apply by default, e.g. only used as a dependency
165
25
        None => quote! { |_: &::conjure_cp::solver::SolverFamily| false },
166
28
        Some(func) => quote! { #func },
167
    };
168

            
169
53
    let expanded = quote! {
170
        use ::conjure_cp::rule_engine::_dependencies::*; // ToDo idk if we need to explicitly do that?
171
        #[::conjure_cp::rule_engine::_dependencies::distributed_slice(::conjure_cp::rule_engine::RULE_SETS_DISTRIBUTED_SLICE)]
172
        pub static #static_ident: ::conjure_cp::rule_engine::RuleSet<'static> =
173
            ::conjure_cp::rule_engine::RuleSet::new(#name, &[#dependencies], #applies_to_family);
174
    };
175

            
176
53
    TokenStream::from(expanded)
177
53
}