1use proc_macro::TokenStream;
2use proc_macro2::Ident;
3use quote::quote;
4use syn::{Data, DataEnum, DeriveInput, Type, parse_macro_input};
5
6#[proc_macro_derive(Screens)]
28pub fn screen(input: proc_macro::TokenStream) -> TokenStream {
29 let input = parse_macro_input!(input as DeriveInput);
30
31 match screens_derive(&input) {
32 Ok(tokens) => tokens,
33 Err(tokens) => tokens,
34 }
35}
36
37fn screens_derive(input: &DeriveInput) -> Result<TokenStream, TokenStream> {
38 let r#enum = get_enum(input)?;
39 let variants = get_screens_variants(r#enum)?;
40
41 let screen_id_tokens = generate_screen_id(&variants);
42 let screen_state_impl = generate_screen_state_impl(&input.ident, &variants);
43
44 Ok(quote! {
45 #screen_id_tokens
46
47 #screen_state_impl
48 }
49 .into())
50}
51
52fn get_enum(input: &DeriveInput) -> Result<&syn::DataEnum, proc_macro::TokenStream> {
53 match &input.data {
54 Data::Enum(data_enum) => Ok(data_enum),
55 _ => Err(quote! {
56 compile_error!("#[derive(ratapp::Screens)] can only be used on enums. Check out the ratapp documentation for more information.");
57 }
58 .into()),
59 }
60}
61
62fn get_screens_variants(input: &DataEnum) -> Result<Vec<(&Ident, &Type)>, proc_macro::TokenStream> {
63 let mut result = Vec::new();
64
65 for variant in &input.variants {
66 let name = &variant.ident;
67 let ty = match &variant.fields {
68 syn::Fields::Unnamed(fields) if fields.unnamed.len() == 1 => &fields.unnamed[0].ty,
69 _ => {
70 return Err(quote! {
71 compile_error!("#[derive(ratapp::Screens)] can only be used on enums with single unnamed field variants (i.e. `Variant(YourScreenType)`). Check out the ratapp documentation for more information.");
72 }.into());
73 }
74 };
75 result.push((name, ty));
76 }
77
78 Ok(result)
79}
80
81fn generate_screen_id(variants: &[(&Ident, &Type)]) -> proc_macro2::TokenStream {
83 let ids = variants.iter().map(|(name, _)| name);
84
85 quote! {
86 #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
87 pub enum ScreenID {
88 #(#ids),*
89 }
90 }
91}
92
93fn generate_screen_state_impl(
94 enum_name: &Ident,
95 variants: &[(&Ident, &Type)],
96) -> proc_macro2::TokenStream {
97 let where_bounds = variants.iter().map(|(_, ty)| {
98 quote! {
99 #ty : ratapp::ScreenWithState<ScreenID, S>
100 }
101 });
102
103 let match_new = variants.iter().map(|(name, ty)| {
104 quote! {
105 ScreenID::#name => #enum_name::#name(#ty::default()),
106 }
107 });
108
109 let match_draw = variants.iter().map(|(name, _)| {
110 quote! {
111 #enum_name::#name(screen) => ScreenWithState::draw(screen, frame, state),
112 }
113 });
114
115 let match_on_event = variants.iter().map(|(name, _)| {
116 quote! {
117 #enum_name::#name(screen) => ScreenWithState::on_event(screen, event, navigator, state).await,
118 }
119 });
120
121 let match_on_enter = variants.iter().map(|(name, _)| {
122 quote! {
123 #enum_name::#name(screen) => ScreenWithState::on_enter(screen, navigator, state).await,
124 }
125 });
126
127 let match_on_exit = variants.iter().map(|(name, _)| {
128 quote! {
129 #enum_name::#name(screen) => ScreenWithState::on_exit(screen, navigator, state).await,
130 }
131 });
132
133 let match_on_pause = variants.iter().map(|(name, _)| {
134 quote! {
135 #enum_name::#name(screen) => ScreenWithState::on_pause(screen, navigator, state).await,
136 }
137 });
138
139 let match_on_resume = variants.iter().map(|(name, _)| {
140 quote! {
141 #enum_name::#name(screen) => ScreenWithState::on_resume(screen, navigator, state).await,
142 }
143 });
144
145 let match_task = variants.iter().map(|(name, _)| {
146 quote! {
147 #enum_name::#name(screen) => ScreenWithState::task(screen, navigator, state).await,
148 }
149 });
150
151 let screen_state_impl = quote! {
152 impl<S> ratapp::ScreenState<S> for #enum_name
153 where
154 #( #where_bounds, )*
155 {
156 type ID = ScreenID;
157
158 fn new(id: Self::ID) -> Self {
159 match id {
160 #(#match_new)*
161 }
162 }
163
164 fn draw(&mut self, frame: &mut ratatui::Frame, state: &S) {
165 use ratapp::ScreenWithState;
166
167 match self {
168 #(#match_draw)*
169 }
170 }
171
172 async fn on_event(&mut self, event: ratatui::crossterm::event::Event, navigator: ratapp::Navigator<Self::ID>, state: &mut S) {
173 use ratapp::ScreenWithState;
174
175 match self {
176 #(#match_on_event)*
177 }
178 }
179
180 async fn on_enter(&mut self, navigator: ratapp::Navigator<Self::ID>, state: &mut S) {
181 use ratapp::ScreenWithState;
182
183 match self {
184 #(#match_on_enter)*
185 }
186 }
187
188 async fn on_exit(&mut self, navigator: ratapp::Navigator<Self::ID>, state: &mut S) {
189 use ratapp::ScreenWithState;
190
191 match self {
192 #(#match_on_exit)*
193 }
194 }
195
196 async fn on_pause(&mut self, navigator: ratapp::Navigator<Self::ID>, state: &mut S) {
197 use ratapp::ScreenWithState;
198
199 match self {
200 #(#match_on_pause)*
201 }
202 }
203
204 async fn on_resume(&mut self, navigator: ratapp::Navigator<Self::ID>, state: &mut S) {
205 use ratapp::ScreenWithState;
206
207 match self {
208 #(#match_on_resume)*
209 }
210 }
211
212 async fn task(&mut self, navigator: ratapp::Navigator<Self::ID>, state: &mut S) {
213 use ratapp::ScreenWithState;
214
215 match self {
216 #(#match_task)*
217 }
218 }
219 }
220 };
221
222 screen_state_impl
223}