1
//! A macro to document enum variants with the things that they are compatible with.
2
//!
3
//!
4
//! As well as documenting each variant, this macro also generates lists of all compatible variants
5
//! for each "thing".
6
//!
7
//! # Motivation
8
//!
9
//! This macro is used in Conjure-Oxide, a constraint modelling tool with support for multiple
10
//! backend solvers (e.g. Minion, SAT).
11
//!
12
//! The Conjure-Oxide AST is used as the singular representation for constraints models throughout
13
//! its crate. A consequence of this is that the AST must contain all possible supported
14
//! expressions for all solvers, as well as the high level Essence language it takes as input.
15
//! Therefore, only a small subset of AST nodes are useful for a particular solver.
16
//!
17
//! The documentation this generates helps rewrite rule implementers determine which AST nodes are
18
//! used for which backends by grouping AST nodes per solver.
19

            
20
#![allow(clippy::unwrap_used)]
21
#![allow(unstable_name_collisions)]
22

            
23
use proc_macro::TokenStream;
24
use std::collections::HashMap;
25

            
26
use itertools::Itertools;
27
use quote::quote;
28
use syn::{
29
    Attribute, Fields, ItemEnum, Meta, Token, Variant, parse_macro_input, parse_quote,
30
    punctuated::Punctuated, visit_mut::VisitMut,
31
};
32

            
33
// A nice S.O answer that helped write the syn code :)
34
// https://stackoverflow.com/a/65182902
35

            
36
struct RemoveCompatibleAttrs;
37
impl VisitMut for RemoveCompatibleAttrs {
38
694
    fn visit_variant_mut(&mut self, i: &mut Variant) {
39
        // 1. generate docstring for variant
40
        // Supported by: minion, sat ...
41
        //
42
        // 2. delete #[compatible] attributes
43

            
44
694
        let mut solvers: Vec<String> = vec![];
45
2576
        for attr in i.attrs.iter() {
46
2576
            if !attr.path().is_ident("compatible") {
47
2056
                continue;
48
520
            }
49
520
            let nested = attr
50
520
                .parse_args_with(Punctuated::<Meta, Token![,]>::parse_terminated)
51
520
                .unwrap();
52
681
            for arg in nested {
53
681
                let ident = arg.path().require_ident().unwrap();
54
681
                let solver_name = ident.to_string();
55
681
                solvers.push(solver_name);
56
681
            }
57
        }
58

            
59
694
        if !solvers.is_empty() {
60
518
            let solver_list: String = solvers.into_iter().intersperse(", ".into()).collect();
61
518
            let doc_string: String = format!("**Supported by:** {solver_list}.\n");
62
518
            let doc_attr: Attribute = parse_quote!(#[doc = #doc_string]);
63
518
            i.attrs.push(doc_attr);
64
518
        }
65

            
66
3094
        i.attrs.retain(|attr| !attr.path().is_ident("compatible"));
67
694
    }
68
}
69

            
70
/// A macro to document enum variants by the things that they are compatible with.
71
///
72
/// # Examples
73
///
74
/// ```
75
/// use conjure_cp_enum_compatibility_macro::document_compatibility;
76
///
77
/// #[document_compatibility]
78
/// pub enum Expression {
79
///    #[compatible(Minion)]
80
///    ConstantInt(i32),
81
///    // ...
82
///    #[compatible(Chuffed)]
83
///    #[compatible(Minion)]
84
///    Sum(Vec<Expression>)
85
///    }
86
/// ```
87
///
88
/// The Expression type will have the following lists appended to its documentation:
89
///
90
///```text
91
/// ## Supported by `minion`
92
///    ConstantInt(i32)
93
///    Sum(Vec<Expression>)
94
///
95
///
96
/// ## Supported by `chuffed`
97
///    ConstantInt(i32)
98
///    Sum(Vec<Expression>)
99
/// ```
100
///
101
/// Two equivalent syntaxes exist for specifying supported solvers:
102
///
103
/// ```
104
///# use conjure_cp_enum_compatibility_macro::document_compatibility;
105
///#
106
///# #[document_compatibility]
107
///# pub enum Expression {
108
///#    #[compatible(Minion)]
109
///#    ConstantInt(i32),
110
///#    // ...
111
///     #[compatible(Chuffed)]
112
///     #[compatible(Minion)]
113
///     Sum(Vec<Expression>)
114
///#    }
115
/// ```
116
///
117
/// ```
118
///# use conjure_cp_enum_compatibility_macro::document_compatibility;
119
///#
120
///# #[document_compatibility]
121
///# pub enum Expression {
122
///#    #[compatible(Minion)]
123
///#    ConstantInt(i32),
124
///#    // ...
125
///     #[compatible(Minion,Chuffed)]
126
///     Sum(Vec<Expression>)
127
///#    }
128
/// ```
129
///
130
#[proc_macro_attribute]
131
11
pub fn document_compatibility(_attr: TokenStream, input: TokenStream) -> TokenStream {
132
    // Parse the input tokens into a syntax tree
133
11
    let mut input = parse_macro_input!(input as ItemEnum);
134
11
    let mut nodes_supported_by_solver: HashMap<String, Vec<syn::Ident>> = HashMap::new();
135

            
136
    // process each item inside the enum.
137
694
    for variant in input.variants.iter() {
138
694
        let variant_ident = variant.ident.clone();
139
2576
        for attr in variant.attrs.iter() {
140
2576
            if !attr.path().is_ident("compatible") {
141
2056
                continue;
142
520
            }
143

            
144
520
            let nested = attr
145
520
                .parse_args_with(Punctuated::<Meta, Token![,]>::parse_terminated)
146
520
                .unwrap();
147
681
            for arg in nested {
148
681
                let ident = arg.path().require_ident().unwrap();
149
681
                let solver_name = ident.to_string();
150
681
                match nodes_supported_by_solver.get_mut(&solver_name) {
151
38
                    None => {
152
38
                        nodes_supported_by_solver.insert(solver_name, vec![variant_ident.clone()]);
153
38
                    }
154
643
                    Some(a) => {
155
643
                        a.push(variant_ident.clone());
156
643
                    }
157
                };
158
            }
159
        }
160
    }
161

            
162
    // we must remove all references to #[compatible] before we finish expanding the macro,
163
    // as it does not exist outside of the context of this macro.
164
11
    RemoveCompatibleAttrs.visit_item_enum_mut(&mut input);
165

            
166
    // Build the doc string.
167

            
168
    // Note that quote wants us to build the doc message first, as it cannot interpolate doc
169
    // comments well.
170
    // https://docs.rs/quote/latest/quote/macro.quote.html#interpolating-text-inside-of-doc-comments
171
11
    let mut doc_msg: String = "# Compatability\n".into();
172
38
    for solver in nodes_supported_by_solver.keys() {
173
        // a nice title
174
38
        doc_msg.push_str(&format!("## {solver}\n"));
175

            
176
        // list all the ast nodes for this solver
177
681
        for node in nodes_supported_by_solver
178
38
            .get(solver)
179
38
            .unwrap()
180
38
            .iter()
181
681
            .map(|x| x.to_string())
182
38
            .sorted()
183
681
        {
184
681
            doc_msg.push_str(&format!("* [`{node}`]({}::{node})\n", input.ident));
185
681
        }
186

            
187
        // end list
188
38
        doc_msg.push('\n');
189
    }
190

            
191
11
    input.attrs.push(parse_quote!(#[doc = #doc_msg]));
192
11
    let expanded = quote! {
193
        #input
194
    };
195

            
196
11
    TokenStream::from(expanded)
197
11
}
198

            
199
/// Generate a `discriminant_from_name!` macro and `discriminant_from_value` function for an enum.
200
///
201
/// The generated discriminants are `1..=variant_count`, in source definition order.
202
///
203
/// # Caveat
204
///
205
/// This emits items with fixed names into the surrounding module, so using it for multiple enums
206
/// in the same module will cause name collisions.
207
#[proc_macro_attribute]
208
8
pub fn generate_discriminants(attr: TokenStream, input: TokenStream) -> TokenStream {
209
8
    let _ = parse_macro_input!(attr as syn::parse::Nothing);
210

            
211
8
    let input = parse_macro_input!(input as ItemEnum);
212
8
    let enum_ident = input.ident.clone();
213
8
    let fn_generics = input.generics.clone();
214
8
    let (_, ty_generics, where_clause) = input.generics.split_for_impl();
215
8
    let fn_vis = input.vis.clone();
216

            
217
688
    let name_arms = input.variants.iter().enumerate().map(|(index, variant)| {
218
688
        let discriminant = index + 1;
219
688
        let variant_ident = &variant.ident;
220
688
        quote! {
221
            (#variant_ident) => { #discriminant };
222
        }
223
688
    });
224

            
225
688
    let value_arms = input.variants.iter().enumerate().map(|(index, variant)| {
226
688
        let discriminant = index + 1;
227
688
        let variant_ident = &variant.ident;
228
688
        let pattern = match &variant.fields {
229
            Fields::Unit => quote!(#enum_ident::#variant_ident),
230
688
            Fields::Unnamed(_) => quote!(#enum_ident::#variant_ident(..)),
231
            Fields::Named(_) => quote!(#enum_ident::#variant_ident { .. }),
232
        };
233

            
234
688
        quote! {
235
            #pattern => #discriminant,
236
        }
237
688
    });
238

            
239
8
    let expanded = quote! {
240
        #input
241

            
242
        #[macro_export]
243
        macro_rules! discriminant_from_name {
244
            #(#name_arms)*
245
        }
246

            
247
        #[allow(dead_code)]
248
        #fn_vis fn discriminant_from_value #fn_generics (
249
            value: &#enum_ident #ty_generics
250
        ) -> usize
251
        #where_clause
252
        {
253
            match value {
254
                #(#value_arms)*
255
            }
256
        }
257
    };
258

            
259
8
    TokenStream::from(expanded)
260
8
}