yaga_derive/
lib.rs

1extern crate proc_macro;
2
3use proc_macro::TokenStream;
4use proc_macro2::Ident;
5use proc_macro2::{Group, TokenStream as TokenStream2};
6use proc_macro_error::abort_call_site;
7use quote::quote;
8use std::collections::HashMap;
9use syn::parse::{Parse, ParseStream};
10use syn::punctuated::Punctuated;
11use syn::{
12    parse_macro_input, parse_quote, Attribute, Data, DataEnum, DataStruct, DeriveInput, Expr,
13    Fields, FieldsNamed, Type,
14};
15use syn::{FieldsUnnamed, Token, Variant};
16
17#[derive(Default)]
18struct DialogueDefs {
19    map: HashMap<String, Expr>,
20}
21
22struct DialogueDef {
23    key: Ident,
24    value: Expr,
25}
26
27#[derive(Default)]
28struct DialogueDefsParenthesized {
29    inner: DialogueDefs,
30}
31
32impl Parse for DialogueDef {
33    fn parse(input: ParseStream) -> syn::Result<Self> {
34        let key: Ident = input.parse()?;
35        let _: Token![=] = input.parse()?;
36        let value: Expr = input.parse()?;
37
38        Ok(DialogueDef { key, value })
39    }
40}
41
42impl Parse for DialogueDefs {
43    fn parse(input: ParseStream) -> syn::Result<Self> {
44        let items: Punctuated<DialogueDef, Token![,]> = Punctuated::parse_terminated(input)?;
45
46        let map = items
47            .into_iter()
48            .map(|dd| (dd.key.to_string(), dd.value))
49            .collect();
50
51        Ok(DialogueDefs { map })
52    }
53}
54
55impl Parse for DialogueDefsParenthesized {
56    fn parse(input: ParseStream) -> syn::Result<Self> {
57        let group: Group = input.parse()?;
58
59        let defs: DialogueDefs = syn::parse2(group.stream())?;
60
61        Ok(DialogueDefsParenthesized { inner: defs })
62    }
63}
64
65#[proc_macro_derive(Dialogue, attributes(dialogue))]
66pub fn derive_dialogue(input: TokenStream) -> TokenStream {
67    let DeriveInput { ident, data, .. } = parse_macro_input!(input as DeriveInput);
68
69    match data {
70        Data::Enum(DataEnum { variants, .. }) => derive_on_enum(variants.iter(), ident),
71        Data::Struct(DataStruct {
72            fields: Fields::Named(ref fields),
73            ..
74        }) => derive_on_struct(fields, ident),
75        _ => {
76            abort_call_site!("Dialogue must be either Enum or Struct")
77        }
78    }
79    .into()
80}
81
82fn derive_on_enum<'a>(variants: impl Iterator<Item = &'a Variant>, ident: Ident) -> TokenStream2 {
83    let variants: Vec<_> = variants
84        .map(|v| (&v.ident, extract_type(&v.fields), extract_prompt(&v.attrs)))
85        .collect();
86
87    let mut opts = Vec::new();
88    let mut names = Vec::new();
89
90    for (i, (field, typ, prompt)) in variants.into_iter().enumerate() {
91        opts.push(quote! {
92            #i => #ident::#field(<#typ as yaga::Dialogue>::compose(#prompt)?),
93        });
94
95        names.push(field);
96    }
97
98    let opts: TokenStream2 = opts.into_iter().collect();
99
100    quote! {
101        impl yaga::Dialogue for #ident {
102            fn compose(prompt: &str) -> std::io::Result<Self> {
103                use dialoguer::Select;
104                use dialoguer::theme::ColorfulTheme;
105
106                let selections = [#(stringify!(#names)),*];
107
108                let idx = Select::with_theme(&ColorfulTheme::default())
109                    .with_prompt(prompt)
110                    .default(0)
111                    .items(&selections[..])
112                    .interact()?;
113
114                Ok(match idx {
115                    #opts
116                    _ => unreachable!(),
117                })
118            }
119        }
120    }
121}
122
123fn extract_prompt(attributes: &Vec<Attribute>) -> Expr {
124    let defs: DialogueDefsParenthesized = match attributes
125        .iter()
126        .find(|attr| attr.path == parse_quote!(dialogue))
127    {
128        Some(attr) => syn::parse(attr.tokens.clone().into()).expect("Expected Def"),
129        None => DialogueDefsParenthesized::default(),
130    };
131
132    defs.inner.map.get("prompt").unwrap().clone()
133}
134
135fn extract_type(fields: &Fields) -> &Type {
136    match fields {
137        Fields::Unnamed(FieldsUnnamed { ref unnamed, .. }) => {
138            if unnamed.len() != 1 {
139                abort_call_site!("Only enum variants with single field are supported at a time")
140            } else {
141                let field = unnamed.first().unwrap();
142                &field.ty
143            }
144        }
145        _ => abort_call_site!("Only unnamed fields are supported for enums at the time"),
146    }
147}
148
149fn derive_on_struct(fields: &FieldsNamed, ident: Ident) -> TokenStream2 {
150    let mut acc = Vec::new();
151
152    for field in &fields.named {
153        let prompt = extract_prompt(&field.attrs);
154        let ident = field.ident.as_ref().unwrap();
155        let typ = &field.ty;
156
157        acc.push(quote! {
158            #ident: <#typ as yaga::Dialogue>::compose(#prompt)?,
159        })
160    }
161
162    let fields: TokenStream2 = acc.into_iter().collect();
163
164    quote! {
165        impl yaga::Dialogue for #ident {
166            fn compose(prompt: &str) -> std::io::Result<Self> {
167                println!("{}", prompt);
168
169                let result = #ident {
170                    #fields
171                };
172
173                Ok(result)
174            }
175        }
176    }
177}