Browse Source

Derive ToPrimitive for enums

Lee Bousfield 7 years ago
parent
commit
5f3a3b0004

+ 48 - 0
derive/src/lib.rs

@@ -68,3 +68,51 @@ pub fn from_primitive(input: TokenStream) -> TokenStream {
 
     res.to_string().parse().unwrap()
 }
+
+#[proc_macro_derive(ToPrimitive)]
+pub fn to_primitive(input: TokenStream) -> TokenStream {
+    let source = input.to_string();
+
+    let ast = syn::parse_macro_input(&source).unwrap();
+    let name = &ast.ident;
+
+    let variants = match ast.body {
+        Enum(ref variants) => variants,
+        _ => panic!("`ToPrimitive` can be applied only to the enums, {} is not an enum", name)
+    };
+
+    let mut idx = 0;
+    let variants: Vec<_> = variants.iter()
+        .map(|variant| {
+            let ident = &variant.ident;
+            match variant.data {
+                Unit => (),
+                _ => {
+                    panic!("`ToPrimitive` can be applied only to unitary enums, {}::{} is either struct or tuple", name, ident)
+                },
+            }
+            if let Some(val) = variant.discriminant {
+                idx = val.value;
+            }
+            let tt = quote!(#name::#ident => #idx);
+            idx += 1;
+            tt
+        })
+        .collect();
+
+    let res = quote! {
+        impl ::num::traits::ToPrimitive for #name {
+            fn to_i64(&self) -> Option<i64> {
+                self.to_u64().map(|x| x as i64)
+            }
+
+            fn to_u64(&self) -> Option<u64> {
+                Some(match *self {
+                    #(variants,)*
+                })
+            }
+        }
+    };
+
+    res.to_string().parse().unwrap()
+}

+ 0 - 0
derive/tests/compile-fail/derive_on_struct.rs → derive/tests/compile-fail/from-primitive/derive_on_struct.rs


+ 0 - 0
derive/tests/compile-fail/enum_with_associated_data.rs → derive/tests/compile-fail/from-primitive/enum_with_associated_data.rs


+ 22 - 0
derive/tests/compile-fail/to-primitive/derive_on_struct.rs

@@ -0,0 +1,22 @@
+// Copyright 2013-2015 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+extern crate num;
+#[macro_use]
+extern crate num_derive;
+
+#[derive(Debug, PartialEq, ToPrimitive)] //~ ERROR
+struct Color {
+    r: u8,
+    g: u8,
+    b: u8,
+}
+
+fn main() {}

+ 21 - 0
derive/tests/compile-fail/to-primitive/enum_with_associated_data.rs

@@ -0,0 +1,21 @@
+// Copyright 2013-2015 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+extern crate num;
+#[macro_use]
+extern crate num_derive;
+
+#[derive(Debug, PartialEq, ToPrimitive)] //~ ERROR
+enum Color {
+    Rgb(u8, u8, u8),
+    Hsv(u8, u8, u8),
+}
+
+fn main() {}

+ 1 - 1
derive/tests/empty_enum.rs

@@ -12,7 +12,7 @@ extern crate num;
 #[macro_use]
 extern crate num_derive;
 
-#[derive(Debug, PartialEq, FromPrimitive)]
+#[derive(Debug, PartialEq, FromPrimitive, ToPrimitive)]
 enum Color {}
 
 #[test]

+ 22 - 1
derive/tests/trivial.rs

@@ -12,7 +12,7 @@ extern crate num;
 #[macro_use]
 extern crate num_derive;
 
-#[derive(Debug, PartialEq, FromPrimitive)]
+#[derive(Debug, PartialEq, FromPrimitive, ToPrimitive)]
 enum Color {
     Red,
     Blue,
@@ -29,3 +29,24 @@ fn test_from_primitive_for_trivial_case() {
     assert_eq!(v,
                [Some(Color::Red), Some(Color::Blue), Some(Color::Green), None]);
 }
+
+#[test]
+fn test_to_primitive_for_trivial_case() {
+    let v: [Option<u64>; 3] = [num::ToPrimitive::to_u64(&Color::Red),
+                               num::ToPrimitive::to_u64(&Color::Blue),
+                               num::ToPrimitive::to_u64(&Color::Green)];
+
+    assert_eq!(v, [Some(0), Some(1), Some(2)]);
+}
+
+#[test]
+fn test_reflexive_for_trivial_case() {
+    let before: [u64; 3] = [0, 1, 2];
+    let after: Vec<Option<u64>> = before.iter()
+        .map(|&x| -> Option<Color> { num::FromPrimitive::from_u64(x) })
+        .map(|x| x.and_then(|x| num::ToPrimitive::to_u64(&x)))
+        .collect();
+    let before = before.into_iter().cloned().map(Some).collect::<Vec<_>>();
+
+    assert_eq!(before, after);
+}