Merge pull request #165 from nvzqz/master

Change derive behavior for enums with less than three variants

Author: alexcrichton

rand-derive/src/lib.rs

@@ -63,12 +63,11 @@ fn impl_rand_derive(ast: &syn::MacroInput) -> quote::Tokens {
             }
 
             let len = body.len();
-            let mut variants = body
+            let mut arms = body
                 .iter()
-                .enumerate()
-                .map(|(index, variant)| {
+                .map(|variant| {
                     let ident = &variant.ident;
-                    let arm = match variant.data {
+                    match variant.data {
                         syn::VariantData::Struct(ref body) => {
                             let fields = body
                                 .iter()
@@ -85,13 +84,24 @@ fn impl_rand_derive(ast: &syn::MacroInput) -> quote::Tokens {
                             quote! { #name::#ident (#(#fields),*) }
                         },
                         syn::VariantData::Unit => quote! { #name::#ident }
-                    };
+                    }
+                });
 
-                    quote! { #index => #arm }
-                })
-                .collect::<Vec<_>>();
-            variants.push(quote! { _ => unreachable!() });
-            quote! { match __rng.gen_range(0, #len) { #(#variants,)* } }
+            match len {
+                1 => quote! { #(#arms)* },
+                2 => {
+                    let (a, b) = (arms.next(), arms.next());
+                    quote! { if __rng.gen() { #a } else { #b } }
+                },
+                _ => {
+                    let mut variants = arms
+                        .enumerate()
+                        .map(|(index, arm)| quote! { #index => #arm })
+                        .collect::<Vec<_>>();
+                    variants.push(quote! { _ => unreachable!() });
+                    quote! { match __rng.gen_range(0, #len) { #(#variants,)* } }
+                },
+            }
         }
     };
 

rand-derive/tests/rand_macros.rs

@@ -19,7 +19,23 @@ struct Tuple(i16, Option<f64>);
 struct Unit;
 
 #[derive(Rand)]
-enum Enum {
+enum EnumUnit {
+    X,
+}
+
+#[derive(Rand)]
+enum Enum1 {
+    X(u8, f32),
+}
+
+#[derive(Rand)]
+enum Enum2 {
+    X(bool),
+    Y,
+}
+
+#[derive(Rand)]
+enum Enum3 {
     X { x: u8, y: isize },
     Y([bool; 4]),
     Z,
@@ -34,6 +50,9 @@ fn smoke() {
         let _ = rng.gen::<Struct>();
         let _ = rng.gen::<Tuple>();
         let _ = rng.gen::<Unit>();
-        let _ = rng.gen::<Enum>();
+        let _ = rng.gen::<EnumUnit>();
+        let _ = rng.gen::<Enum1>();
+        let _ = rng.gen::<Enum2>();
+        let _ = rng.gen::<Enum3>();
     }
 }