rotate_enum/
lib.rs

1//! # rotate-enum crate
2//!
3//! This crate provides simple macros that implements `prev()` and `next()` methods to an enum.
4//!
5//! ## Motivation
6//!
7//! Sometimes you define an enum like this
8//!
9//! ```
10//! enum Direction {
11//!     Up,
12//!     Left,
13//!     Down,
14//!     Right,
15//! }
16//! ```
17//!
18//! and you want to rotate them in some logic,
19//!
20//! ```
21//! # use rotate_enum::RotateEnum;
22//! # #[derive(RotateEnum, PartialEq, Clone, Copy)]
23//! # enum Direction {
24//! #     Up,
25//! #     Left,
26//! #     Down,
27//! #     Right,
28//! # }
29//! let up = Direction::Up;
30//! let left = Direction::Left;
31//! let down = Direction::Down;
32//! let right = Direction::Right;
33//!
34//! assert!(up.next() == left);
35//! assert!(left.next() == down);
36//! assert!(down.next() == right);
37//! assert!(right.next() == up);
38//!
39//! assert!(up.prev() == right);
40//! assert!(left.prev() == up);
41//! assert!(down.prev() == left);
42//! assert!(right.prev() == down);
43//! ```
44//!
45//! You can of course implement these methods manually, but it's repetitive and error prone.
46//! Don't you think it should be automated?
47//! This crate provides a `RotateEnum` derive macro to just do this.
48//!
49//! ## Shifting
50//!
51//! This crate also provides [`ShiftEnum`], which will exhaust at the end of the enum list,
52//! rather than rotating.
53//!
54//! ```
55//! # use rotate_enum::ShiftEnum;
56//! # #[derive(ShiftEnum, PartialEq, Clone, Copy)]
57//! # enum Direction {
58//! #     Up,
59//! #     Left,
60//! #     Down,
61//! #     Right,
62//! # }
63//! let up = Direction::Up;
64//! let left = Direction::Left;
65//! let down = Direction::Down;
66//! let right = Direction::Right;
67//!
68//! assert!(up.next() == Some(left));
69//! assert!(left.next() == Some(down));
70//! assert!(down.next() == Some(right));
71//! assert!(right.next() == None);
72//!
73//! assert!(up.prev() == None);
74//! assert!(left.prev() == Some(up));
75//! assert!(down.prev() == Some(left));
76//! assert!(right.prev() == Some(down));
77//! ```
78//!
79//! Note that you can only derive either one of `RotateEnum` or `ShiftEnum`, but not both, because their semantics conflict.
80//!
81//! ## Iterating
82//!
83//! This crate also provides [`IterEnum`], which will implement [`Iterator`] object
84//! that yields enum variants in sequence. The first yield result will be the same
85//! variant as the one started the iterator, i.e. `Direction::Up.iter().next() == Some(Direction::Up)`.
86//!
87//! ```
88//! # use rotate_enum::IterEnum;
89//! # #[derive(IterEnum, PartialEq, Clone, Copy, Debug)]
90//! # enum Direction {
91//! #     Up,
92//! #     Left,
93//! #     Down,
94//! #     Right,
95//! # }
96//! let up = Direction::Up;
97//! let left = Direction::Left;
98//! let down = Direction::Down;
99//! let right = Direction::Right;
100//!
101//! let mut iter = up.iter();
102//! assert!(iter.next() == Some(up));
103//! assert!(iter.next() == Some(left));
104//! assert!(iter.next() == Some(down));
105//! assert!(iter.next() == Some(right));
106//! assert!(iter.next() == None);
107//!
108//! assert_eq!(up.iter().collect::<Vec<_>>(), vec![up, left, down, right]);
109//! ```
110//!
111//! Or, you could start from `"YourEnum"Iterator::new()`.
112//!
113//! ```
114//! # use rotate_enum::IterEnum;
115//! # #[derive(IterEnum, PartialEq, Clone, Copy, Debug)]
116//! # enum Direction {
117//! #     Up,
118//! #     Left,
119//! #     Down,
120//! #     Right,
121//! # }
122//! assert_eq!(DirectionIterator::new().collect::<Vec<_>>(), vec![
123//!     Direction::Up, Direction::Left, Direction::Down, Direction::Right,
124//! ]);
125//! ```
126//!
127//! Note that it is not the same as `ShiftEnum` in the sense that the iterator is one-directional, which means you can go only forward and not `prev()`.
128//! It can also be used with iterator methods like `collect()`.
129//!
130//!
131//! `IterEnum` also requires deriving `Clone`.
132//!
133//!
134//! ## Usage
135//!
136//! Use `#[derive(...)]` macro to annotate your enum.
137//!
138//! ```rust
139//! use rotate_enum::RotateEnum;
140//!
141//! #[derive(RotateEnum)]
142//! enum Direction {
143//!     Up,
144//!     Left,
145//!     Down,
146//!     Right,
147//! }
148//! ```
149//!
150//!
151//! ## Note
152//!
153//! These macros seem trivial, but it's only possible with procedural macros!
154
155use core::panic;
156
157use proc_macro::TokenStream;
158use quote::quote;
159use syn::{parse_macro_input, Data, DeriveInput};
160
161/// This derive macro will implement `next()` and `prev()` methods that rotates
162/// the variant to the annotated enum.
163///
164/// For code examples, see [module-level docs](index.html).
165///
166/// # Requirements
167///
168/// * It must be applied to an enum. Structs are not supported or won't make sense.
169/// * Enums with any associated data are not supported.
170///
171/// # Generated methods
172///
173/// For example, this macro will implement functions like below for
174/// `enum Direction`.
175///
176/// ```
177/// # enum Direction {
178/// #     Up,
179/// #     Left,
180/// #     Down,
181/// #     Right,
182/// # }
183/// impl Direction {
184///     fn next(self) -> Self {
185///         match self {
186///             Self::Up => Self::Left,
187///             Self::Left => Self::Down,
188///             Self::Down => Self::Right,
189///             Self::Right => Self::Up,
190///         }
191///     }
192///
193///     fn prev(self) -> Self {
194///         match self {
195///             Self::Up => Self::Right,
196///             Self::Left => Self::Up,
197///             Self::Down => Self::Left,
198///             Self::Right => Self::Down,
199///         }
200///     }
201/// }
202/// ```
203#[proc_macro_derive(RotateEnum)]
204pub fn rotate_enum(input: TokenStream) -> TokenStream {
205    let input = parse_macro_input!(input as DeriveInput);
206    let name = input.ident;
207
208    let variants = if let Data::Enum(data) = &input.data {
209        data.variants.iter().collect::<Vec<_>>()
210    } else {
211        panic!("derive(RotateEnum) must be applied to an enum");
212    };
213
214    let nexts = variants
215        .iter()
216        .skip(1)
217        .chain(variants.get(0))
218        .map(|v| (&v.ident))
219        .collect::<Vec<_>>();
220
221    let tokens = quote! {
222        impl #name{
223            pub fn next(self) -> Self {
224                match self {
225                    #(Self::#variants => Self::#nexts, )*
226                }
227            }
228            pub fn prev(self) -> Self {
229                match self {
230                    #(Self::#nexts => Self::#variants, )*
231                }
232            }
233        }
234    };
235
236    tokens.into()
237}
238
239/// This derive macro will implement `next()` and `prev()` methods that shifts
240/// the variant to the annotated enum.
241///
242/// * `next()` will return `Some(Variant)` where `Variant` is next one in the enum, or `None` if it was the last variant of the enum.
243/// * `prev()` will return `Some(Variant)` where `Variant` is previous one in the enum, or `None` if it was the first variant of the enum.
244///
245/// For code examples, see [module-level docs](index.html).
246///
247/// # Requirements
248///
249/// * It must be applied to an enum. Structs are not supported or won't make sense.
250/// * Enums with any associated data are not supported.
251///
252/// # Generated methods
253///
254/// For example, this macro will implement functions like below for
255/// `enum Direction`.
256///
257/// ```
258/// # enum Direction {
259/// #     Up,
260/// #     Left,
261/// #     Down,
262/// #     Right,
263/// # }
264/// impl Direction {
265///     fn next(self) -> Option<Self> {
266///         match self {
267///             Self::Up => Some(Self::Left),
268///             Self::Left => Some(Self::Down),
269///             Self::Down => Some(Self::Right),
270///             Self::Right => None,
271///         }
272///     }
273///
274///     fn prev(self) -> Option<Self> {
275///         match self {
276///             Self::Up => None,
277///             Self::Left => Some(Self::Up),
278///             Self::Down => Some(Self::Left),
279///             Self::Right => Some(Self::Down),
280///         }
281///     }
282/// }
283/// ```
284#[proc_macro_derive(ShiftEnum)]
285pub fn shift_enum(input: TokenStream) -> TokenStream {
286    let input = parse_macro_input!(input as DeriveInput);
287    let name = input.ident;
288
289    let variants = if let Data::Enum(data) = &input.data {
290        data.variants.iter().collect::<Vec<_>>()
291    } else {
292        panic!("derive(RotateEnum) must be applied to an enum");
293    };
294
295    let nexts = variants
296        .iter()
297        .skip(1)
298        .map(|v| quote! { Some(Self::#v) })
299        .chain(Some(quote! { None }))
300        .collect::<Vec<_>>();
301
302    let none_quote = Some(quote! { None });
303    let prevs = variants
304        .iter()
305        .take(variants.len() - 1)
306        .map(|v| quote! { Some(Self::#v) })
307        .collect::<Vec<_>>();
308
309    let prevs = none_quote.iter().chain(&prevs).collect::<Vec<_>>();
310
311    let tokens = quote! {
312        impl #name{
313            pub fn next(self) -> Option<Self> {
314                match self {
315                    #(Self::#variants => #nexts, )*
316                }
317            }
318            pub fn prev(self) -> Option<Self> {
319                match self {
320                    #(Self::#variants => #prevs, )*
321                }
322            }
323        }
324    };
325
326    tokens.into()
327}
328
329/// This derive macro will implement `iter()` method to the annotated enum that sequentially
330/// yield the variant of the enum.
331///
332/// For code examples, see [module-level docs](index.html).
333///
334/// # Requirements
335///
336/// * It must be applied to an enum. Structs are not supported or won't make sense.
337/// * Enums with any associated data are not supported.
338/// * Enum also needs to derive [`Clone`].
339///
340/// # Generated methods
341///
342/// For example, this macro will implement an iterator and methods like below for
343/// `enum Direction`.
344///
345/// ```
346/// # #[derive(Clone, Debug)]
347/// # enum Direction {
348/// #     Up,
349/// #     Left,
350/// #     Down,
351/// #     Right,
352/// # }
353/// struct DirectionIterator(Option<Direction>);
354///
355/// impl Iterator for DirectionIterator {
356///     type Item = Direction;
357///     fn next(&mut self) -> Option<Self::Item> {
358///         let ret = self.0.clone();
359///         self.0 = match self.0 {
360///             Some(Direction::Up) => Some(Direction::Left),
361///             Some(Direction::Left) => Some(Direction::Down),
362///             Some(Direction::Down) => Some(Direction::Right),
363///             Some(Direction::Right) => None,
364///             None => None,
365///         };
366///         ret
367///     }
368/// }
369/// ```
370#[proc_macro_derive(IterEnum)]
371pub fn iter_enum(input: TokenStream) -> TokenStream {
372    let input = parse_macro_input!(input as DeriveInput);
373    let name = input.ident;
374
375    let variants = if let Data::Enum(data) = &input.data {
376        data.variants.iter().collect::<Vec<_>>()
377    } else {
378        panic!("derive(RotateEnum) must be applied to an enum");
379    };
380
381    let first_variant = variants
382        .first()
383        .expect("derive(IterEnum) expects at least one variant in enum");
384
385    let nexts = variants
386        .iter()
387        .skip(1)
388        .map(|v| quote! { Some(#name::#v) })
389        .chain(Some(quote! { None }))
390        .collect::<Vec<_>>();
391
392    let iterator_name = syn::Ident::new(&(name.to_string() + "Iterator"), name.span());
393
394    let tokens = quote! {
395
396        struct #iterator_name(Option<#name>);
397
398        impl #iterator_name {
399            fn new() -> Self {
400                Self(Some(#name::#first_variant))
401            }
402        }
403
404        impl Iterator for #iterator_name {
405            type Item = #name;
406            fn next(&mut self) -> Option<Self::Item> {
407                let ret = self.0.clone();
408                self.0 = match self.0 {
409                    #(Some(#name::#variants) => #nexts, )*
410                    None => None,
411                };
412                ret
413            }
414        }
415
416        impl #name {
417            fn iter(&self) -> #iterator_name {
418                #iterator_name(Some(self.clone()))
419            }
420        }
421    };
422
423    tokens.into()
424}