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}