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
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
//! A macro to document enum variants with the things that they are compatible with.
//!
//!
//! As well as documenting each variant, this macro also generates lists of all compatible variants
//! for each "thing".
//!
//! # Motivation
//!
//! This macro is used in Conjure-Oxide, a constraint modelling tool with support for multiple
//! backend solvers (e.g. Minion, SAT).
//!
//! The Conjure-Oxide AST is used as the singular representation for constraints models throughout
//! its crate. A consequence of this is that the AST must contain all possible supported
//! expressions for all solvers, as well as the high level Essence language it takes as input.
//! Therefore, only a small subset of AST nodes are useful for a particular solver.
//!
//! The documentation this generates helps rewrite rule implementers determine which AST nodes are
//! used for which backends by grouping AST nodes per solver.

#![allow(clippy::unwrap_used)]
#![allow(unstable_name_collisions)]

use proc_macro::TokenStream;
use std::collections::HashMap;

use itertools::Itertools;
use quote::quote;
use syn::{
    parse_macro_input, parse_quote, punctuated::Punctuated, visit_mut::VisitMut, Attribute,
    ItemEnum, Meta, Token, Variant,
};

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

struct RemoveCompatibleAttrs;
impl VisitMut for RemoveCompatibleAttrs {
    fn visit_variant_mut(&mut self, i: &mut Variant) {
        // 1. generate docstring for variant
        // Supported by: minion, sat ...
        //
        // 2. delete #[compatible] attributes

        let mut solvers: Vec<String> = vec![];
        for attr in i.attrs.iter() {
            if !attr.path().is_ident("compatible") {
                continue;
            }
            let nested = attr
                .parse_args_with(Punctuated::<Meta, Token![,]>::parse_terminated)
                .unwrap();
            for arg in nested {
                let ident = arg.path().require_ident().unwrap();
                let solver_name = ident.to_string();
                solvers.push(solver_name);
            }
        }

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

        i.attrs.retain(|attr| !attr.path().is_ident("compatible"));
    }
}

/// A macro to document enum variants by the things that they are compatible with.
///
/// # Examples
///
/// ```
/// use enum_compatability_macro::document_compatibility;
///
/// #[document_compatibility]
/// pub enum Expression {
///    #[compatible(Minion)]
///    ConstantInt(i32),
///    // ...
///    #[compatible(Chuffed)]
///    #[compatible(Minion)]
///    Sum(Vec<Expression>)
///    }
/// ```
///
/// The Expression type will have the following lists appended to its documentation:
///
///```text
/// ## Supported by `minion`
///    ConstantInt(i32)
///    Sum(Vec<Expression>)
///
///
/// ## Supported by `chuffed`
///    ConstantInt(i32)
///    Sum(Vec<Expression>)
/// ```
///
/// Two equivalent syntaxes exist for specifying supported solvers:
///
/// ```
///# use enum_compatability_macro::document_compatibility;
///#
///# #[document_compatibility]
///# pub enum Expression {
///#    #[compatible(Minion)]
///#    ConstantInt(i32),
///#    // ...
///     #[compatible(Chuffed)]
///     #[compatible(Minion)]
///     Sum(Vec<Expression>)
///#    }
/// ```
///
/// ```
///# use enum_compatability_macro::document_compatibility;
///#
///# #[document_compatibility]
///# pub enum Expression {
///#    #[compatible(Minion)]
///#    ConstantInt(i32),
///#    // ...
///     #[compatible(Minion,Chuffed)]
///     Sum(Vec<Expression>)
///#    }
/// ```
///
#[proc_macro_attribute]
pub fn document_compatibility(_attr: TokenStream, input: TokenStream) -> TokenStream {
    // Parse the input tokens into a syntax tree
    let mut input = parse_macro_input!(input as ItemEnum);
    let mut nodes_supported_by_solver: HashMap<String, Vec<syn::Ident>> = HashMap::new();

    // process each item inside the enum.
    for variant in input.variants.iter() {
        let variant_ident = variant.ident.clone();
        for attr in variant.attrs.iter() {
            if !attr.path().is_ident("compatible") {
                continue;
            }

            let nested = attr
                .parse_args_with(Punctuated::<Meta, Token![,]>::parse_terminated)
                .unwrap();
            for arg in nested {
                let ident = arg.path().require_ident().unwrap();
                let solver_name = ident.to_string();
                match nodes_supported_by_solver.get_mut(&solver_name) {
                    None => {
                        nodes_supported_by_solver.insert(solver_name, vec![variant_ident.clone()]);
                    }
                    Some(a) => {
                        a.push(variant_ident.clone());
                    }
                };
            }
        }
    }

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

    // Build the doc string.

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

        // list all the ast nodes for this solver
        for node in nodes_supported_by_solver
            .get(solver)
            .unwrap()
            .iter()
            .map(|x| x.to_string())
            .sorted()
        {
            doc_msg.push_str(&format!("* [`{}`]({}::{})\n", node, input.ident, node));
        }

        // end list
        doc_msg.push('\n');
    }

    input.attrs.push(parse_quote!(#[doc = #doc_msg]));
    let expanded = quote! {
        #input
    };

    TokenStream::from(expanded)
}