rand_enum 0.1.0

Implement rand Distribution for an enum
Documentation 
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

#![doc = include_str!("../README.md")]
#![warn(missing_docs, clippy::all, clippy::pedantic)]

#[macro_use]
extern crate proc_macro_error;

use proc_macro::TokenStream;
use proc_macro2::Ident;
use proc_macro_crate::FoundCrate;
use syn::{spanned::Spanned, DeriveInput, LitInt};

/// Implements [`rand::distributions::Distribution`](https://rust-random.github.io/rand/rand/distributions/distribution/trait.Distribution.html) for a given enum
///
/// Information on weights is as follows. Note that this is not much, and much more information can be found at [`rand::distributions::WeightedIndex`](https://rust-random.github.io/rand/rand/distributions/struct.WeightedIndex.html)
/// - According to rand, a weight cannot be negative
/// - 0: A weight of `0` means it will never be selected
/// - 1: A weight of `1` is the default weight, and is thus redundant
/// - The higher the weight is, the more likely it will be selected
/// - If a type is specified, (i.e 1_u8), that same type must be obeyed for all the other weights
/// - By default the type is i32, as this is Rust's default number type
/// - Weights cannot all be `0`, as there will be no value chosen, ever
#[proc_macro_error]
#[proc_macro_derive(Distribution, attributes(weight))]
pub fn derive_distribute(input: TokenStream) -> TokenStream {
    let rand_crate = match proc_macro_crate::crate_name("rand") {
        Ok(found) => {
            let name = match found {
                FoundCrate::Itself => "crate".to_string(),
                FoundCrate::Name(name) => name,
            };

            Ident::new(&name, name.span())
        }
        Err(_) => abort_call_site!("could not find `rand` crate"),
    };

    let ast: DeriveInput = syn::parse(input).unwrap();
    let enum_ident = &ast.ident;

    let mut variants: Vec<(Ident, LitInt)> = vec![];

    match ast.data {
        syn::Data::Enum(data) => {
            for var in &data.variants {
                let mut variant_weight = None;
                for attr in &var.attrs {
                    match &attr.meta {
                        syn::Meta::List(meta) =>
                        {
                            #[allow(clippy::manual_let_else)]
                            if meta.path.is_ident("weight") {
                                let weight = match meta.parse_args::<LitInt>() {
                                    Ok(weight) => weight,
                                    _ => {
                                        abort!(
                                            meta.tokens.span(),
                                            "could not parse weight. expected a integer literal"
                                        );
                                    }
                                };

                                if let Ok(weight_value) = weight.base10_parse::<u128>() {
                                    if weight_value == 1 {
                                        emit_warning!(
                                            meta.tokens.span(),
                                            "weight of 1 is not recommended. this is the default value and will effectively be ignored"
                                        );
                                    }
                                }

                                variant_weight = Some(weight);
                            }
                        }
                        _ => {
                            abort!(
                                attr.span(),
                                "`#[weight]` attribute requires an integer argument"
                            )
                        }
                    }
                }

                variants.push((
                    var.ident.clone(),
                    variant_weight.unwrap_or_else(|| LitInt::new("1", "1".span())),
                ));
            }
        }
        _ => abort_call_site!("should only be derived for enums"),
    }

    let (variant, weight): (Vec<Ident>, Vec<LitInt>) = variants.into_iter().unzip();

    if weight.iter().all(|weight| {
        let parsed: u128 = weight.base10_parse().unwrap();
        parsed == 0
    }) {
        abort_call_site!("all weights are zero");
    }

    // todo!()

    quote! {
        impl #rand_crate::distributions::Distribution<Colours> for #rand_crate::distributions::Standard {
            fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> Colours {
                let mut items = vec![
                    #((#enum_ident::#variant, #weight)),*
                ];
                let weight_dist = #rand_crate::distributions::WeightedIndex::new(items.iter().map(|(_, weight)| weight)).unwrap();

                let item = items.swap_remove(weight_dist.sample(rng));

                item.0
            }
        }
    }.into()
}